Predicting the Next Frame

It is March 2018. Reinforcement-learning agents that act from raw pixels have spent several years getting good at reacting to one frame at a time, but none of them can close their eyes and picture what comes next. A short paper by two researchers proposes training a network to dream — to compress an environment into a latent model good enough that an agent can practice inside the model instead of the real world. Over the following twenty-one months, Hafner and collaborators at Google Brain and DeepMind turn that proof of concept into a benchmark-beating recipe, and set the template that action-conditioned video generation would rediscover, at far larger scale, six years later.

Where We Were

By the mid-2010s, deep reinforcement learning had shown that a neural network could map raw pixels directly to actions and learn a competent policy purely from trial and error and a scalar reward. The catch was sample efficiency: model-free methods — policy gradients and actor-critic algorithms that never build an explicit model of how the environment evolves — typically needed millions of environment interactions to converge, because every gradient step was informed only by reward, not by any internal picture of cause and effect. For anything expensive to interact with (a robot, a vehicle, a human-in-the-loop simulation) that cost was often prohibitive.

The alternative, model-based reinforcement learning, is older than deep learning itself: learn a model of the environment’s dynamics, then plan or optimize a policy against the model instead of the world. Classical control had used this successfully for decades with hand-built dynamics equations. What was missing was a way to learn those dynamics directly from images. Predicting a whole future frame in pixel space is expensive and brittle — small pixel-level errors compound over a rollout, and a model trained to minimize pixel error tends to collapse uncertain futures into a single blurry average rather than committing to one plausible outcome.

Two ingredients from the generative-modeling side of the field supplied the missing piece. The variational autoencoder, dated to December 2013, had shown how to compress high-dimensional data into a compact latent code trained end-to-end by backpropagation — the ELBO objective and reparameterization trick that make this possible are taught in full in image generation, Chapter 2. And recurrent networks, the LSTM lineage traced in language models, Chapter 3, had already demonstrated that a hidden state could carry information forward through a sequence. Nobody had yet combined the two into a single recipe: compress each observed frame into a small latent code, and train a recurrent model to predict how that code evolves under an agent’s own actions — a model small and fast enough to run many times faster than the real environment, so an agent could rehearse inside it.

The Breakthrough

“World Models” (Ha & Schmidhuber, 2018), posted to arXiv in March 2018, split the agent into three small, separately trained pieces. A Vision component (V) is a convolutional variational autoencoder that compresses each observed frame into a low-dimensional latent vector — 4,348,547 parameters, trained exactly as Chapter 2 of Image Generation describes. A Memory component (M) is a recurrent mixture-density network (MDN-RNN) — 422,368 parameters — that takes the current latent code and the agent’s action and predicts a distribution over the next latent code, rather than a single point estimate, so it can represent genuinely uncertain or multimodal futures instead of averaging them into a blur. A Controller (C) is a tiny linear policy — only 867 parameters — mapping the latent code and the recurrent hidden state directly to an action.

The controller is so small that it can be trained with an evolutionary strategy, CMA-ES, rather than gradient descent, sidestepping the need to backpropagate through the recurrent rollout at all. The full V and M pipeline trained in under an hour of computation on a single GPU; the controller’s evolutionary search ran on a single multi-core CPU machine. The striking demonstration was training the controller entirely inside “hallucinated” rollouts generated by M — the agent never touched the real environment during that phase — and then deploying the resulting policy back in the actual environment. On the CarRacing-v0 driving task the agent scored 906 ± 21 over 100 trials, clearing the 900-point threshold that defines the task as solved and beating the best prior leaderboard entry (838 ± 11) and standard deep RL baselines (591–652). On the VizDoom “take cover” task the policy trained inside the dream survived 918 ± 546 steps in the dream itself and 1,092 ± 556 steps back in the real game, against a roughly 750-step bar for solving the task.

The evolutionary controller and the separately-trained V and M pieces made World Models a proof of concept rather than a general recipe — CMA-ES does not scale past a few thousand parameters, and V and M were trained on data from a fixed, already-competent policy rather than improving together with the agent. The Deep Planning Network, or PlaNet (Hafner et al., 2018), posted to arXiv that November, replaced the hand-off between separately trained modules with one model trained end-to-end and one search procedure. PlaNet learns a single Recurrent State-Space Model (RSSM) that splits the latent state at each timestep into a deterministic path, carried by a recurrent network, and a stochastic path, sampled from a learned distribution conditioned on that recurrent state — combining the sharp long-range memory of a deterministic RNN with the ability to represent genuine uncertainty that a purely stochastic model gives up too easily. Rather than training a separate controller network, PlaNet chooses actions by planning: at every real timestep it uses the cross-entropy method to search over many candidate action sequences, evaluates each by rolling it out through the learned latent dynamics, and executes the first action of the best sequence found — never touching a policy network at all.

