Publish 7.6K parameter neural MCTS self-play agent

README.md

@@ -0,0 +1,45 @@
+---
+license: apache-2.0
+tags:
+- reinforcement-learning
+- self-play
+- monte-carlo-tree-search
+- policy-value-network
+- gradio
+---
+
+# MicroZero
+
+MicroZero is a tiny AlphaZero-style Tic-Tac-Toe agent. A shared neural trunk predicts
+both move priors and position value. Monte Carlo Tree Search converts those estimates
+into stronger self-play targets, while terminal outcomes supervise value learning.
+
+The training loop uses:
+
+- neural-guided PUCT search;
+- root Dirichlet exploration;
+- policy targets from MCTS visit counts;
+- terminal win/draw/loss value targets;
+- all eight square-board symmetries;
+- a bounded replay buffer.
+
+Evaluation alternates playing first and second against random and exact minimax
+opponents.
+
+## Verified results
+
+- 7,626 trainable parameters;
+- 980 neural-MCTS self-play games;
+- 24,000-example bounded replay buffer after symmetry augmentation;
+- 384 wins, 16 draws, and zero losses over 400 games against random play;
+- 200 draws and zero losses over 200 games against exact minimax.
+
+The exact opponent is used only for final evaluation, not as a source of training
+targets. Full metrics and the per-iteration learning history are stored in
+`artifacts/microzero/evaluation.json`.
+
+## Reproduce
+
+```powershell
+uv run python projects/microzero/train.py
+```

evaluation.json

