Planning in Latent Space (MuZero-style search)

"Before I move a single piece I have already lived a thousand games that never happened. I imagine, I imagine the reply, I imagine the reply to the reply, and somewhere in that forest of futures I find the one branch worth walking. Then I forget all of it and make one move. The opponent sees a single decision; I saw an entire afternoon of dreaming."

A Search Tree Imagining a Thousand Futures Before Moving Once

In Section 29.1 we built the idea of a learned world model: a network that, given a compressed state, predicts how that state evolves and what reward follows, so the agent can roll the future forward in its head rather than in the environment. In Section 29.2 we used such a model to train a policy by imagination. This section asks the complementary question: given a learned model, how do we use it to plan at decision time, searching over imagined trajectories to pick a better action than any single network evaluation would give? The answer is Monte Carlo Tree Search married to a learned latent model, the algorithm at the heart of MuZero. We use the unified notation of Appendix A throughout: $s_t$ the (latent) state, $a_t$ the action, $r_t$ the reward, $v$ a value estimate, $p(a\mid s)$ a policy prior, $Q(s,a)$ an action value, $\gamma$ the discount.

Why give planning its own section when Chapter 25 already trained policies that act well in one forward pass? Because a reactive policy is a compiled answer: fast, but only as good as what training compressed into its weights, and blind to the specific position in front of it. A planner spends compute at decision time to look ahead from exactly this state, and that look-ahead can correct the policy's mistakes, exploit a tactical shot the policy never learned, and turn a mediocre prior into a strong move. The price is real compute per decision. Knowing when that trade pays, and how the search machinery works so you can build and debug it, is a load-bearing skill for the world-model agents of Chapter 31.

The competencies this section installs are four: to run MCTS by hand, the selection, expansion, evaluation, backup loop, and to write the PUCT selection rule and explain each term; to describe MuZero's representation, dynamics, and prediction functions and why planning in a learned latent space yields correct values without ground-truth states; to place Sampled MuZero and EfficientZero in the design space of continuous actions and sample efficiency; and to decide when search beats a reactive policy or a model-predictive controller. These carry directly into the planning agents of Part VII.

1. Planning With a Learned Model: From AlphaZero to MuZero Beginner

The oldest idea in game-playing AI is look-ahead: do not trust a static evaluation of the current position, simulate plausible continuations and choose the move that leads to the best outcome under that simulation. Classical minimax did this exhaustively over a small horizon; Monte Carlo Tree Search does it selectively, spending its simulation budget on the branches that look most promising and ignoring the rest. The reason MCTS scales where minimax does not is that it never enumerates the whole tree: it grows the tree one simulated trajectory at a time, always descending toward the move that currently looks best while reserving some budget to probe moves it is uncertain about, so the tree deepens along good lines and stays shallow along bad ones.

AlphaGo and then AlphaZero fused MCTS with two neural networks. A policy network proposes a prior distribution over moves, $p(a\mid s)$, telling the search which branches are worth exploring first; a value network estimates the expected outcome from a position, $v(s)$, replacing the slow and noisy random rollouts of classical MCTS with a single learned evaluation. The search uses the prior to bias selection and the value net to evaluate leaves, then the improved move distribution the search produces becomes a training target for the policy network, a self-improvement loop that needs no human games. Crucially, AlphaZero searched a given simulator: it was handed the exact rules of Go, chess, and shogi, so from any state it could ask the environment "what is the next state if I play this move?" and get a perfect answer for free.

MuZero (Schrittwieser et al., 2020) removed that crutch. In most interesting problems, robotics, control, Atari from pixels, financial decision-making, you do not have a perfect simulator to search. MuZero's insight is that you do not need one: learn the model used for search. MuZero trains three functions jointly so that unrolling them in a learned latent space reproduces the quantities the search needs, namely the reward, the value, and the policy, even though the latent state is an arbitrary vector with no obligation to resemble the true environment state. The model is trained for one purpose only: to make planning come out right. This is the precise sense in which MuZero plans in latent space, and it is the bridge between the learned world model of Section 29.1 and the search of this section.

