Upload 5 files

best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa1c91b6e90fa7fe706b664c5ca0a6dc1533e5d48346e62b87d91da51cc4ad7a
+size 348079598

common_chess.py

@@ -0,0 +1,124 @@
+"""
+Common chess encoding utilities shared across training, evaluation, and UCI engine.
+
+Must match the encoding used in `scripts/generate_data/label_stockfish_dataset.py`.
+"""
+import chess
+import numpy as np
+
+NUM_ACTIONS = 64 * 64 * 5  # 20480: canonical 64×64 squares × 5 promotion types
+NUM_BOARD_PLANES = 18      # 6 own pieces + 6 opponent + 4 castling + 1 ep + 1 side
+
+
+def encode_board(board: chess.Board) -> np.ndarray:
+    """Encode a chess.Board into 18 canonical planes, shape (18, 8, 8), uint8.
+
+    If black is to move, squares are mirrored vertically so that the side
+    to move is always at the bottom (ranks 0-3 own, ranks 4-7 opponent).
+
+    Planes:
+        0-5  : own pawn, knight, bishop, rook, queen, king
+        6-11 : opponent pawn, knight, bishop, rook, queen, king
+        12   : own kingside  castling right (all-1 plane)
+        13   : own queenside castling right (all-1 plane)
+        14   : opponent kingside  castling right (all-1 plane)
+        15   : opponent queenside castling right (all-1 plane)
+        16   : en-passant target square (1-hot)
+        17   : original side-to-move (1 = white, 0 = black)
+    """
+    planes = np.zeros((18, 8, 8), dtype=np.uint8)
+    turn = board.turn
+    mirror = not turn  # mirror squares if black to move
+
+    # Piece planes (own = 0-5, opponent = 6-11)
+    for sq in chess.SQUARES:
+        piece = board.piece_at(sq)
+        if piece is None:
+            continue
+        csq = chess.square_mirror(sq) if mirror else sq
+        row, col = divmod(csq, 8)
+        if piece.color == turn:
+            idx = piece.piece_type - 1          # 0-5 own
+        else:
+            idx = piece.piece_type - 1 + 6      # 6-11 opponent
+        planes[idx, row, col] = 1
+
+    # Castling rights (own / opponent perspective)
+    if board.has_kingside_castling_rights(turn):
+        planes[12, :, :] = 1
+    if board.has_queenside_castling_rights(turn):
+        planes[13, :, :] = 1
+    if board.has_kingside_castling_rights(not turn):
+        planes[14, :, :] = 1
+    if board.has_queenside_castling_rights(not turn):
+        planes[15, :, :] = 1
+
+    # En-passant
+    ep = board.ep_square
+    if ep is not None:
+        cep = chess.square_mirror(ep) if mirror else ep
+        row, col = divmod(cep, 8)
+        planes[16, row, col] = 1
+
+    # Original side-to-move indicator
+    if turn == chess.WHITE:
+        planes[17, :, :] = 1
+
+    return planes
+
+
+def move_to_action_id(move: chess.Move, turn: chess.Color) -> int:
+    """Convert a chess.Move to canonical action id [0, 20480).
+
+    If black to move, from_square and to_square are mirrored vertically
+    so the encoding is invariant under board orientation.
+    """
+    if turn == chess.BLACK:
+        from_sq = chess.square_mirror(move.from_square)
+        to_sq = chess.square_mirror(move.to_square)
+    else:
+        from_sq = move.from_square
+        to_sq = move.to_square
+
+    promo = move.promotion
+    if promo is None:
+        pid = 0
+    elif promo == chess.QUEEN:
+        pid = 1
+    elif promo == chess.ROOK:
+        pid = 2
+    elif promo == chess.BISHOP:
+        pid = 3
+    elif promo == chess.KNIGHT:
+        pid = 4
+    else:
+        pid = 0  # should never happen
+
+    return (from_sq * 64 + to_sq) * 5 + pid
+
+
+def action_id_to_move(action_id: int, turn: chess.Color) -> chess.Move:
+    """Inverse of move_to_action_id."""
+    pid = action_id % 5
+    raw = action_id // 5
+    to_sq = raw % 64
+    from_sq = raw // 64
+
+    if turn == chess.BLACK:
+        from_sq = chess.square_mirror(from_sq)
+        to_sq = chess.square_mirror(to_sq)
+
+    promo_map = {0: None, 1: chess.QUEEN, 2: chess.ROOK,
+                 3: chess.BISHOP, 4: chess.KNIGHT}
+    return chess.Move(from_sq, to_sq, promotion=promo_map[pid])
+
+
+def legal_action_ids(board: chess.Board):
+    """Return (action_ids, moves) for all legal moves on *board*.
+
+    action_ids: list of int length len(moves)
+    moves: list of chess.Move
+    """
+    moves = list(board.legal_moves)
+    action_ids = [move_to_action_id(m, board.turn) for m in moves]
+    return action_ids, moves