Tested on six continuous-control tasks from the DeepMind Control Suite — cartpole swing-up, reacher, cheetah running, finger spinning, ball-in-cup, and the walker — PlaNet outperformed the policy-gradient method A3C (trained on the underlying physical state, not pixels, for 100,000 episodes) within only 100 episodes trained from pixels alone, and after 500 episodes reached performance close to D4PG, a strong model-free actor-critic trained from images for 100,000 episodes — roughly a 200-fold reduction in the number of episodes needed, with a 26% relative improvement over D4PG’s final score on the cheetah task.

Planning through the model at every timestep is itself expensive, though — PlaNet took 10–20 hours of training on a single Nvidia V100 GPU to reach its reported results. Dreamer (Hafner et al., 2019), posted in December 2019, kept PlaNet’s RSSM but removed search from the loop entirely. Dreamer trains an actor network and a value network purely by imagining rollouts inside the latent dynamics model — because the RSSM’s transitions are differentiable, gradients from a predicted future value can be backpropagated straight through an imagined trajectory to update the policy, no environment interaction and no search required at either training or action time. Evaluated on 20 visual control tasks from the same suite, Dreamer reached an average score of 823 using only 5×10⁶ environment steps, matching or beating D4PG’s score of 786 at 10⁸ steps — a 20-fold reduction in environment interaction — while PlaNet, given the identical 5×10⁶-step budget, reached only 332. Dreamer also trained faster in wall-clock terms: about 3 hours per million environment steps, against roughly 11 hours for PlaNet’s online planning and 24 hours for D4PG to reach comparable performance.

The Mathematics

All three papers train the same underlying object: a latent state-space model of the environment, conditioned on the agent’s actions. PlaNet’s Recurrent State-Space Model, which Dreamer also uses, represents the state at each step as a pair — a deterministic part $h_t$ and a stochastic part $s_t$, following the owning paper’s own notation (PlaNet never introduces a separate symbol for the combined state):

$$\begin{aligned} h_t &= f(h_{t-1}, s_{t-1}, a_{t-1}) \\ s_t &\sim p(s_t \mid h_t) \\ \hat{o}_t &\sim p(\hat{o}_t \mid h_t, s_t) \end{aligned}$$

Symbol	Meaning
$h_t$	deterministic recurrent hidden state at time $t$
$s_t$	stochastic latent state at time $t$
$a_{t-1}$	action taken at the previous step
$f$	deterministic recurrent transition, an RNN or GRU cell
$p(s_t \mid h_t)$	learned prior over the next stochastic state given recurrent memory
$\hat{o}_t$	the model’s predicted or reconstructed observation

The model is trained with an encoder $q(s_t \mid h_t, o_t)$ that, given the actual next frame, infers a better estimate of $s_t$ than the prior could predict in advance — the same encoder-plus-prior structure as the single-step VAE, whose ELBO objective is owned by image generation, Chapter 2, extended here across a whole trajectory:

$$\mathcal{L} = \mathbb{E}_q\left[\sum_{t=1}^{T} \Big(\ln p(\hat{o}_t \mid h_t, s_t) - \beta \, D_{\mathrm{KL}}\big[q(s_t \mid h_t, o_t) \,\Vert\, p(s_t \mid h_t)\big]\Big)\right]$$

Symbol	Meaning
$q(s_t \mid h_t, o_t)$	encoder’s posterior over $s_t$, given the real frame
$p(\hat{o}_t \mid h_t, s_t)$	decoder likelihood of the observation
$D_{\mathrm{KL}}$	Kullback-Leibler divergence between encoder posterior and prior
$\beta$	weight balancing reconstruction against the KL penalty
$T$	length of the training trajectory

Training on this objective alone only ever asks the prior to predict one step ahead accurately. PlaNet’s distinguishing contribution, latent overshooting, extends the penalty to every planning horizon the model will later be asked to imagine over:

$$\mathcal{L}_{\text{overshoot}} = \sum_{d=1}^{D} \mathbb{E}_q\left[D_{\mathrm{KL}}\big[q(s_t \mid h_t, o_t) \,\Vert\, p_d(s_t \mid h_{t-d})\big]\right]$$