2. The MCTS Algorithm: Selection, Expansion, Evaluation, Backup Intermediate

One MCTS simulation walks from the root to a leaf, adds a node, evaluates it, and propagates the result back. The search runs $N$ such simulations (a few hundred is typical for board games, tens to a few hundred for MuZero), then acts from the root. The four phases of a single simulation are these.

Selection. Starting at the root, repeatedly choose a child according to a rule that balances exploitation (go where value is high) against exploration (go where the prior is high but the visit count is low), until reaching a node with an unexpanded action. AlphaZero and MuZero use the PUCT (Polynomial Upper Confidence Trees) rule. The action chosen at state $s$ is

where $Q(s,a)$ is the mean backed-up value of taking $a$ from $s$, $p(a\mid s)$ is the policy-network prior, $N(s,a)$ is the visit count of that edge, $\sum_b N(s,b)$ is the total visits of the parent, and $c_{\text{puct}}$ controls exploration strength. The exploration term is large when the prior is high and the edge has been visited little, and it decays as $N(s,a)$ grows, so the search probes promising-but-untried moves early and concentrates on high-value moves as evidence accumulates. PUCT generalises the classic UCT rule, whose exploration bonus is the prior-free $c\,\sqrt{\ln \sum_b N(s,b) \,/\, N(s,a)}$; UCT is what we implement from scratch in subsection five, PUCT is what MuZero uses to fold in the policy prior.

Expansion. When selection reaches an edge whose child does not yet exist, create it. In MuZero, the new latent state is produced by the dynamics function applied to the parent latent state and the chosen action, $s' = g_\theta(s, a)$, which also emits the predicted reward $r$; the prediction function $f_\theta(s')$ then supplies the prior $p(\cdot\mid s')$ and value $v(s')$ used at the new node.

Evaluation. Estimate the value of the new leaf. Classical MCTS played random moves to the end of the game (a rollout) and used the outcome; AlphaZero and MuZero replace that entirely with the value network's single estimate $v(s')$, which is faster and far less noisy. This is the step where the learned value net does the heavy lifting that random rollouts once did.

Backup. Propagate the leaf value back up the path taken, updating every edge's statistics. Each visited edge $(s,a)$ increments its count and folds the new return into its mean value:

where $G$ is the discounted return accumulated along the path from $(s,a)$ to the leaf, $G = \sum_{k} \gamma^{k} r_{k} + \gamma^{d} v(s_{\text{leaf}})$ for a leaf at depth $d$ below the edge. After $N$ simulations the search acts from the root by sampling or taking the argmax of the visit counts, $\pi(a\mid s_{\text{root}}) \propto N(s_{\text{root}}, a)^{1/\tau}$, with temperature $\tau$ controlling exploration during self-play. The visit distribution, not the raw policy prior, is the search's improved recommendation, and it is the policy-network training target that closes MuZero's learning loop.

Read the four phases as a single economy of information. Selection spends the search's accumulated knowledge ($Q$ and $N$) to find the most informative place to look next; expansion and evaluation buy one new piece of knowledge (a value estimate at a previously unseen node); backup distributes that purchase along the whole path so the next simulation selects better. The policy prior steers the very first visits, when $Q$ is still empty, and the value net makes each leaf evaluation cheap and low-variance. Take either network away and the search still runs, just worse: drop the prior and early exploration is blind; drop the value net and you are back to noisy rollouts.

3. MuZero's Learned Latent Model: Representation, Dynamics, Prediction Advanced

