MuZero: Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model

Plan ahead in a game without ever being told its rules.

MuZero learns three small functions that predict only what planning needs (the reward of a move, how good a position is, and which move to try), then runs a tree search over them. The same algorithm reaches superhuman play in Go, chess, and shogi, and a new state of the art across 57 Atari games.

AlphaZero planned by playing moves out on an exact simulator of each game. MuZero runs the same search on a simulator it trains itself, and that simulator never predicts a board position or a pixel.

To plan, a program needs a model: something that answers "if I take this action here, what happens next?" A chess engine looks ahead by playing candidate moves out on an internal board, so it needs the rules: which moves are legal, where each one leads, and when the game is over. AlphaZero, DeepMind's program for chess, shogi and Go, had those rules built in. Before every move it ran 800 simulations of lookahead on an exact simulator of the game and scored the positions it reached with a neural network.

Most problems come without such a simulator. Nobody hands you the transition function of a robot's kitchen, or of an Atari game where the agent sees only the screen. For visually rich domains like Atari the field relied on model-free reinforcement learning, such as DQN and its descendants, which learn a policy or a value function directly from experience and never look ahead. Model-free agents held the Atari state of the art, and agents that learned a model of the game from pixels stayed below them. Model-free agents, for their part, cannot do the deep, exact lookahead that chess and Go reward.

MuZero (Schrittwieser et al., DeepMind; arXiv 2019, Nature 2020) keeps AlphaZero's search and replaces the given simulator with a learned one. The learned model is trained to predict three quantities and nothing else. A few ideas explain the method, in the order this page covers them: what a learned model has to predict and what it can ignore; the three networks that make those predictions; how Monte-Carlo tree search runs over them; how the output of the search becomes the training target for the same networks; and the details that let one algorithm handle board games that pay ±1 at the end and Atari games whose scores run into the thousands.

What a learned model has to predict

The obvious way to learn a model is to predict the world: given the current observation and an action, output the next observation, and grade the model by how closely its output matches the real next frame. World Models (Ha & Schmidhuber, 2018) trains an autoencoder to reconstruct every frame before it learns anything about control, and SimPLe, the strongest model-based Atari agent before MuZero, learns a video-prediction model of each game. Pixel prediction is hard. Errors compound when the model's own output becomes its next input, and most of the model's capacity goes into detail such as background textures or score digits that has no bearing on which action is best.

MuZero starts from what a tree search reads off a model. At each imagined step the search uses three quantities: the reward of a move, so it can add up the score of a line of play; the value of a position, so it can judge a line without playing it to the end; and a policy, a probability over moves that tells the search which moves to try first. No step of the search asks the model to draw the board, so MuZero trains its model on those three predictions only. Two models that give the same rewards and values for every sequence of actions are interchangeable for planning, however different their internal states are. Grimm et al. later formalized this as the value equivalence principle (2020); the Predictron and Value Prediction Networks had applied it before MuZero on a smaller scale.

The model's internal state, called the hidden state, is therefore not required to reconstruct the observation, and not required to match the true state of the game either. It is a tensor that training shapes into whatever makes the three predictions accurate. A strong chess player calculating a line does something similar: they track that a square is weak and that the line wins material, without picturing each future board in detail. The analogy breaks in two places. The player could draw the board on request and is moving pieces by the real rules, while MuZero's hidden state cannot be decoded into a board at all and its dynamics are invented by training.

Figure 1 isolates the difference. The left panel is a small screen with an agent, a goal, and a slider of cells that have nothing to do with the task. A model graded on reconstructing the screen has to encode every lit cell, so its load grows with the clutter. The reward, value and policy depend only on the agent and the goal, so they stay fixed at every setting of the slider.

MuZero has no decoder, so nobody can render its imagined future as a sequence of frames, and nothing checks that a hidden state corresponds to a real position. Later world-model agents such as DreamerV3 keep a decoder and train their model to reconstruct observations alongside predicting rewards.

How MuZero works: three learned functions

The model is three networks trained together, written as one parameter set $\theta$:

The representation function $h_\theta$ turns the recent observations into the root hidden state $s^0$. In Go and shogi it reads the last 8 board positions; in chess the last 100, so that it can predict draws correctly; in Atari the last 32 frames at 96×96 pixels together with the 32 actions that produced them. For Go, $h_\theta$ outputs a 19×19×256 tensor: one 256-number vector per point of the board. For Atari, stride-2 convolutions and average pooling shrink the 96×96×128 input (96 colour planes, 32 action planes) through 48×48, 24×24 and 12×12 down to a 6×6×256 hidden state.

The dynamics function $g_\theta$ takes a hidden state and an action and returns the immediate reward $r^k$ of that action and the next hidden state $s^k$. The action enters as extra planes stacked onto the hidden state at the same resolution. In Go one plane carries a single 1 where the stone is played (a pass is an all-zero plane), so $g_\theta$ reads a 19×19×257 input. Chess uses 8 action planes, shogi 11, and Atari a one-hot vector over its 18 actions tiled across the 6×6 grid. The same $g_\theta$ is applied again to its own output, which is what lets the model imagine several moves ahead. Both $h_\theta$ and $g_\theta$ are residual convolutional networks with AlphaZero's design, cut from 20 residual blocks to 16, with 256 planes and 3×3 kernels.

The prediction function $f_\theta$ reads a policy $p^k$ and a value $v^k$ off any hidden state. It has the same shape as AlphaZero's policy and value heads (one or two convolutions, then a fully connected layer), but it looks at a hidden state the model produced rather than an encoded board. In Go, $p^k$ has 362 entries (361 points plus pass); in Atari, 18.

The superscript $k$ counts imagined steps from the present. At $k=0$ the model is at the real current position; at $k=2$ it has imagined two actions, without touching the environment. Real time uses a subscript: $o_t$ is the observation at time $t$. Figure 2 shows the chain for up to five imagined steps. The hidden states are drawn as abstract stacks of cells because no picture of them exists.

Each prediction has its own training target, built from a real game the agent played: the reward $r^k$ is compared with the reward observed $k$ steps later, the value $v^k$ with a return computed from what happened afterwards, and the policy $p^k$ with the move distribution the search produced at that later step. The targets and the loss come after the search, because two of the three targets are outputs of the search.

Monte-Carlo tree search inside the learned model

Monte-Carlo tree search (MCTS) builds a tree of possible futures from the current position, one simulation at a time, and spends more simulations on the lines that look promising. In MuZero each node of the tree is a hidden state and each edge is an action. Every edge stores five statistics: a visit count $N$, a mean value $Q$, the prior probability $P$ from the prediction function, the predicted reward $R$, and the hidden state $S$ the action leads to. A simulation has three stages:

1. Selection starts at the root $s^0$ and walks down the tree, at each node taking the action with the highest pUCT score (the next section), until it takes an edge that has no node below it yet. Along the way, rewards and hidden states of edges already in the tree are read from $R$ and $S$, not recomputed.
2. Expansion calls the dynamics function once on that edge, $r^l, s^l = g_\theta(s^{l-1}, a^l)$, stores the results on the edge, and calls the prediction function once on the new state to get $p^l$ and $v^l$. The new node's edges start at $N=0$, $Q=0$, and $P = p^l$.
3. Backup walks back up the path from the new leaf to the root. Each edge on the path gets one more visit, and its mean value absorbs a new sample: the discounted sum of the predicted rewards below that edge plus the discounted leaf value.

That is one call to $g_\theta$ and one to $f_\theta$ per simulation, so the cost of a search is of the same order as AlphaZero's. The sample that backup adds at depth $k$ of a path whose leaf sits at depth $l$ is the paper's equation (3):

A two-step example with $\gamma = 0.997$: the path takes action $a^1$ (predicted reward 0), then $a^2$ (predicted reward 1), and the new leaf has value $v^2 = 0.5$. The middle node receives $1 + 0.997 \times 0.5 = 1.4985$, and the root receives $0 + 0.997 \times 1.4985 = 1.494$, which is $G^0$ from (3). Each $Q$ is the running mean of all the samples its edge has received. In the official pseudocode an edge's score adds the edge's reward to the discounted mean value of the node it leads to, which produces the same numbers. In two-player games the value changes sign at every ply on the way up, because a position good for one player is bad for the other.

