MuZero and Value-Equivalent Models › Representation, Dynamics, and Prediction Networks

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.