Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

MuZero and Value-Equivalent Models › Planning Without Reconstructing Observations

Michael BrenndoerferJuly 16, 202613 min read

World Models HandbookAll chapters

Part 8 of 12Decision-Centric Research Lineages

Chapter 46MuZero and Value-Equivalent Models › Planning Without Reconstructing Observations

Reading section 3 of 8 · Planning Without Reconstructing Observations · 13 minSection 3 of 8 · 13 min

Section 3 of 8 · 13 min read

MuZero and Value-Equivalent Models › Planning Without Reconstructing ObservationsOne section · Sections in this chapter3 / 8 ↓

Sections in this chapter8

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

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.

The search treeLink Copied

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.

In[16]:

Code

puct_c = 1.5
puct_total_visits = 20.0
puct_visits = np.arange(1, 21)
puct_priors = [0.25, 0.5, 0.75]
puct_bonus_curves = [
    puct_c * prior * np.sqrt(puct_total_visits) / (1.0 + puct_visits)
    for prior in puct_priors
]

Out[17]:

Visualization

Line chart of PUCT exploration bonus decaying with visit count for three priors. — PUCT exploration bonus as a function of an action's visit count, for three prior probabilities. The bonus shrinks rapidly as an action is tried and starts higher for larger priors, showing how the network's prior shapes early exploration within a fixed simulation budget.

The prior $p(a)$ directs finite simulations toward actions favored by the network. A good prior can save search effort; a mistaken one can delay exploration of a strong action. Thus the learned policy shapes search, but mutual improvement must be measured rather than presumed.

When a simulation reaches a node that has never been expanded, we call the prediction network to get a value estimate, and we call the dynamics network once per action to materialize the children. Then the value is backed up along the path, with the immediate reward predicted by the dynamics on each transition folded in. If $z_{i+1}$ follows action $a_i$ from $z_i$ , the state-return recursion is:

\begin{aligned} G(z_i) &= r(z_i,a_i) + \gamma\,G(z_{i+1}), \\ G(z_0) &= \sum_{i=0}^{H-1}\gamma^i r(z_i,a_i) + \gamma^H v(z_H), \end{aligned}

where:

$G(z_i)$ : the backed-up discounted return from the node holding latent state $z_i$
$r(z_i,a_i)$ : the reward predicted by dynamics on the edge from $z_i$ to $z_{i+1}$
$\gamma$ : the discount factor, which weights nearer rewards more heavily
$H$ : the number of simulated transitions before the value bootstrap
$v(z_H)$ : the prediction network's value estimate at the search horizon

The code stores each child's incoming-edge reward and its estimate of the return from the child state. A parent adds that reward to the discounted child-state return, matching the recursion above. This has the same discounted, additive form as a Monte Carlo return, but its rewards are predicted rather than observed, so reliability depends on those predictions and on the horizon value estimate.

After a fixed number of simulations, the tutorial search returns a root visit-count distribution $\pi(a) = N(a) / \sum_b N(b)$ and the largest estimated root action value. These become direct targets only in the later toy comparison. Original MuZero trains its policy head on search visit counts, while its value target incorporates outcomes or $n$ -step returns.

Model-based policy improvement

Exact policy improvement in a known tabular MDP can produce a non-worse policy under its assumptions. This finite-budget learned-model search is only an approximation: biased rewards, values, or latent transitions can make its chosen action worse. We therefore compare search decisions with the exact gridworld solution rather than presuming improvement.

Our fixed-horizon gridworld has an absorbing goal rather than a terminating episode, so the tutorial does not need a termination head. Original MuZero likewise does not use an explicit continuation head. Its search can proceed past a terminal state and relies on absorbing-state training and value predictions rather than learned termination. A model for environments with terminal states can predict continuation or discount explicitly, but neither this tutorial nor the cited MuZero results establish that such a head generally improves learning.

In[18]:

Code

ACTION_ONEHOT = torch.eye(N_ACTIONS)


class MCTSNode:
    __slots__ = ("z", "reward", "prior", "N", "W", "Q", "children", "expanded")

    def __init__(self, z=None, reward=0.0, prior=1.0):
        self.z = z
        self.reward = reward
        self.prior = prior
        self.N = 0
        self.W = 0.0
        self.Q = 0.0
        self.children = {}
        self.expanded = False


@torch.no_grad()
def expand_node(node, model):
    """Materialize all children of a leaf using the learned dynamics."""
    logits, value = model.predict(node.z)
    priors = torch.softmax(logits, dim=-1).squeeze(0).numpy()

    z_rep = node.z.expand(N_ACTIONS, -1)
    z_next, rewards = model.next_state(z_rep, ACTION_ONEHOT)
    _, child_values = model.predict(z_next)

    for a in range(N_ACTIONS):
        child = MCTSNode(
            z=z_next[a : a + 1],
            reward=float(rewards[a].item()),
            prior=float(priors[a]),
        )
        child.Q = float(child_values[a].item())
        node.children[a] = child
    node.expanded = True
    return float(value.item())


