wtedw/nanoAlphaZero

AlphaZero in a single file. Contribute to wtedw/nanoAlphaZero development by creating an account on GitHub.

12 stars · 1 forks

README

nanoAlphaZero nanoAlphaZero is a game-agnostic, high-performance implementation of AlphaZero. It reaches perfect play in games like Hex, and grandmaster-level strength in chess. demo: play against the models

How is this different from other implementations?

It scales to chess. Not a toy AlphaZero implementation. It can train a grandmaster-level chess model in under 24h on a TPU. Genuinely game-agnostic. We validate the core logic across Hex, Connect4, Go, and Chess, and demonstrate how to train AlphaZero on custom games of your own using a Colab notebook. Training is one JAX function. Self-play, MCTS, and training are fused into a single jitted call (run_fn). It's dead simple to run. Clone the repo, then uv run train --env chess. It's fast. Our custom, TPU-native JAX environments (pgx1) run orders of magnitude faster than the reference implementation. For MCTS, we parallelize the sequential halving algorithm from Gumbel MuZero via mctx.

Environment step benchmarks:

env pgx pgx1 speedup env/s (batch 4096)

go_9x9 80.5 ms 0.535 ms 150x 7.7M

go_19x19 656 ms 1.595 ms 411x 2.6M

chess 904 ms 0.832 ms 1087x 4.9M

note: all code is optimized for TPUs; correctness on GPUs isn't guaranteed.

Setup On a fresh TPU VM:

1. Install uv

2. Clone the repo

3. Train (first run resolves + installs deps)

Primarily developed and maintained with coding agents

Train Train a model, then drop into an interactive game against it. The trained params and resolved model configuration auto-save to artifacts/alphazero_<env>.safetensors. uv run train --env ttt uv run train --env connect4 uv run train --env hex5 uv run train --env chess Supported games:

game envs status

Tic-Tac-Toe ttt solid, reaches perfect play

Hex hex4–hex9 solid up to 8x8 (hex9 is less tested)

Chess chess solid, reaches strong play given enough compute

Connect4 connect4 reaches perfect play outcomes, can struggle to maintain it

Go go3–go9 trains well up to 9x9

Options:

flag effect

--save PATH custom save path for the checkpoint

--no-save train only, don't write a checkpoint

--no-play train only, skip the interactive game afterward

--enable-wandb log metrics and upload versioned model artifacts

Watch it train (all in the terminal) Metrics are logged straight to the terminal. It periodically prints an ASCII loss curve and other env-specific diagnostics: ➜ ~ uv run train --env connect4 Warmup finished in 26.1s. Model has 9,458,308 parameters. Starting training for 2000 cycles... Cycle 1/2000 | 1.74s phase1 selfplay | p1_wins=317 p2_wins=222 ties=0 n_legal_avg_mid=6.77 phase2 drain | consumable=71585 slices=8 phase3 train | loss=0.8216 loss_v=0.5008 loss_pi=0.3209 batch[r+=4359 r-=3826 r0=7 valid=8192 invalid=0] Cycle 2/2000 | 1.74s phase1 selfplay | p1_wins=297 p2_wins=204 ties=0 n_legal_avg_mid=6.79 phase2 drain | consumable=70214 slices=8 phase3 train | loss=0.8335 loss_v=0.4998 loss_pi=0.3337 batch[r+=4407 r-=3784 r0=1 valid=8192 invalid=0] Cycle 3/2000 | 1.74s phase1 selfplay | p1_wins=299 p2_wins=229 ties=0 n_legal_avg_mid=6.78 phase2 drain | consumable=69670 slices=8 phase3 train | loss=0.8200 loss_v=0.4980 loss_pi=0.3220 batch[r+=4383 r-=3804 r0=5 valid=8192 invalid=0]

...

