World Models for Control
How learning a latent dynamics model lets an agent train inside its own imagination, why that is the answer to robotics' sample-efficiency problem, and where the learned simulator's errors get exploited.
Model-free reinforcement learning needs millions of environment interactions. On a simulated humanoid that is a weekend; on a real robot arm it is years of wall-clock time and several replaced gearboxes. The oldest idea in the field for closing that gap is to learn a model of the environment and train the policy against the model, so that the expensive resource, real interaction, is spent on improving the model rather than on gradient steps.
The modern form of this is a latent world model. Rather than predicting future pixels, the agent learns a compact recurrent state, predicts forward in that state space, and trains an actor and critic entirely on imagined rollouts. DreamerV3 does exactly this and reports outperforming specialised methods across over 150 tasks with a single fixed hyperparameter configuration, including collecting diamonds in Minecraft from scratch with no human data and no curriculum (Hafner et al., 2023, arXiv:2301.04104). The fixed-configuration claim matters more than the Minecraft one: robustness techniques based on normalisation, balancing and transformations are what let one recipe hold across domains that previously each needed their own tuning.
The three parts
A world model of this kind has a representation, a dynamics head and a decoder. The representation encodes observations into a stochastic latent \(z_t\) alongside a deterministic recurrent state \(h_t\). The dynamics head predicts \(z_{t+1}\) from \((h_t, z_t, a_t)\) without seeing the next observation, which is the part that must be right, since imagination runs on it alone. The decoder reconstructs observations and rewards, and exists mainly to force the latent to retain information.
Training then splits cleanly. The world model is fit by supervised prediction on replayed real experience. The actor and critic are trained on rollouts imagined from replayed latents, typically a horizon of 15 to 20 steps, never touching the environment. The environment supplies data; it does not supply gradients for the policy.
Learned simulators and latent actions
A different lineage drops the reward and asks whether a controllable environment can be learned from video alone. Genie trains on unlabelled internet gameplay video and learns a latent action model, inferring what action was taken between consecutive frames without any ground-truth action labels, producing an 11B-parameter foundation world model that can be acted in frame by frame from a text prompt, a photograph or a sketch (Bruce et al., ICML 2024, arXiv:2402.15391).
For robotics that is the interesting direction, because action-labelled robot data is the scarce resource and video is not. If latent actions inferred from human video can be aligned with a robot's real action space, then video becomes usable pretraining data for control rather than only for perception.
Where it breaks
Model exploitation. A policy optimised against a learned model will find the model's errors. Classic symptoms are physically impossible exploits: the agent discovers a latent trajectory the dynamics head assigns high reward to and that the real environment cannot produce. Short imagination horizons, ensembles, and penalising uncertainty are the standard defences, and all of them trade optimism for conservatism.
Compounding error over the horizon. One-step prediction error compounds multiplicatively across an imagined rollout, which is the reason horizons stay near 15 steps rather than 500. Long-horizon planning through a learned model remains largely unsolved.
Reconstruction spends capacity on the wrong things. A decoder trained to reproduce pixels allocates capacity by pixel count, so a large moving background outweighs a small task-critical object. Reward-relevant detail is systematically under-represented; this is the motivation for reconstruction-free variants that predict in representation space instead.
Real robots do not reset. Sample efficiency is not the only obstacle. Even a perfect world model needs data from a real robot that must be manually reset after every failure, and no amount of imagination removes the human in that loop.
12 flashcards for this concept
Click a card to reveal the answer.