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()