@@ -0,0 +1,167 @@
+{
+  "model": "MicroZero",
+  "parameters": 7626,
+  "self_play_iterations": 14,
+  "self_play_games": 980,
+  "final_replay_examples": 24000,
+  "training_history": [
+    {
+      "iteration": 1,
+      "replay_examples": 4128,
+      "x_win_rate": 0.5428571428571428,
+      "o_win_rate": 0.14285714285714285,
+      "draw_rate": 0.3142857142857143,
+      "average_game_moves": 7.371428571428571,
+      "policy_loss": 2.1499555636854732,
+      "value_loss": 0.5162482239744243
+    },
+    {
+      "iteration": 2,
+      "replay_examples": 8128,
+      "x_win_rate": 0.5142857142857142,
+      "o_win_rate": 0.0,
+      "draw_rate": 0.4857142857142857,
+      "average_game_moves": 7.142857142857143,
+      "policy_loss": 1.9038925133645535,
+      "value_loss": 0.3203776921145618
+    },
+    {
+      "iteration": 3,
+      "replay_examples": 12328,
+      "x_win_rate": 0.42857142857142855,
+      "o_win_rate": 0.014285714285714285,
+      "draw_rate": 0.5571428571428572,
+      "average_game_moves": 7.5,
+      "policy_loss": 1.4543613067695074,
+      "value_loss": 0.26955098012576295
+    },
+    {
+      "iteration": 4,
+      "replay_examples": 16456,
+      "x_win_rate": 0.5142857142857142,
+      "o_win_rate": 0.0,
+      "draw_rate": 0.4857142857142857,
+      "average_game_moves": 7.371428571428571,
+      "policy_loss": 1.229265012878638,
+      "value_loss": 0.22285609325537314
+    },
+    {
+      "iteration": 5,
+      "replay_examples": 20952,
+      "x_win_rate": 0.2857142857142857,
+      "o_win_rate": 0.0,
+      "draw_rate": 0.7142857142857143,
+      "average_game_moves": 8.028571428571428,
+      "policy_loss": 1.1138035529633847,
+      "value_loss": 0.1854344406960214
+    },
+    {
+      "iteration": 6,
+      "replay_examples": 24000,
+      "x_win_rate": 0.3142857142857143,
+      "o_win_rate": 0.014285714285714285,
+      "draw_rate": 0.6714285714285714,
+      "average_game_moves": 7.957142857142857,
+      "policy_loss": 1.0146886479664357,
+      "value_loss": 0.15539901554980812
+    },
+    {
+      "iteration": 7,
+      "replay_examples": 24000,
+      "x_win_rate": 0.2571428571428571,
+      "o_win_rate": 0.0,
+      "draw_rate": 0.7428571428571429,
+      "average_game_moves": 8.085714285714285,
+      "policy_loss": 0.8795330280319174,
+      "value_loss": 0.11288299490796759
+    },
+    {
+      "iteration": 8,
+      "replay_examples": 24000,
+      "x_win_rate": 0.24285714285714285,
+      "o_win_rate": 0.0,
+      "draw_rate": 0.7571428571428571,
+      "average_game_moves": 8.314285714285715,
+      "policy_loss": 0.7982094690520712,
+      "value_loss": 0.09204166978993948
+    },
+    {
+      "iteration": 9,
+      "replay_examples": 24000,
+      "x_win_rate": 0.3,
+      "o_win_rate": 0.014285714285714285,
+      "draw_rate": 0.6857142857142857,
+      "average_game_moves": 7.985714285714286,
+      "policy_loss": 0.7609958442601752,
+      "value_loss": 0.09048273768077822
+    },
+    {
+      "iteration": 10,
+      "replay_examples": 24000,
+      "x_win_rate": 0.2857142857142857,
+      "o_win_rate": 0.0,
+      "draw_rate": 0.7142857142857143,
+      "average_game_moves": 8.17142857142857,
+      "policy_loss": 0.7380844823857571,
+      "value_loss": 0.08676699136799955
+    },
+    {
+      "iteration": 11,
+      "replay_examples": 24000,
+      "x_win_rate": 0.18571428571428572,
+      "o_win_rate": 0.0,
+      "draw_rate": 0.8142857142857143,
+      "average_game_moves": 8.371428571428572,
+      "policy_loss": 0.7269294456915653,
+      "value_loss": 0.074583658938633
+    },
+    {
+      "iteration": 12,
+      "replay_examples": 24000,
+      "x_win_rate": 0.3142857142857143,
+      "o_win_rate": 0.0,
+      "draw_rate": 0.6857142857142857,
+      "average_game_moves": 8.142857142857142,
+      "policy_loss": 0.7146875767948779,
+      "value_loss": 0.07327091548235175
+    },
+    {
+      "iteration": 13,
+      "replay_examples": 24000,
+      "x_win_rate": 0.22857142857142856,
+      "o_win_rate": 0.0,
+      "draw_rate": 0.7714285714285715,
+      "average_game_moves": 8.17142857142857,
+      "policy_loss": 0.7104355310823055,
+      "value_loss": 0.07133727120433717
+    },
+    {
+      "iteration": 14,
+      "replay_examples": 24000,
+      "x_win_rate": 0.2,
+      "o_win_rate": 0.0,
+      "draw_rate": 0.8,
+      "average_game_moves": 8.371428571428572,
+      "policy_loss": 0.7074488862714869,
+      "value_loss": 0.06895742299867437
+    }
+  ],
+  "versus_random": {
+    "games": 400,
+    "wins": 384,
+    "draws": 16,
+    "losses": 0,
+    "win_rate": 0.96,
+    "non_loss_rate": 1.0,
+    "average_moves": 6.065
+  },
+  "versus_exact_minimax": {
+    "games": 200,
+    "wins": 0,
+    "draws": 200,
+    "losses": 0,
+    "win_rate": 0.0,
+    "non_loss_rate": 1.0,
+    "average_moves": 9.0
+  }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:702c351788afad7d80eb3b55117230644708b5a46952b84f44806476e976af2c
+size 31392

source/app.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+from game import is_terminal, legal_actions, play, winner
+from mcts import NeuralMCTS
+from model import MicroZeroNet
+from safetensors.torch import load_file
+
+ARTIFACT = Path(__file__).resolve().parent / "artifacts" / "microzero"
+MODEL = MicroZeroNet()
+MODEL.load_state_dict(load_file(ARTIFACT / "model.safetensors"))
+MODEL.eval()
+SYMBOLS = {0: ".", 1: "X", -1: "O"}
+
+
+def neural_action(
+    model: MicroZeroNet,
+    board: np.ndarray,
+    player: int,
+    simulations: int,
+    seed: int,
+) -> int:
+    search = NeuralMCTS(model, simulations=simulations)
+    policy = search.policy(
+        board,
+        player,
+        temperature=0.0,
+        add_noise=False,
+        rng=np.random.default_rng(seed),
+    )
+    return int(np.argmax(policy))
+
+
+def render(board: list[int], message: str) -> str:
+    cells = [SYMBOLS[int(value)] for value in board]
+    rows = [" | ".join(cells[index : index + 3]) for index in range(0, 9, 3)]
+    return "```\n" + "\n---------\n".join(rows) + f"\n```\n{message}"
+
+
+def new_game() -> tuple[list[int], str]:
+    board = [0] * 9
+    return board, render(board, "You are X. Choose a cell from 0 to 8.")
+
+
+def move(board: list[int], action: int) -> tuple[list[int], str]:
+    array = np.asarray(board, dtype=np.int8)
+    action = int(action)
+    if action not in legal_actions(array):
+        return board, render(board, "That cell is occupied.")
+    array = play(array, action, 1)
+    if is_terminal(array):
+        message = "You win." if winner(array) == 1 else "Draw."
+        return array.tolist(), render(array.tolist(), message)
+    response = neural_action(MODEL, array, -1, simulations=96, seed=7000 + action)
+    array = play(array, response, -1)
+    result = winner(array)
+    if result == -1:
+        message = f"MicroZero played {response} and won."
+    elif is_terminal(array):
+        message = f"MicroZero played {response}. Draw."
+    else:
+        message = f"MicroZero played {response}. Your turn."
+    return array.tolist(), render(array.tolist(), message)
+
+
+with gr.Blocks(title="MicroZero") as demo:
+    gr.Markdown("# MicroZero\nPlay against the neural MCTS agent.")
+    state = gr.State([0] * 9)
+    board = gr.Markdown(render([0] * 9, "You are X. Choose a cell from 0 to 8."))
+    action = gr.Slider(0, 8, value=4, step=1, label="Your move")
+    with gr.Row():
+        play_button = gr.Button("Play", variant="primary")
+        reset_button = gr.Button("New game")
+    play_button.click(move, inputs=[state, action], outputs=[state, board])
+    reset_button.click(new_game, outputs=[state, board])
+
+
+if __name__ == "__main__":
+    demo.launch()

source/game.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from functools import cache
+
+import numpy as np
+
+WIN_LINES = [
+    (0, 1, 2),
+    (3, 4, 5),
+    (6, 7, 8),
+    (0, 3, 6),
+    (1, 4, 7),
+    (2, 5, 8),
+    (0, 4, 8),
+    (2, 4, 6),
+]
+
+
+def winner(board: np.ndarray) -> int:
+    for left, middle, right in WIN_LINES:
+        value = int(board[left])
+        if value and value == board[middle] == board[right]:
+            return value
+    return 0
+
+
+def is_terminal(board: np.ndarray) -> bool:
+    return winner(board) != 0 or not np.any(board == 0)
+
+
+def legal_actions(board: np.ndarray) -> np.ndarray:
+    return np.flatnonzero(board == 0)
+
+
+def play(board: np.ndarray, action: int, player: int) -> np.ndarray:
+    if board[action] != 0:
+        raise ValueError(f"Cell {action} is occupied.")
+    next_board = board.copy()
+    next_board[action] = player
+    return next_board
+
+
+def canonical(board: np.ndarray, player: int) -> np.ndarray:
+    return board.astype(np.float32) * player
+
+
+def symmetries(
+    state: np.ndarray,
+    policy: np.ndarray,
+) -> list[tuple[np.ndarray, np.ndarray]]:
+    board = state.reshape(3, 3)
+    probabilities = policy.reshape(3, 3)
+    transformed = []
+    for rotations in range(4):
+        rotated_board = np.rot90(board, rotations)
+        rotated_policy = np.rot90(probabilities, rotations)
+        transformed.append(
+            (rotated_board.reshape(-1).copy(), rotated_policy.reshape(-1).copy())
+        )
+        transformed.append(
+            (
+                np.fliplr(rotated_board).reshape(-1).copy(),
+                np.fliplr(rotated_policy).reshape(-1).copy(),
+            )
+        )
+    return transformed
+
+
+@cache
+def minimax_value(board_key: tuple[int, ...], player: int) -> int:
+    board = np.asarray(board_key, dtype=np.int8)
+    result = winner(board)
+    if result:
+        return 1 if result == player else -1
+    legal = legal_actions(board)
+    if not len(legal):
+        return 0
+    return max(
+        -minimax_value(tuple(play(board, int(action), player)), -player)
+        for action in legal
+    )
+
+
+def minimax_action(board: np.ndarray, player: int) -> int:
+    legal = legal_actions(board)
+    values = [
+        -minimax_value(tuple(play(board, int(action), player)), -player)
+        for action in legal
+    ]
+    best = max(values)
+    return int(legal[values.index(best)])

source/mcts.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+from game import canonical, is_terminal, legal_actions, play, winner
+from model import MicroZeroNet
+
+
+class NeuralMCTS:
+    def __init__(
+        self,
+        model: MicroZeroNet,
+        simulations: int,
+        exploration: float = 1.35,
+    ) -> None:
+        self.model = model
+        self.simulations = simulations
+        self.exploration = exploration
+        self.priors: dict[tuple, np.ndarray] = {}
+        self.state_visits: dict[tuple, int] = {}
+        self.edge_visits: dict[tuple, int] = {}
+        self.edge_value: dict[tuple, float] = {}
+
+    @staticmethod
+    def key(board: np.ndarray, player: int) -> tuple:
+        return (*board.tolist(), player)
+
+    @torch.inference_mode()
+    def expand(self, board: np.ndarray, player: int) -> float:
+        key = self.key(board, player)
+        logits, value = self.model(torch.from_numpy(canonical(board, player))[None])
+        logits = logits[0].numpy()
+        legal = legal_actions(board)
+        masked = np.full(9, -1e9, dtype=np.float64)
+        masked[legal] = logits[legal]
+        probabilities = np.exp(masked - np.max(masked))
+        probabilities /= probabilities.sum()
+        self.priors[key] = probabilities
+        self.state_visits[key] = 0
+        return float(value[0])
+
+    def search(self, board: np.ndarray, player: int) -> float:
+        result = winner(board)
+        if result:
+            return 1.0 if result == player else -1.0
+        if is_terminal(board):
+            return 0.0
+        state_key = self.key(board, player)
+        if state_key not in self.priors:
+            return self.expand(board, player)
+
+        best_score = -float("inf")
+        best_action = -1
+        root_visits = max(1, self.state_visits[state_key])
+        for action in legal_actions(board):
+            edge = (*state_key, int(action))
+            visits = self.edge_visits.get(edge, 0)
+            value = self.edge_value.get(edge, 0.0)
+            prior = self.priors[state_key][action]
+            score = value + self.exploration * prior * math.sqrt(root_visits) / (
+                1 + visits
+            )
+            if score > best_score:
+                best_score = score
+                best_action = int(action)
+
+        child = play(board, best_action, player)
+        value = -self.search(child, -player)
+        edge = (*state_key, best_action)
+        visits = self.edge_visits.get(edge, 0) + 1
+        previous = self.edge_value.get(edge, 0.0)
+        self.edge_visits[edge] = visits
+        self.edge_value[edge] = previous + (value - previous) / visits
+        self.state_visits[state_key] += 1
+        return value
+
+    def policy(
+        self,
+        board: np.ndarray,
+        player: int,
+        *,
+        temperature: float,
+        add_noise: bool,
+        rng: np.random.Generator,
+    ) -> np.ndarray:
+        key = self.key(board, player)
+        if key not in self.priors:
+            self.expand(board, player)
+        if add_noise:
+            legal = legal_actions(board)
+            noise = rng.dirichlet(np.full(len(legal), 0.30))
+            priors = self.priors[key].copy()
+            priors[legal] = 0.75 * priors[legal] + 0.25 * noise
+            self.priors[key] = priors
+        for _ in range(self.simulations):
+            self.search(board, player)
+        visits = np.zeros(9, dtype=np.float64)
+        for action in legal_actions(board):
+            visits[action] = self.edge_visits.get((*key, int(action)), 0)
+        if temperature <= 1e-6:
+            policy = np.zeros(9, dtype=np.float64)
+            policy[int(np.argmax(visits))] = 1.0
+            return policy
+        adjusted = visits ** (1 / temperature)
+        if adjusted.sum() == 0:
+            adjusted[legal_actions(board)] = 1
+        return adjusted / adjusted.sum()

source/model.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class MicroZeroNet(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.trunk = nn.Sequential(
+            nn.Linear(9, 64),
+            nn.LayerNorm(64),
+            nn.SiLU(),
+            nn.Linear(64, 64),
+            nn.SiLU(),
+        )
+        self.policy = nn.Linear(64, 9)
+        self.value = nn.Sequential(
+            nn.Linear(64, 32),
+            nn.SiLU(),
+            nn.Linear(32, 1),
+            nn.Tanh(),
+        )
+
+    def forward(self, board: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        hidden = self.trunk(board)
+        return self.policy(hidden), self.value(hidden).squeeze(-1)
+
+
+def parameter_count(model: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in model.parameters())

source/requirements.txt

@@ -0,0 +1,4 @@
+gradio>=5,<7
+numpy>=2,<3
+safetensors>=0.6,<1
+torch>=2.7,<3

source/train.py

@@ -0,0 +1,265 @@
+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()