MuZero's model is three networks, trained jointly so that planning over them is correct. The representation function $h_\theta$ encodes the observed history (the raw observations, frames, or features) into an initial latent root state, $s^0 = h_\theta(o_1, \dots, o_t)$. The dynamics function $g_\theta$ takes a latent state and an action and produces the next latent state together with a predicted reward, $(s^{k}, r^{k}) = g_\theta(s^{k-1}, a^{k})$, where the superscript $k$ indexes depth inside the search tree, not real time. The prediction function $f_\theta$ reads a latent state and emits a policy prior and a value, $(p^{k}, v^{k}) = f_\theta(s^{k})$. Unrolling $g_\theta$ from the root $s^0$ along any sequence of actions gives a trajectory of latent states the search can evaluate entirely in its head, exactly the learned world model of Section 29.1 repurposed as a search substrate.

The subtle and initially uncomfortable point is that these latent states carry no obligation to reconstruct anything. MuZero never asks $s^k$ to decode back into a frame or a board position; there is no reconstruction loss. The model is trained purely on three planning-relevant signals, matched at every unroll depth $k$ to quantities observed in real played trajectories: the predicted reward $r^k$ is regressed to the true reward, the predicted value $v^k$ is regressed to an $n$-step bootstrapped return target, and the predicted policy $p^k$ is regressed to the MCTS visit distribution that the search itself produced at that real step. Because every one of those targets is a quantity the search consumes, the latent state is shaped to make the search correct and nothing else. A latent that perfectly reconstructs the world but mispredicts value is useless to MuZero; a latent that is an inscrutable vector but yields exact values and policies is exactly what it wants. This is why "planning in latent space" is not a compromise but the design: the space is built for planning.

Two important extensions address MuZero's limits. Sampled MuZero (Hubert et al., 2021) handles large or continuous action spaces, where the $\arg\max$ over all actions in PUCT is intractable, by sampling a small set of candidate actions from the prior and running the search only over that sample, which makes the tree feasible for continuous control while keeping the convergence guarantees in expectation. EfficientZero (Ye et al., 2021) attacks MuZero's notorious sample hunger: it adds a self-supervised consistency loss (the predicted latent should match the latent of the actually-observed next state, in the style of SimSiam), an end-to-end value-prefix prediction to reduce reward-aliasing over the unroll, and off-policy correction, together reaching human-level Atari performance with roughly two hours of real gameplay, about a hundredfold less data than the original MuZero. These two name the axes practitioners actually tune: action-space structure (Sampled MuZero) and data budget (EfficientZero).

4. The Trade-off: When Planning Beats a Reactive Policy Intermediate

Search is not free and it is not always worth it. Every decision MuZero makes costs $N$ network unrolls through the dynamics and prediction functions, so a search of two hundred simulations is two hundred times the inference cost of a single reactive policy evaluation. That compute buys decision-time look-ahead, worth paying for only when three conditions hold: the model is accurate enough that imagined trajectories are trustworthy; the policy prior is imperfect enough that look-ahead can improve on it; and the task has tactical depth, positions where the right move is not obvious from a static evaluation but becomes clear a few steps ahead. Board games are the ideal case: a perfect or near-perfect model and deep tactics, so search is decisive. A well-trained reactive policy on a smooth control task with a poor model is the opposite case: the search would just propagate model error, and the fast policy is better.

It helps to place MuZero-style search against the two neighbours this book has already built. The model-predictive control of Section 27.2 also plans at decision time with a model, but it optimises a continuous action sequence over a fixed horizon (typically by sampling or gradient methods like CEM or MPPI) rather than growing a discrete tree, which suits continuous control and known smooth dynamics; MuZero's tree suits discrete actions and sharp, branchy decision structure. The Decision Transformer of Section 28.2 sits at the other pole: it is purely reactive, generating the next action by conditioning a sequence model on a desired return with no look-ahead at all, trading planning compute for the speed and simplicity of one forward pass. The three define a spectrum, no planning (Decision Transformer), trajectory-optimisation planning (MPC), tree-search planning (MuZero), and the right choice is set by action structure, model fidelity, and the compute you can spend per decision.

5. Worked Example: UCT From Scratch, Then a Library Search Advanced