Cycle 48/2000 | 1.75s phase1 selfplay | p1_wins=167 p2_wins=129 ties=25 n_legal_avg_mid=6.11 phase2 drain | consumable=110212 slices=13 phase3 train | loss=0.9877 loss_v=0.3518 loss_pi=0.6359 batch[r+=3923 r-=3513 r0=756 valid=8192 invalid=0] Cycle 49/2000 | 1.75s phase1 selfplay | p1_wins=162 p2_wins=120 ties=17 n_legal_avg_mid=6.09 phase2 drain | consumable=108795 slices=13 phase3 train | loss=0.9803 loss_v=0.3480 loss_pi=0.6323 batch[r+=3792 r-=3632 r0=768 valid=8192 invalid=0] Cycle 50/2000 | 1.75s phase1 selfplay | p1_wins=182 p2_wins=122 ties=20 n_legal_avg_mid=5.94 phase2 drain | consumable=104636 slices=12 phase3 train | loss=0.9567 loss_v=0.3485 loss_pi=0.6083 batch[r+=3882 r-=3546 r0=764 valid=8192 invalid=0]

── loss over last 50 cycles ── 1.3409 ┤ ╭╮ 1.2984 ┤ ╭─╯╰──╮ 1.2559 ┤ ╭─╯ ╰╮ 1.2134 ┤ ╭──╯ ╰──╮ 1.1709 ┤ ╭╯ ╰╮ 1.1284 ┤ ╭─╯ ╰╮ 1.0859 ┤ ╭╯ ╰───╮ 1.0434 ┤ ╭╯ ╰──╮ 1.0009 ┤ ╭─╯ ╰──╮ 0.9584 ┤ ╭─╯ ╰── 0.9159 ┤ ╭──╯ ╭╮ 0.8734 ┤ ╭─╯ ╭─╯╰──╮ 0.8309 ┤╭╮ ╭─╯ ╭─╯ ╰─╮ 0.7884 ┤╯╰─╯ ╭──╯ ╰─╮ 0.7459 ┤ ╭╯ ╰─╮ 0.7034 ┤ ╭─╯ ╰─╮╭╮ ╭╮ 0.6609 ┤ ╭╯ ╰╯╰─╯│ 0.6184 ┤ ╭╯ ╰────╮ 0.5759 ┤ ╭╯ ╰ 0.5334 ┤ ╭╯ 0.4909 ┤───────╮ ╭─╯ 0.4484 ┤ ╰╭─╯──╮ 0.4059 ┤ ╭──╯ ╰──────────────────────╮ 0.3634 ┤ ╭╯ ╰─────────╮ 0.3209 ┤────╯ ╰── ● total ● value ● policy

...

Cycle 1998/2000 | 1.75s phase1 selfplay | p1_wins=172 p2_wins=116 ties=40 n_legal_avg_mid=6.04 phase2 drain | consumable=105882 slices=12 phase3 train | loss=0.0731 loss_v=0.0186 loss_pi=0.0545 batch[r+=3409 r-=3007 r0=1776 valid=8192 invalid=0] Cycle 1999/2000 | 1.75s phase1 selfplay | p1_wins=174 p2_wins=116 ties=21 n_legal_avg_mid=6.05 phase2 drain | consumable=109770 slices=13 phase3 train | loss=0.0647 loss_v=0.0153 loss_pi=0.0494 batch[r+=3400 r-=3085 r0=1707 valid=8192 invalid=0] Cycle 2000/2000 | 1.75s phase1 selfplay | p1_wins=147 p2_wins=127 ties=33 n_legal_avg_mid=6.02 phase2 drain | consumable=106517 slices=13 phase3 train | loss=0.0692 loss_v=0.0164 loss_pi=0.0528 batch[r+=3268 r-=3120 r0=1804 valid=8192 invalid=0]

── ladder Elo over cycles 50..2000 ── 2137.5776┤ ╭───────╮╭──╮╭───╮╭╮╭───────────── 1959.4461┤ ╭─╯ ╰╯ ╰╯ ╰╯╰╯ 1781.3146┤ ╭─╯ 1603.1832┤╭─╯ 1425.0517┤│ 1246.9202┤│ 1068.7888┤│ 890.6573 ┤│ 712.5259 ┤│ 534.3944 ┤│ 356.2629 ┤│ 178.1315 ┤│ 0.0000 ┤╯ ● total

Training finished in 3627.1s. ✅ Saved model params to artifacts/alphazero_connect4.safetensors

