DOCSEARCHbyAlgolia
Submit Docs

Maia-2: Unified Human-Like Chess Model

repository·main·Indexed 41 days ago

https://github.com/csslab/maia2

Official implementation of Maia-2, a skill-aware chess model designed for human-AI alignment. It predicts human moves and game outcomes across various Elo levels using a multi-task PyTorch architecture combining ChessResNet and Transformers. Supports batch and position-wise inference, training on Lichess rated Rapid and Blitz data, and runs on CUDA, Apple MPS, or CPU.

Tokens
9K
Snippets
33
Records
37
Agent score
99%

What's inside Maia-2

  1. Install maia2

    main

    You can install the package directly via pip. Maia-2 supports Python 3.10–3.12 and runs on CUDA, Apple MPS, or CPU.

    For development or to use the validated dependency baseline, it is recommended to use a fresh Python 3.12 environment via Conda.

    # Standard installation
pip install maia2

# Development installation
conda create -n maia2 python=3.12 -y
conda activate maia2
git clone https://github.com/CSSLab/maia2.git
cd maia2
python -m pip install -r maia2/requirements.txt
python -m pip install -e . --no-deps

# Install contributor tools
python -m pip install -e ".[dev]"

  2. Download Lichess data for training

    main

    Maia-2 uses .pgn.zst archives from the Lichess standard-rated database. To be included, a game's PGN Event must contain the exact, case-sensitive Rated marker and either the Rapid or Blitz marker. Games with BOT players or missing per-move clock annotations are excluded.

    Use the fetch_data.sh script to download data. Note that December 2019 is skipped to match the original training pipeline.

    # Download one month for a local test
./maia2/fetch_data.sh /path/to/lichess_data 2023-01 2023-01

# Download the released training range
./maia2/fetch_data.sh /path/to/lichess_data 2018-05 2023-11

  3. Configure and run maia2 training

    main

    Training is performed using train.run(cfg, device="auto"). You should use one of the maintained YAML presets for new training runs:

    • maia2-training-rapid.yaml: For games with Rated and Rapid markers.
    • maia2-training-blitz.yaml: For games with Rated and Blitz markers.

    Important Notes:

    • Main Guard: You must use a if __name__ == "__main__": guard because training uses spawned preprocessing workers.
    • Checkpoint Paths: Checkpoints are saved under <save_root>/<lr>_<batch_size>_<weight_decay>/.
    • Device Selection: "cuda" (unindexed) uses all visible CUDA devices via DataParallel; "cuda:N" selects a specific device.
    • Resuming: To resume training, set from_checkpoint, checkpoint_epoch, checkpoint_year, and checkpoint_month in your config, while keeping the original full date range.
    from importlib.resources import as_file, files
from maia2 import train, utils


def main():
    game_type = "rapid"  # Use "blitz" for rated Blitz training.
    config_resource = files("maia2.configs").joinpath(
        f"maia2-training-{game_type}.yaml"
    )
    with as_file(config_resource) as config_path:
        cfg = utils.parse_args(config_path)

    cfg.data_root = "/path/to/lichess_data"
    cfg.save_root = f"/path/to/checkpoints/{game_type}"
    train.run(cfg, device="auto")


if __name__ == "__main__":
    main()

  4. Validate optimizer state during checkpoint loading

    main

    The training system performs strict validation on the optimizer state when loading checkpoints to prevent silent corruption or mismatched training states.

    Key validations include:

    • Mapping Check: Ensures optimizer_state_dict and its internal state are valid mappings.
    • Step Consistency: Every parameter in the AdamW optimizer must have a step value. The _validate_optimizer_step_count function ensures that the steps recorded in the optimizer's parameter states match the optimizer_steps value stored in the training_metadata.
    • Hyperparameter Check: _validate_optimizer_hyperparameters ensures the optimizer configuration in the checkpoint is compatible with the current cfg.

  5. Understand the training schedule and epoch logic

    main

    The training schedule is determined by _training_schedule(cfg, pgn_paths).

    • Fresh Run: If cfg.from_checkpoint is False, the schedule consists of cfg.max_epochs, where each epoch covers the full range of provided pgn_paths.
    • Resumed Run: If cfg.from_checkpoint is True, the schedule starts at checkpoint_epoch - 1. For the first epoch of the resumed run, it only processes PGN files that appear after the checkpoint_month in the provided list. Subsequent epochs cover the full range of pgn_paths.

    Note: The checkpoint_month must be present in the configured training range. If it is not found, a ValueError is raised.

  6. Resume training from a checkpoint

    main

    To resume a training run, configure your configuration object (cfg) with the following keys:

    • from_checkpoint: Set to True to enable resuming.
    • checkpoint_epoch: The epoch number of the last completed month.
    • checkpoint_month: The month (as an integer) that the checkpoint represents.
    • checkpoint_year: The year of the checkpoint.

    When from_checkpoint is enabled, the training schedule will skip all months prior to the checkpoint's month and resume training from the next available month in the configured range. The system also validates that the checkpoint's optimizer state and metadata match the current configuration to ensure reproducibility.

  7. Perform position-wise inference with maia2

    main

    For processing individual positions, use inference.prepare() once to prepare mappings, then iterate through your data using inference.inference_each.

    inference.inference_each returns two values:

    1. move_probs: A dictionary of move probabilities.
    2. white_expected_score: A White-perspective expected score (0 = White loss, 0.5 = draw, 1 = White win). Note that this is not a calibrated win probability for the active player.
    prepared = inference.prepare()