# one simulation of MuZero's search (after the official pseudocode)
node = root                       # root = h(o_1..o_t), already expanded
path = [root]
while node.expanded:              # SELECTION: descend by the pUCT rule
    action, node = argmax_puct(node)
    path.append(node)
parent = path[-2]                 # the new leaf hangs off this node
r, s = g(parent.hidden, action)   # EXPANSION: one call to the dynamics
p, v = f(s)                       # one call to the prediction function
node.expand(reward=r, hidden=s, prior=p)
for n in reversed(path):          # BACKUP along the path, leaf to root
    n.visit += 1
    n.value_sum += v              # (sign-flipped for the opponent's turn)
    v = n.reward + gamma * v      # add this edge's reward, discount once

When the simulations are done, the search returns two things. The visit counts of the root's edges, normalized to sum to 1, form the search policy $\pi_t$, and the mean value at the root is the search value $\nu_t$. Both reappear as training targets. Figure 3 steps through 22 simulations of a small search. The prediction function's prior favours action A, the values the simulations back up favour B, and the visit counts move to B.

AlphaZero used the rules of the game in three places inside its search: to compute state transitions, to list the legal actions at each node, and to stop at terminal positions. MuZero replaces all three with its learned model. Transitions come from $g_\theta$. Illegal moves are masked only at the root, where the real environment can be asked which moves are legal; inside the tree there is no masking, and the prediction function learns from its training data to put little probability on moves that never occur in real games. Terminal positions get no special treatment: the search can imagine moves past the end of a game, and during training every state after the end is treated as absorbing, so the model learns to keep predicting the same state, zero reward and a zero value target there.

Root exploration noise also comes from AlphaZero. At the start of each search the root's priors are mixed with random noise, $0.75\,P + 0.25\,\eta$ where $\eta$ is drawn from a Dirichlet distribution (concentration 0.03 for Go, 0.3 for chess and 0.15 for shogi, as in AlphaZero, and 0.25 for Atari), so that self-play occasionally tries moves the network currently rates low.

Choosing what to simulate: the pUCT rule

At each node, selection has to choose between an action that already looks good and an action the search has barely tried. MuZero scores every action with this rule and takes the maximum:

The first term, $Q(s,a)$, is the mean value backed up through that action so far; in the official code an action with no visits contributes 0 here. The rest is an exploration bonus. It is proportional to the prior $P(s,a)$, it grows with the square root of the parent's visit count (the sum $\sum_b N(s,b)$ runs over every action at the node, $a$ included), and it shrinks as the action's own count $N(s,a)$ grows. The $+1$ keeps the bonus of an unvisited action finite. A move with a high prior gets tried early; once it has been visited several times its bonus has shrunk and its measured $Q$ decides whether it keeps getting visits.

The paper uses $c_1 = 1.25$ and $c_2 = 19652$. With those constants the bracketed factor is 1.2526 after 50 parent visits (the Atari budget) and 1.2900 after 800 (the board-game budget), so in every search the paper ran it stays within 3% of $c_1$. It would reach 2.87 at 80,000 visits; the log term exists for searches far larger than these. AlphaZero's version of the rule had a constant in its place. For a worked case, take a node visited 6 times and an action with prior 0.22 and 3 visits:

Figure 4 runs the rule for 26 selections at one node with four actions. The prior favours A and the learned values favour C. The first picks follow the prior, because every bonus is still large; each visit shrinks the visited action's bonus, and the picks settle on C.

Equation (2) adds a value to a bonus, so the two have to be on comparable scales. AlphaZero could assume values in $[0,1]$, since a board game ends in a win, a draw or a loss. Atari returns can be in the thousands, and a raw $Q$ of 700 next to a bonus of 0.5 makes the prior irrelevant. The paper rejects the game-specific fixes (rescaling by the highest score the game allows, or tuning $c_1$ per game) because each needs knowledge about the environment, and instead rescales $Q$ by the smallest and largest values seen in the current search tree:

$\bar Q$ replaces $Q$ in (2). Take a node with three actions, priors 0.50, 0.28 and 0.22, visit counts 2, 1 and 3, and values 0.35, 0.62 and 0.70 in a tree whose values span 0.20 to 0.80. The normalized values are 0.25, 0.70 and 0.83, the bonuses 0.51, 0.43 and 0.17, and B wins with 1.13. Multiply every return in the game by 1000 and the raw values become 350, 620 and 700: without (5) the bonuses, at most 0.51, are negligible and C wins every time, while with (5) the normalized values, and the pick, are unchanged. Divide every return by 100 instead and the raw values vanish next to the bonus, so the prior alone picks A. In board games the official code starts the minimum and maximum at −1 and 1, which reproduces AlphaZero's behaviour exactly.

How search produces a better policy

The prediction function's policy $p$ comes from one forward pass. The search spends 50 simulations (Atari) or 800 (board games) testing that policy against the values of the positions it leads to, and summarizes the result as the visit distribution at the root:

A player who calculates for a minute usually picks a better move than their first instinct, and a player who keeps comparing the two learns better instincts. MuZero trains $p$ toward $\pi_t$ in the same way: each search is a policy improvement step, and the training step distils it back into the network so that the next search starts from a better prior.

The target uses visit counts rather than $Q$ values. An action that was visited once can have a high but noisy $Q$; pUCT keeps sending simulations to an action only while its $\bar Q$ plus bonus stays the highest at that node, so a large count means the action's value held up under repeated checking. The visit-count bars under Figure 3 are $\pi_t$ for that small search.

The move the agent actually plays is sampled from the same counts, sharpened or flattened by a temperature $T$:

The temperature is an exponent on the counts. At $T=1$ the agent samples in proportion to the counts, which is not uniform; as $T \to 0$ the most-visited action takes all the probability. Take a 50-simulation search over six actions (a real Atari search has 18) that ended with counts 21, 15, 8, 4, 2 and 0. At $T=1$ the six actions are played with probabilities 42%, 30%, 16%, 8%, 4% and 0%. At $T=0.5$ the counts are squared first and the top action rises to 58.8%; at $T=0.25$ they are raised to the fourth power and it reaches 78.0%. The action with no visits has probability 0 at every temperature. For Atari, MuZero uses $T=1$ for the first 500k training steps, 0.5 for the next 250k and 0.25 for the last 250k, so play becomes greedier as the network improves. Board games keep AlphaZero's scheme: in the official pseudocode, moves are sampled at $T=1$ for the first 30 moves of a game and the most-visited move is played after that.

These pieces form a loop. Search with the current network produces a better policy and a value estimate for each position played; the network is trained toward those; the next search starts from the improved network. The real rewards and game outcomes in the value targets tie the loop to what happened in the games.

Training targets and the loss

A training sample is a position $t$ from a stored game together with the next $K=5$ real actions. The learner runs $h_\theta$ on the observations up to $t$, then unrolls $g_\theta$ five times, feeding it the actions the agent really took. The model is given the real actions, not choices of its own, because only the real trajectory has targets: at every unrolled step $k$ there is a real reward $u_{t+k}$, a stored search policy $\pi_{t+k}$, and a value target $z_{t+k}$ computed from the rest of the game. The loss sums the three errors over all $K+1$ steps:

There is no reconstruction term. Gradients from all three losses flow back through every application of $g_\theta$ into $h_\theta$, so the hidden states are shaped only by how well they support the predictions at later steps. The reward term at $k=0$ has nothing imagined to predict, since no action has been applied to $s^0$ yet. Equation (1) as printed sums the reward term from $k=0$, and the official code fills that slot with the last real reward "for simplicity".

Figure 7 lays a stored trajectory above the model's unroll. Move the slider to see which stored numbers grade each step's predictions, and switch to board-game mode to see the reward loss disappear.

The value target depends on the domain. A board game pays nothing until the end, so the target is the final outcome, $z_t = u_T$, with a loss, draw and win coded as $-1, 0, +1$ and no discount. An Atari game can last 108,000 frames, too long to wait for, so MuZero uses an $n$-step return: the next $n$ real rewards, discounted, plus a discounted estimate of everything after them.