Symbol	Meaning
$D$	maximum overshoot distance trained for
$d$	a specific overshoot distance, $1 \le d \le D$
$p_d(s_t \mid h_{t-d})$	the state predicted $d$ steps ahead, without seeing intervening frames
$h_{t-d}$	the recurrent state $d$ steps before $t$

Dreamer discards PlaNet’s search entirely and instead trains a value function and a policy to act inside imagined rollouts. It learns a value network $V_\psi$ that, from any imagined state $s_\tau$, estimates discounted future imagined rewards up to the imagination horizon $H$ plus its own bootstrapped estimate of everything beyond that horizon — then trains a policy network $\pi_\phi$ to maximize that value by differentiating straight through the imagined trajectory:

$$\begin{aligned} V_\psi(s_\tau) &\approx \mathbb{E}_q\left[\sum_{n=\tau}^{t+H-1} \gamma^{\,n-\tau} r_n + \gamma^{\,t+H-\tau} V_\psi(s_{t+H})\right] \\ \max_{\phi} \;&\; \mathbb{E}_q\left[\sum_{\tau=t}^{t+H} V_\psi(s_\tau)\right] \end{aligned}$$

Symbol	Meaning
$V_\psi$	learned value network, parameters $\psi$
$\pi_\phi$	learned policy network, parameters $\phi$
$\gamma$	discount factor for future rewards
$r_n$	imagined reward at step $n$
$H$	imagination horizon, the number of steps rolled out before bootstrapping
$\tau$	an imagined timestep between $t$ and $t+H$

What It Enabled

The three papers form a clean lineage: a proof of concept with hand-off between separately trained pieces, then an end-to-end model with expensive planning, then an end-to-end model with cheap, fast policy execution learned purely in imagination.

	World Models, 2018	PlaNet, 2018	Dreamer, 2019
Latent dynamics model	VAE plus MDN-RNN, trained separately from the controller	Single RSSM, trained end to end	Same RSSM, trained end to end
How actions are chosen	Tiny linear controller, evolved with CMA-ES	Online planning, cross-entropy search over the model at every step	Actor network, trained by backprop through imagined rollouts
Test-time action cost	One forward pass through a near-trivial network	A full search over many candidate rollouts, every step	One forward pass through the policy network
Reported data efficiency	Solved CarRacing-v0 at 906 ± 21 over 100 trials	~200× fewer episodes than D4PG on DeepMind Control Suite tasks	20× fewer environment steps than D4PG on 20 visual control tasks

The core idea — compress observations into a latent space, learn how that latent space evolves under an agent’s actions, then use the learned dynamics instead of (or alongside) the real environment — is the world-model concept this book keeps returning to. It resurfaces directly, and at a completely different scale, in Genie (2024-02) and GameNGen (2024-08), where the “latent dynamics model” becomes a large video generative model conditioned on user actions, capable of rendering an interactive game frame by frame rather than a low-dimensional control vector. It resurfaces again, scaled to general-purpose robotics and embodied simulation, in the world-foundation-model wave covered in Chapter 16 — NVIDIA’s Cosmos platform and DeepMind’s Genie 2 and 3.

What did not resurface, at least not directly, was the actor-critic-in-imagination training loop itself scaling to photorealistic video. Dreamer’s imagination happens in a latent space of a few hundred dimensions; rendering a plausible next video frame at the resolution and fidelity later chapters demand required a different generative backbone entirely — diffusion, covered starting at Chapter 5. What this chapter’s lineage supplies to that later work is not the architecture but the framing: that an environment can be represented as a compressed, learnable, action-conditioned latent process, worth predicting on its own terms rather than only as a byproduct of policy learning.

Key Papers

Paper	Authors	Year	Link	Why it matters
World Models	Ha & Schmidhuber	2018	arXiv:1803.10122	First demonstration of training a control policy entirely inside a learned, hallucinated latent environment; introduces the Vision-Memory-Controller split
Learning Latent Dynamics for Planning from Pixels	Hafner et al.	2018	arXiv:1811.04551	Introduces the Recurrent State-Space Model and latent overshooting; end-to-end model-based agent that plans directly in latent space
Dream to Control: Learning Behaviors by Latent Imagination	Hafner et al.	2019	arXiv:1912.01603	Replaces online planning with an actor-critic trained purely on imagined latent rollouts, removing search from both training and action selection

PM Lens