Overview
A physicist simulating a plasma instability on a supercomputer waits three days for a single trajectory. A chemist screening reaction conditions runs hundreds of wet-lab experiments to map a yield surface. In both cases, the bottleneck is the same: the real environment (or its faithful numerical simulator) is too expensive to query at the scale that modern optimization and planning algorithms demand. World models offer an escape. They are learned approximations of environment dynamics that trade a small amount of accuracy for orders-of-magnitude speedup, letting an agent "imagine" thousands of future trajectories before committing to a single real action.
This chapter builds world models from the ground up. We start with latent dynamics models, in particular the Recurrent State-Space Model (RSSM) architecture behind DreamerV3, and derive the Evidence Lower Bound (ELBO) objective that trains them (Section 44.1). We then use trained world models for counterfactual reasoning: asking "what would have happened if I had taken a different action?" and quantifying the confidence of those answers (Section 44.2). Finally, we build a complete model-based planner that selects actions by rolling out imagined trajectories inside the world model, and we diagnose the compound error that accumulates over long planning horizons (Section 44.3).
World models sit at the intersection of several threads in this book. They extend the differentiable programming ideas of Chapter 42 to sequential decision problems, complement the numerical simulators of Chapter 43 with learned surrogates, and provide the dynamics backbone for the optimization methods of Chapter 45 and the automated experiment design of Chapter 46.
Prerequisites
This chapter assumes familiarity with PyTorch (tensor operations, autograd, training loops), variational inference at the level of Chapter 32: Bayesian Discovery and Uncertainty, and the basics of reinforcement learning (states, actions, rewards, policies). Prior exposure to recurrent neural networks (GRUs or LSTMs) is helpful but not required; we introduce the relevant mechanics as needed. Readers who have worked through Chapter 42: Differentiable Programming will find the gradient-through-time derivations familiar.
Learning Outcomes
- Explain why learned dynamics models accelerate scientific discovery by reducing the cost of trajectory simulation.
- Implement a Recurrent State-Space Model (RSSM) with separate deterministic and stochastic state components.
- Derive the ELBO objective for training latent dynamics models and connect each term to reconstruction, dynamics, and representation learning.
- Use a trained world model for counterfactual reasoning: intervening on past actions and propagating the consequences.
- Build a rollout-based model predictive control (MPC) planner that selects actions by imagining future trajectories.
- Diagnose and mitigate compound error accumulation over multi-step rollouts.
Sections
44.1 Latent Dynamics Models
Recurrent State-Space Models (RSSM), the DreamerV3 architecture, transformer world models, the ELBO training objective, and the distinction between deterministic and stochastic state components.
44.2 Counterfactual Reasoning with World Models
Interventions in latent space, do-calculus over learned dynamics, confidence-calibrated counterfactual predictions, and applications to scientific "what-if" questions.
44.3 Building a World Model Planner
Recipe: build an RSSM world model, train it on Gymnasium trajectories, implement 5-step rollout planning with the Cross-Entropy Method, and diagnose compound error accumulation.
Bibliography
Foundational Papers
DreamerV3: the state-of-the-art RSSM-based world model agent that masters Atari, DMLab, and Minecraft without domain-specific tuning. The architecture and training objective we implement in this chapter.
The original RSSM paper (PlaNet), introducing the split deterministic-stochastic state and the latent overshooting objective.
A seminal paper coining the modern usage of "world models," combining a VAE with an RNN to learn compact environment representations for planning.
MuZero: learns a world model without access to environment rules, combining learned dynamics with Monte Carlo Tree Search.
Transformer World Models
IRIS: a transformer-based world model that achieves human-level Atari performance with 100k environment interactions, demonstrating the viability of attention-based dynamics models.
TWM: a transformer world model that further improves sample efficiency, comparing favorably to RSSM-based approaches.
Books and Surveys
A comprehensive survey covering world model architectures, training objectives, and applications, with a taxonomy that extends beyond driving to scientific domains.
The standard RL textbook; Chapter 8 covers model-based planning and the Dyna architecture that prefigures modern world model approaches.
Tools and Libraries
The primary deep learning framework used throughout this chapter for implementing RSSM components and training world models.
The standard reinforcement learning environment interface, used here to generate training trajectories and evaluate planners.
Community PyTorch implementation of DreamerV3, useful as a reference for production-grade RSSM training.
A reliable set of RL algorithm implementations in PyTorch, useful for generating baseline policies and comparison trajectories.