columns = ["board", "move", "active_elo", "opponent_elo"]

for fen, move, elo_self, elo_oppo in data.loc[:, columns].head(10).itertuples(
    index=False, name=None
):
    move_probs, white_expected_score = inference.inference_each(
        maia2_model,
        prepared,
        fen,
        elo_self,
        elo_oppo,
    )
    predicted_move = max(move_probs, key=move_probs.get)
    print(
        f"Move: {move}; predicted: {predicted_move}; "
        f"White expected score: {white_expected_score}"
    )

  8. Perform batch inference with maia2

    main

    Use inference.inference_batch to run a pre-trained model on a dataset. You can load models using model.from_pretrained.

    Set device to "auto" to automatically select CUDA, then MPS, then CPU. You can also specify "cuda", "mps", or "cpu" explicitly. Adjust batch_size and num_workers based on your available memory.

    from maia2 import dataset, inference, model

maia2_model = model.from_pretrained(type="rapid", device="auto")
data = dataset.load_example_test_dataset()

data, accuracy = inference.inference_batch(
    data,
    maia2_model,
    verbose=1,
    batch_size=1024,
    num_workers=0,
)
print(accuracy)

  9. Preprocessing chess positions for the model

    main

    The preprocessing function converts a FEN string and Elo ratings into tensors suitable for the model.

    Key behaviors:

    • Board Mirroring: If the FEN indicates it is Black's turn ('b'), the board is automatically mirrored to maintain a consistent perspective for the model.
    • Elo Mapping: Elo ratings are mapped to categories using the provided elo_dict.
    • Legal Move Masking: It generates a boolean tensor representing all possible legal moves for the current position.

    Note: If the position has no legal moves, it will raise a ValueError.

    from maia2.inference import preprocessing

# Requires dictionaries generated by prepare()
all_moves_dict, elo_dict, _ = prepare()

board_input, elo_self, elo_oppo, legal_moves = preprocessing(
    fen='rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR b KQkq - 0 1',
    elo_self=1500,
    elo_oppo=1500,
    elo_dict=elo_dict,
    all_moves_dict=all_moves_dict
)

  10. Process PGN files in parallel chunks

    main

    The process_chunks function handles the high-level orchestration of reading and preprocessing PGN files. It splits the work into chunks and uses multiprocessing to process them in parallel.

    Arguments:

    • cfg: Configuration object. Must have game_type (default 'rapid') and verbose (boolean).
    • pgn_path: Path to the PGN file.
    • pgn_chunks: A list of (start_pos, end_pos) tuples representing byte offsets in the file.
    • elo_dict: A dictionary used for mapping Elo ratings to categories.

    Returns:

    • ret: A list of processed game data tuples.
    • count: Total number of games processed.
    • len(pgn_chunks): The number of chunks processed.
    def process_chunks(cfg, pgn_path, pgn_chunks, elo_dict):
    # ... implementation ...
    return ret, count, len(pgn_chunks)

  11. Calculate number of worker processes

    main

    Use get_num_processes(num_cpu_left) to determine how many CPU cores can be dedicated to preprocessing without starving the main training process. It calculates max(1, total_cpu_count - num_cpu_left).

    from maia2.train import get_num_processes

# If you want to leave 2 CPUs free for other tasks
n_workers = get_num_processes(num_cpu_left=2)

  12. Convert a chess board to a tensor

    main

    The board_to_tensor(board) function converts a chess.Board object into a PyTorch tensor of shape (18, 8, 8).

    Tensor Structure:

    • Channels 0-5: White pieces (Pawn, Knight, Bishop, Rook, Queen, King)
    • Channels 6-11: Black pieces (Pawn, Knight, Bishop, Rook, Queen, King)
    • Channel 12: Player's turn (1.0 for White, 0.0 for Black)
    • Channel 13: Castling rights (White Kingside, White Queenside, Black Kingside, Black Queenside)
    • Channel 17: En passant target square
    import chess
import torch
from maia2.utils import board_to_tensor

board = chess.Board()
tensor = board_to_tensor(board)
# tensor shape: [18, 8, 8]