kb://library/muzero-and-learned-model-planningstable2026-06-16

MuZero & Learned-Model Planning

muzeroalphazeromctsmodel-based-rllearned-modelvalue-equivalenceplanningtree-searchlatent-dynamicsefficientzerotest-time-computefrontier

MuZero & Learned-Model Planning

One-line summary. MuZero (Schrittwieser et al. 2019) is what you get when you take AlphaZero's tree search but throw away the rulebook: instead of being handed a perfect simulator, it learns its own latent dynamics model and runs Monte-Carlo Tree Search inside that learned model. The radical, easily-missed move is that the learned model is never asked to reconstruct the real world — only to predict the three quantities planning actually uses (reward, policy, value). This is "value equivalence," and it's why the same algorithm mastered Go, chess, shogi (matching rule-aware AlphaZero) and 57 Atari games from pixels, with no game rules supplied.

What it is (intuition first)

The companion article [[world-models]] is about learning a simulator of the world. This article is about the other half of the loop: once you have a learned model, how do you actually plan with it? The cleanest, most influential answer is MuZero — and the surprising lesson is that planning works best when the learned model is deliberately not a faithful world model at all.

Start with chess. AlphaZero ([[reinforcement-learning-fundamentals]]) plays superhuman chess by combining a neural net (which scores positions and suggests moves) with Monte-Carlo Tree Search (MCTS) — it imagines move sequences, building a search tree, and picks the move that looks best after deep look-ahead. But AlphaZero has an enormous cheat: it is given the rules of chess. It has a perfect, hand-coded simulator. To search the tree it just applies the rules: "if I move the bishop here, the board becomes exactly that." That's fine for board games where the rules are known and exact. It is useless for Atari from pixels, or robotics, or the real world — there is no rulebook to hand the agent.

MuZero's question: what if the agent had to learn its own simulator, and that simulator only had to be good enough to plan well? A naïve approach would learn a model that predicts the next screen (the next pixels), then plan by rolling that forward — that's the generative [[world-models]] route, and it's hard (you waste capacity predicting every pixel, and small errors compound into garbage). MuZero refuses to predict pixels. Instead it learns an abstract latent state — an internal vector with no required correspondence to the real game state — under exactly one constraint: when you plan with it, the reward, value, and policy predictions come out right.

That is the whole trick. The model is a fiction the agent tells itself; the fiction is allowed to be wrong about everything except the things planning consumes. DeepMind's framing: MuZero "learns a model that, when applied iteratively, predicts the quantities most directly relevant to planning: the reward, the action-selection policy, and the value function." Not the world. The planning-relevant projection of the world.

This page assumes you're comfortable with the RL basics ([[reinforcement-learning-fundamentals]] — MDPs, value functions, policies) and have at least skimmed [[world-models]] (learning the dynamics) and [[test-time-compute-reasoning]] (spending compute at inference to think). MuZero sits exactly at the intersection of those three.

Why it matters

  • It removed the "you must know the rules" assumption. AlphaZero needed a perfect simulator. MuZero proved you can keep AlphaZero's planning power while learning the simulator from experience — the same algorithm, one set of ideas, going from rule-based board games to pixel-based Atari. That generality is the headline.
  • It is the cleanest statement of "plan in a learned latent space." Every model-based agent that imagines futures in an abstract latent — Dreamer ([[world-models]]), EfficientZero, robotics MPC — is a cousin. MuZero made the value-equivalence design choice explicit and showed it scales.
  • It reframed what a model is *for*. A model isn't a truth-telling oracle about the world; it's a compute substrate for search. This is a genuinely different philosophy from the "render reality faithfully" school, and the tension between them is live (see [[world-models]] caveats).
  • It's the RL ancestor of LLM test-time search. The pattern "neural net proposes + tree search disposes + the search result trains the net" is exactly the loop now being ported to language-model reasoning ([[test-time-compute-reasoning]], [[self-play-and-verifier-driven-training]], [[process-outcome-rewards]]). Understanding MuZero is understanding where "let the model search before answering" came from.

How it works (real mechanics)

0. The base you must know: AlphaZero = net + MCTS + self-play