================================================== Playing connect_four. You are 'X' (player 1). Enter a column number (1-7) to drop your piece. Commands: undo, restart, quit ==================================================

. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1 2 3 4 5 6 7

Your move (X):

Train a custom game in Colab We can train AlphaZero in Colab on a completely new game by hoisting the core logic from this package. Run the notebook in Colab on a TPU.

Eval Chess Run head-to-head matches or round-robin tournaments between nanoAlphaZero checkpoints, DeepMind's Searchless Chess (search-free transformer models), and Stockfish. uv sync --group dev

fetch openings, reference models, and BayesElo

reproduce the 400-simulation tournament between a 24h trained model, 48h trained model and searchless270m

score an existing games.pgn with BayesElo, e.g. reproducing the Searchless

Chess paper's all-models-vs-Stockfish matchup

they're JAX/Haiku based we can run batched tournaments against them on TPUs they don't require search, so one model inference should reproduce the same playing ability

Here's how our 10x256nbt models do against Searchless Chess 270M.

Checkpoint Training Search W-D-L vs 270M Score Relative Elo

model34400 24h 400 sims 205-144-163 54.1% +27

model34400 24h 800 sims 256-131-125 62.8% +90

model68800 48h 400 sims 261-131-120 63.8% +99

model68800 48h 800 sims 325-103-84 73.5% +177

For more details, see the evaluation report. Hex On a solved game like hex we can watch the value head acquire perfect play directly. The trainer periodically prints its verdict on every Black opening move next to the known perfect-play outcome. As training proceeds, MSE vs perfect falls toward 0 and sign accuracy climbs to 1.000. For example, these are the value head outputs for each opening once training finishes. hex6 Perfect play (B=Black wins / .=Black loses), [x]=model sign mismatch: . . . . . B . B B B B B B B B B B B B B B B B B B B B B B . B . . . . . MSE vs perfect = 0.0004 | sign accuracy = 1.000

hex7 Perfect play (B=Black wins / .=Black loses), [x]=model sign mismatch: . . . . . . B . . B . B B B . B B B B B . B B B B B B B . B B B B B . B B B . B . . B . . . . . . MSE vs perfect = 0.0082 | sign accuracy = 1.000

Compare this to the known perfect play outcomes in Hex

hex6 bottom row, right hex7 bottom row, middle

MSE vs perfect is a good proxy for progress, but a near-zero MSE does not guarantee the model can reliably beat a strong opponent like MoHex, so we set configs to keep training well past that point.

hex7: value head progression Cycle 1 — value head is flat ~0, it knows nothing yet. MSE 0.9996 | sign acc 0.673 Cycle 1/5000 | 8.40s

--- Hex value-head after each Black opening (value = White-to-move perspective; negative => Black-winning) --- -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 -0.00 +0.00 -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 +0.00 -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 +0.00 -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 -0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 +0.00 Perfect play (B=Black wins / .=Black loses), [x]=model sign mismatch: [.][.][.][.][.][.] B . . [B] . [B][B] B . B B B B B [.] [B] B B B B B B . B B B B B [.] B B B [.] B [.][.] [B] . . . . . . MSE vs perfect = 0.9996 | sign accuracy = 0.673

Cycle 100 — coarse map forming, many sign errors remain. MSE 0.6118 | sign acc 0.837 Cycle 100/5000 | 8.39s

--- Hex value-head after each Black opening (value = White-to-move perspective; negative => Black-winning) --- +0.46 +0.43 +0.42 +0.59 +0.61 +0.65 -0.59 +0.05 +0.21 +0.41 +0.49 +0.15 -0.36 -0.18 +0.16 +0.19 +0.20 +0.45 -0.78 +0.06 +0.47 -0.14 -0.29 -0.59 -0.73 -0.74 -0.39 -0.08 +0.59 -0.04 -0.70 +0.62 -0.28 -0.22 +0.24 -0.31 -0.67 -0.02 +0.04 +0.04 +0.18 +0.22 -0.57 +0.53 +0.46 +0.54 +0.44 +0.35 +0.39 Perfect play (B=Black wins / .=Black loses), [x]=model sign mismatch: . . . . . . B . . [B] . [B] B B . [B][B][B] B [B] . B B B B B B B . B B [B] B B . B B B . [B] . . B . . . . . . MSE vs perfect = 0.6118 | sign accuracy = 0.837

