Update train.py

train.py

@@ -1,555 +1,626 @@
-def train():
-    """
-    Chess AI (AlphaZero-style) — tối ưu cho Google Colab
-    =====================================================
-    Cách dùng trên Colab:
-      !pip install chess tinygrad
-      from google.colab import drive; drive.mount('/content/drive')
-      !python chess_ai_colab.py
-    """
-    import os
-    import time
-
-    # ==================== ENV ====================
-    os.environ["DEV"] = "CUDA"
-
-    import chess
-    import chess.pgn
-    import numpy as np
-    import pickle
-    from tinygrad import Tensor, nn
-    from tinygrad.nn.state import get_state_dict, load_state_dict, get_parameters
-    from tinygrad.nn.optim import Adam
-
-    def log(msg):
-        print(msg, flush=True)
-
-    # =============================================================
-    # CONFIG
-    # =============================================================
-    CFG = {
-        "mcts_sims":       16,
-        "buffer_maxlen":   52_000,
-        "min_buffer":      1,
-        "batch_size":      512,
-        "lr":              1e-4,
-        "train_steps":     30,
-        "save_every":      10,
-        "pgn_limit":       2000,
-        "res_blocks":      6,
-        "max_moves":       81,
-        "checkpoint_path": "dataset stockfish/latest.pkl",
-        "pgn_path":        "dataset stockfish/million_games.pgn",
-        "log_path":        "dataset stockfish/training_log.txt",
-    }
-    CFG["pgn_chunk"]      = 500
-    CFG["pgn_offset_file"] = "dataset stockfish/pgn_offset.pkl"
-    CFG["channels"]       = 256
-    TOTAL_STEPS_TARGET    = CFG["train_steps"]
-    os.makedirs(os.path.dirname(CFG["checkpoint_path"]), exist_ok=True)
-
-    # =============================================================
-    # HUGGING FACE CONFIG
-    # =============================================================
-    HF_REPO  = "bobbydeveloper2014/bybychessengine"
-    HF_TOKEN = "yourtoken"
-    HF_SYNC_FILES = {
-        "latest.pkl":       CFG["checkpoint_path"],
-        "pgn_offset.pkl":   CFG["pgn_offset_file"],
-        "training_log.txt": CFG["log_path"],
+!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
+"""
+Chess AI — AlphaZero-style
+Backend : PyTorch CUDA  (thay thế TinyGrad)
+Checkpoint : pickle  (giữ nguyên)
+Deps : torch, chess, numpy, tqdm
+"""
+
+EVAL_CACHE = {}
+CACHE_MAX  = 100_000
+
+import os, pickle, time, traceback
+import chess, chess.pgn
+import numpy as np
+from tqdm.auto import tqdm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim import Adam
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"[Device] {DEVICE}  |  "
+      f"{torch.cuda.get_device_name(0) if DEVICE.type == 'cuda' else 'CPU'}")
+
+# =============================================================
+# CONFIG
+# =============================================================
+CFG = {
+    "mcts_sims":       16,
+    "buffer_maxlen":   50_000,
+    "min_buffer":      2_000,
+    "batch_size":      512,
+    "lr":              1e-4,
+    "train_steps":     30,
+    "save_every":      1,
+    "pgn_limit":       2000,
+    "res_blocks":      6,
+    "channels":        256,
+    "max_moves":       81,
+    "checkpoint_path": "/content/drive/MyDrive/dataset stockfish/latest.pkl",
+    "pgn_path":        "/content/drive/MyDrive/dataset stockfish/million_games.pgn",
+    "log_path":        "/content/drive/MyDrive/dataset stockfish/training_log.txt",
+}
+CFG["pgn_chunk"]       = 500
+CFG["pgn_offset_file"] = "/content/drive/MyDrive/dataset stockfish/pgn_offset.pkl"
+TOTAL_STEPS_TARGET     = CFG["train_steps"]
+os.makedirs(os.path.dirname(CFG["checkpoint_path"]), exist_ok=True)
+
+# =============================================================
+# ENCODING  (NumPy — CPU)
+# =============================================================
+MCTS_BATCH = 16
+
+PIECE_IDX = {
+    chess.PAWN: 0, chess.KNIGHT: 1, chess.BISHOP: 2,
+    chess.ROOK: 3, chess.QUEEN:  4, chess.KING:   5,
+}
+PROMOTION_OFFSET = 4096
+PROMOTION_MAP    = {chess.QUEEN: 0, chess.ROOK: 1,
+                    chess.BISHOP: 2, chess.KNIGHT: 3}
+ACTION_SIZE = 4672
+
+ENCODE_CACHE = {}
+
+def encode_board(board: chess.Board) -> np.ndarray:
+    key = board.fen()
+    if key in ENCODE_CACHE:
+        return ENCODE_CACHE[key]
+    x = np.zeros((12, 8, 8), dtype=np.float32)
+    for sq, piece in board.piece_map().items():
+        row, col = divmod(sq, 8)
+        offset = 0 if piece.color == chess.WHITE else 6
+        x[PIECE_IDX[piece.piece_type] + offset][row][col] = 1.0
+    ENCODE_CACHE[key] = x
+    return x
+
+def move_to_idx(move: chess.Move) -> int:
+    if move.promotion is not None:
+        return PROMOTION_OFFSET + move.to_square * 4 + PROMOTION_MAP[move.promotion]
+    return move.from_square * 64 + move.to_square
+
+# =============================================================
+# MODEL  (PyTorch)
+# =============================================================
+class ResBlock(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+        self.c1  = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(channels)
+        self.c2  = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(channels)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = F.relu(self.bn1(self.c1(x)))
+        x = self.bn2(self.c2(x))
+        return F.relu(x + residual)
+
+
+class ChessNet(nn.Module):
+    def __init__(self, channels: int = 256, num_blocks: int = 6):
+        super().__init__()
+        self.stem_conv = nn.Conv2d(12, channels, 3, padding=1, bias=False)
+        self.stem_bn   = nn.BatchNorm2d(channels)
+        self.blocks    = nn.ModuleList([ResBlock(channels) for _ in range(num_blocks)])
+
+        # Policy head
+        self.policy_conv = nn.Conv2d(channels, 32, 1, bias=False)
+        self.policy_bn   = nn.BatchNorm2d(32)
+        self.policy_fc   = nn.Linear(32 * 8 * 8, ACTION_SIZE)
+
+        # Value head
+        self.value_conv  = nn.Conv2d(channels, 8, 1, bias=False)
+        self.value_bn    = nn.BatchNorm2d(8)
+        self.value_fc1   = nn.Linear(8 * 8 * 8, 256)
+        self.value_fc2   = nn.Linear(256, 1)
+
+    def forward(self, x: torch.Tensor):
+        x = F.relu(self.stem_bn(self.stem_conv(x)))
+        for block in self.blocks:
+            x = block(x)
+
+        p = F.relu(self.policy_bn(self.policy_conv(x)))
+        p = self.policy_fc(p.flatten(1))
+
+        v = F.relu(self.value_bn(self.value_conv(x)))
+        v = F.relu(self.value_fc1(v.flatten(1)))
+        v = torch.tanh(self.value_fc2(v))
+        return p, v
+
+
+model     = ChessNet(CFG["channels"], CFG["res_blocks"]).to(DEVICE)
+optimizer = Adam(model.parameters(), lr=CFG["lr"])
+
+# =============================================================
+# REPLAY BUFFER
+# =============================================================
+class ReplayBuffer:
+    def __init__(self, maxlen: int):
+        self.maxlen   = maxlen
+        self.states   = np.zeros((maxlen, 12, 8, 8), dtype=np.float32)
+        self.policies = np.zeros((maxlen, ACTION_SIZE), dtype=np.float32)
+        self.values   = np.zeros((maxlen, 1), dtype=np.float32)
+        self.ptr  = 0
+        self.size = 0
+
+    def extend(self, samples):
+        for state, pi, z in samples:
+            self.states[self.ptr]    = state
+            self.policies[self.ptr]  = pi
+            self.values[self.ptr, 0] = z
+            self.ptr  = (self.ptr + 1) % self.maxlen
+            self.size = min(self.size + 1, self.maxlen)
+
+    def sample_batch(self, batch_size: int):
+        idx = np.random.randint(0, self.size, batch_size)
+        return self.states[idx], self.policies[idx], self.values[idx]
+
+    def __len__(self):
+        return self.size
+
+# =============================================================
+# CHECKPOINT  (pickle — giữ nguyên format)
+# =============================================================
+def save_checkpoint(games_played: int):
+    data = {
+        "games":       games_played,
+        "total_steps": total_steps,
+        "weights":     {k: v.cpu() for k, v in model.state_dict().items()},
+        "optimizer":   optimizer.state_dict(),
     }
+    with open(CFG["checkpoint_path"], "wb") as f:
+        pickle.dump(data, f)
+    tqdm.write(f"  [CKPT] saved → games={games_played}, steps={total_steps:,}")
+
+
+def load_checkpoint():
+    path = CFG["checkpoint_path"]
+    if not os.path.exists(path):
+        tqdm.write("  [CKPT] không tìm thấy checkpoint, bắt đầu từ đầu.")
+        return 0, 0
+    with open(path, "rb") as f:
+        data = pickle.load(f)
+    model.load_state_dict(data["weights"])
+    optimizer.load_state_dict(data["optimizer"])
+    games = data.get("games", 0)
+    steps = data.get("total_steps", 0)
+    tqdm.write(f"  [CKPT] loaded → games={games}, steps={steps:,}")
+    return games, steps
+
+# =============================================================
+# PGN LOADING
+# =============================================================
+def load_pgn_offset():
+    if os.path.exists(CFG["pgn_offset_file"]):
+        with open(CFG["pgn_offset_file"], "rb") as f:
+            return pickle.load(f)
+    return 0
+
+
+def load_pgn_chunk(buf: ReplayBuffer, path: str, max_games: int = 500):
+    if not os.path.exists(path):
+        return 0, None
+    offset = load_pgn_offset()
+    count  = 0
+    with open(path, "r") as f:
+        f.seek(offset)
+        while count < max_games:
+            game = chess.pgn.read_game(f)
+            if game is None:
+                break
+            result  = game.headers.get("Result", "1/2-1/2")
+            white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0
+            board   = game.board()
+            samples = []
+            for move in game.mainline_moves():
+                pi = np.zeros(ACTION_SIZE, dtype=np.float32)
+                pi[move_to_idx(move)] = 1.0
+                z = white_z if board.turn == chess.WHITE else -white_z
+                samples.append((encode_board(board), pi, np.float32(z)))
+                board.push(move)
+            buf.extend(samples)
+            count += 1
+        new_offset = f.tell()
+    with open(CFG["pgn_offset_file"], "wb") as f:
+        pickle.dump(new_offset, f)
+    tqdm.write(f"[PGN] loaded chunk={count}, buffer={len(buf):,}")
+    return count, new_offset
+
+
+def load_pgn(buf: ReplayBuffer, path: str, limit: int = 2000):
+    if not os.path.exists(path):
+        tqdm.write(f"  [PGN] không tìm thấy {path}, bỏ qua.")
+        return
+    count = 0
+    with open(path) as f:
+        pbar = tqdm(total=limit, desc="Loading PGN")
+        while count < limit:
+            game = chess.pgn.read_game(f)
+            if game is None:
+                break
+            result  = game.headers.get("Result", "1/2-1/2")
+            white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0
+            board   = game.board()
+            samples = []
+            for move in game.mainline_moves():
+                pi = np.zeros(ACTION_SIZE, dtype=np.float32)
+                pi[move_to_idx(move)] = 1.0
+                z = white_z if board.turn == chess.WHITE else -white_z
+                samples.append((encode_board(board), pi, np.float32(z)))
+                board.push(move)
+            buf.extend(samples)
+            count += 1
+            pbar.update(1)
+        pbar.close()
+    tqdm.write(f"  [PGN] loaded {count} games, buffer={len(buf):,}")
+
+# =============================================================
+# INFERENCE  (torch.no_grad)
+# =============================================================
+def _softmax_masked(logits_np: np.ndarray, mask: np.ndarray) -> np.ndarray:
+    p = logits_np.astype(np.float64)
+    p = np.nan_to_num(p, nan=0.0, posinf=0.0, neginf=0.0)
+    p -= p.max()
+    np.exp(p, out=p)
+    p *= mask
+    total = p.sum()
+    if total > 0:
+        p /= total
+    else:
+        p = mask / mask.sum()
+    return p
+
+
+def predict(board: chess.Board):
+    fen    = board.fen()
+    cached = EVAL_CACHE.get(fen)
+    if cached is not None:
+        return cached
+
+    x = torch.from_numpy(encode_board(board)[None]).to(DEVICE)  # (1,12,8,8)
+    model.eval()
+    with torch.no_grad():
+        logits, v = model(x)
+
+    p_np     = logits[0].cpu().numpy()
+    v_scalar = float(v[0, 0].cpu().item())
+
+    mask = np.zeros(ACTION_SIZE, dtype=np.float64)
+    for m in board.legal_moves:
+        mask[move_to_idx(m)] = 1.0
+
+    p = _softmax_masked(p_np, mask)
+
+    result = (p, v_scalar)
+    if len(EVAL_CACHE) >= CACHE_MAX:
+        EVAL_CACHE.clear()
+    EVAL_CACHE[fen] = result
+    return result
+
+
+def predict_batch(boards):
+    n      = len(boards)
+    states = np.empty((n, 12, 8, 8), dtype=np.float32)
+    for i, b in enumerate(boards):
+        states[i] = encode_board(b)
+
+    x = torch.from_numpy(states).to(DEVICE)
+    model.eval()
+    with torch.no_grad():
+        logits, values = model(x)
+
+    logits_np = logits.cpu().numpy().astype(np.float64, copy=False)
+    values_np = values.cpu().numpy()
+
+    outputs = []
+    for i, board in enumerate(boards):
+        mask = np.zeros(ACTION_SIZE, dtype=np.float64)
+        for m in board.legal_moves:
+            mask[move_to_idx(m)] = 1.0
+        p = _softmax_masked(logits_np[i], mask)
+        outputs.append((p.copy(), float(values_np[i, 0])))
+    return outputs
+
+# =============================================================
+# MCTS  (CPU — giữ nguyên)
+# =============================================================
+class MCTSNode:
+    __slots__ = ("board", "prior", "parent", "children", "visits", "total_value")
+
+    def __init__(self, board: chess.Board, prior: float = 0.0, parent=None):
+        self.board       = board
+        self.prior       = prior
+        self.parent      = parent
+        self.children    = {}
+        self.visits      = 0
+        self.total_value = 0.0
+
+    @property
+    def q(self) -> float:
+        return self.total_value / self.visits if self.visits > 0 else 0.0
+
+    def uct_score(self) -> float:
+        c_puct        = 1.5
+        parent_visits = self.parent.visits if self.parent else 1
+        u = c_puct * self.prior * (parent_visits ** 0.5) / (1 + self.visits)
+        return self.q + u
+
+
+def terminal_value(board: chess.Board) -> float:
+    outcome = board.outcome()
+    if outcome is None or outcome.winner is None:
+        return 0.0
+    return -1.0
+
+
+def expand_node(node: MCTSNode, policy: np.ndarray):
+    for move in node.board.legal_moves:
+        if move not in node.children:
+            child_board = node.board.copy(stack=False)
+            child_board.push(move)
+            node.children[move] = MCTSNode(
+                child_board, prior=float(policy[move_to_idx(move)]), parent=node)
+
+
+def backup_path(path, value: float):
+    for n in reversed(path):
+        n.visits      += 1
+        n.total_value += value
+        value = -value
+
+
+def mcts_search(root: MCTSNode, num_sims: int = 64):
+    remaining = num_sims
+    while remaining > 0:
+        batch_size = min(MCTS_BATCH, remaining)
+        leaves, paths = [], []
+
+        for _ in range(batch_size):
+            node = root
+            path = [node]
+            while node.children and not node.board.is_game_over():
+                node = max(node.children.values(), key=lambda n: n.uct_score())
+                path.append(node)
+            leaves.append(node)
+            paths.append(path)
+
+        boards_to_eval, eval_indices = [], []
+        values   = [0.0] * batch_size
+        policies = [None] * batch_size
+
+        for i, node in enumerate(leaves):
+            if node.board.is_game_over():
+                values[i] = terminal_value(node.board)
+            else:
+                boards_to_eval.append(node.board)
+                eval_indices.append(i)
+
+        if boards_to_eval:
+            outputs = predict_batch(boards_to_eval)
+            for idx, (policy, value) in zip(eval_indices, outputs):
+                policies[idx] = policy
+                values[idx]   = value
+
+        for i, node in enumerate(leaves):
+            if not node.board.is_game_over():
+                policy = policies[i]
+                if node == root:
+                    legal = list(node.board.legal_moves)
+                    if legal:
+                        noise = np.random.dirichlet([0.3] * len(legal))
+                        for j, move in enumerate(legal):
+                            k = move_to_idx(move)
+                            policy[k] = 0.75 * policy[k] + 0.25 * noise[j]
+                expand_node(node, policy)
+            backup_path(paths[i], values[i])
+
+        remaining -= batch_size
+
+# =============================================================
+# SELF-PLAY
+# =============================================================
+def play_one_game(temperature: float = 1.0, game_id: int = 0):
+    board    = chess.Board()
+    history  = []
+    move_num = 0
+
+    move_bar = tqdm(
+        total=CFG["max_moves"], desc=f"  ♟ Game #{game_id}",
+        leave=False, dynamic_ncols=True)
 
-    hf_api = None
     try:
-        from huggingface_hub import HfApi, create_repo
-        hf_api = HfApi(token=HF_TOKEN)
-        create_repo(HF_REPO, token=HF_TOKEN, exist_ok=True, repo_type="model")
-        log(f"  [HF] sẵn sàng sync với {HF_REPO}")
-    except Exception as e:
-        log(f"  [HF] lỗi khởi tạo (bỏ qua): {e}")
-
-    def push_to_hf(games_played, total_steps):
-        if hf_api is None:
-            return
-        for remote_name, local_path in HF_SYNC_FILES.items():
-            if not os.path.exists(local_path):
-                continue
-            try:
-                hf_api.upload_file(
-                    path_or_fileobj=local_path,
-                    path_in_repo=remote_name,
-                    repo_id=HF_REPO,
-                    repo_type="model",
-                    commit_message=f"checkpoint @ game={games_played} steps={total_steps}",
-                )
-            except Exception as e:
-                log(f"  [HF][ERROR] upload {remote_name}: {e}")
-        log(f"  [HF] uploaded -> games={games_played}, steps={total_steps:,}")
-
-    # =============================================================
-    # ENCODING
-    # =============================================================
-    MCTS_BATCH       = 16
-    ACTION_SIZE      = 4672
-    PROMOTION_OFFSET = 4096
-    PROMOTION_MAP    = {chess.QUEEN: 0, chess.ROOK: 1, chess.BISHOP: 2, chess.KNIGHT: 3}
-    PIECE_IDX        = {chess.PAWN: 0, chess.KNIGHT: 1, chess.BISHOP: 2,
-                        chess.ROOK: 3, chess.QUEEN: 4, chess.KING: 5}
-    ENCODE_CACHE     = {}
-    EVAL_CACHE       = {}
-    CACHE_MAX        = 100_000
-
-    def encode_board(board):
-        key = board.fen()
-        if key in ENCODE_CACHE:
-            return ENCODE_CACHE[key]
-        x = np.zeros((12, 8, 8), dtype=np.float32)
-        for sq, piece in board.piece_map().items():
-            r, c = divmod(sq, 8)
-            offset = 0 if piece.color == chess.WHITE else 6
-            x[PIECE_IDX[piece.piece_type] + offset][r][c] = 1.0
-        ENCODE_CACHE[key] = x
-        return x
-
-    def move_to_idx(move):
-        if move.promotion is not None:
-            return PROMOTION_OFFSET + move.to_square * 4 + PROMOTION_MAP[move.promotion]
-        return move.from_square * 64 + move.to_square
-
-    # =============================================================
-    # MODEL
-    # =============================================================
-    class ResBlock:
-        def __init__(self, channels):
-            self.c1  = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
-            self.bn1 = nn.BatchNorm(channels)
-            self.c2  = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
-            self.bn2 = nn.BatchNorm(channels)
-        def __call__(self, x):
-            r = x
-            x = self.bn1(self.c1(x)).relu()
-            x = self.bn2(self.c2(x))
-            return (x + r).relu()
-
-    class ChessNet:
-        def __init__(self, channels=256, num_blocks=6):
-            self.stem_conv   = nn.Conv2d(12, channels, 3, padding=1, bias=False)
-            self.stem_bn     = nn.BatchNorm(channels)
-            self.blocks      = [ResBlock(channels) for _ in range(num_blocks)]
-            self.policy_conv = nn.Conv2d(channels, 32, 1, bias=False)
-            self.policy_bn   = nn.BatchNorm(32)
-            self.policy_fc   = nn.Linear(32 * 8 * 8, ACTION_SIZE)
-            self.value_conv  = nn.Conv2d(channels, 8, 1, bias=False)
-            self.value_bn    = nn.BatchNorm(8)
-            self.value_fc1   = nn.Linear(8 * 8 * 8, 256)
-            self.value_fc2   = nn.Linear(256, 1)
-        def __call__(self, x):
-            x = self.stem_bn(self.stem_conv(x)).relu()
-            for b in self.blocks: x = b(x)
-            p = self.policy_bn(self.policy_conv(x)).relu()
-            p = self.policy_fc(p.flatten(1))
-            v = self.value_bn(self.value_conv(x)).relu()
-            v = self.value_fc1(v.flatten(1)).relu()
-            v = self.value_fc2(v).tanh()
-            return p, v
-
-    model     = ChessNet(CFG["channels"], CFG["res_blocks"])
-    optimizer = Adam(get_parameters(model), lr=CFG["lr"])
-
-    # =============================================================
-    # CHECKPOINT — weights luôn lưu numpy, tương thích Metal + CUDA
-    # =============================================================
-    def _tensor_to_numpy(v):
-        """Convert bất kỳ kiểu weight nào → numpy float32."""
-        if isinstance(v, np.ndarray):
-            return v.astype(np.float32)
-        if isinstance(v, Tensor):
-            try:
-                return v.numpy().astype(np.float32)
-            except Exception:
-                pass
-        if hasattr(v, 'numpy'):
-            try:
-                return np.array(v.numpy(), dtype=np.float32)
-            except Exception:
-                pass
-        if hasattr(v, 'data'):
+        while not board.is_game_over() and move_num < CFG["max_moves"]:
+            root = MCTSNode(board.copy(stack=False))
+            mcts_search(root, CFG["mcts_sims"])
+
+            moves_list = list(root.children.keys())
+            if not moves_list:
+                break
+
+            visits = np.array(
+                [root.children[m].visits for m in moves_list], dtype=np.float32)
+            s = visits.sum()
+            if s <= 0 or np.isnan(s):
+                visits[:] = 1.0
+
+            idxs = np.array([move_to_idx(m) for m in moves_list], dtype=np.int32)
+            pi   = np.zeros(ACTION_SIZE, dtype=np.float32)
+            pi[idxs] = visits
+            s = pi.sum()
+            if s > 0 and not np.isnan(s):
+                pi /= s
+
+            history.append((encode_board(board), pi, board.turn))
+
+            if move_num < 30 and temperature > 0:
+                probs = visits ** (1.0 / temperature)
+                s = probs.sum()
+                if s <= 0 or np.isnan(s):
+                    probs = np.ones(len(visits), dtype=np.float32)
+                probs /= probs.sum()
+                move = moves_list[np.random.choice(len(moves_list), p=probs)]
+            else:
+                move = moves_list[int(np.argmax(visits))]
+
+            board.push(move)
+            move_num += 1
+            move_bar.update(1)
+            move_bar.set_postfix(
+                turn="W" if board.turn == chess.BLACK else "B",
+                move=move.uci(), sims=CFG["mcts_sims"])
+    finally:
+        move_bar.close()
+
+    result  = board.result() if board.is_game_over() else "1/2-1/2"
+    white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0
+    samples = [
+        (state, pi, np.float32(white_z if color == chess.WHITE else -white_z))
+        for state, pi, color in history
+    ]
+    return samples, result
+
+# =============================================================
+# TRAINING STEP
+# =============================================================
+def train_step(buf: ReplayBuffer):
+    if len(buf) < CFG["min_buffer"]:
+        return None
+
+    model.train()
+    train_bar = tqdm(
+        total=CFG["train_steps"], desc="  🔧 Training",
+        leave=False, dynamic_ncols=True)
+
+    loss = loss_policy = loss_value = 0.0
+
+    for _ in range(CFG["train_steps"]):
+        s, p, z = buf.sample_batch(CFG["batch_size"])
+
+        S = torch.from_numpy(np.ascontiguousarray(s)).to(DEVICE)   # (B,12,8,8)
+        P = torch.from_numpy(np.ascontiguousarray(p)).to(DEVICE)   # (B,ACTION_SIZE)
+        Z = torch.from_numpy(np.ascontiguousarray(z)).to(DEVICE)   # (B,1)
+
+        logits, v = model(S)
+
+        # Policy loss: stable cross-entropy via log_softmax
+        log_prob   = F.log_softmax(logits, dim=1)
+        loss_policy = -(P * log_prob).sum(dim=1).mean()
+
+        # Value loss: MSE
+        loss_value = F.mse_loss(v, Z)
+
+        loss = loss_policy + loss_value
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        train_bar.update(1)
+        train_bar.set_postfix(loss=f"{loss.item():.4f}")
+
+    train_bar.close()
+
+    return {
+        "loss":        loss.item(),
+        "loss_policy": loss_policy.item(),
+        "loss_value":  loss_value.item(),
+        "steps":       CFG["train_steps"],
+    }
+
+# =============================================================
+# MAIN
+# =============================================================
+if __name__ == "__main__":
+    step_since_pgn = 0
+
+    print("=" * 60)
+    print("  Chess AI — AlphaZero-style  |  PyTorch CUDA")
+    print(f"  Device : {DEVICE}")
+    print("=" * 60)
+
+    buf = ReplayBuffer(CFG["buffer_maxlen"])
+    games_played, total_steps = load_checkpoint()
+
+    if games_played == 0:
+        load_pgn(buf, CFG["pgn_path"], CFG["pgn_limit"])
+    load_pgn_chunk(buf, CFG["pgn_path"], max_games=500)
+
+    print(f"\n  Bắt đầu training từ game #{games_played + 1}")
+    print(f"  Buffer: {len(buf):,} | Min train: {CFG['min_buffer']:,}\n")
+
+    game_bar = tqdm(initial=games_played, unit="game",
+                    desc="🎮 Games", dynamic_ncols=True)
+    step_bar = tqdm(
+        total=TOTAL_STEPS_TARGET, initial=total_steps,
+        unit="step", desc="Steps", dynamic_ncols=True,
+        bar_format="{desc}: {n_fmt}/{total_fmt} |{bar}| {elapsed}<{remaining}")
+
+    while True:
+        t0 = time.time()
+        step_since_pgn += CFG["train_steps"]
+        if step_since_pgn >= CFG["pgn_chunk"]:
+            load_pgn_chunk(buf, CFG["pgn_path"], CFG["pgn_chunk"])
+            step_since_pgn = 0
+
+        # ── Self-play ────────────────────────────────────────
+        try:
+            samples, result = play_one_game(temperature=1.0, game_id=games_played + 1)
+        except Exception as e:
+            ts  = time.strftime("%Y-%m-%d %H:%M:%S")
+            msg = f"[ERROR] self-play: {e}\n{traceback.format_exc()}"
+            tqdm.write(f"[{ts}] {msg}")
             try:
-                return np.array(v.data, dtype=np.float32)
+                with open(CFG["log_path"], "a") as f:
+                    f.write(f"[{ts}] {msg}\n")
             except Exception:
                 pass
+            continue
+
+        buf.extend(samples)
+        games_played += 1
+        elapsed_game  = time.time() - t0
+
+        # ── Train ────────────────────────────────────────────
+        t1    = time.time()
+        info  = train_step(buf)
+        elapsed_train = time.time() - t1
+
+        if info:
+            step_bar.update(1)
+            total_steps += 1
+            step_bar.set_postfix(
+                progress=f"{total_steps}/{TOTAL_STEPS_TARGET}",
+                loss=f"{info['loss']:.4f}",
+                p=f"{info['loss_policy']:.4f}",
+                v=f"{info['loss_value']:.4f}",
+                t=f"{elapsed_train:.1f}s",
+            )
+
+        game_bar.update(1)
+        game_bar.set_postfix(
+            result=result, moves=len(samples),
+            buf=f"{len(buf):,}",
+            loss=f"{info['loss']:.4f}" if info else "warmup",
+            t=f"{elapsed_game:.1f}s",
+        )
+
+        # ── Log ──────────────────────────────────────────────
+        ts   = time.strftime("%Y-%m-%d %H:%M:%S")
+        line = (
+            f"game={games_played:5d} | {result} | moves={len(samples):3d} | "
+            f"buf={len(buf):,} | steps={total_steps:,} | "
+            + (f"loss={info['loss']:.4f} "
+               f"(p={info['loss_policy']:.4f} v={info['loss_value']:.4f})"
+               if info else "warmup")
+            + f" | game={elapsed_game:.1f}s train={elapsed_train:.1f}s"
+        )
         try:
-            return np.array(v, dtype=np.float32)
+            with open(CFG["log_path"], "a") as f:
+                f.write(f"[{ts}] {line}\n")
         except Exception:
-            return None
-
-    def save_checkpoint(games_played):
-        # Luôn save numpy để macOS Metal load được
-        numpy_weights = {}
-        for k, v in get_state_dict(model).items():
-            arr = _tensor_to_numpy(v)
-            if arr is not None:
-                numpy_weights[k] = arr
-
-        data = {
-            "games":       games_played,
-            "total_steps": total_steps,
-            "weights":     numpy_weights,
-        }
-        with open(CFG["checkpoint_path"], "wb") as f:
-            pickle.dump(data, f)
-        log(f"  [CKPT] saved → games={games_played}, steps={total_steps:,}")
-
-    def load_checkpoint():
-        path = CFG["checkpoint_path"]
-        if not os.path.exists(path):
-            log("  [CKPT] không tìm thấy checkpoint, bắt đầu từ đầu.")
-            return 0, 0
-
-        with open(path, "rb") as f:
-            data = pickle.load(f)
-
-        # Convert tất cả về Tensor để load vào model
-        clean = {}
-        nan_keys = []
-        for k, v in data["weights"].items():
-            arr = _tensor_to_numpy(v)
-            if arr is not None:
-                if not np.isfinite(arr).all():
-                    nan_keys.append(k)
-                clean[k] = Tensor(arr)
-        load_state_dict(model, clean)
-
-        if nan_keys:
-            log(f"  [CKPT][CẢNH BÁO] checkpoint chứa NaN/Inf ở {len(nan_keys)} tensor: {nan_keys[:5]}{'...' if len(nan_keys) > 5 else ''}")
-            log("  [CKPT] -> checkpoint có thể đã hỏng. Kiểm tra lại trước khi train tiếp.")
-
-        games = data.get("games", 0)
-        steps = data.get("total_steps", 0)
-        log(f"  [CKPT] loaded → games={games}, steps={steps:,}")
-        return games, steps
-
-    # =============================================================
-    # REPLAY BUFFER
-    # =============================================================
-    class ReplayBuffer:
-        def __init__(self, maxlen):
-            self.maxlen   = maxlen
-            self.states   = np.zeros((maxlen, 12, 8, 8), dtype=np.float32)
-            self.policies = np.zeros((maxlen, ACTION_SIZE), dtype=np.float32)
-            self.values   = np.zeros((maxlen, 1), dtype=np.float32)
-            self.ptr  = 0
-            self.size = 0
-
-        def extend(self, samples):
-            for state, pi, z in samples:
-                self.states[self.ptr]      = state
-                self.policies[self.ptr]    = pi
-                self.values[self.ptr, 0]   = z
-                self.ptr  = (self.ptr + 1) % self.maxlen
-                self.size = min(self.size + 1, self.maxlen)
-
-        def sample_batch(self, batch_size):
-            idx = np.random.randint(0, self.size, batch_size)
-            return self.states[idx], self.policies[idx], self.values[idx]
-
-        def __len__(self):
-            return self.size
-
-    # =============================================================
-    # PGN
-    # =============================================================
-    def load_offset():
-        if os.path.exists(CFG["pgn_offset_file"]):
-            with open(CFG["pgn_offset_file"], "rb") as f:
-                return pickle.load(f)
-        return 0
-
-    def load_pgn_chunk(buf, path, max_games=500):
-        if not os.path.exists(path):
-            return 0, None
-        offset = load_offset()
-        count  = 0
-        with open(path, "r") as f:
-            f.seek(offset)
-            while count < max_games:
-                game = chess.pgn.read_game(f)
-                if game is None: break
-                result  = game.headers.get("Result", "1/2-1/2")
-                white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0
-                board   = game.board()
-                samples = []
-                for move in game.mainline_moves():
-                    pi = np.zeros(ACTION_SIZE, dtype=np.float32)
-                    pi[move_to_idx(move)] = 1.0
-                    z = white_z if board.turn == chess.WHITE else -white_z
-                    samples.append((encode_board(board), pi, np.float32(z)))
-                    board.push(move)
-                buf.extend(samples)
-                count += 1
-            new_offset = f.tell()
-        with open(CFG["pgn_offset_file"], "wb") as f:
-            pickle.dump(new_offset, f)
-        log(f"[PGN] loaded chunk={count}, buffer={len(buf):,}")
-        return count, new_offset
-
-    def load_pgn(buf, path, limit=2000):
-        if not os.path.exists(path):
-            log(f"  [PGN] không tìm thấy {path}, bỏ qua.")
-            return
-        count = 0
-        with open(path) as f:
-            while count < limit:
-                game = chess.pgn.read_game(f)
-                if game is None: break
-                result  = game.headers.get("Result", "1/2-1/2")
-                white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0
-                board   = game.board()
-                samples = []
-                for move in game.mainline_moves():
-                    pi = np.zeros(ACTION_SIZE, dtype=np.float32)
-                    pi[move_to_idx(move)] = 1.0
-                    z = white_z if board.turn == chess.WHITE else -white_z
-                    samples.append((encode_board(board), pi, np.float32(z)))
-                    board.push(move)
-                buf.extend(samples)
-                count += 1
-        log(f"  [PGN] loaded {count} games, buffer={len(buf):,}")
-
-    # =============================================================
-    # INFERENCE
-    # =============================================================
-    def predict_batch(boards):
-        n      = len(boards)
-        states = np.empty((n, 12, 8, 8), dtype=np.float32)
-        for i, b in enumerate(boards):
-            states[i] = encode_board(b)
-        x        = Tensor(states)
-        p_logits, values = model(x)
-        p_logits = p_logits.numpy().astype(np.float64, copy=False)
-        values   = values.numpy()
-        outputs  = []
-        for i, board in enumerate(boards):
-            p = np.nan_to_num(p_logits[i], nan=0.0, posinf=0.0, neginf=0.0, copy=False)
-            mask = np.zeros(ACTION_SIZE, dtype=np.float64)
-            mask[[move_to_idx(m) for m in board.legal_moves]] = 1.0
-            p -= p.max(); np.exp(p, out=p); p *= mask
-            s = p.sum()
-            p = p / s if s > 0 else mask / mask.sum()
-            outputs.append((p.copy(), float(values[i, 0])))
-        return outputs
-
-    # =============================================================
-    # MCTS
-    # =============================================================
-    class MCTSNode:
-        __slots__ = ("board", "prior", "parent", "children", "visits", "total_value")
-        def __init__(self, board, prior=0.0, parent=None):
-            self.board       = board
-            self.prior       = prior
-            self.parent      = parent
-            self.children    = {}
-            self.visits      = 0
-            self.total_value = 0.0
-        @property
-        def q(self):
-            return self.total_value / self.visits if self.visits > 0 else 0.0
-        def uct_score(self):
-            pv = self.parent.visits if self.parent else 1
-            return self.q + 1.5 * self.prior * (pv ** 0.5) / (1 + self.visits)
-
-    def terminal_value(board):
-        out = board.outcome()
-        if out is None or out.winner is None: return 0.0
-        return -1.0
-
-    def expand_node(node, policy):
-        for move in node.board.legal_moves:
-            if move not in node.children:
-                cb = node.board.copy(stack=False); cb.push(move)
-                node.children[move] = MCTSNode(cb, prior=float(policy[move_to_idx(move)]), parent=node)
-
-    def backup_path(path, value):
-        for n in reversed(path):
-            n.visits += 1; n.total_value += value; value = -value
-
-    def mcts_search(root, num_sims=64):
-        remaining = num_sims
-        while remaining > 0:
-            batch_size = min(MCTS_BATCH, remaining)
-            leaves, paths = [], []
-            for _ in range(batch_size):
-                node = root; path = [node]
-                while node.children and not node.board.is_game_over():
-                    node = max(node.children.values(), key=lambda n: n.uct_score())
-                    path.append(node)
-                leaves.append(node); paths.append(path)
-
-            boards_to_eval, eval_indices = [], []
-            values   = [0.0] * batch_size
-            policies = [None] * batch_size
-            for i, node in enumerate(leaves):
-                if node.board.is_game_over(): values[i] = terminal_value(node.board)
-                else: boards_to_eval.append(node.board); eval_indices.append(i)
-
-            if boards_to_eval:
-                for idx, (pol, val) in zip(eval_indices, predict_batch(boards_to_eval)):
-                    policies[idx] = pol; values[idx] = val
-
-            for i, node in enumerate(leaves):
-                if not node.board.is_game_over():
-                    policy = policies[i]
-                    if node == root:
-                        legal = list(node.board.legal_moves)
-                        if legal:
-                            noise = np.random.dirichlet([0.3] * len(legal))
-                            for j, move in enumerate(legal):
-                                k = move_to_idx(move)
-                                policy[k] = 0.75 * policy[k] + 0.25 * noise[j]
-                    expand_node(node, policy)
-                backup_path(paths[i], values[i])
-            remaining -= batch_size
-
-    # =============================================================
-    # SELF-PLAY
-    # =============================================================
-    def play_one_game(temperature=1.0, game_id=0):
-        board    = chess.Board()
-        history  = []
-        move_num = 0
-        Tensor.training = False
-        try:
-            while not board.is_game_over() and move_num < CFG["max_moves"]:
-                root = MCTSNode(board.copy(stack=False))
-                mcts_search(root, CFG["mcts_sims"])
-                moves_list = list(root.children.keys())
-                if not moves_list: break
-                visits = np.array([root.children[m].visits for m in moves_list], dtype=np.float32)
-                s = visits.sum()
-                if s <= 0 or np.isnan(s): visits[:] = 1.0
-                idxs = np.array([move_to_idx(m) for m in moves_list], dtype=np.int32)
-                pi = np.zeros(ACTION_SIZE, dtype=np.float32); pi[idxs] = visits
-                s = pi.sum()
-                if s > 0 and not np.isnan(s): pi /= s
-                history.append((encode_board(board), pi, board.turn))
-                if move_num < 30 and temperature > 0:
-                    probs = visits ** (1.0 / temperature)
-                    s = probs.sum()
-                    if s <= 0 or np.isnan(s): probs = np.ones(len(visits), dtype=np.float32)
-                    probs /= probs.sum()
-                    move = moves_list[np.random.choice(len(moves_list), p=probs)]
-                else:
-                    move = moves_list[int(np.argmax(visits))]
-                board.push(move); move_num += 1
-        finally:
-            Tensor.training = True
-
-        result  = board.result() if board.is_game_over() else "1/2-1/2"
-        white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0
-        samples = [(s, p, np.float32(white_z if c == chess.WHITE else -white_z)) for s, p, c in history]
-        return samples, result
-
-    # =============================================================
-    # TRAINING STEP
-    # =============================================================
-    def train_step(buf):
-        if len(buf) < CFG["min_buffer"]:
-            return None
-        loss_policy = loss_value = loss = 0.0
-        skipped = 0
-        for _ in range(CFG["train_steps"]):
-            s, p, z = buf.sample_batch(CFG["batch_size"])
-            S = Tensor(np.ascontiguousarray(s))
-            P = Tensor(np.ascontiguousarray(p))
-            Z = Tensor(np.ascontiguousarray(z))
-            logits, v   = model(S)
-            # clamp logits trước softmax để tránh exp() overflow -> inf
-            logits      = logits.clip(-30, 30)
-            logits      = logits - logits.max(axis=1).reshape(-1, 1)
-            exp_l       = logits.exp()
-            prob        = exp_l / exp_l.sum(axis=1).reshape(-1, 1)
-            loss_policy = -(P * (prob + 1e-8).log()).sum(axis=1).mean()
-            loss_value  = (v - Z).square().mean()
-            loss        = loss_policy + loss_value
-
-            loss_val = float(loss.item())
-            if not np.isfinite(loss_val):
-                # loss NaN/Inf -> bỏ qua step này, KHÔNG cập nhật weights
-                skipped += 1
-                log(f"  [WARN] loss={loss_val} (NaN/Inf) -> skip optimizer.step()")
-                continue
-
-            optimizer.zero_grad(); loss.backward()
-            # gradient clipping: chặn gradient explosion gây NaN lan vào weights
-            for t in get_parameters(model):
-                if t.grad is not None:
-                    t.grad = t.grad.clip(-1.0, 1.0)
-            optimizer.step()
-        return {"loss": float(loss.item()), "loss_policy": float(loss_policy.item()),
-                "loss_value": float(loss_value.item()), "steps": CFG["train_steps"]}
-
-    # =============================================================
-    # MAIN
-    # =============================================================
-    if __name__ == "__main__":
-        step_since_pgn = 0
-        print("=" * 60)
-        print("  Chess AI — AlphaZero-style, tối ưu Colab")
-        print("=" * 60)
-        buf = ReplayBuffer(CFG["buffer_maxlen"])
-        games_played, total_steps = load_checkpoint()
-        if games_played == 0:
-            load_pgn(buf, CFG["pgn_path"], CFG["pgn_limit"])
-        load_pgn_chunk(buf, CFG["pgn_path"], max_games=500)
-        print(f"\n  Bắt đầu training từ game #{games_played + 1}")
-        print(f"  Buffer: {len(buf):,} | Min train: {CFG['min_buffer']:,}\n")
-
-        while True:
-            t0 = time.time()
-            step_since_pgn += CFG["train_steps"]
-            if step_since_pgn >= CFG["pgn_chunk"]:
-                load_pgn_chunk(buf, CFG["pgn_path"], CFG["pgn_chunk"]); step_since_pgn = 0
-
-            try:
-                samples, result = play_one_game(temperature=1.0, game_id=games_played + 1)
-            except Exception as e:
-                import traceback
-                ts = time.strftime("%Y-%m-%d %H:%M:%S")
-                msg = f"[ERROR] self-play: {e}\n{traceback.format_exc()}"
-                log(f"[{ts}] {msg}")
-                try:
-                    with open(CFG["log_path"], "a") as f: f.write(f"[{ts}] {msg}\n")
-                except Exception: pass
-                continue
-
-            buf.extend(samples); games_played += 1
-            elapsed_game = time.time() - t0
-            t1   = time.time()
-            info = train_step(buf)
-            elapsed_train = time.time() - t1
-
-            if info:
-                total_steps += 1
-                log(f"  Steps: {total_steps}/{TOTAL_STEPS_TARGET} "
-                    f"loss={info['loss']:.4f} p={info['loss_policy']:.4f} v={info['loss_value']:.4f} "
-                    f"t={elapsed_train:.1f}s")
-
-            loss_str = f"{info['loss']:.4f}" if info else "warmup"
-            log(f"🎮 Game {games_played}: result={result} moves={len(samples)} "
-                f"buf={len(buf):,} loss={loss_str} t={elapsed_game:.1f}s")
-
-            ts   = time.strftime("%Y-%m-%d %H:%M:%S")
-            line = (f"game={games_played:5d} | {result} | moves={len(samples):3d} | buf={len(buf):,} | "
-                    f"steps={total_steps:,} | "
-                    + (f"loss={info['loss']:.4f} (p={info['loss_policy']:.4f} v={info['loss_value']:.4f})" if info else "warmup")
-                    + f" | game={elapsed_game:.1f}s train={elapsed_train:.1f}s")
-            try:
-                with open(CFG["log_path"], "a") as f: f.write(f"[{ts}] {line}\n")
-            except Exception: pass
-
-            if games_played % CFG["save_every"] == 0:
-                save_checkpoint(games_played)
-    return push_to_hf(games_played, total_steps)
-
+            pass
 
-train()
+        if games_played % CFG["save_every"] == 0:
+            save_checkpoint(games_played)