Model-Based Deep RL: World Models, Dreamer, and MuZero

This builds on Model-Based Deep RL: Sample Efficiency and PETS, which made the case that a good model turns experience into imagined planning, showed how a learned model's errors compound over a rollout, and built PETS — a probabilistic ensemble planned online with model-predictive control.

PETS still predicts full next states in the environment's native coordinates. That is the assumption the two methods here relax. The first learns a compact latent world and imagines inside it; the second predicts only the handful of quantities planning actually consumes. Both keep the same trade — cheap imagined computation for expensive real interaction — but change what the model has to get right.

Latent-space world models

PETS plans in the environment's native state space. That works when the state is a clean vector of joint angles, but not when it is a stream of images: a pixel-space dynamics model must reconstruct every future frame, spending its capacity predicting textures and shadows that are irrelevant to control. The alternative is to learn a latent state — a compact code $z$ that captures what matters for control — and do all the modeling and planning there.

World Models, Ha and Schmidhuber (2018), was the clean demonstration.¹ It splits the agent into three parts. A vision model (V), a variational autoencoder, compresses each high-dimensional frame $o_t$ into a small latent $z_t$. A memory model (M), a recurrent network with a mixture-density output (an MDN-RNN), models the dynamics in latent space: given $z_t$ and action $a_t$ it predicts a distribution over the next latent $z_{t+1}$. A tiny controller (C) — a linear policy — maps the latent state and the RNN's hidden state to an action. Almost all the parameters live in V and M, which are trained unsupervised on collected rollouts; the controller is small enough to optimize with a black-box search.

The central result was that the controller can be trained entirely inside the model's dream. Once M has learned the latent dynamics, it is itself a simulated environment: sample a latent, feed it and an action to M, get the next latent and a predicted reward, repeat. Ha and Schmidhuber trained a car-racing and a VizDoom controller by evolution against M's hallucinated rollouts and then transferred the policy back to the real environment, where it performed well — learning to act without touching the environment during policy optimization at all. A temperature knob on M's sampling controlled how noisy the dream was, which mattered: too deterministic a dream let the controller exploit the model's flaws, the same failure PETS guards against with ensembles.

The Dreamer line

World Models optimized a fixed learned model, then a controller, in separate phases. The Dreamer line, Hafner et al. (2020–2023), made the whole thing one end-to-end, continually-improving loop and turned it into a top general agent.² The core is a recurrent state-space model (RSSM) that learns a latent MDP: a latent state combining a deterministic recurrent part and a stochastic part, with heads that predict, from the latent, the next latent, the reward, the discount (episode continuation), and, for training the representation, a reconstruction of the observation.

What makes Dreamer more than a bigger World Model is how the policy is learned. The agent learns an actor ${\pi }_{\varphi }$ and critic $v_{\psi }$ purely from imagined latent trajectories. From latent states drawn from real experience it rolls the learned model forward a short horizon $H$, entirely in latent space, producing imagined states, actions, and predicted rewards ${\hat{r}}_k$. The critic regresses onto a $\lambda$-return over the imagined rollout, bootstrapped past the horizon by its own value,

and the actor maximizes the imagined return ${\mathbb{E}}_{{\pi }_{\varphi }}\left[{\sum }_k{V}_k^{\lambda }\right]$. Because the entire imagined rollout is differentiable through the learned model, the actor is trained by backpropagating value gradients straight through the dynamics — a signal no model-free method has access to.

Successive versions sharpened it. DreamerV2 discretized the stochastic latent into categorical variables and was the first world-model agent to reach human-level on the Atari 200M benchmark. DreamerV3 fixed a single set of hyperparameters and, without tuning, learned across more than 150 tasks spanning continuous and discrete control, proprioception and pixels — and collected diamonds in Minecraft from scratch, a long-horizon exploration task that had resisted every prior method. The idea is constant: learn a latent world, then do almost all the learning by imagining inside it, spending real interaction only to keep the model accurate.

MuZero: planning with a learned latent model

AlphaZero reached superhuman play in Go, chess, and shogi by combining MCTS with a learned policy-and-value network — but it planned with the real game rules, a perfect simulator handed to it. That is exactly the assumption the tabular chapter's MCTS made, and it fails the moment the rules are unknown or the rules are the physics of an Atari game. MuZero, Schrittwieser et al. (2020), removes it: it learns a model and plans with MCTS against that learned model, matching AlphaZero in board games and DQN-class methods in Atari, all with no given rules.³

MuZero's decisive idea is what its model predicts. A pixel-accurate model like World Models reconstructs future observations; MuZero does not — it never predicts an observation at all. Its learned model predicts only the three quantities that MCTS actually consumes: the reward, the value, and the policy. This is the payoff of the compounding-error discussion taken to its conclusion: if planning only ever reads reward, value, and policy off the model, then the model only needs to get those right, and the enormous burden of predicting irrelevant visual detail disappears.

