In the search-and-learning lesson we sketched the recipe: a network gives MCTS a prior, MCTS gives the network a target, the loop closes. Here we open the hood. We derive the PUCT formula AlphaZero actually uses, write down its loss term by term, and then watch MuZero strip away the one assumption AlphaZero still depends on — the simulator. What's left is the same recipe applied to Atari, Go, chess, shogi, and (in 2024) Olympiad math. Three networks, three losses, one self-play loop.
Learning Objectives
After this lesson, you will be able to:
Derive PUCT — the AlphaZero variant of UCB1 — and explain exactly what term the neural policy adds to the standard exploration bonus
Write down the AlphaZero loss (value MSE + policy cross-entropy + L2) term by term and understand why each piece is needed
Distinguish AlphaZero (needs a perfect simulator) from MuZero (learns a latent dynamics), and trace MuZero's three networks: representation, dynamics, prediction
Read MuZero's K-step unrolled loss and explain why K > 1 is required to keep the learned dynamics from collapsing
Connect the AlphaZero recipe to modern reasoning systems — AlphaProof, AlphaGeometry, and the o1/R1 lineage of math-and-code RL
Don't be intimidated by the names. AlphaZero and MuZero are just MCTS with a smarter prior, trained by self-play. Once you've seen the three equations — PUCT, AZ loss, MuZero loss — the rest is engineering.
where s is the raw board state (no human features), π is a probability distribution over legal actions, and v ∈ [-1, 1] is the predicted value of s from the current player's perspective. One forward pass, two outputs. Silver et al. 2017 ("Mastering the Game of Go without Human Knowledge") used a ResNet with 20–40 blocks; newer reproductions often swap in a transformer.
Standard MCTS uses UCB1 to choose which child to descend during selection:
a∗=argamax[Q(s,a)+cN(s,a)lnN(s)]
UCB1 treats every move equally before it has data. For a 19x19 Go board with 361 possible moves, that's a problem — most moves are obviously bad and shouldn't get any visits at all. AlphaZero fixes this by inserting the network's policy prior P(s, a) directly into the selection rule. This is the PUCT formula:
a∗=argamax[Q(s,a)+cpuct⋅P(s,a)⋅1+N(s,a)N(s)]
Three things change between UCB1 and PUCT:
The denominator becomes 1 + N(s, a) instead of N(s, a) — no special-case for unvisited children, the prior already biases toward them.
The numerator becomes √N(s) instead of √(ln N(s)) — grows faster, gives the prior longer-lasting influence.
The bonus is multiplied by P(s, a) — the network's prior on this move.
If the network is uniform (every move equally likely), PUCT is roughly UCB1 with a different schedule. If the network is sharp, PUCT spends almost all the simulation budget on the moves the network believes in — which is what you want with a finite simulation count and a 250-wide branching factor.
The four MCTS phases (selection, expansion, simulation, backup) survive into AlphaZero, but two of them change in important ways:
Selection. Walk down the tree from the root, picking the child that maximizes PUCT, until you reach a node that hasn't been expanded.
Expansion. Add the new node. Query the network: (P, v) = f_θ(s_new). Initialize children's priors to P, set the new node's value to v.
Simulation.No random rollouts. Classic MCTS plays random moves from the leaf until the game ends — slow and noisy on Go. AlphaZero replaces this entire phase with the network's value estimate v. One forward pass instead of a full random game.
Backup. Walk back to the root, incrementing N(s, a) and updating Q(s, a) as a running mean of the value estimates seen along the path. Crucially, when backing up across a turn change in a two-player game, flip the sign: a position that's good for me is bad for you.
After a fixed number of simulations (800 for Go, 50–400 for Atari), the agent picks its actual move from the visit-count distribution at the root:
π_MCTS(a | s) ∝ N(s, a)^(1/τ)
where τ is a temperature. During the first 30 moves of training games, AlphaZero uses τ = 1 for exploration; after that, τ → 0 (greedy on visit count). Visit count — not Q — because the most-visited child is the one MCTS has the most evidence about, while a high-Q but rarely-visited child might be an optimistic outlier.
Quick check
Why does AlphaZero use the visit-count distribution π_MCTS(a|s) ∝ N(s,a)^(1/τ) as its policy target instead of the network's own output π?
AlphaZero starts with random weights θ_0 — no pretraining, no human games, no opening book. Then:
The current network θ_t plays many games against itself. Each move uses MCTS guided by θ_t.
For every visited state s, record the MCTS visit distribution π_MCTS(a | s) — this becomes the policy target.
When the game ends with result z ∈ {-1, 0, +1} (from the player-to-move's perspective in each position), assign z as the value target for every state in the game.
Train θ_{t+1} by gradient descent on the loss below using a buffer of recent self-play data.
Replace θ_t ← θ_{t+1}. Repeat for millions of games.
That's the entire training loop. No human expert data, no curriculum, no demonstrations. The only inputs are the rules of the game (for the simulator) and the initial state.
The value MSE forces the network to be calibrated. If the network says v = 0.8 for a position and the game ends in a loss (z = -1), the gradient pulls v down. Over millions of self-play games, v_θ(s) converges to the true expected outcome under the search-augmented policy.
The cross-entropy is KL divergence in disguise. Up to constants in π_MCTS, minimizing -π_MCTS · log p is the same as minimizing KL(π_MCTS || p_θ). The network is being pulled to match the MCTS distribution.
There is no separate "policy gradient" term. AlphaZero is not REINFORCE. It's pure supervised learning, where the labels come from MCTS rollouts.
L2 is essential. Without it, the network overfits the recent replay buffer and self-play diversity collapses — every game becomes a copy of the last.
What Do You Think?
AlphaProof — DeepMind's 2024 IMO silver-medalist — is best described as AlphaZero applied to what domain?
AlphaZero has one assumption that limits its reach: it needs a perfect simulator. To expand a node during MCTS, you have to know what state results from taking an action. For Go, chess, and shogi this is fine — the rules give you the transition function for free. For Atari from pixels, for a real robot, or for any problem where the dynamics are unknown, AlphaZero simply cannot run.
MuZero (Schrittwieser et al. 2019/2020 "Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model") removes this dependency. The trick: do MCTS in a learned latent space, where the "states" are just internal representations the network is free to define however helps the loss go down. There are three networks:
The key observation: the latent state s_k has no fixed meaning. It's whatever vector the dynamics network finds useful for predicting reward and value. It doesn't have to match the actual Atari frame, the actual board position, or anything physically interpretable. The training loss is the only thing that constrains it.
To prevent the latent dynamics from collapsing to a degenerate solution, MuZero unrolls K steps from each training position and demands consistency across all of them:
Why K > 1 is essential. If K = 0, the dynamics network never has to produce a useful next latent — only the representation head matters. If K = 1, the dynamics is constrained at one step out, but nothing prevents it from "forgetting" the action by step 2. K = 5 (the value used in MuZero) forces five consecutive single-step predictions to all be consistent with real observed rewards and bootstrapped returns. The dynamics has to encode information about the action long enough to predict rewards five steps later.
z_k is a bootstrapped n-step return. Specifically, z_k = ∑_{i=0}^{n-1} γ^i r_{t+k+i+1} + γ^n ν_{t+k+n}, where ν is the value network's estimate at the n-step horizon. This is the same n-step TD target used elsewhere in deep RL — it interpolates between Monte Carlo (high variance, unbiased) and one-step bootstrapping (low variance, biased).
The policy target π_MCTS,k comes from MCTS run in real environment time, not in latent space. During self-play, you run MCTS at every real timestep and record its visit distribution. During training, those distributions become targets at the corresponding latent step.
What Do You Think?
In MuZero's K-step loss, why is K > 1 essential? What would break if you trained with K = 1?
Quick check
MuZero uses three separate networks. Match them to their roles:
AlphaZero needs the rules of the game. MuZero needs only four things:
A stream of observationso_0, o_1, o_2, ... — pixels, sensor readings, raw text tokens, whatever.
A discrete action space — must be enumerable (so MCTS can branch over it).
A scalar reward signalr_t — even sparse rewards (zero everywhere except at game-end) work.
A done flag — when each episode terminates.
That's a much weaker requirement than "give me the transition function." It's the same interface OpenAI Gym exposes for every Atari game, every robotics simulation, every gridworld. Which is why a single MuZero implementation, with a single set of hyperparameters, beats AlphaZero at Go and sets state-of-the-art on Atari from pixels in the same paper.
The generality cuts both ways: in domains where the rules are available cheaply (board games), AlphaZero's perfect simulator gives it a small edge. In domains where they aren't (Atari, robotics, real-world control), MuZero is the only one of the two that can be applied at all.
#The Recipe Outside Games: AlphaProof and Reasoning RL
The same loop now powers reasoning models. The mapping is direct:
AlphaZero on Go
AlphaProof / Reasoning RL
Board state s
Partial proof / partial reasoning trace
Legal moves
Tactic applications / next-token candidates
Network f_θ(s) = (π, v)
LLM head — policy over next tokens, value of partial solution
Game result z ∈ {-1, +1}
Verifier verdict — proof checker, code executor, math grader
Self-play
Self-generation of attempted solutions
MCTS over moves
Search over reasoning steps / tree-of-thought
This is exactly the training loop in DeepMind's AlphaProof (Trinh et al. 2024) and the o1 / R1 lineage of math-and-code RL modelsReasoning ModelsReasoning models are LLMs trained to perform extended chain-of-thought reasoning before producing a final answer, improving performance on complex tasks.Learn more →. The key ingredient that makes it work in these domains is a verifiable reward: a Lean proof checker, a Python test runner, a math grader. Without verifiable reward, MCTS has nothing to back up, and the AlphaZero loop reduces to RLHF, which is noisier and gameable.
Reproductions of AlphaZero and MuZero converge on a similar set of numbers, which is useful both for sanity-checking your own implementation and for understanding the compute budget.
Simulations per move at training time. Go: 800. Chess/shogi: 800. Atari (MuZero): 50–100. The training process is robust to lowering this; the played strength at deployment scales roughly linearly with simulations per move.
c_puct. Typically 1.0–2.5. Schrittwieser et al. use 1.25. Larger values bias toward exploration; smaller values trust the network's prior more.
Dirichlet noise at the root. Add α · Dir(α_0) to the prior at the root for exploration. α_0 ≈ 0.3 for Go, scaled inversely with branching factor for other games. Without it, self-play collapses to a single deterministic policy and stops generating diverse data.
Replay buffer size. Order of 1M positions (Go) to 10M (Atari). Older self-play data ages out.
Network depth. Go uses 20–40 ResNet blocks. Modern reproductions use transformers of comparable parameter count.
Training scale. Original AlphaZero Go: ~5M games of self-play across thousands of TPUs. MuZero on Atari: ~1B environment steps. Educational reproductions on toy games (tic-tac-toe, Connect 4) converge in minutes on a laptop.
The playground below is a complete, runnable AlphaZero implementation on a 3x3 tic-tac-toe board. It uses a tiny two-layer policy/value network, runs MCTS with PUCT, generates self-play games, and trains the network on the AlphaZero loss. After 100 self-play games and 200 gradient steps, the win rate against a random opponent climbs from ~33% (chance for the first player on tic-tac-toe) to near-perfect.
Edit anything. Try:
Setting C_PUCT = 0.1 (network trusted too much) or C_PUCT = 10 (search too random) and watch training stall.
Setting N_SIMS = 5 and noting how the learned policy gets noisy.
Setting SELFPLAY_GAMES = 5 and watching the value head fail to converge for lack of data.
Loading visualization...
What to look for in the output:
The first eval line (before any training) is close to random — roughly the chance a first-mover gets on tic-tac-toe.
By 60–80 self-play games, losses should drop near zero. Tic-tac-toe is a draw under optimal play, so a fully-trained agent draws or wins but never loses against random.
The losses-going-to-zero behavior is the AlphaZero loop visibly working in your browser. There is no domain knowledge in this script — only the rules of tic-tac-toe in step() and winner().
The frontier in 2025 is applying the AlphaZero recipe to domains that aren't board games. The recipe needs three ingredients:
A model. AlphaZero needs a simulator; MuZero needs only an observation stream. For LLM reasoning, the model is the language model itself acting as both policy and dynamics — every generated token is both a "move" and the next "state."
A verifier. Game outcome for Go; Lean proof checker for AlphaProof; unit tests for code RL; math graders for o1/R1. Without a verifier, MCTS has nothing to back up. RLHF's learned reward model is a noisy substitute that introduces gaming.
A search procedure over a discrete-enough action space. AlphaZero searches over board moves; AlphaProof over tactic applications; reasoning RL over CoT branches or tool calls. Continuous high-dimensional actions (robotics) push toward MCTS variants or policy gradient methods instead.
When all three are present — a model, a verifier, and a search — the AlphaZero loop applies. When any is missing, you get something weaker (RLHF, behavior cloning, vanilla policy gradient). The bet behind reasoning RL is that more and more domains will become verifier-rich as auto-grading, formal methods, and tool execution mature.
The leap from AlphaGo (2015) to AlphaProof (2024) is the same algorithm with the simulator stripped away and the verifier swapped out — and that closes out the Reinforcement Learning track. Eighteen lessons from the Bellman equation to MuZero-in-latent-space; you can now read the policy-gradient and search-and-learning literature on its own terms. Next, Track 8 (RAG & Knowledge Systems) shifts to a different problem entirely: how do you give a frozen LLM access to a knowledge base it was not trained on? The retrieval, ranking, and evaluation patterns there are the bridge into agent systems in Track 9, where reasoning modelsReasoning ModelsReasoning models are LLMs trained to perform extended chain-of-thought reasoning before producing a final answer, improving performance on complex tasks.Learn more → trained with the very recipes you just learned start to plan and act.