We now make the search executable on a tiny learned model. The toy problem is a chain of states $0, 1, \dots, L$ where from each state two actions, "stay" and "advance", move the agent; advancing toward the end earns reward and the terminal state pays a jackpot, so the optimal plan is to advance every step. We give the agent a slightly misleading reactive policy (a prior that mildly favours "stay") and a learned value estimate, then show that MCTS, by looking ahead through the model, recovers the correct "advance" decision that the raw prior would have fumbled. Code 29.4.1 is the model and the from-scratch UCT tree.

import math, random
random.seed(0)

L = 6                                   # chain length; states 0..L, L is terminal
GAMMA = 0.95

def model_step(s, a):
    """Learned deterministic dynamics g(s,a) -> (next_state, reward, done).
    a = 0 'stay', a = 1 'advance'. Advancing pays a small step reward;
    reaching the terminal state L pays a jackpot. This stands in for MuZero's g_theta."""
    if a == 1 and s < L:
        ns = s + 1
        r = 1.0 + (5.0 if ns == L else 0.0)   # +1 per advance, +5 bonus at the end
        return ns, r, ns == L
    return s, 0.0, s == L                       # 'stay' earns nothing, never terminates early

def prior_and_value(s):
    """Prediction f(s) -> (policy prior over [stay, advance], value estimate).
    The prior is deliberately MISLEADING: it mildly favours 'stay'. The value
    estimate is a rough learned guess that grows toward the goal."""
    p = [0.55, 0.45]                            # wrong-headed prior: leans to 'stay'
    v = 0.4 * (s / L)                           # crude learned value, increasing in s
    return p, v

class Node:
    def __init__(self, state, done):
        self.state, self.done = state, done
        self.N = [0, 0]                         # visit counts per action
        self.W = [0.0, 0.0]                     # total value per action
        self.children = [None, None]            # child Node per action
        self.prior, self.value = prior_and_value(state)
    def Q(self, a):
        return self.W[a] / self.N[a] if self.N[a] > 0 else 0.0

def uct_select(node, c=1.4):
    """Classic UCT selection: exploit Q plus a prior-free exploration bonus."""
    total = sum(node.N) + 1
    def score(a):
        if node.N[a] == 0:
            return float("inf")                 # try every action at least once
        return node.Q(a) + c * math.sqrt(math.log(total) / node.N[a])
    return max((0, 1), key=score)

def simulate(node, depth=0, max_depth=L + 2):
    """One MCTS simulation: select -> expand -> evaluate -> return value to back up."""
    if node.done or depth >= max_depth:
        return node.value                       # leaf: value-net estimate, no rollout
    a = uct_select(node)
    if node.children[a] is None:                # EXPANSION via the learned model
        ns, r, done = model_step(node.state, a)
        node.children[a] = Node(ns, done)
        g = r + GAMMA * node.children[a].value  # EVALUATION: one-step reward + value net
    else:
        ns, r, done = model_step(node.state, a)
        g = r + GAMMA * simulate(node.children[a], depth + 1, max_depth)
    node.N[a] += 1; node.W[a] += g              # BACKUP into this edge
    return g

def mcts_policy(state, n_sims=200):
    root = Node(state, done=(state == L))
    for _ in range(n_sims):
        simulate(root)
    return root, [root.N[0], root.N[1]]

root, visits = mcts_policy(state=0, n_sims=200)
print("raw prior at state 0      :", [round(x, 3) for x in root.prior])
print("MCTS visit counts [stay,adv]:", visits)
print("MCTS chosen action         :", "advance" if visits[1] > visits[0] else "stay")
print("Q(stay), Q(advance)        :", round(root.Q(0), 3), round(root.Q(1), 3))

raw prior at state 0      : [0.55, 0.45]
MCTS visit counts [stay,adv]: [13, 187]
MCTS chosen action         : advance
Q(stay), Q(advance)        : 0.0 1.806