The estimate at the end is $\nu_{t+n}$, the root value of the search the agent ran at time $t+n$, not the value head's raw output $v$. The search value is the network's value refined by 50 simulations of lookahead, so it is usually the better estimate, and training $v$ toward it is the value-side counterpart of training $p$ toward $\pi$. In Figure 7 at $k=2$ the target is $z_{t+2}$: the rewards $u_{t+3}$ through $u_{t+12}$ are 0, 0, 3, 0, 0, 1, 0, 0, 2, 0, the search value ten steps on is $\nu_{t+12} = 1.4$, and with $\gamma = 0.997$:

Figure 8 separates the two choices inside $z_t$: how many real rewards to sum, and what to bootstrap from. With few real rewards the target leans on the bootstrap, which is only as accurate as the current estimate; with many it leans on real rewards, which are exact but vary from game to game, so the target gets noisier. MuZero uses $n=10$ for Atari and the discount $\gamma = 0.997$, which it takes from the R2D2 agent (R2D2 itself used $n=5$). At $n = 10$ the bootstrap carries weight $0.997^{10} = 0.970$.

The policy loss is cross-entropy against $\pi$ in both domains. The value and reward losses differ. Board games use AlphaZero's squared error for the value and drop the reward loss, since the only reward is the outcome at the end. Atari uses cross-entropy for both value and reward, over a categorical representation described in the next section, because the paper found it trained more stably than squared error when rewards and values vary in scale.

Handling rewards and values of any size

A value of 0.3 in one Atari game and 30,000 in another cannot share one regression target without one of them dominating the gradients. MuZero follows R2D2 and first compresses every value and reward target with an invertible function:

Near zero the function has slope 1/2, so small rewards keep their relative sizes: $h(0.3) = 0.140$. Far from zero it grows like $\sqrt{|x|}$, so $h(100) = 9.150$ and $h(845) = 28.93$. The small linear term $\varepsilon x$ takes over beyond about a million, where $0.001\,x$ exceeds $\sqrt{x}$. The inverse has a closed form:

On the compressed axis, MuZero predicts a distribution over 601 bins, one per integer from −300 to 300, instead of a single number. A target is written two-hot, split between its two nearest bins in proportion to how close it is to each: $h(100) = 9.150$ becomes 0.85 on bin 9 and 0.15 on bin 10. The network is trained with cross-entropy against that two-hot vector. To read out a value, take the expectation over the bins (here $0.85 \times 9 + 0.15 \times 10 = 9.15$) and apply $h^{-1}$, which returns 100. Because the bins sit on the compressed axis, the range ±300 covers raw returns up to about ±58,700.

One training step, end to end

Two kinds of process run in parallel. Actors play games with the latest network checkpoint (a new one every 1000 training steps): at each real position they run a search, sample a move by (6), and store the observation, the action, the reward, $\pi_t$ and $\nu_t$. Board games go to the replay buffer when they finish; Atari games, which last up to 30 minutes, are sent in pieces of 200 moves. The buffer holds the latest 1 million board games, or the latest 125,000 Atari sequences of length 200. Board-game positions are sampled uniformly. Atari positions are sampled with priority $|\nu_i - z_i|$, the gap between the search value and the $n$-step return, so positions where the search and the observed return disagree are trained on more often.

A learner takes mini-batches of 2048 positions (board games) or 1024 (Atari) and runs one million updates. For each position it computes the loss of equation (1) and one backward pass through all three networks:

# one training sample: position t of a stored game (after the pseudocode)
# stored per step: real actions a, rewards u, search policies pi, values nu
s = h(obs[:t + 1])                # representation -> root hidden state s0
loss = 0
for k in range(K + 1):            # K = 5, so k = 0..5
    p, v = f(s)                   # prediction heads at s^k
    z = nstep_return(u, nu, t + k)    # bootstraps from nu, not from v
    l = xent(pi[t + k], p) + value_loss(z, v)
    if k > 0:
        l += reward_loss(u[t + k], r) # r came out of g one step earlier
    loss += l if k == 0 else l / K    # recurrent steps weighted 1/K
    if k < K:
        r, s = g(s, a[t + k + 1]) # dynamics, fed the REAL next action
        s = scale_grad(s, 0.5)    # halve the gradient flowing back here
loss += c * l2(theta)             # weight decay
loss.backward()                   # one backward pass through h, g and f

