"For years they fed me the world one frame at a time and asked me only to act, and I obliged, blindly, the way a swimmer obeys a current he cannot see. Then one day they let me close my eyes and dream the next frame myself, and the dream was wrong in a hundred small ways, but it was mine, and I could run it forward a thousand steps in the time it once took to take one. I planned inside that imagined world, made my mistakes where mistakes cost nothing, and woke to act on a single confident move. They call it a world model, as if I had swallowed the world whole. I have swallowed only a small latent shadow of it, just enough to rehearse tomorrow before tomorrow arrives, which, I have come to suspect, is all that anyone who plans has ever really had."
A Latent Dynamics Model Dreaming One Step Ahead of the Truth
Chapter Overview
Every agent in this part so far has learned to act inside a world it could not see the inside of. The model-free methods of Chapter 25 mapped observations straight to values or actions and threw the dynamics away, paying for that simplicity in millions of environment steps. This chapter is the culmination of Part VI, where the agent stops treating the environment as a black box and instead learns a compact internal model of how the world evolves, then plans or learns inside that model rather than in the costly real world. The move is the oldest idea in the book wearing new clothes: a world model is a temporal predictor, a learned state-space model whose latent state summarizes the past well enough to forecast the future, and whose forecasts are accurate enough to act on. The lineage runs straight back to the Kalman filter and the structured state-space models of earlier parts; here the dynamics are learned end to end and the prediction target is the agent's own future experience.
The chapter builds that idea in two movements. The first learns the model. A latent environment model encodes high-dimensional observations into a low-dimensional latent state and learns a transition that rolls that state forward, so the agent carries a small differentiable simulator of its world. Dreamer is the headline method: it learns a recurrent latent world model from pixels and then improves its policy entirely inside the model, in imagination, never touching the real environment during the policy-improvement step, which is where its extraordinary sample efficiency comes from. Predictive state representations offer a different and elegant answer to what the latent state should even be, defining it not as a hidden cause to be inferred but directly as predictions about future observations, sidestepping the latent-variable inference that the Kalman and HMM lineage required.
The second movement plans with the model. Once the agent holds a learned simulator, it can search: roll the model forward under candidate action sequences and pick the branch that pays. MuZero is the landmark result, learning its model and planning with Monte Carlo tree search over that learned model, mastering board games and Atari without ever being told the rules, because the rules are exactly what the learned dynamics recover. Planning in a learned latent space is what separates these methods from classical search, which needed a hand-given simulator. The chapter closes on the hardest problem in planning, the long horizon, where a flat one-step model and a flat search both drown in the number of steps. Temporal abstraction, planning over learned subgoals and extended options rather than primitive actions, is the lever that lets a learned model reach across hundreds or thousands of steps.
The temporal-thread payoff here is the strongest in the book. The world model is the learned successor to the state-space model of Chapter 7 and the structured state-space networks of Chapter 13: a recurrent latent state, a learned transition, an observation model, and a filtering step that updates belief as new observations arrive. What was once a tool for forecasting a univariate series becomes, with a reward head and an action input, the engine an agent rehearses its future on. Forecasting and decision making, the two halves of this book, meet in the same equations.
Prerequisites
This chapter sits at the confluence of two threads, so it draws on both. From Chapter 26: Advanced Reinforcement Learning you need the model-based reinforcement-learning material of Section 26.1: the distinction between model-free and model-based agents, why a learned model promises sample efficiency, and how model error compounds when the agent plans on an imperfect simulator, since every method here lives or dies by how it manages that compounding error. From Chapter 13: State-Space and Continuous-Time Neural Models you need the structured state-space view of a sequence as a latent state rolled forward by a learned transition, because a world model is exactly that object with actions and rewards attached. The deepest root is Chapter 7: State-Space Models and Filtering: the Kalman filter, the predict-update cycle, and the belief state are the classical skeleton that the learned latent dynamics of this chapter put muscle on, and recognizing Dreamer's recurrent state-space model as a neural Kalman filter is the single most clarifying connection in the chapter. Readers wanting to refresh state-space filtering, model-based RL, or the structured state-space networks before diving in will find each through the Table of Contents.
If you keep one idea from this chapter, keep this: a world model is a learned temporal predictor of the environment, the state-space model of Chapters 7 and 13 fitted with actions and a reward head, and once an agent holds that compact latent simulator it can improve its policy by learning in imagination (Dreamer) or by searching over the learned model (MuZero), trading expensive real-world interaction for cheap rehearsal inside its own forecasts. Latent environment models learn the simulator, predictive state representations redefine the latent state as predictions about the future, planning in latent space turns the learned model into searchable structure, and temporal abstraction over subgoals and options is what lets that planning reach across long horizons where flat one-step search drowns.
Chapter Roadmap
- 29.1 Latent Environment Models Learning a compact latent state that summarizes high-dimensional observations and a transition that rolls it forward, the differentiable simulator an agent carries inside it, recognized as the learned successor to the Kalman filter and structured state-space models of Chapters 7 and 13.
- 29.2 Dreamer and Recurrent World Models Learning a recurrent latent world model from pixels and improving the policy entirely in imagination: the recurrent state-space model as a neural filtering loop, the PlaNet and Dreamer line through DreamerV3, and why rehearsing inside the model rather than the world is where the sample efficiency comes from.
- 29.3 Predictive State Representations Defining the latent state directly as predictions about future observations rather than as a hidden cause to be inferred, the predictive-state and successor-representation view that sidesteps latent-variable inference while keeping the temporal summary the agent needs to act.
- 29.4 Planning in Latent Space Turning a learned model into searchable structure: rolling the latent dynamics forward under candidate actions and running Monte Carlo tree search over the learned model, the MuZero and AlphaZero line that masters games and Atari without ever being given the rules.
- 29.5 Temporal Abstraction and Long-Horizon Planning Planning over learned subgoals and extended options rather than primitive actions, the temporal-abstraction lever that lets a learned model and its search reach across hundreds or thousands of steps where flat one-step planning drowns in the horizon.
Once you have worked through the five sections, the Hands-On Lab below chains them into one progression on a single task. Each lab step maps to a section, so the lab doubles as a review: you will learn the latent model of Section 29.1, improve a policy in imagination as in Section 29.2, plan with search over the learned model from Section 29.4, and add a subgoal abstraction from Section 29.5, so that by the end the agent both dreams and plans inside a world it taught itself.
Hands-On Lab: Dream and Plan
Objective
Build an agent that learns a model of its world and then both dreams and plans inside it. You will first train a small latent world model from the latent-environment-model machinery of Section 29.1, an encoder, a learned latent transition, and reward and decoder heads, on a control task, so the agent ends up holding a compact differentiable simulator of its environment. You will then improve a policy entirely in imagination, Dreamer-style, following Section 29.2: roll the learned model forward under the current policy, compute returns inside those imagined trajectories, and backpropagate to improve the actor without touching the real environment during the improvement step. Next you will plan with Monte Carlo tree search over the same learned model, the MuZero-style search of Section 29.4, comparing a planned action against the policy's reflexive one. Finally you will add a subgoal abstraction from Section 29.5 and watch the agent reach a long-horizon goal that flat one-step search could not. The single thread is that a learned world model is one object that supports two complementary uses, fast amortized learning in imagination and deliberate search at decision time, and a capable temporal agent wants both.
What You'll Practice
- Training a latent world model with an encoder, a learned transition, and reward and observation heads, the simulator of Section 29.1.
- Improving a policy by learning in imagination over model rollouts, the Dreamer-style loop of Section 29.2.
- Running Monte Carlo tree search over a learned model to plan an action, the MuZero-style search of Section 29.4.
- Adding a learned subgoal abstraction to extend planning across a long horizon, following Section 29.5.
- Reading a world model as a temporal predictor, the learned descendant of the state-space filter of Chapters 7 and 13.
Setup
You need a Python environment with torch for the model and policy, gymnasium for the environment, and numpy for bookkeeping. A pixel or low-dimensional control task with a clear goal works best, since the long-horizon stage needs a goal far enough away that one-step search struggles. The sketch below is the heart of Section 29.1: a recurrent latent transition that takes the previous latent state and the action and predicts the next latent state, the learned analogue of the Kalman predict step, with a reward head reading off the latent so the agent can score imagined futures without rendering them. Everything else in the lab hangs off this one learned dynamics function.
import torch
import torch.nn as nn
class LatentDynamics(nn.Module):
"""A learned transition: the neural successor to the Kalman predict step."""
def __init__(self, latent_dim=32, action_dim=4, hidden=256):
super().__init__()
self.gru = nn.GRUCell(latent_dim + action_dim, hidden)
self.to_latent = nn.Linear(hidden, latent_dim) # next latent state
self.reward_head = nn.Linear(hidden, 1) # score an imagined step
def forward(self, latent, action, h):
x = torch.cat([latent, action], dim=-1)
h = self.gru(x, h) # roll the state forward
next_latent = self.to_latent(h)
reward = self.reward_head(h)
return next_latent, reward, h
Steps
Step 1: Train the latent world model
Collect experience from the environment and fit the latent world model of Section 29.1: encode observations into latents, train the transition above to predict the next latent, and train the reward and decoder heads to reconstruct reward and observation. Validate the model by rolling it forward open-loop and comparing its imagined trajectory against a real one, making the compounding model error of model-based RL concrete.
Step 2: Improve a policy in imagination
Freeze the world model and improve a policy entirely inside it, the Dreamer-style loop of Section 29.2: starting from encoded real states, roll the model forward under the current actor, compute imagined returns, and backpropagate through the differentiable rollout to update the actor and a critic. The real environment is touched only to gather data and seed the rollouts, never during the improvement step itself.
Step 3: Plan with search over the learned model
At decision time, run Monte Carlo tree search over the learned model from Section 29.4: expand candidate action sequences by rolling the latent dynamics forward, score leaves with the critic, and return the best first action. Compare the planned action against the policy's reflexive action from Step 2 and measure where deliberate search beats the amortized policy.
Step 4: Add a subgoal abstraction for the long horizon
Extend the agent to a long-horizon goal where flat one-step search from Step 3 drowns in the number of steps. Following Section 29.5, introduce a learned subgoal or option, plan over subgoals rather than primitive actions, and let a low-level controller reach each subgoal, so the search horizon shrinks from raw steps to a handful of abstract moves.
Step 5: Compare dreaming against planning
Put the three modes side by side on the same task: the reflexive policy from Step 2, the search-augmented agent from Step 3, and the temporally abstract planner from Step 4. Report sample efficiency and final performance for each, and discuss when learning in imagination suffices and when deliberate search over the learned model earns its cost, the central trade-off of the chapter.
Expected Output
The lab produces one learned world model and three ways of using it. Step 1 yields a latent simulator whose open-loop rollouts track reality for a useful horizon before drifting, the visible signature of model error that every later step must respect. Step 2 yields a policy improved entirely in imagination, reaching competence with far fewer real environment steps than a model-free agent would need, the sample-efficiency payoff of Dreamer. Step 3 shows search over the learned model outperforming the reflexive policy on states where a few steps of lookahead matter, the MuZero-style case for planning. Step 4 shows the temporally abstract planner solving a long-horizon goal that the flat planner could not reach, making concrete why temporal abstraction is the long-horizon lever. The reader finishes able to build a world model and to use it both ways, learning in imagination for cheap broad competence and searching at decision time for sharp deliberate moves, and, most of all, able to see the world model for what it is: the state-space temporal predictor of the earlier parts of this book, turned into the engine of an agent that rehearses its future before living it.
The from-scratch latent dynamics and tree search of this lab are worth building once so the imagined rollout and the search over a learned model stop being abstractions, but you should rarely reimplement the search core in production. DeepMind's mctx library provides batched, JAX-native Monte Carlo tree search, the exact MuZero-style and Gumbel-AlphaZero search of Section 29.4, behind a single function call that plans over any learned model you supply, collapsing the hand-rolled tree of Step 3 into a few lines while running thousands of simulations in parallel on accelerator hardware. The reference Dreamer implementations similarly package the recurrent state-space model and the imagination loop of Section 29.2 into a configurable agent. The discipline the chapter teaches still governs the libraries: a search that plans flawlessly over a model whose rollouts diverge after ten steps will plan confidently into a future that never arrives. The library makes the search fast; it does not make the world model accurate.
Stretch Goals
- Replace the recurrent transition of Section 29.1 with a structured state-space layer from Chapter 13 and measure whether the longer effective memory improves open-loop rollout accuracy, making the world-model-as-state-space-model connection literal.
- Implement a predictive state representation from Section 29.3 as an alternative latent and compare its rollout fidelity and planning performance against the inferred-latent model, contrasting predictions-as-state with hidden-cause-as-state.
- Stress the agent by degrading the world model (fewer training steps, a smaller latent) and measure how planning and imagination each fail as model error grows, exposing which use of the model is more robust to an imperfect simulator.
What's Next?
This chapter completes Part VI: Sequential Decision Making. Across eight chapters the agent grew from a Markov decision process it could solve exactly, through partial observability and bandits, the deep and advanced reinforcement learning that scaled it, optimal control and imitation, and sequences reframed as decisions, to this final step, an agent that learns a model of its world and plans inside it. The thread that ran through it all was temporal: a decision is a forecast of consequences, and a world model is the forecaster made explicit. Part VII: Building Intelligent Temporal Systems turns from learning to act toward reasoning about why. Chapter 30: Temporal Reasoning and Causality asks what it means for a temporal system to understand cause and effect rather than merely predict correlation, the foundation an agent needs to plan in a world it can intervene on rather than only observe, and the natural sequel to a chapter about planning inside a learned model. The full path through Part VII and the rest of the book is laid out in the Table of Contents.
Bibliography & Further Reading
World Models and Latent Dynamics
Ha, D., Schmidhuber, J. "World Models." NeurIPS, 2018. arXiv:1803.10122. arxiv.org/abs/1803.10122
The paper that named the world-model agenda, training a variational autoencoder and a recurrent dynamics model so an agent could learn a policy inside its own learned latent simulation, the conceptual root of Section 29.1.
Hafner, D., et al. "Learning Latent Dynamics for Planning from Pixels (PlaNet)." ICML, 2019. arXiv:1811.04551. arxiv.org/abs/1811.04551
The paper introducing the recurrent state-space model and planning from pixels by latent rollouts, the latent-dynamics-and-planning bridge between Sections 29.1 and 29.4.
Dreamer and Learning in Imagination
Hafner, D., Lillicrap, T., Ba, J., Norouzi, M. "Dream to Control: Learning Behaviors by Latent Imagination (Dreamer)." ICLR, 2020. arXiv:1912.01603. arxiv.org/abs/1912.01603
The paper that improved a policy entirely in imagination by backpropagating through learned latent rollouts, the headline learning-in-imagination method of Section 29.2.
Hafner, D., Lillicrap, T., Norouzi, M., Ba, J. "Mastering Atari with Discrete World Models (DreamerV2)." ICLR, 2021. arXiv:2010.02193. arxiv.org/abs/2010.02193
The paper that brought discrete latent world models to Atari, the first imagination-trained agent to match top model-free methods on the benchmark, deepening Section 29.2.
Hafner, D., Pasukonis, J., Ba, J., Lillicrap, T. "Mastering Diverse Domains through World Models (DreamerV3)." 2023. arXiv:2301.04104. arxiv.org/abs/2301.04104
The paper that made one fixed-hyperparameter world-model agent work across more than a hundred diverse tasks, including collecting diamonds in Minecraft from scratch, the maturation of the Dreamer line in Section 29.2.
Predictive State Representations
Littman, M. L., Sutton, R. S., Singh, S. "Predictive Representations of State." NeurIPS, 2001. proceedings.neurips.cc/paper/2001
The paper defining a predictive state representation, a state expressed directly as predictions about future observations rather than as a hidden cause to be inferred, the foundation of Section 29.3.
Dayan, P. "Improving Generalization for Temporal Difference Learning: The Successor Representation." Neural Computation, 5(4), 613-624, 1993. gatsby.ucl.ac.uk/~dayan/papers/d93b.pdf
The paper introducing the successor representation, a predictive encoding of state as expected future occupancy, the conceptual cousin of predictive state representations that recurs in Section 29.3.
Planning with Learned Models and Search
Schrittwieser, J., et al. "Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model (MuZero)." Nature, 588, 604-609, 2020. arXiv:1911.08265. arxiv.org/abs/1911.08265
The paper that mastered board games and Atari by learning its own model and planning with Monte Carlo tree search over it, without ever being given the rules, the landmark of Section 29.4.
Silver, D., et al. "A General Reinforcement Learning Algorithm that Masters Chess, Shogi, and Go through Self-Play (AlphaZero)." Science, 362, 1140-1144, 2018. arXiv:1712.01815. arxiv.org/abs/1712.01815
The paper that unified search and learning through self-play with a known simulator, the immediate predecessor whose given model MuZero replaced with a learned one, framing Section 29.4.
Ye, W., Liu, S., Kurutach, T., Abbeel, P., Gao, Y. "Mastering Atari Games with Limited Data (EfficientZero)." NeurIPS, 2021. arXiv:2111.00210. arxiv.org/abs/2111.00210
The paper that made MuZero-style learned-model search dramatically sample efficient, reaching human-level Atari from two hours of data, the data-efficient frontier of Section 29.4.
Temporal Abstraction and Long-Horizon Planning
Hafner, D., Lee, K.-H., Fischer, I., Abbeel, P. "Deep Hierarchical Planning from Pixels (Director)." NeurIPS, 2022. arXiv:2206.04114. arxiv.org/abs/2206.04114
The paper that planned over learned subgoals inside a world model, letting a manager set goals in latent space for a worker to reach, the long-horizon temporal-abstraction method of Section 29.5.
Tools and Libraries
DeepMind. "mctx: Monte Carlo Tree Search in JAX." 2021. github.com/google-deepmind/mctx
The library of batched, JAX-native Monte Carlo tree search, including MuZero and Gumbel-AlphaZero search over any learned model, the off-the-shelf planning core for Section 29.4 and the lab.