engine.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+cd "$(dirname "$0")"
+exec /root/miniconda3/envs/lichessbot/bin/python -u uci_engine.py --ckpt best.pt --device cpu

model.py

@@ -0,0 +1,176 @@
+"""
+ChessResNet: a ResNet-style policy-value network for chess.
+
+Architecture:
+  Stem: Conv3x3 18→channels, GroupNorm, GELU
+  Tower: N residual blocks (Conv3x3→GN→GELU→Conv3x3→GN→+→GELU)
+  Policy head: spatial Conv1x1 → 320 channels → reshape to [B, 20480]
+  Value head:  Conv1x1 → 32 → Flatten → Linear 256 → Linear 1 → tanh
+
+Default config (channels=256, blocks=24) yields ~29M parameters.
+"""
+import math
+from typing import Optional
+
+import torch
+from torch import nn
+
+
+class ResidualBlock(nn.Module):
+    """Pre-activation residual block with GroupNorm."""
+
+    def __init__(self, channels: int, norm_groups: int = 32):
+        super().__init__()
+        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
+        self.norm1 = nn.GroupNorm(norm_groups, channels)
+        self.act1 = nn.GELU()
+        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
+        self.norm2 = nn.GroupNorm(norm_groups, channels)
+        self.act2 = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        out = self.conv1(x)
+        out = self.norm1(out)
+        out = self.act1(out)
+        out = self.conv2(out)
+        out = self.norm2(out)
+        out = out + residual
+        out = self.act2(out)
+        return out
+
+
+class ChessResNet(nn.Module):
+    """Policy-value ResNet for chess.
+
+    Args:
+        channels: Number of filters in the residual tower (default: 256).
+        blocks:   Number of residual blocks (default: 20).
+        num_actions: Size of action space (default: 20480).
+        norm_groups: Number of groups for GroupNorm (default: 32).
+    """
+
+    def __init__(
+        self,
+        channels: int = 256,
+        blocks: int = 24,
+        num_actions: int = 20480,
+        norm_groups: int = 32,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.blocks = blocks
+        self.num_actions = num_actions
+
+        # ---- Stem ----
+        self.stem = nn.Sequential(
+            nn.Conv2d(18, channels, 3, padding=1, bias=False),
+            nn.GroupNorm(norm_groups, channels),
+            nn.GELU(),
+        )
+
+        # ---- Residual tower ----
+        tower = []
+        for _ in range(blocks):
+            tower.append(ResidualBlock(channels, norm_groups))
+        self.tower = nn.Sequential(*tower)
+
+        # ---- Policy head (spatial): linear logits, no activation ----
+        # 320 = 64 destination squares × 5 promotion types
+        self.policy_head = nn.Conv2d(channels, 320, 1, bias=True)
+
+        # ---- Value head ----
+        self.value_head = nn.Sequential(
+            nn.Conv2d(channels, 32, 1, bias=False),
+            nn.GroupNorm(8, 32),
+            nn.GELU(),
+            nn.Flatten(),
+            nn.Linear(32 * 8 * 8, 256),
+            nn.GELU(),
+            nn.Linear(256, 1),
+            nn.Tanh(),
+        )
+
+        self._init_weights()
+
+    def _init_weights(self):
+        """Initialize weights with scaled normal for stability."""
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out',
+                                        nonlinearity='relu')
+            elif isinstance(m, nn.Linear):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+    def forward(self, boards: torch.Tensor):
+        """
+        Args:
+            boards: [B, 18, 8, 8] float tensor (values 0.0 or 1.0)
+
+        Returns:
+            policy_logits: [B, 20480]  raw logits for all action IDs
+            value:         [B]         tanh-squashed scalar [-1, 1]
+        """
+        x = self.stem(boards)
+        x = self.tower(x)
+
+        # Policy head: [B, 320, 8, 8] (NCHW) -> [B, 20480]
+        # Action ID encoding:
+        #   action_id = ((from_sq * 64) + to_sq) * 5 + promo_id
+        # Spatial (h,w) = from_square = h*8 + w
+        # Channel c = to_sq * 5 + promo_id  (0-319)
+        pol = self.policy_head(x)               # [B, 320, 8, 8] NCHW
+        B = pol.shape[0]
+        pol = pol.permute(0, 2, 3, 1)           # [B, 8, 8, 320] NHWC
+        policy_logits = pol.reshape(B, 8 * 8 * 320)  # [B, 20480]
+        # After permute+reshape:
+        #   flat_idx = (h*8+w) * 320 + c
+        #            = from_sq * 320 + to_sq * 5 + promo_id
+        #            = ((from_sq * 64) + to_sq) * 5 + promo_id  ✓
+
+        # Value head
+        value = self.value_head(x).squeeze(-1)  # [B]
+
+        return policy_logits, value
+
+
+def count_parameters(model: nn.Module) -> int:
+    """Return total number of trainable parameters."""
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+def get_model_config(model: ChessResNet) -> dict:
+    """Return model hyperparameters for checkpoint saving."""
+    return dict(
+        channels=model.channels,
+        blocks=model.blocks,
+        num_actions=model.num_actions,
+    )
+
+
+def create_model_from_config(config: dict) -> ChessResNet:
+    """Create a model from a config dict (as stored in checkpoints)."""
+    return ChessResNet(
+        channels=config.get("channels", 256),
+        blocks=config.get("blocks", 24),
+        num_actions=config.get("num_actions", 20480),
+    )
+
+
+if __name__ == "__main__":
+    m = ChessResNet(channels=256, blocks=24)
+    n_params = count_parameters(m)
+    print(f"ChessResNet(channels=256, blocks=24): {n_params:,} parameters")
+    # ~29.0M expected
+
+    m = ChessResNet(channels=128, blocks=4)
+    n_params = count_parameters(m)
+    print(f"ChessResNet(channels=128, blocks=4):  {n_params:,} parameters")
+
+    # Test forward
+    x = torch.randn(4, 18, 8, 8)
+    pol, val = m(x)
+    print(f"Policy logits shape: {pol.shape}  (expected [4, 20480])")
+    print(f"Value shape:         {val.shape}  (expected [4])")

