erlangzhang/alan-aumaton

1"""Board / move encoding for behavioral cloning.2 3Everything is expressed from the **side-to-move's perspective** ("my" pieces always4occupy planes 0-5 and move *up* the board). When it is Black to move we mirror the5board vertically so the network sees one consistent point of view regardless of the6colour you actually played. This is what lets a single policy capture *your* style.7 8Board tensor : (12, 8, 8) float32 -- 6 piece types x {mine, theirs}.9Move label   : an int in [0, 4096) encoding (from_square, to_square) in the10               canonical frame: index = from * 64 + to. Promotions are collapsed to11               the from->to pair (queen is assumed at inference).12"""13 14import chess15import numpy as np16 17NUM_MOVES = 64 * 64  # 409618BOARD_PLANES = 1219 20# piece_type (1..6) -> plane offset (0..5)21_PIECE_TO_PLANE = {22    chess.PAWN: 0,23    chess.KNIGHT: 1,24    chess.BISHOP: 2,25    chess.ROOK: 3,26    chess.QUEEN: 4,27    chess.KING: 5,28}29 30 31def _canonical_square(square, turn):32    """Map a real square into the side-to-move frame (vertical mirror for Black)."""33    return square if turn == chess.WHITE else chess.square_mirror(square)34 35 36def board_to_tensor(board):37    """Return a (12, 8, 8) float32 tensor from the side-to-move's perspective."""38    tensor = np.zeros((BOARD_PLANES, 8, 8), dtype=np.float32)39    turn = board.turn40    for square, piece in board.piece_map().items():41        csq = _canonical_square(square, turn)42        rank, file = divmod(csq, 8)43        is_mine = piece.color == turn44        plane = _PIECE_TO_PLANE[piece.piece_type] + (0 if is_mine else 6)45        tensor[plane, rank, file] = 1.046    return tensor47 48 49def move_to_index(move, board):50    """Encode a move as from*64 + to in the canonical frame."""51    turn = board.turn52    cf = _canonical_square(move.from_square, turn)53    ct = _canonical_square(move.to_square, turn)54    return cf * 64 + ct55 56 57def is_underpromotion(move):58    """True for knight/bishop/rook promotions (collapsed away by our encoding)."""59    return move.promotion is not None and move.promotion != chess.QUEEN60 61 62def legal_move_index_map(board):63    """Map {canonical_index: chess.Move} for the current legal moves.64 65    When several legal moves share a from->to pair (only under-promotions), the66    queen promotion wins so inference never emits an under-promotion.67    """68    mapping = {}69    for move in board.legal_moves:70        idx = move_to_index(move, board)71        if idx in mapping and is_underpromotion(move):72            continue  # keep the already-stored (queen / non-promo) move73        mapping[idx] = move74    return mapping75