train.py

from __future__ import annotations

import json
import random
from collections import deque
from pathlib import Path

import numpy as np
import torch
import trackio
from game import (
    canonical,
    is_terminal,
    legal_actions,
    minimax_action,
    play,
    symmetries,
    winner,
)
from mcts import NeuralMCTS
from model import MicroZeroNet, parameter_count
from safetensors.torch import save_file
from torch.nn import functional as F
from torch.utils.data import DataLoader, TensorDataset

PROJECT_DIR = Path(__file__).resolve().parent
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "microzero"


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def self_play_game(
    model: MicroZeroNet,
    rng: np.random.Generator,
    simulations: int,
) -> tuple[list[tuple[np.ndarray, np.ndarray, float]], int, int]:
    board = np.zeros(9, dtype=np.int8)
    player = 1
    history = []
    mcts = NeuralMCTS(model, simulations=simulations)
    moves = 0
    while not is_terminal(board):
        temperature = 1.0 if moves < 4 else 0.15
        policy = mcts.policy(
            board,
            player,
            temperature=temperature,
            add_noise=True,
            rng=rng,
        )
        history.append((canonical(board, player), policy.copy(), player))
        action = int(rng.choice(9, p=policy))
        board = play(board, action, player)
        player = -player
        moves += 1
    result = winner(board)
    examples = []
    for state, policy, state_player in history:
        outcome = float(result * state_player)
        for transformed_state, transformed_policy in symmetries(state, policy):
            examples.append((transformed_state, transformed_policy, outcome))
    return examples, result, moves


def train_replay(
    model: MicroZeroNet,
    replay: deque,
    epochs: int,
) -> dict:
    states = torch.tensor(np.stack([item[0] for item in replay]), dtype=torch.float32)
    policies = torch.tensor(
        np.stack([item[1] for item in replay]),
        dtype=torch.float32,
    )
    outcomes = torch.tensor([item[2] for item in replay], dtype=torch.float32)
    loader = DataLoader(
        TensorDataset(states, policies, outcomes),
        batch_size=256,
        shuffle=True,
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0015, weight_decay=0.001)
    model.train()
    policy_losses = []
    value_losses = []
    for _ in range(epochs):
        for state, target_policy, target_value in loader:
            logits, value = model(state)
            policy_loss = -(
                target_policy * F.log_softmax(logits, dim=1)
            ).sum(dim=1).mean()
            value_loss = F.mse_loss(value, target_value)
            loss = policy_loss + value_loss
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            policy_losses.append(policy_loss.item())
            value_losses.append(value_loss.item())
    return {
        "policy_loss": float(np.mean(policy_losses)),
        "value_loss": float(np.mean(value_losses)),
    }


def neural_action(
    model: MicroZeroNet,
    board: np.ndarray,
    player: int,
    simulations: int,
    seed: int,
) -> int:
    mcts = NeuralMCTS(model, simulations=simulations)
    policy = mcts.policy(
        board,
        player,
        temperature=0.0,
        add_noise=False,
        rng=np.random.default_rng(seed),
    )
    return int(policy.argmax())


def arena(
    model: MicroZeroNet,
    opponent: str,
    games: int,
    simulations: int,
    seed: int,
) -> dict:
    rng = np.random.default_rng(seed)
    wins = 0
    draws = 0
    losses = 0
    lengths = []
    for game in range(games):
        agent_player = 1 if game % 2 == 0 else -1
        board = np.zeros(9, dtype=np.int8)
        player = 1
        moves = 0
        while not is_terminal(board):
            if player == agent_player:
                action = neural_action(
                    model,
                    board,
                    player,
                    simulations,
                    seed + game * 11 + moves,
                )
            elif opponent == "random":
                action = int(rng.choice(legal_actions(board)))
            else:
                action = minimax_action(board, player)
            board = play(board, action, player)
            player = -player
            moves += 1
        result = winner(board)
        if result == agent_player:
            wins += 1
        elif result == 0:
            draws += 1
        else:
            losses += 1
        lengths.append(moves)
    return {
        "games": games,
        "wins": wins,
        "draws": draws,
        "losses": losses,
        "win_rate": wins / games,
        "non_loss_rate": (wins + draws) / games,
        "average_moves": float(np.mean(lengths)),
    }


def main() -> None:
    seed_everything(2043)
    rng = np.random.default_rng(2043)
    model = MicroZeroNet()
    replay: deque = deque(maxlen=24_000)
    iterations = 14
    games_per_iteration = 70
    trackio.init(
        project="microzero",
        name="neural-mcts-self-play-v1",
        config={
            "parameters": parameter_count(model),
            "iterations": iterations,
            "games_per_iteration": games_per_iteration,
            "self_play_simulations": 32,
            "symmetry_augmentation": 8,
        },
    )
    history = []
    for iteration in range(1, iterations + 1):
        outcomes = []
        moves = []
        for _ in range(games_per_iteration):
            examples, result, game_moves = self_play_game(
                model,
                rng,
                simulations=32,
            )
            replay.extend(examples)
            outcomes.append(result)
            moves.append(game_moves)
        training = train_replay(model, replay, epochs=4)
        record = {
            "iteration": iteration,
            "replay_examples": len(replay),
            "x_win_rate": float(np.mean(np.asarray(outcomes) == 1)),
            "o_win_rate": float(np.mean(np.asarray(outcomes) == -1)),
            "draw_rate": float(np.mean(np.asarray(outcomes) == 0)),
            "average_game_moves": float(np.mean(moves)),
            **training,
        }
        history.append(record)
        trackio.log(record)

    random_arena = arena(
        model,
        opponent="random",
        games=400,
        simulations=64,
        seed=3043,
    )
    minimax_arena = arena(
        model,
        opponent="minimax",
        games=200,
        simulations=96,
        seed=4043,
    )
    results = {
        "model": "MicroZero",
        "parameters": parameter_count(model),
        "self_play_iterations": iterations,
        "self_play_games": iterations * games_per_iteration,
        "final_replay_examples": len(replay),
        "training_history": history,
        "versus_random": random_arena,
        "versus_exact_minimax": minimax_arena,
    }
    trackio.log(
        {
            "random_win_rate": random_arena["win_rate"],
            "random_non_loss_rate": random_arena["non_loss_rate"],
            "minimax_non_loss_rate": minimax_arena["non_loss_rate"],
        }
    )
    trackio.finish()
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    save_file(model.state_dict(), ARTIFACT_DIR / "model.safetensors")
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()