| """ |
| Board and move encoding for TinyChess. |
| |
| Design notes |
| ------------ |
| * Positions are *canonicalised* to side-to-move = White by vertically mirroring |
| the board and swapping piece colours when Black is to move. This is a |
| lossless symmetry of chess (ignoring nothing: castling rights and en-passant |
| are mirrored too), and it is a standard, non-cheating canonicalisation. |
| * A move is packed into a single int32 with structural fields only. We |
| deliberately store *rules-derived* structure (from/to/promo/piece/capture/ |
| castle/ep) and no heuristic evaluation. |
| * Per-square features are kept separate (64 vectors); nothing is collapsed to a |
| single vector at encoding time. |
| """ |
| from __future__ import annotations |
|
|
| import numpy as np |
| import chess |
|
|
| |
| |
| |
| |
| _F_FROM = 0 |
| _F_TO = 6 |
| _F_PROMO = 12 |
| _F_PIECE = 15 |
| _F_CAP = 18 |
| _F_CASTLE = 19 |
| _F_EP = 20 |
|
|
| PROMO_CODES = {None: 0, chess.KNIGHT: 1, chess.BISHOP: 2, chess.ROOK: 3, chess.QUEEN: 4} |
| PROMO_INV = {v: k for k, v in PROMO_CODES.items()} |
| N_PROMO = 5 |
| N_PIECE_TYPES = 6 |
|
|
| MOVE_FIELD_SIZES = dict(from_sq=64, to_sq=64, promo=N_PROMO, piece=N_PIECE_TYPES, |
| capture=2, castle=2, ep=2) |
|
|
|
|
| def pack_move(board: chess.Board, move: chess.Move) -> int: |
| """Pack a legal move on `board` into an int32 of structural fields.""" |
| piece = board.piece_at(move.from_square) |
| ptype = (piece.piece_type - 1) if piece is not None else 0 |
| is_ep = board.is_en_passant(move) |
| is_cap = board.is_capture(move) |
| is_castle = board.is_castling(move) |
| v = (move.from_square << _F_FROM) | (move.to_square << _F_TO) |
| v |= PROMO_CODES.get(move.promotion, 0) << _F_PROMO |
| v |= ptype << _F_PIECE |
| v |= int(is_cap) << _F_CAP |
| v |= int(is_castle) << _F_CASTLE |
| v |= int(is_ep) << _F_EP |
| return v |
|
|
|
|
| def unpack_move(v: int) -> dict: |
| return dict( |
| from_sq=(v >> _F_FROM) & 63, |
| to_sq=(v >> _F_TO) & 63, |
| promo=(v >> _F_PROMO) & 7, |
| piece=(v >> _F_PIECE) & 7, |
| capture=(v >> _F_CAP) & 1, |
| castle=(v >> _F_CASTLE) & 1, |
| ep=(v >> _F_EP) & 1, |
| ) |
|
|
|
|
| def unpack_fields(v: np.ndarray) -> np.ndarray: |
| """Vectorised unpack -> int64 array [..., 7] in MOVE_FIELD order.""" |
| v = np.asarray(v, dtype=np.int64) |
| return np.stack([ |
| (v >> _F_FROM) & 63, |
| (v >> _F_TO) & 63, |
| (v >> _F_PROMO) & 7, |
| (v >> _F_PIECE) & 7, |
| (v >> _F_CAP) & 1, |
| (v >> _F_CASTLE) & 1, |
| (v >> _F_EP) & 1, |
| ], axis=-1) |
|
|
|
|
| MOVE_FIELD_ORDER = ["from_sq", "to_sq", "promo", "piece", "capture", "castle", "ep"] |
|
|
|
|
| def packed_to_uci(v: int) -> str: |
| d = unpack_move(v) |
| m = chess.Move(d["from_sq"], d["to_sq"], promotion=PROMO_INV.get(d["promo"])) |
| return m.uci() |
|
|
|
|
| def packed_to_move(v: int) -> chess.Move: |
| d = unpack_move(v) |
| return chess.Move(d["from_sq"], d["to_sq"], promotion=PROMO_INV.get(d["promo"])) |
|
|
|
|
| |
| |
| |
|
|
| def canonical_board(board: chess.Board) -> tuple[chess.Board, bool]: |
| """Return (board with White to move, flipped?). |
| |
| Uses python-chess `mirror()` which flips vertically and swaps colours, |
| correctly transforming castling rights and the en-passant square. |
| """ |
| if board.turn == chess.WHITE: |
| return board, False |
| return board.mirror(), True |
|
|
|
|
| def flip_square(sq: int) -> int: |
| return sq ^ 56 |
|
|
|
|
| def flip_move(move: chess.Move) -> chess.Move: |
| return chess.Move(flip_square(move.from_square), flip_square(move.to_square), |
| promotion=move.promotion) |
|
|
|
|
| |
| |
| |
| |
| N_PIECE_TOKENS = 13 |
| |
| SQ_EXTRA = ["last_from", "last_to", "ep_square", "own_backrank_castle"] |
| N_SQ_EXTRA = len(SQ_EXTRA) |
|
|
| GLOBAL_FEATS = [ |
| "own_oo", "own_ooo", "opp_oo", "opp_ooo", |
| "has_ep", |
| "halfmove_clock_n", "fullmove_n", "repetition_n", |
| "own_pawns_n", "own_knights_n", "own_bishops_n", "own_rooks_n", "own_queens_n", |
| "opp_pawns_n", "opp_knights_n", "opp_bishops_n", "opp_rooks_n", "opp_queens_n", |
| "in_check", "material_diff_n", |
| ] |
| N_GLOBAL = len(GLOBAL_FEATS) |
|
|
| _PT_ORDER = [chess.PAWN, chess.KNIGHT, chess.BISHOP, chess.ROOK, chess.QUEEN, chess.KING] |
| _PIECE_VAL = {chess.PAWN: 1, chess.KNIGHT: 3, chess.BISHOP: 3, chess.ROOK: 5, |
| chess.QUEEN: 9, chess.KING: 0} |
|
|
|
|
| def encode_board(cboard: chess.Board, last_move: chess.Move | None = None, |
| repetition: int = 0): |
| """Encode a *canonical* (White-to-move) board. |
| |
| Returns (squares int8[64], extras int8[64,N_SQ_EXTRA], glob float32[N_GLOBAL]). |
| """ |
| sq = np.zeros(64, dtype=np.int8) |
| for s, piece in cboard.piece_map().items(): |
| idx = _PT_ORDER.index(piece.piece_type) + 1 |
| if piece.color == chess.BLACK: |
| idx += 6 |
| sq[s] = idx |
|
|
| extras = np.zeros((64, N_SQ_EXTRA), dtype=np.int8) |
| if last_move is not None: |
| extras[last_move.from_square, 0] = 1 |
| extras[last_move.to_square, 1] = 1 |
| if cboard.ep_square is not None: |
| extras[cboard.ep_square, 2] = 1 |
| if cboard.has_kingside_castling_rights(chess.WHITE): |
| extras[chess.H1, 3] = 1 |
| extras[chess.E1, 3] = 1 |
| if cboard.has_queenside_castling_rights(chess.WHITE): |
| extras[chess.A1, 3] = 1 |
| extras[chess.E1, 3] = 1 |
|
|
| g = np.zeros(N_GLOBAL, dtype=np.float32) |
| g[0] = cboard.has_kingside_castling_rights(chess.WHITE) |
| g[1] = cboard.has_queenside_castling_rights(chess.WHITE) |
| g[2] = cboard.has_kingside_castling_rights(chess.BLACK) |
| g[3] = cboard.has_queenside_castling_rights(chess.BLACK) |
| g[4] = cboard.ep_square is not None |
| g[5] = min(cboard.halfmove_clock, 100) / 100.0 |
| g[6] = min(cboard.fullmove_number, 120) / 120.0 |
| g[7] = min(repetition, 3) / 3.0 |
| mat = 0 |
| for i, pt in enumerate(_PT_ORDER[:5]): |
| nw = len(cboard.pieces(pt, chess.WHITE)) |
| nb = len(cboard.pieces(pt, chess.BLACK)) |
| g[8 + i] = nw / 8.0 |
| g[13 + i] = nb / 8.0 |
| mat += _PIECE_VAL[pt] * (nw - nb) |
| g[18] = cboard.is_check() |
| g[19] = np.tanh(mat / 10.0) |
| return sq, extras, g |
|
|
|
|
| def encode_position(board: chess.Board, last_move: chess.Move | None = None, |
| repetition: int = 0): |
| """Canonicalise then encode. Returns (squares, extras, glob, flipped).""" |
| cb, flipped = canonical_board(board) |
| lm = last_move |
| if flipped and lm is not None: |
| lm = flip_move(lm) |
| sq, ex, g = encode_board(cb, lm, repetition) |
| return sq, ex, g, flipped |
|
|
|
|
| def legal_packed(board: chess.Board): |
| """Return (packed candidate array int32[L], canonical board, flipped).""" |
| cb, flipped = canonical_board(board) |
| return np.array([pack_move(cb, m) for m in cb.legal_moves], dtype=np.int32), cb, flipped |
|
|