Three learned functions do the work, all operating on an abstract latent state with no required meaning beyond being useful for planning.

Inside MCTS these functions replace the game rules. The representation function turns the current observation into the root latent state. Then the search expands the tree entirely in latent space: to expand a node, the dynamics function produces the child's latent state and the reward on that edge, and the prediction function supplies a value to back up and a policy prior to bias which children to explore. The tree search is the same selection–expansion–backup loop as AlphaZero's; only the transitions and evaluations now come from learned functions instead of a simulator.

Training ties the three functions together through the search. The network is unrolled $K$ steps from a real state, applying $g_{\theta }$ to the actions the agent actually took, and at each unrolled step $k$ three losses are summed: the predicted policy $p^k$ matches the MCTS-improved policy ${\pi }_{t+k}$ (as in AlphaZero, the search is a policy-improvement operator), the predicted value $v^k$ matches the observed $n$-step return or game outcome $z_{t+k}$, and the predicted reward $r^k$ matches the real reward $u_{t+k}$,

There is no reconstruction loss — the latent states are shaped only by these three prediction targets, so the model learns whatever internal representation makes reward, value, and policy predictable, and nothing more.

The result placed MuZero as the natural successor to AlphaZero. It matched AlphaZero's superhuman strength in Go, chess, and shogi without being told the rules, and on the Atari benchmark it set a new state of the art among methods of its class, learning to play from pixels with the same learned-model-plus-MCTS architecture. It is the tabular chapter's closing promise made good at scale: MCTS gives the deep, focused lookahead; the learned representation, dynamics, and prediction functions supply the model MCTS needs — a model that predicts only what planning reads, and so has nothing to get wrong except the things that matter.

Beyond the three: short rollouts, latent planning, and sample-efficient MuZero

PETS, Dreamer, and MuZero mark out the design space along two axes — what the model represents, and how planning consumes it:

The work after them mostly finds better points inside this space; three lines are worth naming because they became defaults.

MBPO: short model rollouts inside a model-free learner. The compounding-error arithmetic above says long imagined rollouts are untrustworthy — so MBPO (Janner et al., 2019) uses very short ones and lets a model-free learner do the rest.⁴ It trains an ensemble model as in PETS, but instead of planning with MPC it generates short (often single-step) imagined transitions branched from real states in the replay buffer and feeds them, alongside real data, to an off-policy SAC. Branching from real states keeps every imagined transition close to the data the model was fit on, so ${\sigma }_H$ never has room to blow up, while the imagined transitions still multiply the effective sample count. Janner et al. gave a monotonic-improvement analysis bounding the model-induced error by the rollout length, formalizing the keep $H$ short instinct, and MBPO matched or beat PETS and SAC on the MuJoCo benchmarks. It is the clean hybrid: model-based sample efficiency, model-free robustness, with the model trusted only one step at a time.

TD-MPC: plan in a latent value model. TD-MPC (Hansen et al., 2022) fuses the PETS and Dreamer ideas.⁵ It learns a latent dynamics model like Dreamer, but rather than reconstruct observations it trains the latent purely from a value signal (a TD loss on a learned latent Q), the MuZero move of predicting only what planning reads. At decision time it plans with MPPI (a CEM relative) over short latent rollouts, and it bootstraps the finite-horizon plan with the learned value at the end — so a short plan sees past its own horizon through the value function. That combination — learned latent value model, short latent planning, value bootstrap — set strong sample-efficiency results on continuous control and carried directly to real robots. It reads as the natural join of this lesson's three lines: a latent world model (Dreamer), planned online by MPC (PETS), with a learned value that predicts what matters (MuZero).

EfficientZero: MuZero made sample-efficient. MuZero's headline results used enormous amounts of data. EfficientZero (Ye et al., 2021) added three targeted fixes — a self-supervised consistency loss tying the latent dynamics to the encoder's own next-state features, a learned prediction of the sum of rewards over the search horizon, and a correction for off-policy staleness in the value target — and reached human-level median performance on the Atari 100k benchmark, i.e. after only about two hours of gameplay.⁶ It was the first time a MuZero-style learned-model planner was competitive in the extreme low-data regime, closing much of the gap between model-based planning and the sample budgets real applications allow.

The one idea

Strip the three systems to their common shape and they answer one question three ways: what must a learned model predict for planning to be worth it? PETS predicts full next states but distrusts itself through an ensemble and replans every step, so error never accumulates. World Models and Dreamer predict a compact latent state and do their learning inside it, imagining thousands of trajectories per real one. MuZero predicts neither states nor pixels but only reward, value, and policy — the minimal quantities MCTS reads — and plans against those. Across all three the same sentence holds: a good model turns experience into imagined planning, and the work is in building a model that is trustworthy exactly where planning queries it.