--- title: "Model-Based Deep RL: World Models, Dreamer, and MuZero" module: Modern Deep Reinforcement Learning moduleNumber: 5 lessonNumber: 6 order: 506 summary: > A companion to the PETS lesson. PETS plans in the environment's native state space; these methods change what the model represents. World Models and Dreamer learn a compact latent state and do almost all their learning by imagining inside it, with value gradients flowing through the differentiable dynamics. MuZero predicts neither states nor pixels — only the reward, value, and policy that MCTS reads — and plans with search against that learned model, AlphaZero without the rules. We close with MBPO, TD-MPC, and EfficientZero. topics: [Deep RL] sources: - book: Grokking Deep RL ref: "Ch. 12 — model-based methods; sample efficiency and learned models" - book: Sutton & Barto ref: "Ch. 8 — Planning and Learning (Dyna, MCTS), scaled to function approximation" --- This builds on [Model-Based Deep RL: Sample Efficiency and PETS](/reinforcement-learning/modern-deep-rl/model-based-rl), 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.[^worldmodels] 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. $$ % caption: The World Models loop. The vision model V encodes an observation into a % latent z; the memory model M (a recurrent network) predicts the next latent from % z and the action; the controller C maps the latent state to an action. Because M % predicts z-to-z, the agent can roll the loop forward with no images at all — % imagining, or "dreaming," entire trajectories in latent space. \begin{tikzpicture}[>=stealth, font=\small, box/.style={draw, minimum width=20mm, minimum height=10mm, align=center, font=\footnotesize}] \definecolor{acc}{HTML}{2348F2} \definecolor{red}{HTML}{C0392B} \node[box] (obs) at (0,0) {observation}; \node[box, draw=acc, text=acc, thick] (enc) at (3.3,0) {encode V\\latent z}; \node[box] (mem) at (6.8,0) {memory M\\next z}; \node[box] (ctrl) at (6.8,-2.3) {controller C\\action}; \draw[->, acc, thick] (obs) -- (enc); \draw[->, acc, thick] (enc) -- (mem); \draw[->, black, thick] (enc.south) .. controls (3.3,-2.3) .. (ctrl.west); \draw[->, black, thick] (mem) -- (ctrl) node[midway, right, font=\scriptsize, text=black] {latent state}; % imagine loop: next z feeds back as z_t \draw[->, red, thick] (mem.north) .. controls (6.8,1.6) and (3.3,1.6) .. (enc.north) node[midway, above, font=\scriptsize, text=red] {imagine: feed next z back, no image}; \end{tikzpicture} $$ 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.[^dreamer] 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_\phi$ 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, $$ V^\lambda_k = \hat r_k + \gamma\big[(1-\lambda)\,v_\psi(\hat z_{k+1}) + \lambda\,V^\lambda_{k+1}\big], \qquad V^\lambda_H = v_\psi(\hat z_H), $$ and the actor maximizes the imagined return $\mathbb{E}_{\pi_\phi}\!\big[\sum_{k} V^\lambda_k \big]$. 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. $$ % caption: Dreamer's imagination loop. Real experience is encoded into latent % states and used to fit the world model (left). Learning happens on the right, % entirely in latent imagination: from a real latent the model rolls forward H % steps predicting reward and the next latent; the critic fits the imagined % returns and the actor is trained to maximize them, with gradients flowing back % through the differentiable model. Only occasionally does the improved actor act % in the real environment to gather more data. \begin{tikzpicture}[>=stealth, font=\small, box/.style={draw, minimum width=22mm, minimum height=10mm, align=center, font=\footnotesize}] \definecolor{acc}{HTML}{2348F2} \definecolor{red}{HTML}{C0392B} \node[box] (env) at (0,0) {real experience\\(replay)}; \node[box, draw=acc, text=acc, thick] (wm) at (0,-2.3) {world model\\(latent MDP)}; \node[box] (imag) at (4.6,-2.3) {imagine H steps\\in latent space}; \node[box] (ac) at (9.0,-2.3) {actor + critic\\on imagined returns}; \draw[->, black, thick] (env) -- (wm) node[midway, right, font=\scriptsize, text=black] {f\/it model}; \draw[->, acc, thick] (wm) -- (imag); \draw[->, acc, thick] (imag) -- (ac); % gradient back through the model \draw[->, red, thick] (ac.north) .. controls (9.0,-0.6) and (4.6,-0.6) .. (imag.north) node[midway, above, font=\scriptsize, text=red] {value gradients through model}; % improved actor collects new real data \draw[->, black, thick] (ac.north) .. controls (9.0,1.4) and (0,1.4) .. (env.north) node[midway, above, font=\scriptsize, text=black] {act occasionally, gather data}; \end{tikzpicture} $$ ```algorithm caption: $\textsc{Dreamer}$ — learn a latent world model, act in imagination initialize world model, actor $\pi_\phi$, critic $v_\psi$, replay buffer $\mathcal{D}$ repeat // model learning sample a sequence batch from $\mathcal{D}$ update the RSSM to predict next latent, reward, discount, and reconstruct observations // behavior learning, entirely in latent imagination for each latent $z$ from the batch do imagine a rollout $\hat z_0{=}z, \hat a_k \sim \pi_\phi, \; \hat z_{k+1}, \hat r_k$ for $k < H$ compute $\lambda$-returns $V^\lambda_k$ over the imagined rollouts update critic $v_\psi$ to regress $V^\lambda_k$; update actor $\pi_\phi$ to maximize $V^\lambda_k$ // data collection act with $\pi_\phi$ in the real environment; add transitions to $\mathcal{D}$ until converged ``` 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](/reinforcement-learning/deep-rl/case-studies) 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] 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. > **Definition (MuZero's three functions).** A **representation** function > $h_\theta$ maps the observed history to an initial latent state $s^0 = h_\theta(o_{\le t})$. > A **dynamics** function $g_\theta$ maps a latent state and action to the next > latent state and an immediate reward, $(s^{k}, r^{k}) = g_\theta(s^{k-1}, a^k)$. > A **prediction** function $f_\theta$ maps a latent state to a policy and value, > $(p^k, v^k) = f_\theta(s^k)$. The latent states are never decoded back to > observations; they exist only to make the reward, value, and policy predictions > accurate. 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. $$ % caption: MuZero inside MCTS. The representation function h encodes the real % observation into a root latent state; the dynamics function g unrolls the tree % in latent space, producing each child's latent state and the edge reward r; the % prediction function f scores each latent state with a value v and a policy prior % p. No node is ever decoded to pixels — the tree plans over abstract states % chosen only to predict reward, value, and policy well. \begin{tikzpicture}[>=stealth, font=\small, st/.style={circle, draw, fill=white, minimum size=6.5mm, inner sep=0pt}, box/.style={draw, minimum width=17mm, minimum height=7mm, align=center, font=\scriptsize}] \definecolor{acc}{HTML}{2348F2} \definecolor{red}{HTML}{C0392B} % observation -> root via h \node[box] (obs) at (-3.4,2.6) {observation}; \node[st, draw=acc, thick] (root) at (0,2.6) {s0}; \draw[->, acc, thick] (obs) -- (root) node[midway, above, font=\scriptsize, text=black] {rep. h}; % tree in latent space \node[st] (c1) at (-1.6,0.9) {s1}; \node[st] (c2) at (1.6,0.9) {s1}; \node[st] (g1) at (-2.6,-0.8) {s2}; \node[st] (g2) at (-0.6,-0.8) {s2}; \draw[->, black] (root) -- (c1) node[midway, left=1mm, font=\scriptsize, text=red] {r}; \draw[->, black] (root) -- (c2) node[midway, right=1mm, font=\scriptsize, text=red] {r}; \draw[->, black] (c1) -- (g1) node[midway, left=1mm, font=\scriptsize, text=red] {r}; \draw[->, black] (c1) -- (g2); \node[acc, anchor=west, font=\scriptsize] at (1.9,0.0) {dynamics g: (s, a) to (next s, r)}; % prediction f on a node \node[box] (pred) at (3.0,-1.7) {pred. f: v, p}; \draw[->, acc, thick] (g2) .. controls (1.0,-1.7) .. (pred.west); \node[anchor=north, font=\scriptsize, text=black] at (-1.0,-1.6) {latent tree, no pixels}; \end{tikzpicture} $$ 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}$, $$ \mathcal{L}(\theta) = \sum_{k=0}^{K} \Big[ \ell^r(u_{t+k}, r^k) + \ell^v(z_{t+k}, v^k) + \ell^p(\pi_{t+k}, p^k) \Big] + c\,\|\theta\|^2. $$ 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. $$ % caption: MuZero's training unroll. From a real state the network unrolls K steps % along the actually-taken actions; at each step the reward head is trained to the % real reward, the value head to the observed return, and the policy head to the % MCTS-improved policy. There is no pixel reconstruction: the abstract states are % shaped only by these three prediction targets. \begin{tikzpicture}[>=stealth, font=\small, st/.style={circle, draw, fill=white, minimum size=7mm, inner sep=0pt}] \definecolor{acc}{HTML}{2348F2} \definecolor{red}{HTML}{C0392B} \node[st, draw=acc, thick] (s0) at (0,0) {s0}; \node[st] (s1) at (3.0,0) {s1}; \node[st] (s2) at (6.0,0) {s2}; \draw[->, acc, thick] (s0) -- (s1) node[midway, above, font=\scriptsize, text=black] {g, action a0}; \draw[->, acc, thick] (s1) -- (s2) node[midway, above, font=\scriptsize, text=black] {g, action a1}; % targets under each node \foreach \n/\x in {s0/0, s1/3.0, s2/6.0} { \node[anchor=north, font=\scriptsize, align=center, text=red] at (\x,-0.55) {reward r}; \node[anchor=north, font=\scriptsize, align=center, text=black] at (\x,-1.15) {value v}; \node[anchor=north, font=\scriptsize, align=center, text=acc] at (\x,-1.75) {policy p}; } \node[anchor=west, font=\scriptsize, text=black] at (6.4,-1.15) {targets: real reward, return, MCTS policy}; \end{tikzpicture} $$ 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: | Method | Model represents | Planning | Trust mechanism | | --- | --- | --- | --- | | PETS | full next state (native coords) | online MPC, replan each step | probabilistic ensemble | | Dreamer | compact latent state | actor-critic on imagined rollouts | short horizon $H$, reconstruction | | MuZero | reward, value, policy only | MCTS over learned model | no reconstruction; three prediction targets | 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.[^mbpo] 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. ```algorithm caption: $\textsc{MBPO}$ — model-generated short rollouts into a model-free learner initialize policy $\pi$, model ensemble $\{\hat p_i\}$, env buffer $\mathcal{D}_{env}$, model buffer $\mathcal{D}_{model}$ repeat take an action with $\pi$ in the real environment; add the transition to $\mathcal{D}_{env}$ fit the ensemble $\{\hat p_i\}$ on $\mathcal{D}_{env}$ for several branch rollouts do sample a start state $s$ from $\mathcal{D}_{env}$ roll $\hat p_i$ forward $k$ steps under $\pi$ ($k$ small); add imagined transitions to $\mathcal{D}_{model}$ update $\pi$ with SAC on batches drawn from $\mathcal{D}_{env} \cup \mathcal{D}_{model}$ until converged ``` **TD-MPC: plan in a latent value model.** **TD-MPC** (Hansen et al., 2022) fuses the PETS and Dreamer ideas.[^tdmpc] 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.[^efficientzero] 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. $$ % caption: Three post-PETS/MuZero systems located by the design axes of this lesson. % MBPO keeps native-state modeling but only one-step rollouts fed to a model-free % learner; TD-MPC plans in a latent value model with a value bootstrap; EfficientZero % keeps MuZero's latent-model MCTS but adds losses that make it data-efficient. \begin{tikzpicture}[>=stealth, font=\footnotesize, box/.style={draw, minimum width=32mm, minimum height=15mm, align=center, font=\scriptsize}] \definecolor{acc}{HTML}{2348F2} \definecolor{red}{HTML}{C0392B} \node[box, draw=acc, text=acc] (mbpo) at (0,0) {MBPO\\short rollouts\\into SAC}; \node[box, draw=acc, text=acc] (tdmpc) at (4.4,0) {TD-MPC\\latent value model\\+ MPC + bootstrap}; \node[box, draw=red, text=red] (ez) at (8.8,0) {Ef\/f\/icientZero\\MuZero + consistency\\for low data}; \node[anchor=north, font=\scriptsize, text=black] at (0,-1.1) {2019}; \node[anchor=north, font=\scriptsize, text=black] at (4.4,-1.1) {2022}; \node[anchor=north, font=\scriptsize, text=black] at (8.8,-1.1) {2021}; \end{tikzpicture} $$ ## 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. [^worldmodels]: **Ha & Schmidhuber (2018)**, "World Models," _NeurIPS_ (and arXiv:1803.10122). The V (variational autoencoder), M (mixture-density recurrent network modeling latent dynamics), and C (small linear controller) decomposition; the controller trained by evolution inside M's "dream" and transferred to the real car-racing and VizDoom environments, with a sampling-temperature control on the dream. [^dreamer]: **Hafner, Lillicrap, Ba, Norouzi (2020)**, "Dream to Control: Learning Behaviors by Latent Imagination," _ICLR_ (Dreamer); **Hafner et al. (2021)**, "Mastering Atari with Discrete World Models," _ICLR_ (DreamerV2); **Hafner, Pasukonis, Ba, Lillicrap (2023)**, "Mastering Diverse Domains through World Models," arXiv:2301.04104 (DreamerV3). A recurrent state-space model learns a latent MDP; actor and critic are learned purely from imagined latent rollouts, with value gradients backpropagated through the differentiable dynamics. DreamerV3 with fixed hyperparameters spans 150+ tasks and collects diamonds in Minecraft from scratch. [^mbpo]: **Janner, Fu, Zhang, Levine (2019)**, "When to Trust Your Model: Model-Based Policy Optimization," _NeurIPS_. MBPO: an ensemble dynamics model generates short (often one-step) rollouts branched from real replay states, fed with real data to an off-policy SAC; a monotonic-improvement bound ties the model error to the rollout length, and the method matches or exceeds PETS and SAC on MuJoCo. [^tdmpc]: **Hansen, Wang, Su (2022)**, "Temporal Difference Learning for Model Predictive Control," _ICML_. TD-MPC: a latent dynamics model trained from a TD/value loss (no observation reconstruction), planned online with MPPI over short latent rollouts and bootstrapped by a learned terminal value; strong sample efficiency on continuous control and transfer to real robots. [^efficientzero]: **Ye, Liu, Kurutach, Abbeel, Gao (2021)**, "Mastering Atari Games with Limited Data," _NeurIPS_. EfficientZero: MuZero plus a self-supervised latent-consistency loss, a predicted value-prefix (reward-sum) head, and an off-policy value correction, reaching human-level median performance on the Atari 100k (two-hour) benchmark. [^muzero]: **Schrittwieser, Antonoglou, Hubert, Simonyan, Sifre, Schmitt, Guez, Lockhart, Hassabis, Graepel, Lillicrap, Silver (2020)**, "Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model," _Nature_ 588. MuZero's representation ($h$), dynamics ($g$), and prediction ($f$) functions operate on abstract latent states with no reconstruction loss; MCTS plans over the learned model, and training matches the reward, value, and MCTS-improved policy at each unrolled step; matched AlphaZero in Go/chess/shogi without given rules and set a new state of the art on Atari.