MuZero and Value-Equivalent Models

Picture a video game. The screen is a torrent of pixels: explosions, particle effects, a scrolling scoreboard, a blinking cursor in the corner. Now picture the decision the game asks of you: move left, jump, or wait. Much of that pixel torrent may not matter for the decision. What matters is a task-relevant summary: where am I, where is the enemy, how much health do I have, is the platform solid? A game-playing agent must learn which details to retain for its decisions, and its training objective influences what its internal representation keeps.

A prominent visual-model lineage learns to reconstruct observations alongside its decision model. PlaNet and the Dreamer family, discussed in Part VIII: Decision-Centric Research Lineages, use a reconstruction objective to help train their latent states. That objective can spend capacity on observation details that do not change an action. This is one approach to model-based reinforcement learning, not its historical definition; classical control models often operate directly on compact physical states.

MuZero, introduced by Schrittwieser and colleagues in 2019 and refined in 2020, asked a sharper question. What if we skip the reconstruction entirely? What if the only requirement on the learned model is that it supports good decisions?

This decision-centric answer became highly influential. MuZero learns three connected functions:

• a representation function that turns an observation into a latent state,
• a dynamics function that advances a latent state under an action and predicts the reward,
• a prediction function that reads a policy and a value out of a latent state.

Here, a "latent state" means an internal representation the network learns, with no requirement that it reconstruct the raw observation. The reward is scalar feedback from the environment, and the policy is a distribution over actions. These three functions form MuZero's learned model; self-play data-collection workers, a replay buffer, and tree search are also essential parts of the agent.

There is no observation decoder or reconstruction loss in the original MuZero objective. Planning uses Monte Carlo tree search over the learned latent dynamics. Schrittwieser et al. reported strong results in Go, chess, shogi, and Atari without giving the model the games' transition rules. The paper's large-data Atari comparison used 20 billion frames for MuZero and 37.5 billion for R2D2; its 200-million-frame MuZero Reanalyze setting is a separate experiment. Neither result supports a blanket claim of a hundredfold or thousandfold sample-efficiency improvement.

This chapter builds a MuZero-style tutorial model on a small gridworld and tests latent prediction, tree search, and search-generated targets. The measurements also show limits to what these experiments establish: more search can worsen the planner's Q-greedy action agreement, while better value targets alone do not establish better learned-policy decisions. We will connect this to value equivalence, a formal criterion based on selected Bellman updates rather than pixel-level prediction.

MuZero uses a learned latent model for planning instead of querying the environment during search. It does not plan in observation space. It plans in a representation learned from reward, value, and policy targets. This is a substantial shift in perspective. In a classical state-space model, the state variables may be supplied by the problem: the position of a chess piece, the reading of a sensor. In MuZero, the latent state is learned from data, and its usefulness must be evaluated through predictions and decisions. If planning fails, training and evaluation must diagnose and improve the model; MuZero does not automatically discard and relearn the entire state space.

Formally, in a Markov decision process the agent interacts with an environment over discrete time steps $t = 0, 1, 2, \ldots$. At each step the environment is in some state, the agent selects an action, and the environment returns a reward and a new state. The learned model is called latent because its state $s^k$ is not the environment's state; it is a vector of the network's own invention. Throughout this chapter we distinguish:

• the observation $o_t$ (what the agent sees at wall-clock time $t$);
• the action $a_t$ (what the agent chooses at wall-clock time $t$);
• the reward $u_t$ (what the environment returns at wall-clock time $t$);
• the policy $\pi$ (a distribution over actions);
• the transition $p(s_{t+1}\mid s_t, a_t)$ (how the environment's state evolves);
• the observation model $p(o_t\mid s_t)$ (how observations are generated from states);
• the latent state $s^k$ and unrolled horizon $k$ used inside the learned model.

Keeping these separate matters, because the central confusion you may run into is mixing up the unrolled time inside the learned model with the wall-clock time of the real environment. The superscript $k$ counts steps taken inside the imagination of the network. The subscript $t$ counts steps taken in the world. In our tutorial, a single observation $o_t$ is encoded into the initial latent state $s^0$. Published MuZero's representation instead takes an observation-and-action history. After the root is encoded, search rolls forward through learned transitions rather than querying new environment transitions, though it may query the environment for legal actions at the root.

Write the environment's observation at time $t$ as $o_t$ (what the agent sees), the action as $a_t$ (what the agent chooses), the reward as $u_t$ (what the environment returns), and the latent state as $s^k$ (the compressed summary the network invents at model step $k$). The superscript $k$ indexes steps after the encoded root, not wall-clock time. In the training equations below, $K$ is the replay-supervised unroll length; search depth is a separate choice. MuZero parameterizes three functions with a shared network $\theta$ (layers within each function have their own parameters, conventionally grouped as $\theta$). The representation function encodes available history into an initial latent state; our compact tutorial uses only the current observation. The dynamics function advances that latent state under an action while predicting the reward, and the prediction function reads a policy and a value out of a latent state:

$$\begin{aligned} s^0 &= h_\theta(o_t), \\ (s^{k+1}, r^{k+1}) &= g_\theta(s^k, a^k), \\ (p^k, v^k) &= f_\theta(s^k). \end{aligned}$$

where:

• $o_t$: the observation the agent sees at time $t$
• $h_\theta$: the representation function, which encodes available history into an initial latent state $s^0$; the displayed $h_\theta(o_t)$ is our single-observation tutorial specialization
• $a^k$: the action taken at unrolled step $k$ (encoded as a one-hot or embedding)
• $g_\theta$: the dynamics function, which maps a latent state and an action to the next latent state $s^{k+1}$ and the reward $r^{k+1}$ on the transition into it
• $s^k$: the latent state at unrolled step $k$; the tutorial implementation below bounds it with a tanh activation
• $f_\theta$: the prediction function, which reads a latent state $s^k$ and emits a policy $p^k$ (a distribution over actions, used as a prior during search) and a value $v^k$ (the expected discounted return from that latent state)
• $k \in \{0, 1, \ldots, K\}$: the model-step index within the training unroll, not wall-clock time
• $K$: the number of model steps used for a replay-supervised training unroll; it is not the tree-search depth

Read the three equations top to bottom and the whole algorithm is visible. The first line says "look at the world and compress it." The second says "given a situation and an action, imagine the resulting situation and the reward it brings." The third says "given a situation, what would you do and how good is it?" Everything else in the chapter is a consequence of these three lines plus a search procedure.

Notice what is missing: there is no inverse dynamics function, no decoder, and no reconstruction head. The latent state $s^k$ is not constrained to be interpretable and is not compared to a real observation. The states produced by the dynamics function are trained to support reward, value, and policy predictions useful for planning. The model need not reproduce every observation detail, although its latent state can still retain irrelevant information.

The original MuZero dynamics function is deterministic: given a latent state and an action, it returns one next latent state rather than a sampled distribution. Our gridworld's position dynamics are deterministic, although its observation includes independent noise. A deterministic model can still predict useful expected values in a stochastic environment, but a single imagined branch does not represent the full outcome distribution. Search quality must be checked against real-environment outcomes.

Suppose we sample a trajectory from the replay buffer and pick a starting index $t$. We run the network forward from $s^0 = h_\theta(o_t)$, unrolling $K$ steps with the actual actions the agent took:

$$\begin{aligned} s^0 &\xrightarrow{\;a_t\;} (s^1, r^1) \\ &\xrightarrow{\;a_{t+1}\;} (s^2, r^2) \\ &\xrightarrow{\;a_{t+2}\;} \cdots \\ &\xrightarrow{\;a_{t+K-1}\;} (s^K, r^K). \end{aligned}$$

At the root and each unrolled state we predict value $v^k$ and policy $p^k$; each modeled transition predicts reward $r^{k+1}$. The reward target is the observed reward on that transition. The policy target $\pi_{t+k}$ comes from search at the corresponding real state. In the original MuZero training objective, the value target $z_{t+k}$ is the final outcome for board games or an observed $n$-step return bootstrapped from a later search value for Atari. With our reward indexing, a schematic loss is:

$$\mathcal{L}(\theta) = \mathbb{E}_{\tau \sim \mathcal{D}} \left[\sum_{k=1}^{K} \ell^{r}(u_{t+k}, r^k) + \sum_{k=0}^{K} \bigl(\ell^{v}(z_{t+k}, v^k) + \ell^{p}(\pi_{t+k}, p^k)\bigr)\right] + c \lVert \theta \rVert_2^2$$

where:

• $\tau \sim \mathcal{D}$: a trajectory sampled from the replay buffer $\mathcal{D}$, from which the starting index $t$ is drawn
• $\mathcal{L}(\theta)$: the total training loss, minimized over the shared parameters $\theta$
• $K$: the number of model steps in this training unroll
• $u_{t+k}$: the observed reward on the $k$th transition after the starting state, for $k=1,\ldots,K$
• $r^k$: the reward predicted by dynamics on that transition; no root reward $r^0$ is predicted
• $z_{t+k}$: the outcome or $n$-step bootstrapped value target appropriate to the domain
• $v^k$: the value predicted by the prediction function at unrolled step $k$
• $\pi_{t+k}$: the policy target produced by the search at time $t+k$
• $p^k$: the policy prior predicted by the prediction function at unrolled step $k$
• $\ell^{r}$: the reward loss, whose form depends on the domain and target representation
• $\ell^{v}$: the value loss, likewise chosen for the domain and target representation
• $\ell^{p}$: the policy loss, a cross-entropy between the search visit-count distribution $\pi_{t+k}$ and the network prior $p^k$
• $c$: the coefficient of the $L_2$ weight-regularization term
• $\lVert \theta \rVert_2^2$: the squared $L_2$ norm of all parameters in the shared network $\theta$, which penalizes large weights

In published MuZero, board games use a squared-error value loss and, because those experiments have no intermediate rewards, omit the reward-prediction loss. Atari uses categorical-support cross-entropy for transformed reward and value targets. The policy loss compares search visit counts with the network prior. The tutorial below uses simple scalar mean-squared reward and value losses for ease of inspection, not the complete published objective.

The expectation is over trajectories in replay. In published MuZero, actors repeatedly generate new trajectories using MCTS; replay is refreshed, not a permanently fixed dataset. Our tutorial instead uses a fixed exploratory dataset so that the small experiment is reproducible.

Search provides policy targets at replayed states; the value target also incorporates actual outcomes or later bootstrapped search estimates. This can improve decision targets, but it does not guarantee that a model trained on them will improve, especially outside the replay distribution. Our fixed-data experiment should be read as a controlled demonstration, not as a reproduction of MuZero's self-play loop.

This objective is unusual and worth examining closely. It does not ask the latent to reconstruct the next observation or resemble the environment state. Reward, value, and policy losses instead constrain the quantities used for decisions. Many different latent encodings can satisfy those targets, and stable learning depends on data, architecture, optimization, and target quality; the losses alone do not guarantee it.

The original MuZero objective supervises unrolled reward, value, and policy predictions, but does not include an explicit loss matching a predicted latent to the representation of the next observation. Correct predictions constrain the dynamics indirectly, while multiple latent representations may remain equivalent for the supervised targets.

An auxiliary consistency objective can encourage the predicted state and the encoding of the next observation to agree. It is an additional inductive bias, not a required component of MuZero or a proof that search will be reliable.

For example, EfficientZero uses a self-supervised temporal consistency component. Our toy model uses a simpler squared-distance variant. Write $\tilde s^{k+1}$ for the state component of $g_\theta(s^k,a^k)$; then:

$$\ell^{\text{cons}} = \big\lVert\, \tilde s^{k+1} - \operatorname{sg}\!\big(h_\theta(o_{t+k+1})\big) \,\big\rVert_2^2$$

where:

• $\ell^{\text{cons}}$: the latent consistency loss at unrolled step $k$
• $\tilde s^{k+1}$: the state component of the dynamics prediction after applying $a^k$ at $s^k$
• $\operatorname{sg}(\cdot)$: the stop-gradient operator, which treats its argument as a constant during backpropagation
• $h_\theta(o_{t+k+1})$: the latent state the representation function would have produced had it seen the observation $k+1$ steps later
• $k$: the unrolled step at which the consistency loss is computed, ranging over $0, \ldots, K-1$

This term compares the predicted state with a stopped-gradient next-observation encoding. It is self-predictive rather than reconstructive: nothing decodes pixels. Stop-gradient blocks the direct gradient through the target occurrence of the encoder. The encoder can still change through the online branch and the reward, value, and policy losses. The squared distance is a tutorial choice, not a latent probability model.

Without a stopped target, both branches can move toward one another. Stopping one branch changes the gradient direction but does not by itself prevent representational collapse: a constant zero latent satisfies the consistency term and survives $\tanh$. Decision losses and training design must preserve useful information.

The toy experiment asks whether this combination of losses reduces the representation of four unpredictable, decision-irrelevant observation bits. We will measure the result rather than assume that the bits disappear.

Consider the kinds of information a world model might be asked to represent:

• Decision-relevant, predictable. Where the agent is, where the goal is, how much health remains. When these factors affect reward or action quality, the learned model needs enough information to predict relevant rewards, values, and policy targets. It need not store each variable literally in the latent; a fixed goal, for example, can be implicit in the network parameters.
• Decision-relevant, unpredictable. An opponent's future move in a stochastic game may affect the return, but its realized outcome is not known at the current state. The representation may need information that predicts the distribution of possible moves or their expected effect on value; a scalar value target alone does not guarantee that it identifies the full distribution.
• Decision-irrelevant, predictable. The scrolling scoreboard animation. The dynamics can predict it, and nothing needs it. It may survive or not; nothing punishes its presence, and nothing rewards it.
• Decision-irrelevant, unpredictable. A random distractor bit. Consistency may discourage encoding it, but finite-capacity training can still leave measurable information in the latent.

These two axes distinguish decision relevance from predictability. A decision-relevant but unpredictable variable may call for uncertainty-aware values; a predictable but irrelevant variable need not be retained. An unpredictable irrelevant variable is a plausible candidate to discard, not something the objective is guaranteed to remove.

A reconstruction loss can reward encoding distractor detail, depending on its weighting and the model's capacity. MuZero's original decision-target objective lacks that direct incentive, but it can still retain distractors incidentally. The later probe measures how much remains in this particular learned latent.

Let us now build the smallest possible environment in which this effect is visible.

The learned value head is trained against behavior-policy Monte Carlo returns. Its RMSE of 0.104 is only marginally below the empirical behavior baseline's 0.105 in this run; the difference does not establish a robust improvement. There is no multi-step unrolled loss in the tutorial training loop, so attributing value behavior to such a loss would be unsupported. Search may use the learned reward and dynamics predictions to change action rankings, which we test next.

The gap to the optimal curve gives room for search or revised targets to help, but their benefit must be measured; neither is guaranteed by the architecture.

Training gives us a latent model. Now we need to use it. MuZero plans by running Monte Carlo tree search directly over latent states, with $g_\theta$ standing in for the environment's transition function.

This is the core of the method, so it is worth being explicit about what has changed relative to classical planning. In the search methods discussed in Part VII: Planning and Agency, search operates over states that mean something: board positions, robot configurations, belief distributions. Here, search operates over vectors in $\mathbb{R}^8$ that have no agreed-upon meaning at all. The tree is built out of network activations. The only thing that makes it a planning tree rather than a random graph is that the reward and value heads read useful quantities off its nodes.

This may feel unusual at first. In a classical planner, if you inspect a node of the search tree, you can describe what it represents: this is the position after moving the knight to f3, this is the belief state after observing the red light. In MuZero, the analogous inspection yields eight floating-point numbers with no built-in board description. We can inspect predicted rewards, values, policies, or a separately trained probe, but any human interpretation of a latent node needs its own validation. The tree can support decisions through its action values and visit counts without each node having a readable label.

A node in the tree stores a latent state $z$, the reward that the dynamics function predicts on the transition into that node, a prior probability $p(a)$ supplied by the prediction network, and running statistics. We write $N(z, a)$ for the number of times action $a$ has been tried at a node holding $z$, $W(z, a)$ for the sum of values backed up through those trials, and $Q(z, a) = W(z, a) / N(z, a)$ for the resulting mean.

When the notation is clear from context we drop the argument $z$ and write $N$, $W$, and $Q$ for a single node. The full form $N(z, a)$ is used in the PUCT rule below, where $z$ identifies the node and $a$ selects an edge.

Each simulation walks down the tree from the root, choosing at each node the action that maximizes the PUCT score. PUCT (Predictor + Upper Confidence bounds applied to Trees), a rule from the AlphaZero lineage, balances exploiting actions with high estimated value against exploring actions with high prior probability:

$$a^\star = \arg\max_a \left[ Q(z, a) + c_{\text{puct}} \cdot p(a) \cdot \frac{\sqrt{\sum_b N(z, b)}}{1 + N(z, a)} \right]$$

where:

• $Q(z, a)$: the mean value of taking action $a$ at latent state $z$, equal to $W(z, a) / N(z, a)$
• $N(z, a)$: the visit count for action $a$ at $z$
• $p(a)$: the prior probability the prediction network assigns to action $a$
• $\sum_b N(z, b)$: the total visits to $z$ across all actions
• $c_{\text{puct}}$: a constant controlling exploration strength

The displayed rule is schematic. Its edge mean $Q(z,a)=W(z,a)/N(z,a)$ is defined after an edge has been visited. The tutorial stores a child-state estimate as child.Q, seeds it from the value head before that child is visited, and scores the parent edge with child.reward + GAMMA * child.Q. It also uses sqrt(node.N + 1) instead of the displayed square root of total action visits, so the exploration bonus is nonzero when the parent has no visits. These choices retain the prior/visit tradeoff without making the equation a line-by-line transcription of the code.

The first term favors actions with high estimated return. The second encourages exploration according to the network prior $p(a)$. At fixed total node visits, its denominator $1 + N(z,a)$ makes the bonus for a particular action fall as that action is tried; a larger prior raises that action's bonus. As total node visits grow, the square-root numerator increases, so the bonus is not constant. Its net behavior depends on how visits are distributed across actions and on the scale of $Q$.

The exploration bonus is a concrete, plottable quantity. Holding the total visit count fixed, the bonus for an action decreases as that action is tried and increases with the network's prior for it.

Reading the search code top to bottom, expand_node applies prediction and then dynamics for every action in a batch. select_action ranks the parent-edge return child.reward + GAMMA * child.Q plus exploration; simulate backs up the same reward/value convention. This compact tutorial eagerly creates all children and initializes their values; it is not the exact optimized tree implementation used in the MuZero paper, which expands a selected leaf per simulation.

Take the cell furthest from the goal, $(3,3)$. Six moves reach the absorbing goal, so the reward arrives on transition six and has discount exponent five: $\gamma^5 = 0.95^5 \approx 0.774$. We can see this by expanding the return:

$$\begin{aligned} G(z_0) &= r(z_0,a_0) + \gamma\, G(z_1) \\ &= 0 + \gamma\,(0 + \gamma\,(0 + \gamma\,(0 + \gamma\,(0 + \gamma\,1)))) \\ &= \gamma^5 = 0.95^5 \approx 0.774. \end{aligned}$$

where the reward is 1 only on the final step into the goal and 0 on every other step. The behaviour policy takes a longer, less direct route, so its average discounted return from that cell is lower than the optimal value. Let us see what search produces.

The six transitions contribute rewards at discount exponents zero through five. Only the final transition has nonzero reward.

The printed visit-count policy chooses left most often, which is one optimal action from $(3,3)$. Its estimated value of 0.741 is below the exact 0.774 but above this dataset's empirical behavior estimate of 0.700. This four-step search cannot reach the goal from the far corner: its estimate combines predicted rewards along searched edges with a learned value bootstrap, rather than observing the sixth-step goal reward. Search can select a good action while remaining imperfectly calibrated. The paired-budget evaluation below tests all grid cells.

The natural question is how this finite search behaves as its simulation budget grows. We compare the Q-greedy root action, $\arg\max_a Q(a)$, with an exact optimal action and compute root-value RMSE against $V^\star$. This diagnostic is distinct from selecting an action from MuZero's root visit-count policy. The same noisy observation for each state is reused across budgets, making their comparison paired. Neither metric is guaranteed to improve monotonically in an approximate learned model.

The bars show Q-greedy optimal-action agreement falling from about 0.92 at one simulation to about 0.91 at 64, while value RMSE improves from 0.087 to 0.075. The intermediate budgets are not monotonic: agreement peaks near 0.98 at four simulations, and RMSE is lowest at 16. More search improved one diagnostic but worsened the other in the one-to-64 comparison. Approximate model/value errors can change Q rankings even when value calibration improves; this does not by itself measure the action selected from root visit counts. These 64 observations are too few to establish a general budget law.

The tutorial's decision-target model illustrates the motivation for value equivalence, but we have not proven that it satisfies a formal value-equivalence criterion. Here is the criterion and the distinction it makes.

Model equivalence asks for a faithful transition/reward description. Value equivalence asks for agreement only on selected decision-relevant Bellman tests. A model can differ in transitions to states that a chosen function set assigns the same value, yet agree on those Bellman updates. Whether that suffices for planning depends on the policies and functions included; a narrow test set gives a weaker guarantee.

Here $T_m^\pi v(s)$ is the expected immediate reward plus discounted expected $v$ at the next state under policy $\pi$ in model $m$. If the selected function set is sufficiently rich and closed under relevant Bellman iterations, agreement can imply agreement of corresponding policy values. Equality of all $V^\pi$ is not the paper's Definition 1.

Grimm and colleagues use this criterion to study model classes sufficient for selected planning calculations. The lesson is not that prediction accuracy is irrelevant, but that the quantities needed by the planner determine which predictive distinctions matter. A model that reconstructs pixels accurately can still misrank actions; one that omits pixels can still support good decisions on its tested distribution.

Consider two models of our gridworld:

• Model A has an observation decoder and is trained to reconstruct all 20 observation coordinates, including the four unpredictable bits.
• Model B produces an 8-dimensional latent with no interpretation whatsoever, and no decoder.

For a decoder of the current encoded observation, retaining the noise bits can reduce reconstruction loss. For a predictor of future independent bits, no deterministic estimate can eliminate expected squared error. These are distinct tasks; the plot below tests how much current-bit information each encoder retains, not whether future noise is forecast.

Model B has no direct reconstruction incentive, but its latent can still encode distractor bits incidentally. Consistency may reduce that information without guaranteeing its removal. We compare trained probes rather than infer what happened from the objectives alone.

This decoder comparison does not test formal value equivalence. For a Bellman example, imagine two gridworld models with the same position transitions and rewards but different distributions for the next observation's independent noise bits. Their full observation-transition kernels differ. Yet for policies and value functions that depend only on position, both models give the same immediate reward and expected next-position value, hence the same selected Bellman updates. That restricted agreement illustrates value equivalence without claiming that either trained network above satisfies the formal criterion.

The practical goal is task sufficiency: the latent should carry enough information for the policies and values the agent needs. In this fixed-goal gridworld, position matters, while the four independent distractor bits do not affect optimal actions. The goal location is constant, so it need not be separately encoded in every observation.

Sufficiency is relative to a task, state distribution, and decision rule. A representation adequate for navigation may omit information needed for contact-rich manipulation. A compact decision-centric latent can be efficient for its training objective, but transfer to a new reward or environment must be tested.

There is a relationship to bisimulation, discussed in Part III: Representing Agents and Worlds. Under usual MDP assumptions, bisimilar states have equal immediate rewards for each action and matching transition probabilities over equivalence classes. Such a quotient preserves values for policies compatible with the abstraction. Value equivalence asks only for selected Bellman tests and can therefore be weaker.

We can test all of this empirically. We train a second model, identical in architecture except that it also has a decoder and is trained with a reconstruction loss instead of the consistency loss. Then we ask two questions of each model:

1. How large is the sampled between-position share of latent variance? A noise-invariant latent would vary little within a position, but eight sampled noise masks per cell also put sampling variation into this statistic. Decision sufficiency alone does not require noise invariance.
2. How accurately can a linear probe read the noise bits out of the latent? If the latent is noise-invariant, a linear probe should be no better than chance.

These two diagnostics are complementary. The finite-sample variance share describes one aspect of latent organization, while the linear probe tests recoverable current-noise information. A model can score differently on the two measures, so we inspect both without treating either as a sufficiency test.