uci_engine.py

@@ -0,0 +1,267 @@
+"""
+Minimal UCI engine wrapper for the trained ChessResNet model.
+
+Can be used by chess GUIs (Arena, cutechess, En Croissant, etc.) or
+test harnesses that speak the UCI protocol.
+
+Usage:
+    python uci_engine.py --ckpt runs/stage1_stockfish_30m/best.pt --device cuda
+    python uci_engine.py --ckpt runs/stage1_stockfish_30m/best.pt --device cpu
+"""
+import argparse
+import sys
+from pathlib import Path
+
+import chess
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from common_chess import encode_board, move_to_action_id, legal_action_ids
+from model import ChessResNet, create_model_from_config
+
+# ── Draw-aware move selection helpers ──────────────────────────────────────
+
+PIECE_VALUES = {
+    chess.PAWN: 100,
+    chess.KNIGHT: 320,
+    chess.BISHOP: 330,
+    chess.ROOK: 500,
+    chess.QUEEN: 900,
+    chess.KING: 0,
+}
+
+
+def material_score_for_side(board: chess.Board, side: chess.Color) -> int:
+    """Return *side*'s material advantage in centipawns (positive = side ahead)."""
+    score = 0
+    for piece_type in chess.PIECE_TYPES:
+        value = PIECE_VALUES[piece_type]
+        score += len(board.pieces(piece_type, side)) * value
+        score -= len(board.pieces(piece_type, not side)) * value
+    return score
+
+
+def move_causes_drawish(board: chess.Board, move: chess.Move) -> bool:
+    """Check whether *move* immediately leads to a drawish outcome."""
+    b = board.copy(stack=True)
+    b.push(move)
+    if b.is_repetition(3):
+        return True
+    if b.can_claim_threefold_repetition():
+        return True
+    if b.is_fifty_moves():
+        return True
+    if b.can_claim_fifty_moves():
+        return True
+    if b.is_stalemate():
+        return True
+    if b.is_insufficient_material():
+        return True
+    return False
+
+
+def choose_move_with_draw_awareness(
+    board: chess.Board,
+    legal_moves: list,
+    legal_action_ids: list[int],
+    policy_logits: torch.Tensor,
+    value_pred,
+    topk: int = 12,
+) -> chess.Move:
+    """Draw-aware top-k rerank of legal moves.
+
+    Parameters
+    ----------
+    board : current python-chess Board (must retain move stack).
+    legal_moves : list of chess.Move, parallel to *legal_action_ids*.
+    legal_action_ids : list of int action IDs parallel to *legal_moves*.
+    policy_logits : full [N_ACTIONS] torch.Tensor (on any device).
+    value_pred : scalar value-head output (can be tensor or float).
+    topk : number of top candidates to consider for reranking.
+
+    Returns a legal ``chess.Move``.
+    """
+    side = board.turn
+    root_value = float(value_pred.squeeze().item() if hasattr(value_pred, "item") else value_pred)
+    material = material_score_for_side(board, side)
+
+    # Material fallback: override value signal when material gap is large
+    if material >= 500:
+        root_value = max(root_value, 0.50)
+    if material <= -500:
+        root_value = min(root_value, -0.50)
+
+    # Sort legal moves by policy logit descending
+    scored = [
+        (float(policy_logits[aid].item() if hasattr(policy_logits, "item") else policy_logits[aid]), move)
+        for aid, move in zip(legal_action_ids, legal_moves)
+    ]
+    scored.sort(key=lambda x: x[0], reverse=True)
+    sorted_moves = [m for _, m in scored]
+    original_top1 = sorted_moves[0]
+
+    # ── Advantage: avoid draws ────────────────────────────────────────
+    if root_value > 0.35:
+        for move in sorted_moves[:topk]:
+            if not move_causes_drawish(board, move):
+                return move
+        return original_top1  # fallback – every top-k move is drawish
+
+    # ── Disadvantage: prefer draws ────────────────────────────────────
+    if root_value < -0.35:
+        for move in sorted_moves[:topk]:
+            if move_causes_drawish(board, move):
+                return move
+        return original_top1  # fallback – no drawish move in top-k
+
+    # ── Near-equality: stick with policy top1 ─────────────────────────
+    return original_top1
+
+
+class UCIEngine:
+    """Minimal UCI chess engine using a trained ChessResNet model."""
+
+    def __init__(self, ckpt_path: str, device: str = "cuda"):
+        self.device = device if torch.cuda.is_available() and device == "cuda" else "cpu"
+        self.board = chess.Board()
+        self.model = self._load_model(ckpt_path)
+        self.model.eval()
+        self._stop_requested = False
+
+    def _load_model(self, ckpt_path: str) -> ChessResNet:
+        ckpt = torch.load(ckpt_path, map_location=self.device, weights_only=True)
+        model_config = ckpt.get("model_config", {})
+        if not model_config:
+            model_config = {
+                "channels": ckpt.get("args", {}).get("channels", 256),
+                "blocks": ckpt.get("args", {}).get("blocks", 20),
+                "num_actions": 20480,
+            }
+        model = create_model_from_config(model_config)
+        model.load_state_dict(ckpt["model"])
+        model.to(self.device)
+        return model
+
+    def uci_new_game(self):
+        """Reset board for a new game."""
+        self.board.reset()
+        self._stop_requested = False
+
+    def set_position(self, fen: str | None = None, moves: list[str] | None = None):
+        """Set up position from FEN and optional move list."""
+        if fen:
+            self.board.set_fen(fen)
+        else:
+            self.board.reset()
+        if moves:
+            for m in moves:
+                self.board.push(chess.Move.from_uci(m))
+
+    def get_best_move(self, movetime_ms: int = 1000) -> tuple[str, float]:
+        """
+        Return (bestmove_uci, top_logit) by evaluating the current board.
+
+        Only considers legal moves.  This is a single-forward-pass evaluator;
+        it does NOT do MCTS or search.
+        """
+        # Encode board
+        planes = encode_board(self.board)
+        inp = torch.from_numpy(planes).unsqueeze(0).float().to(self.device)  # [1,18,8,8]
+
+        with torch.no_grad():
+            with torch.autocast(device_type=self.device, enabled=(self.device == "cuda")):
+                policy_logits, value = self.model(inp)
+
+        policy_logits = policy_logits.squeeze(0)  # [20480]
+
+        # Get legal moves and their action IDs
+        action_ids, moves = legal_action_ids(self.board)
+
+        if not moves:
+            return "0000", float("-inf")
+
+        # Draw-aware top-k rerank (avoids threefold-repetition, 50-move, etc.)
+        best_move_obj = choose_move_with_draw_awareness(
+            self.board, moves, action_ids, policy_logits, value, topk=12
+        )
+        best_move = best_move_obj.uci()
+        best_logit = float(policy_logits[move_to_action_id(best_move_obj, self.board.turn)])
+
+        return best_move, best_logit
+
+    def handle_go(self, tokens: list[str]):
+        """Process 'go' command and output bestmove."""
+        movetime_ms = 1000
+        if "movetime" in tokens:
+            idx = tokens.index("movetime") + 1
+            if idx < len(tokens):
+                movetime_ms = int(tokens[idx])
+
+        best_move, _ = self.get_best_move(movetime_ms)
+        print(f"bestmove {best_move}", flush=True)
+
+    def handle_position(self, tokens: list[str]):
+        """Process 'position' command."""
+        fen = None
+        moves = []
+
+        if "startpos" in tokens:
+            pass  # use starting position (board.reset() already done or standard)
+        elif "fen" in tokens:
+            # Collect FEN string up to "moves" keyword
+            idx = tokens.index("fen") + 1
+            fen_parts = []
+            while idx < len(tokens) and tokens[idx] != "moves":
+                fen_parts.append(tokens[idx])
+                idx += 1
+            fen = " ".join(fen_parts)
+
+        if "moves" in tokens:
+            idx = tokens.index("moves") + 1
+            moves = tokens[idx:]
+
+        self.set_position(fen, moves)
+
+    def run(self):
+        """Main UCI loop: read commands from stdin, respond to stdout."""
+        while True:
+            line = sys.stdin.readline()
+            if not line:
+                break
+            line = line.strip()
+            if not line:
+                continue
+
+            parts = line.split()
+            cmd = parts[0]
+
+            if cmd == "uci":
+                print("id name JoeyStage1StockfishDistill", flush=True)
+                print("id author Joey", flush=True)
+                print("uciok", flush=True)
+            elif cmd == "isready":
+                print("readyok", flush=True)
+            elif cmd == "ucinewgame":
+                self.uci_new_game()
+            elif cmd == "position":
+                self.handle_position(parts[1:])
+            elif cmd == "go":
+                self.handle_go(parts[1:])
+            elif cmd == "stop":
+                self._stop_requested = True
+            elif cmd == "quit":
+                break
+
+
+def main():
+    parser = argparse.ArgumentParser(description="UCI chess engine")
+    parser.add_argument("--ckpt", default='', help="Path to checkpoint .pt file")
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = parser.parse_args()
+
+    engine = UCIEngine(args.ckpt, args.device)
+    engine.run()
+
+
+if __name__ == "__main__":
+    main()