MuZero and AlphaZero

Why Does This Exist?

AlphaZero (Silver and colleagues, 2017) mastered chess, shogi, and Go from scratch by pairing a neural network with Monte Carlo tree search. But it cheated in one respectable way: at every new node it asked the true game engine for the exact next board. Take away the rulebook, put it in front of pixels or a robot, and planning stalls because there is no simulator to query.

Early model-based RL answered by reconstructing observations, predicting next video frames pixel by pixel. That spends capacity on leaves fluttering in the background while errors compound down the imagined rollout. MuZero (Schrittwieser and colleagues, 2020) cut the knot: a planner only needs the three quantities that steer decisions, the policy $\mathbf{p}$, the value $v$, and the reward $r$. It learns a compact latent simulator for exactly those, searches inside it, and matched AlphaZero on board games while topping the Atari benchmark with no rules given. For the sibling that dreams in latent space differently, see Dreamer.

The one idea you need first

Tree search picks moves with a policy guided bonus. AlphaZero and MuZero use PUCT: $Q(s, a)$ plus the prior $P(s, a)$ times an exploration fraction that shrinks as the edge gets visits. The network supplies the prior and the value; search sharpens both into visit counts; self-play turns those visits into the next training targets.

Think of It Like This

How It Actually Works

Three learned functions replace the simulator. The representation $h$ encodes past observations into a root hidden state $s^0$. The dynamics $g$ steps the hidden state forward: $(r^k, s^k) = g(s^{k-1}, a^k)$, predicting the immediate reward and next hidden state with no obligation to decode pixels. The prediction $f$ reads any hidden state out: $(\mathbf{p}^k, v^k) = f(s^k)$, a policy prior plus a value estimate.

Search runs MCTS entirely among hidden states. Selection maximizes the PUCT score:

$Q$ exploits, the prior $P$ steers toward plausible moves, and the fraction explores the unvisited. Because rewards differ wildly between games, MuZero min-max normalizes every $Q$ into $[0, 1]$ from the tree's own observed range before scoring. Expansion calls $g$, evaluation calls $f$, and backup carries $G = r + \gamma v$ up the path into running averages. After search the real move is sampled from visit counts with a temperature $T$.

Learning without reconstructing

Training unrolls the model $K = 5$ steps along real trajectories and matches three targets: the search visit distribution for policy, bootstrapped $n$-step returns for value, and observed rewards for $r$, plus weight regularization. Gradients flow back through $g$ into $h$, so hidden states keep exactly the information the three targets need. No pixel loss exists anywhere, which is why the model stays small and the rollouts stay honest.

Worked example

Root hidden state with $c = 1.0$, $\gamma = 0.9$, parent visits $4$. Edges: $a_1$ has $Q = 1.2$, prior $0.7$, $N = 3$; $a_2$ has $Q = 0.8$, prior $0.3$, $N = 1$. Since $\sqrt{4} = 2$:

Search picks $a_1$. Dynamics predicts $r^1 = 0.5$ with next hidden state, prediction returns leaf value $v^1 = 2.0$. The backup return is $G = 0.5 + 0.9 \times 2.0 = 2.3$, and the edge average moves from $3 \times 1.2 = 3.6$ over $3$ visits to $5.9$ over $4$:

The diagram draws this selection, expansion, and backup with the same numbers.

Code

This snippet replays the worked example: PUCT scoring, the dynamics reward plus leaf value backup, and the averaged update.

1import math2
3
4def puct(q: float, p: float, n_parent: int, n_child: int, c: float = 1.0) -> tuple[float, float]:5    """Exploration bonus and total PUCT score for one latent edge."""6    u = c * p * math.sqrt(n_parent) / (1 + n_child)7    return u, q + u8
9
10u1, s1 = puct(1.2, 0.7, 4, 3)11u2, s2 = puct(0.8, 0.3, 4, 1)12print(f"a1: U={u1:.4f} score={s1:.4f}")  # -> a1: U=0.3500 score=1.550013print(f"a2: U={u2:.4f} score={s2:.4f}")  # -> a2: U=0.3000 score=1.100014g = 0.5 + 0.9 * 2.015print(f"G={g:.4f} Q={((3 * 1.2) + g) / 4:.4f}")  # -> G=2.3000 Q=1.4750

The printout matches the hand computation step for step: $1.55$ against $1.10$, backup $2.3$, updated $Q$ of $1.475$.

Watch Out For

The Quick Version

• AlphaZero plans with a perfect simulator plus search; MuZero learns the simulator from play and searches inside it.
• Three networks split the job: $h$ encodes observations, $g$ steps hidden states with rewards, $f$ reads out policy and value.
• Nothing reconstructs observations; policy, value, and reward targets shape the hidden states alone.
• Latent MCTS uses PUCT selection with min-max normalized $Q$ values and backs $r + \gamma v$ up each path.
• Short $K = 5$ training unrolls and hidden state normalization keep imagined rollouts from diverging.

What to Read Next

• Monte Carlo Tree Search (MCTS)
• Model-Based Deep RL
• Dreamer: World Models for RL