AlphaZero has a single neural network f(state) → (policy p, value v): given a board, output a probability over moves (p, the prior) and an estimate of who's winning (v). At decision time it runs MCTS: from the current position, repeatedly (a) select a path down the tree using a score that balances value (exploit) against prior × (visit-count bonus) (explore — the PUCT rule), (b) expand a new leaf and evaluate it with the net, (c) back-up the value along the path. After hundreds of simulations, the visit-count distribution at the root is a sharpened, search-improved policy π. AlphaZero plays the most-visited move; in training it (i) uses π as the target for the policy head (policy improvement) and (ii) uses the game's final outcome z as the target for the value head. Self-play generates the data; the improved policy from search becomes the next net's training signal. MuZero keeps all of this. The only thing it changes is where the tree's transitions come from.

1. MuZero's three learned functions: h, g, f

AlphaZero applied the real rules to get child states. MuZero replaces the rules with three learned networks:

  • Representation `h` — encodes the observation history into an initial hidden state s⁰ = h(o₁, …, o_t). This is the root of the search tree. (For Atari, o are frames; for Go, the board.)
  • Dynamics `g` — the learned transition: given a hidden state and an action, produce the next hidden state and the predicted immediate reward: s^{k}, r^{k} = g(s^{k-1}, a^{k}). This is the simulator. It runs entirely in latent spaces are abstract vectors, not boards or frames.
  • Prediction `f` — given a hidden state, output policy and value: p^{k}, v^{k} = f(s^{k}). (This is AlphaZero's net, now applied to latent states.)
s⁰              = h(o₁ … o_t)            # encode history → root latent
for each action a in the tree:
    s', r       = g(s, a)                # learned dynamics: next latent + reward
    p, v        = f(s')                  # policy prior + value at that latent
# MCTS uses (p, v, r) exactly like AlphaZero — but every node is a LATENT state

MCTS now runs over the tree of latent states produced by g, scored by f. The agent imagines "if I take this action then that action…" purely as latent rollouts — never decoding back to pixels, never consulting any rules. The search yields the same product as AlphaZero: a visit-count policy π at the root, which is what the agent plays and what trains the policy head.

2. The load-bearing idea: value equivalence (don't model the world, model the value)

Here is the part everyone misses on first read. The hidden states `s^k` are semantically unconstrained. There is no loss telling g to make s^1 correspond to the actual game position after action a. You could not necessarily decode a board out of s^1. MuZero never reconstructs the observation. The only constraints are on the three outputs:

  • the predicted reward r^k must match the real reward observed,
  • the predicted policy p^k must match the MCTS-improved π at that step,
  • the predicted value v^k must match a bootstrapped return target.

This is the value-equivalence principle (Grimm et al. 2020): two models are "value equivalent" if they produce the same Bellman updates / value predictions for the policies you care about, even if their internal states are nothing alike. MuZero learns the simplest model that is value-equivalent over its own trajectory distribution — it spends zero capacity on world-detail that doesn't move reward/value/policy. Grimm's follow-up ("Proper Value Equivalence," 2021) shows MuZero can be read as minimizing an upper bound on a proper-value-equivalence loss. This is the theoretical backbone: MuZero's model is a value-predictor unrolled in a computation graph, not a world simulator.

3. Training: unroll K steps, backprop through time

MuZero is trained end-to-end. Sample a real trajectory; pick a start; unroll the model `K` steps (typically K=5) by feeding the actually-taken actions through g, and align each step's predictions to targets from the real trajectory:

s⁰ = h(o_{1:t})
loss = 0
for k in 0..K:
    p^k, v^k = f(s^k)
    loss += policy_loss(p^k, π_{t+k})          # match MCTS-improved search policy
    loss += value_loss (v^k, z_{t+k})           # match n-step bootstrapped return
    if k>0: loss += reward_loss(r^k, u_{t+k})   # match observed reward
    s^{k+1}, r^{k+1} = g(s^k, a_{t+k})          # advance the learned dynamics
# one big backprop-through-time through h, g, f jointly

Crucially `π` (the policy target) is produced by MCTS, not by f alone — so the search at data-collection time improves the policy, and the net is trained to imitate that improved policy, which makes the next round of search stronger. That virtuous loop — search improves policy → policy trains net → better net makes search stronger — is inherited straight from AlphaZero and is the engine of [[self-play-and-verifier-driven-training]]. Value targets use bootstrapping (n-step return + a value-at-the-end), so MuZero works in long-horizon, sparse-reward Atari, not just terminal-reward board games.

4. The headline results (and what's a claim vs. a fact)

  • Board games: MuZero matched AlphaZero in Go, chess, and shogi without being given the rules — it had to discover a sufficient model from scratch. (Fact, reproduced; the board-game regime is the strongest evidence.)
  • Atari: state-of-the-art across the 57-game ALE suite from pixels at the time of publication. (Reported by DeepMind; later independently reproduced in open frameworks like LightZero, though exact numbers are sensitive to compute and tuning.)

5. The variant lineage (each fixes a real limitation)

Plain MuZero assumes a small discrete action set, a deterministic environment, and a big compute budget. The follow-ups each lift one assumption:

  • EfficientZero (Ye et al. 2021) — sample efficiency. Vanilla MuZero is data-hungry; EfficientZero hits super-human Atari-100k (≈2 hours of game experience) by adding three fixes: a self-supervised consistency loss (SimSiam/SPR-style) that forces g's predicted next-latent to match h of the real next observation (re-introducing a little grounding into the otherwise-unconstrained latent), a value-prefix head (predict the sum of rewards over a window, easing the "exactly when does reward arrive" problem), and off-policy correction for stale replay data. This is the practical "MuZero you can actually afford" and a key bridge back toward [[world-models]]' grounded-latent camp.
  • Sampled MuZero (Hubert et al. 2021) — large / continuous action spaces. You can't enumerate children when actions are continuous (robot torques); Sampled MuZero plans over a sampled subset of actions and corrects for the sampling, extending MuZero to arbitrarily complex action spaces.
  • Gumbel MuZero (Danihelka et al. 2022) — search efficiency. Replaces the heuristic PUCT/visit-count policy improvement with a principled Gumbel-top-k action-selection that guarantees policy improvement even with very few simulations — letting MuZero plan well with a tiny search budget.
  • Stochastic MuZero (Antonoglou et al. 2022) — randomness. Real environments (and games like 2048, backgammon) have chance events; it adds afterstates and chance nodes so the model can represent stochastic transitions instead of pretending the world is deterministic.

These compose: e.g. EfficientZero-V2 (2024) blends Gumbel search with continuous control for sample-efficient robotics.

Key ideas & tradeoffs

AxisAlphaZeroMuZeroEfficientZero / variants
SimulatorGiven (exact rules)Learned latent dynamics gLearned + lightly grounded (consistency)
DomainsKnown-rule board gamesBoard games and pixel Atari+ small-data, continuous, stochastic
What the model predictsn/a (rules are exact)reward, value, policy only+ next-latent consistency, value-prefix
State semanticsReal boardAbstract, unconstrained latentLatent nudged toward real observation
Core betSearch a perfect modelSearch a value-equivalent modelMake value-equivalence data-efficient

Load-bearing ideas:

  • A model for planning ≠ a model of the world. This is the philosophical heart and the cleanest contrast with the generative [[world-models]] camp. MuZero says: don't model what you can't use; model exactly the projection search consumes. It's elegant and it scales — and it has a real cost (next section).
  • Search is policy improvement. MCTS isn't just "look-ahead at test time" — the search-improved policy is the training target. Search and learning are the same flywheel. This is the idea now animating LLM reasoning RL ([[process-outcome-rewards]], [[grpo]], [[rlvr]]).
  • Compounding error is bounded by the policy prior. Because the model is only accurate near the trajectories it trained on, the PUCT prior p biases search toward actions where the model is reliable — search stays in-distribution. This is a feature, but it's also the seed of the central caveat below.
  • Grounding is a dial, not a binary. Pure MuZero grounds the latent not at all (only via reward/value/policy); EfficientZero adds a little observation-consistency; full generative world models ground everything in pixels. The whole [[world-models]] spectrum is "how much do you force the latent to resemble reality?"

Honest caveats & open questions

  1. MuZero's learned model does *not* generalize to unseen policies — and that's a measured result, not a quibble. De Vries, Voelcker et al. ("What model does MuZero learn?", 2023) found the learned model "struggles to generalize when evaluating unseen policies, which limits its capacity for additional policy improvement." In plain terms: MuZero does not learn a robust, reusable world model you can hand to a different planner or policy and trust. It learns a model narrowly accurate along the trajectory distribution its own policy visits — value-equivalence is local. MCTS quietly papers over this by using the policy prior to keep search in regions where the model is accurate. So the romantic reading ("MuZero discovers the rules of chess") overstates it: it discovers a sufficient-for-its-own-search fiction, not the rules.
  1. Value equivalence is a strength *and* a ceiling. The same property that makes the model cheap to learn (ignore world-detail) makes it brittle for transfer, model inspection, or off-policy re-use. If you want a model that supports replanning under a new objective, MuZero's latent is the wrong tool — you'd want a more grounded world model, paying the cost the generative camp pays.
  1. Compute and tuning are heavy; reproduction is non-trivial. Original MuZero used large self-play infrastructure. Open reimplementations exist (EfficientZero's public code, LightZero, muzero-general) and broadly validate the ideas, but exact headline numbers are sensitive to compute, replay ratio, and search budget. Treat specific scores as configuration-dependent, not universal constants.
  1. It assumes a clean reward signal. MuZero is built on rewards (game score, win/loss). It says nothing about where the reward comes from — which is the entire hard problem in open-ended, real-world, or language settings. Porting MuZero-style search to domains without a crisp reward (e.g. "good reasoning") requires a verifier or reward model, which reintroduces all the difficulties of [[rlvr]] and [[process-outcome-rewards]].
  1. The board-game evidence is far stronger than the everything-else evidence. Matching AlphaZero in Go/chess/shogi is rock-solid. The leap from "super-human Atari/2048" to "general model-based planning for robots and agents" is real progress but still narrow and short-horizon. Be excited about the principle; don't over-extrapolate from the demos — the same honesty discipline this library applies to [[world-models]] and [[embodied-ai-and-vla]].

How it connects to OpenAlice

OpenAlice's Alice is a deployed agentic LLM system — a perceive→decide→act loop ([[agentic-loops]]) with tools and memory, layered on frontier models, not a model-based-RL training lab. MuZero connects as a conceptual ancestor and a design lens, and it's worth being precise about which seams are real today vs. aspirational:

  • "Search before you act" is MuZero's central idea, ported to language. The pattern MuZero crystallized — the policy proposes, a search refines it, and the search result becomes the training signal — is exactly what modern LLM reasoning RL is doing ([[test-time-compute-reasoning]], [[self-play-and-verifier-driven-training]], [[grpo]], [[rlvr]], [[process-outcome-rewards]]). When Alice spends extra inference compute exploring candidate plans/tool-calls before committing, it's doing the shape of MuZero's MCTS in a discrete tool/action space — even though there's no learned latent dynamics model under it today. MuZero is the principled, learned version of that look-ahead.
  • Value equivalence is a useful warning for any "world model for the agent" ambition. If Alice ever drives a richly-simulated or embodied surface (the streaming/world-rendering candidate in [[world-models]] and [[embodied-ai-and-vla]]), MuZero's lesson cuts both ways: a model trained only to predict task value will be cheap and brittle; a model you want to replan against needs grounding. Knowing the tradeoff up front is the value here.
  • The verifier bottleneck is the same one Alice faces. MuZero works because Atari/Go hand it a clean reward. Alice's harder problem — and the org's bench rig (bench.blal.pro) — is getting a trustworthy reward/verifier for agentic and coding tasks ([[agentic-rl-long-horizon-coding]]). MuZero is a clean illustration that given a good reward, search + learning is a flywheel; the OpenAlice-relevant frontier is supplying that reward without a game's built-in score.
  • The honesty transfers verbatim. "MuZero discovers a model good enough for its own search, not a true world model" is the model-based-RL version of "a fluent answer is not a correct one." Measure the thing you actually care about ([[agent-evaluation]] is the discipline) — a model that plans well on-distribution can be silently useless off it.

Read [[world-models]] first (learning the dynamics) and [[reinforcement-learning-fundamentals]] for MCTS/value functions; this article is the planning-with-that-model half, and it's the cleanest bridge from classic model-based RL to today's [[test-time-compute-reasoning]].

See also

  • [[world-models]] — the companion: learning the latent/generative dynamics this article plans with (Dreamer, Genie, JEPA, Sora).
  • [[reinforcement-learning-fundamentals]] — MDPs, value functions, policy improvement; the MCTS + PUCT machinery MuZero inherits from AlphaZero.
  • [[test-time-compute-reasoning]] — spending inference compute on search/look-ahead; MuZero is its model-based-RL ancestor.
  • [[self-play-and-verifier-driven-training]] — the search-improves-policy-improves-net flywheel, generalized.
  • [[grpo]] · [[rlvr]] · [[process-outcome-rewards]] — porting "search + reward + policy improvement" to language-model reasoning.
  • [[embodied-ai-and-vla]] · [[agentic-rl-long-horizon-coding]] — where learned-model planning meets robots and long-horizon agents (and the verifier bottleneck).
  • [[agentic-loops]] · [[agent-evaluation]] — Alice's perceive→decide→act loop and the discipline of measuring whether a plan actually worked.