Cycle 250 — magnitudes sharpening toward ±1, only 2 sign errors left. MSE 0.1272 | sign acc 0.959 Cycle 250/5000 | 8.40s

--- Hex value-head after each Black opening (value = White-to-move perspective; negative => Black-winning) --- +0.92 +0.88 +0.80 +0.88 +0.87 +0.82 -0.94 +0.28 -0.11 -0.59 -0.26 -0.96 -1.00 -0.99 +0.24 -0.98 -0.97 -0.99 -1.00 -0.98 +0.54 -0.59 -0.86 -1.00 -1.00 -1.00 -0.95 -0.89 +0.16 -0.98 -1.00 -0.99 -0.99 -0.99 +0.40 -0.95 -1.00 -0.94 +0.63 -0.78 +0.89 +0.53 -0.94 +0.74 +0.94 +0.91 +0.91 +0.91 +0.93 Perfect play (B=Black wins / .=Black loses), [x]=model sign mismatch: . . . . . . B . [.] B [.] B B B . B B B B B . B B B B B B B . B B B B B . B B B . B . . B . . . . . . MSE vs perfect = 0.1272 | sign accuracy = 0.959

Cycle 500 — all signs correct, magnitudes nearly saturated. MSE 0.0260 | sign acc 1.000 Cycle 500/5000 | 8.40s

--- Hex value-head after each Black opening (value = White-to-move perspective; negative => Black-winning) --- +1.00 +0.99 +0.99 +0.99 +0.99 +0.98 -0.99 +0.65 +0.08 -0.89 +0.98 -0.97 -1.00 -0.99 +0.95 -1.00 -1.00 -0.99 -1.00 -0.99 +0.90 -0.87 -0.94 -1.00 -1.00 -1.00 -0.94 -0.53 +0.93 -0.99 -1.00 -1.00 -1.00 -1.00 +0.99 -0.94 -1.00 -0.98 +0.97 -0.90 +0.98 +0.99 -0.97 +0.98 +1.00 +1.00 +1.00 +1.00 +1.00 Perfect play (B=Black wins / .=Black loses), [x]=model sign mismatch: . . . . . . B . . B . B B B . B B B B B . B B B B B B B . B B B B B . B B B . B . . B . . . . . . MSE vs perfect = 0.0260 | sign accuracy = 1.000

Cycle 1000 — pretty much close to perfect play outcomes. MSE 0.0007 | sign acc 1.000 Cycle 1000/5000 | 8.40s

--- Hex value-head after each Black opening (value = White-to-move perspective; negative => Black-winning) --- +1.00 +1.00 +1.00 +0.99 +1.00 +1.00 -0.99 +0.99 +0.99 -0.90 +0.99 -1.00 -1.00 -0.99 +0.98 -1.00 -1.00 -1.00 -1.00 -1.00 +0.91 -0.95 -0.93 -1.00 -1.00 -1.00 -0.97 -0.96 +0.98 -1.00 -1.00 -0.98 -1.00 -1.00 +1.00 -0.99 -1.00 -0.99 +0.97 -0.96 +0.99 +1.00 -0.99 +0.99 +0.99 +0.99 +1.00 +1.00 +1.00 Perfect play (B=Black wins / .=Black loses), [x]=model sign mismatch: . . . . . . B . . B . B B B . B B B B B . B B B B B B B . B B B B B . B B B . B . . B . . . . . . MSE vs perfect = 0.0007 | sign accuracy = 1.000

Todo

Add documentation for mcts.py's Gumbel MuZero implementation Add documentation for how to support custom games Verify larger hex boards still hit perfect play Test Go models against reference opponent

Acknowledgements This project would not have been possible without the amazing work of the following:

MCTX — search algorithm (paper) PGX — game environments Flashbax — replay buffers Scaling Scaling Laws — experiments & model architecture KataGo — model architecture and countless training methods TPU Starter — TPU setup guide Research supported with Cloud TPUs from Google's TPU Research Cloud (TRC)

source