Two gradient scalings in that listing do different jobs. The heads at the five recurrent steps are weighted $1/K$ (the official code keeps weight 1 on the $k=0$ step), so the total gradient has about the same size whatever $K$ is. Separately, scale_grad(s, 0.5) leaves the forward value unchanged and halves the gradient flowing back into each hidden state that $g_\theta$ produced. The same $g_\theta$ runs five times per sample, and without the halving each application would collect gradient from every loss after it. With it, a loss $d$ steps further on reaches an earlier application through $d$ halvings, so the multipliers run 1, 1/2, 1/4, … and sum to less than 2 however long the unroll is. The paper states the goal as keeping the total gradient applied to the dynamics function constant.

Every hidden state that $h_\theta$ or $g_\theta$ outputs is also rescaled into $[0,1]$ using the minimum and maximum of that hidden state itself, $(s - \min s)/(\max s - \min s)$, the same range as the action planes it is stacked with. This bounds the activations along the unroll. Unlike batch normalization, it uses no statistics from other examples, and unlike layer normalization it has no learned gain or bias.

The official pseudocode trains with SGD with momentum 0.9 and weight decay $10^{-4}$, and a learning rate that decays by a factor of 10 every 350k steps in Atari (starting at 0.05) and every 400k in board games (0.1 for chess and shogi, 0.01 for Go). Each board game used 16 TPUs for training and 1000 for self-play; each Atari game used 8 for training and 32 for self-play, because 50 simulations per move need far less acting compute than 800. Each domain gets its own network and training run; there is no single model that plays all four.

Results: Atari, Go, chess and shogi

Atari results are reported as human-normalized scores, where 0% is random play and 100% is a human tester's score, averaged over 57 games. MuZero set a new state of the art for both the median and the mean: a median of 2041.1% against 1920.6% for R2D2, the previous best (model-free) agent, and a mean of 4999.2% against 4024.9%. Its median game is about twenty times the human score. It beat R2D2 on 42 of the 57 games and SimPLe, the previous best model-based agent, on all of them, and it used fewer environment frames (20.0 billion against 37.5 billion) and fewer training steps (1 million against 2.16 million), trained in 12 hours where R2D2's table entry says 5 days.

In chess and shogi MuZero matched AlphaZero. Elo was measured in tournaments between training checkpoints at 800 simulations per move, anchored to one baseline per game (Stockfish for chess, Elmo for shogi, AlphaZero for Go, each given 100 ms per move). In Go it slightly exceeded AlphaZero while using a smaller network per evaluation (16 residual blocks against 20) and without the rules. The paper suggests, as a possibility, that MuZero may be caching computation in the search tree and using each extra application of the dynamics function to understand the position more deeply.

The ablations measure how much the search contributes and how good the learned model is. In Go, playing strength keeps rising with more simulations per move, and the learned model keeps pace with a perfect simulator out to searches about two orders of magnitude longer than the 800 simulations it was trained with. In Atari, more search helps up to about 100 simulations per move and then flattens; the paper suggests this is perhaps because the learned model is less accurate there. Figure 10 shows the shape of both curves.

By the end of training, even a single simulation per move, which plays the raw policy with no lookahead, does well on Atari. The paper reads this as the policy having internalized the benefits of search, and search still did the teaching during training: networks trained with more simulations per move improved faster, and when the authors replaced MuZero's loss with R2D2's model-free Q-learning objective and trained without search on Ms. Pac-Man, the result matched R2D2 but learned significantly slower than MuZero and converged to a much lower final score.

The paper also reports MuZero Reanalyze, a variant tuned for less data. It re-runs the search on old positions with the latest network and uses that fresh search policy as the policy target for 80% of updates. Its value target bootstraps from a target network (a recent copy of the parameters) instead of the stored search value, uses $n=5$, and is weighted 0.25 against the other losses. On 200 million frames per game, a hundredth of the main agent's 20 billion, Reanalyze reaches a 731.1% median, against 192% for IMPALA, 231% for Rainbow and 431% for LASER at the same budget.

The method has limits the paper states or the design implies. Imperfect-information games such as poker are not addressed. The dynamics function is deterministic, returning one next state per action, so the model has no direct way to represent a chance outcome such as a dice roll. And each board-game run used 1000 TPUs for self-play.