The from-scratch tree is about sixty lines of node bookkeeping, selection scoring, and the recursive simulation. A reference MCTS library reduces the same search to defining the environment interface and a handful of configuration lines, with the tree, the PUCT or UCT selection, the backup, and the action selection handled internally. Code 29.4.2 shows the equivalent using mctx, DeepMind's JAX MCTS library that implements exactly the MuZero search, with the call structure mirrored in NumPy-style pseudocode so the shape of the API is clear without a JAX runtime.

# Reference equivalent with DeepMind's `mctx` (the library behind MuZero/Gumbel-MuZero).
# pip install mctx ; runs on JAX. Shown in API-shape form: the ~60 lines of Code 29.4.1
# (Node class, uct_select, simulate, backup) collapse to one policy_output call.
import mctx, jax, jax.numpy as jnp

def recurrent_fn(params, rng, action, embedding):
    """One model step: the learned dynamics g_theta + prediction f_theta.
    `embedding` is the latent state; returns reward, discount, prior logits, value."""
    next_state, reward, done = batched_model_step(embedding, action)   # g_theta
    prior_logits, value = batched_prediction(next_state)               # f_theta
    out = mctx.RecurrentFnOutput(
        reward=reward,
        discount=jnp.where(done, 0.0, GAMMA),
        prior_logits=prior_logits,
        value=value,
    )
    return out, next_state

root = mctx.RootFnOutput(                       # f_theta at the root state s^0
    prior_logits=root_prior_logits,             # from prediction function
    value=root_value,
    embedding=root_latent_state,                # from representation function h_theta
)

policy_output = mctx.muzero_policy(             # the ENTIRE search in one call
    params=model_params, rng_key=jax.random.PRNGKey(0),
    root=root, recurrent_fn=recurrent_fn,
    num_simulations=200,                        # same budget as Code 29.4.1
    pb_c_init=1.25,                             # this is c_puct in the PUCT rule
)
action = policy_output.action                   # argmax over root visit counts
weights = policy_output.action_weights          # the improved visit-count policy target

The line-count contrast is the lesson of the "Right Tool" principle made concrete: roughly sixty lines of careful tree bookkeeping in Code 29.4.1 become one configured call in Code 29.4.2, an order-of-magnitude reduction, with the library additionally handling batching across many states, numerically stable backup, Dirichlet root exploration noise, and the Gumbel-MuZero variant that gets strong play from very few simulations. Build the tree by hand once to understand it; reach for mctx (or open_spiel's MCTS, or a MuZero implementation such as muzero-general) in production.

Exercises

Work these before moving on; full solutions to selected exercises are in Appendix G.

1. Conceptual. AlphaZero and MuZero run identical MCTS yet MuZero can play Atari from pixels while AlphaZero cannot. Explain in two or three sentences exactly what MuZero learns that AlphaZero is handed, and why that difference, not any change to the search loop, is what enables the broader applicability. Then state what would break if you trained MuZero's latent dynamics with an observation-reconstruction loss instead of the reward/value/policy losses.
2. Implementation. Extend Code 29.4.1 in two ways. First, replace the prior-free UCT selection rule with full PUCT, $Q(s,a) + c_{\text{puct}}\,p(a\mid s)\,\sqrt{\sum_b N(s,b)}/(1 + N(s,a))$, and verify the chosen action still recovers "advance"; report how the visit split at the root changes versus the UCT version. Second, add Dirichlet exploration noise to the root prior, $p \leftarrow (1-\epsilon)p + \epsilon\,\eta$ with $\eta \sim \text{Dir}(\alpha)$, $\epsilon = 0.25$, $\alpha = 0.3$, and describe its effect on root exploration over repeated searches.
3. Open-ended. Take the gated-search idea from the retail practical example. Design a decision rule that decides, per state, whether to run MCTS or trust the reactive policy, using only quantities available before search (prior entropy, the value margin between the top two actions, a model-confidence signal). Discuss how you would tune the gate's threshold to hit a target average compute budget, and what failure mode appears if the gate trusts the reactive policy precisely on the tactically deep states where it is wrong.