def select_action(node, c_puct):
    total = np.sqrt(node.N + 1)
    best_a, best_score = 0, -np.inf
    for a, child in node.children.items():
        q_edge = child.reward + GAMMA * child.Q
        score = q_edge + c_puct * child.prior * total / (1 + child.N)
        if score > best_score:
            best_score, best_a = score, a
    return best_a


_selection_probe = MCTSNode()
_selection_probe.N = 1
_selection_probe.children = {
    0: MCTSNode(reward=1.0, prior=0.5),
    1: MCTSNode(reward=0.0, prior=0.5),
}
_selection_probe.children[0].Q = 0.0
_selection_probe.children[1].Q = 0.5
assert select_action(_selection_probe, c_puct=0.1) == 0


def simulate(node, model, c_puct, depth, max_depth):
    if not node.expanded:
        value = expand_node(node, model)
        node.N += 1
        node.W += value
        node.Q = node.W / node.N
        return value
    if depth >= max_depth:
        value = node.Q
        node.N += 1
        node.W += value
        node.Q = node.W / node.N
        return value
    a = select_action(node, c_puct)
    child = node.children[a]
    value = child.reward + GAMMA * simulate(
        child, model, c_puct, depth + 1, max_depth
    )
    node.N += 1
    node.W += value
    node.Q = node.W / node.N
    return value


_depth_cap_probe = MCTSNode()
_depth_cap_probe.expanded = True
_depth_cap_probe.N = 1
_depth_cap_probe.W = 0.5
_depth_cap_probe.Q = 0.5
assert np.isclose(simulate(_depth_cap_probe, None, 1.5, 4, 4), 0.5)
assert _depth_cap_probe.N == 2
assert np.isclose(_depth_cap_probe.W, 1.0)


@torch.no_grad()
def run_mcts(model, obs, n_sims=32, c_puct=1.5, max_depth=4):
    """Plan from a single observation and return (value, action values, policy)."""
    z0 = model.encode(torch.tensor(obs, dtype=torch.float32).unsqueeze(0))
    root = MCTSNode(z=z0)
    expand_node(root, model)
    for _ in range(n_sims):
        simulate(root, model, c_puct, 0, max_depth)

    visits = np.array(
        [root.children[a].N for a in range(N_ACTIONS)], dtype=np.float64
    )
    q_values = np.array(
        [
            root.children[a].reward + GAMMA * root.children[a].Q
            for a in range(N_ACTIONS)
        ]
    )
    policy = visits / max(visits.sum(), 1.0)
    return float(q_values.max()), q_values, policy

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.

Worked example: planning from a single cellLink Copied

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.

In[19]:

Code

goal_steps = np.arange(0, 6)
discount_weights = GAMMA**goal_steps

Out[20]:

Visualization

Bar chart of discount factors gamma^0 through gamma^5; the final goal reward receives gamma^5. — Discount weights for the six transitions from the far corner. The only nonzero reward arrives on transition six, with exponent five and optimal return 0.774.

In[21]:

Code

rng_demo = np.random.default_rng(2)
demo_obs = observe((3, 3), rng_demo)
demo_value, demo_q, demo_policy = run_mcts(
    VE_MODEL, demo_obs, n_sims=64, max_depth=4
)

Out[22]:

Console

Search value from (3,3):  0.741
Optimal value from (3,3): 0.774
Behaviour value from (3,3): 0.700

Root action values: {'up': 0.741, 'right': 0.606, 'down': 0.653, 'left': 0.717}
Search policy:       {'up': 0.219, 'right': 0.0, 'down': 0.016, 'left': 0.766}

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.

Does search get better with more simulations?Link Copied

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.

In[23]:

Code

def evaluate_planning(
    model, sims_list=(1, 4, 16, 64), n_noise=4, max_depth=4, seed=3
):
    rng_local = np.random.default_rng(seed)
    optimal_sets = [
        set(np.flatnonzero(Q_STAR[s] >= Q_STAR[s].max() - 1e-9))
        for s in range(N_STATES)
    ]
    paired_inputs = [
        (s, observe((s // N, s % N), rng_local))
        for s in range(N_STATES)
        for _ in range(n_noise)
    ]
    results = {}
    for n_sims in sims_list:
        matches, values, targets = 0, [], []
        for s, o in paired_inputs:
            root_value, q_vals, _ = run_mcts(
                model, o, n_sims=n_sims, max_depth=max_depth
            )
            matches += int(int(np.argmax(q_vals)) in optimal_sets[s])
            values.append(root_value)
            targets.append(V_STAR[s])
        rmse = float(
            np.sqrt(np.mean((np.array(values) - np.array(targets)) ** 2))
        )
        results[n_sims] = (matches / len(paired_inputs), rmse)
    return results


PLANNING_RESULTS = evaluate_planning(VE_MODEL)
SEARCH_VALUE_RMSE = {k: metrics[1] for k, metrics in PLANNING_RESULTS.items()}

Out[24]:

Visualization

Bar chart of Q-greedy optimal-action agreement across 1, 4, 16, and 64 simulations, with corresponding value RMSE labels. — Q-greedy optimal-action agreement and root-value RMSE over 64 paired noisy observations (four per state) at four search budgets. Agreement uses argmax root Q, not the visit-count policy; these are measured outcomes, not a monotonic-improvement guarantee.

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.