World Models: Latent Dynamics, RSSM & Learning in Imagination

"I do not need to crash a thousand real cars to learn to drive. I crash a thousand imaginary ones, in a dream I built from watching, and wake up cautious."

An Agent That Learned to Plan Inside Its Own Head

Every generative model so far in this chapter produced observations: frames, objects, scenes. A world model produces dynamics. It is the point where generative vision meets decision-making, and it closes an arc that began far back in the book. Chapter 15 introduced state estimation with the Kalman filter: a hidden state, a motion model that predicts how it evolves, and observations that correct it. A world model is that idea made fully learned and deeply nonlinear, with a VAE-style encoder (Chapter 31) for the observation model and a neural network for the transition. The cross-reference map names this exactly: motion models and Kalman state estimation, deepened through video understanding and VAE latents, transformed here into learned latent dynamics.

1. Why Imagine? The Sample-Efficiency Argument Beginner

Reinforcement learning that acts directly in an environment is sample-hungry: a robot or game agent may need millions of real interactions to learn a good policy, and real interactions are slow, expensive, or dangerous. The world-model proposal, articulated in Ha and Schmidhuber's "World Models" (2018), is to learn a model of the environment from logged experience, then train the policy by simulating, or imagining, rollouts inside that model. Imagined steps are cheap (a forward pass, no physics engine, no real robot), so the policy can practice billions of times in its own dream and only occasionally touch reality to refine the model. The illustration below captures the bargain: the agent crashes a thousand imaginary cars in a dream it built from watching, and touches the costly real road only now and then.

This is the central payoff and it is dramatic. Make it concrete. A model-free agent that needs 10 million real robot steps, at one second each, faces over 100 days of continuous, collision-prone real-world acting. A Dreamer agent learns from 1 million real steps plus a billion imagined ones, and the imagined steps run at a millisecond each, so almost all the practice moves into the dream. That turns 100 days of robot time into roughly 11 days of real acting plus an overnight GPU run.

That is the felt meaning of world-model agents such as DreamerV3 (Hafner et al., 2023) reaching strong performance with one to two orders of magnitude fewer real environment steps than model-free baselines, with a single configuration mastering domains from Atari to continuous control to Minecraft. The price is that the dream must be good enough; a policy that learns to exploit a flaw in the model (driving through an imaginary wall the model failed to render) fails in reality. The quality of the world model is the ceiling on the policy, which is why evaluation (Section 36.8) is so central.

2. The RSSM: Deterministic Memory Meets Stochastic State Advanced

The technical heart is the Recurrent State-Space Model. The latent state at time $t$ is split into two parts: a deterministic recurrent state $h_t$ that carries reliable memory of the past, and a stochastic state $z_t$, a sampled latent capturing the irreducible uncertainty of what comes next. This split is the RSSM's signature insight. A purely deterministic model cannot represent genuine randomness (a die roll, an opponent's choice); a purely stochastic model struggles to carry long-range memory. Splitting them gets both.

The deterministic part $h_t$ is the hidden state of a gated recurrent unit (GRU), a small recurrent network whose learned gates decide how much of its memory vector to keep versus overwrite at each step. It is the same recurrent update operator you met inside RAFT in Chapter 26, reused here to give the world model a stable thread of memory through time.

The model has four learned components. The recurrent model updates the deterministic state, $h_t = f(h_{t-1}, z_{t-1}, a_{t-1})$, folding in the previous action $a_{t-1}$. The representation model (the encoder, used during training) infers the stochastic state from the current observation, $z_t \sim q(z_t \mid h_t, x_t)$. The transition (prior) model predicts the stochastic state from memory alone, $\hat{z}_t \sim p(z_t \mid h_t)$, with no access to the observation, this is the predictor used to imagine. Decoder, reward, and continuation heads read out pixels, reward, and episode-end from the full state $(h_t, z_t)$. Training maximizes a variational lower bound (the ELBO you met in Chapter 31), reconstructing observations and rewards while a KL term pulls the observation-informed posterior $q$ toward the observation-free prior $p$, so the prior alone can carry the dream forward.

Figure 36.5.1 unrolls the RSSM and highlights the one path that makes imagination possible: the prior $p(z_t \mid h_t)$, which predicts the next stochastic state from memory alone. During training the posterior $q$ (which peeks at the observation) drives reconstruction; during imagination only the prior runs, so the model generates futures without any real input. The code below implements one RSSM transition step.

# One RSSM transition step: a GRU carries deterministic memory, a prior predicts
# the next stochastic state from memory alone (used to imagine), and a posterior
# uses the observation (used in training). The KL between them teaches the dream.
import torch
import torch.nn as nn

class RSSMCell(nn.Module):
    """One step of a Recurrent State-Space Model: deterministic GRU memory plus
    a stochastic latent drawn from a prior (imagination) or posterior (training)."""
    def __init__(self, stoch=32, deter=256, action_dim=6, obs_feat=1024, hidden=256):
        super().__init__()
        self.gru = nn.GRUCell(stoch + action_dim, deter)
        self.prior_net = nn.Sequential(nn.Linear(deter, hidden), nn.SiLU(),
                                       nn.Linear(hidden, 2 * stoch))   # mean, logstd
        self.post_net = nn.Sequential(nn.Linear(deter + obs_feat, hidden), nn.SiLU(),
                                      nn.Linear(hidden, 2 * stoch))
        self.stoch = stoch

    def _sample(self, params):
        mean, logstd = params.chunk(2, dim=-1)
        std = torch.exp(logstd.clamp(-5, 2))
        eps = torch.randn_like(std)
        return mean + std * eps, mean, std        # reparameterized sample

    def forward(self, prev_stoch, prev_action, prev_deter, obs_feat=None):
        x = torch.cat([prev_stoch, prev_action], dim=-1)
        deter = self.gru(x, prev_deter)            # advance deterministic memory
        prior = self.prior_net(deter)              # predict next state from memory alone
        if obs_feat is None:                       # IMAGINATION: dream forward
            stoch, mean, std = self._sample(prior)
            return stoch, deter, (mean, std), None
        post = self.post_net(torch.cat([deter, obs_feat], dim=-1))  # use observation
        stoch, _, _ = self._sample(post)           # TRAINING: posterior sample
        return stoch, deter, prior, post           # KL(post || prior) trains the dream

cell = RSSMCell()
z, h = torch.zeros(4, 32), torch.zeros(4, 256)
a = torch.zeros(4, 6)
z, h, prior, post = cell(z, a, h, obs_feat=torch.randn(4, 1024))
print(z.shape, h.shape)   # torch.Size([4, 32]) torch.Size([4, 256])

3. Learning in Imagination: The Dreamer Loop Advanced

With a trained RSSM, the Dreamer recipe (Hafner et al., 2020) learns behavior without touching the environment. Starting from latent states encoded from real replayed experience, it rolls the RSSM forward using only the prior and the current policy's actions, generating an imagined trajectory of latent states, predicted rewards, and continuation flags. An actor-critic is then trained on this imagined trajectory: the critic learns to predict the long-horizon value of latent states, and the actor (policy) is updated to maximize that value. Because the entire rollout is differentiable, gradients of value can flow back through the imagined dynamics into the policy, an analytic policy gradient that model-free methods cannot get.

# Train DreamerV3 on a Gym environment with the reference implementation
# (github.com/danijar/dreamerv3). The repo is launched from its main entry
# point with a task and a config name; the schematic call below mirrors that:
#   python dreamerv3/main.py --configs defaults --task gym_CartPole-v1 \
#                            --run.steps 5e4
# Internally this builds the RSSM, the actor-critic, the replay buffer, and the
# imagination loop, all wired together with DreamerV3's stabilization tricks.

4. The Limits and the Road Ahead Intermediate

The RSSM-Dreamer recipe is elegant and sample-efficient but has clear boundaries. Its latent is small and its decoder modest, so it excels on games and control tasks with structured dynamics but does not, by itself, generate the rich, photorealistic, long-horizon video that Sections 36.1 and 36.2 produce. The next two sections take the two roads out of this limitation. Section 36.6 scales the world model up to pixel-space generative simulators (GAIA-1, playable game engines) that fuse the video-generation machinery with action conditioning, trading the RSSM's compact latent for the raw generative power of large video models. Section 36.7 goes the opposite way: it argues that the pixel decoder is wasteful and that prediction should happen purely in representation space, the JEPA philosophy. Both descend from the same insight you just built: a latent state plus a learned transition is a model of the world.

5. The Original World Model: V, M, C Advanced

Before the RSSM unified everything into one trained objective, Ha and Schmidhuber (2018) proposed a deliberately modular agent with three pieces, each trained almost independently, that already contained every idea this section builds on. The decomposition is worth studying precisely because the parts are visible. A Vision module (V) is a variational autoencoder that compresses each frame $o_t$ into a low-dimensional latent code $z_t$ (the same VAE you built in Chapter 31). A Memory module (M) is a recurrent network that predicts the next latent given the current latent and action. A Controller (C) is a tiny linear policy. Compression, prediction, and control are cleanly separated, which makes the agent easy to reason about and, as we will see, easy to train the controller inside the dream.

The Memory module is the conceptual ancestor of the RSSM transition. The next-frame latent is genuinely uncertain (the same action in the same place can lead to different futures), so M does not predict a single $z_{t+1}$; it predicts a full probability distribution. Ha and Schmidhuber use a mixture density network on top of an LSTM, so the next latent is modeled as a mixture of $K$ Gaussians whose mixing weights, means, and variances are all functions of the recurrent hidden state $h_t$:

Each term is one mode of the predicted future. The mixing weights $\pi_k(h_t)$ (a softmax over the $K$ components, so $\sum_k \pi_k = 1$) say how probable each mode is; $\mu_k(h_t)$ is where that mode sits in latent space and $\sigma_k^2(h_t)$ how spread out it is. A single Gaussian would be forced to average over forking futures and would predict the blurry midpoint of two outcomes that never actually occurs; the mixture instead keeps the modes apart. This is the mixture-density-network idea (Bishop, 1994) applied to dynamics, and it is exactly the role the stochastic latent $s_t$ later plays inside the RSSM: a learned representation of irreducible branching in the world.

The Controller is deliberately minuscule. It is a single linear map from the concatenated latent and memory state to an action,

where $[z_t, h_t]$ stacks the VAE latent and the LSTM hidden state into one vector and $(W_c, b_c)$ are the only parameters the controller has. Because V and M already did the hard representational work, C has so few parameters that it can be optimized by a black-box evolution strategy rather than by backpropagation. Ha and Schmidhuber use CMA-ES (covariance matrix adaptation evolution strategy), which maintains a Gaussian distribution over controller parameter vectors, samples a population, scores each by the return of a rollout, and shifts the distribution toward the high-scoring samples. The decisive move, the one that names this whole section, is that those scoring rollouts are run inside M: the controller never touches the real environment during its own optimization. It learns to drive by driving in the dream that M generates, which is why Ha and Schmidhuber titled the relevant experiment "training inside the dream." If the dream is faithful, a controller good in M is good in reality; if M has exploitable flaws, the controller learns to exploit them, the model-exploitation failure flagged in the Fun Note above.

6. PlaNet and the RSSM Objective Advanced

Ha and Schmidhuber's three modules are trained separately, which means M never gets a gradient signal about what V should encode, and the controller cannot send any pressure back through the model. PlaNet (Hafner et al., 2019) fuses the pieces into a single latent-dynamics model trained end-to-end by one variational objective, and in doing so introduces the precise split that defines the RSSM. The state at time $t$ is a pair: a deterministic recurrent state $h_t$ that is a deterministic function of the past, and a stochastic state $s_t$ sampled from a learned distribution. The deterministic part is a GRU update,

which folds the previous stochastic state and action into memory. Conditioned on $h_t$, the model defines four learned distributions: the prior (transition) $p(s_t \mid h_t)$ that predicts the stochastic state from memory alone, the posterior (representation) $q(s_t \mid h_t, o_t)$ that corrects that prediction using the observation, the decoder $p(o_t \mid h_t, s_t)$ that reconstructs pixels, and the reward head $p(r_t \mid h_t, s_t)$ that predicts the scalar reward. Only the prior is needed to imagine; the posterior exists only during training, to supply a target the prior must learn to match.

Training maximizes a variational lower bound on the log-likelihood of the observed sequence, the same ELBO logic as the VAE in Chapter 31 but now summed over time and carrying a reward term:

Read term by term, the objective is intuitive. The first term, $\log p(o_t \mid h_t, s_t)$, is reconstruction: the state must contain enough information to redraw the frame. The second, $\log p(r_t \mid h_t, s_t)$, forces the state to be reward-relevant, so the model does not waste capacity on visually salient but task-irrelevant detail. The third, the KL divergence, is the term that makes imagination possible: it pulls the observation-informed posterior $q$ toward the observation-free prior $p$, so that at imagination time the prior alone produces stochastic states statistically indistinguishable from the ones the posterior would have produced had it seen the frame. Without that KL, the prior would never learn to dream; with it, the model can roll forward with no input.

A single-step ELBO trains the prior only to predict one step ahead, but planning needs accurate predictions many steps out, where one-step errors compound. PlaNet adds latent overshooting: it also requires the prior, run $d$ steps without observations, to match the posterior at the destination, averaged over multiple distances $d$:

Here $p(s_t \mid s_{t-d})$ is the prior advanced $d$ steps forward from the state at $t-d$ using no observations, and the KL drives that multi-step prediction toward the true posterior at $t$. Summing over $d$ from $1$ to $D$ trains the model to be self-consistent over many horizons at once, which is exactly the regime a planner exercises when it imagines long rollouts. The intuition is direct: if you intend to plan twenty steps ahead, train the model to predict twenty steps ahead, not just one.

PlaNet has no learned policy at all. It plans online with CEM (the cross-entropy method) as model-predictive control. At each real step it samples many candidate action sequences from a diagonal Gaussian, imagines each one through the latent dynamics, scores each by the sum of predicted rewards, keeps the top-K highest-scoring sequences (the elites), refits the Gaussian to those elites, and repeats for a few iterations. It then executes only the first action of the best sequence and replans at the next step (receding horizon). Planning replaces the policy entirely; the model is the plan.

7. Dreamer: Actor-Critic in Imagination Advanced

Replanning from scratch at every step, as PlaNet does, is expensive, and it throws away everything the agent learned on the previous step. Dreamer (Hafner et al., 2020) keeps PlaNet's RSSM but replaces online planning with a learned actor-critic trained entirely on imagined trajectories. The critic $v_\psi(s_t)$ estimates the expected long-horizon return from a latent state; the actor $\pi_\theta(a_t \mid s_t)$ is the policy. Both are trained on rollouts the RSSM dreams from real replayed start states, so no environment interaction is spent on behavior learning.

The training target is the $\lambda$-return, a geometric average of $n$-step returns that trades off the bias of the critic's bootstrap against the variance of long imagined rollouts. It is defined recursively over the imagination horizon $H$:

The recursion is the heart of the method, so unpack it. At horizon $H$ the return is just the critic's estimate $v_\psi(s_H)$, the boundary condition. Working backward, each step earns the immediate reward $r_t$ plus the discounted future, and that future is a blend: with weight $(1-\lambda)$ it trusts the critic's one-step bootstrap $v_\psi(s_{t+1})$, and with weight $\lambda$ it trusts the deeper return $V^\lambda_{t+1}$ assembled from the rest of the imagined rollout. Setting $\lambda = 0$ recovers the pure one-step TD target (low variance, high bias if the critic is wrong); $\lambda = 1$ recovers the full Monte-Carlo imagined return (unbiased given the model, high variance). Intermediate $\lambda$ interpolates. The discount $\gamma$ down-weights distant rewards as usual.

The critic regresses onto this target with a squared loss, treating the $\lambda$-return as a fixed value via the stop-gradient $\text{sg}(\cdot)$ so the regression target does not chase the parameters being optimized:

The actor maximizes the same $\lambda$-return, plus an entropy bonus $\mathcal{H}(\pi_\theta(a_t \mid s_t))$ scaled by $\eta$ that keeps the policy exploratory and prevents premature collapse onto a single action:

The decisive advantage over model-free actor-critics lives in how this gradient is computed. Because the entire imagined rollout, the RSSM transitions, the reward head, and the reparameterized stochastic states, is differentiable, the gradient of $V^\lambda_t$ with respect to the actor parameters $\theta$ can flow backward through the imagined dynamics analytically. The actor receives a direct, low-variance signal of how nudging each action changes the predicted future return, rather than the high-variance score-function (REINFORCE) estimate a model-free agent is stuck with. This analytic policy gradient through a learned model is the single most important reason Dreamer is sample-efficient.

8. DreamerV2 and V3: Stabilizing the Recipe Advanced

The original Dreamer worked on continuous control but was fragile on the discrete, visually diverse Atari suite, and it needed per-domain tuning. The V2 and V3 revisions are a sequence of stabilization changes that culminate in one fixed hyperparameter set working across more than 150 tasks. Each change is small and well-motivated, and together they are why DreamerV3 (Hafner et al., 2023; Nature, 2025) is the default modern recipe.

DreamerV2 (Hafner et al., 2021) replaced the Gaussian stochastic latent with a vector of categorical variables. Discrete latents are a better fit for the sharp, multi-modal transitions of games (a screen either changes or it does not), but sampling from a categorical is not differentiable. DreamerV2 uses the straight-through estimator: the forward pass samples a hard one-hot category, while the backward pass pretends the sample was the continuous softmax probability, so gradients still flow into the prior and posterior. It also introduced KL balancing, which recognizes that the prior-posterior KL is doing two jobs at once and should not weight them equally. Training the prior toward the posterior (the dream learning to match reality) matters more than regularizing the posterior toward the prior, so the two directions get different weights:

The first term, with the posterior held fixed by $\text{sg}$, trains the prior to predict; the second, with the prior held fixed, gently regularizes the posterior. The asymmetry $\alpha = 0.8$ puts most of the pressure on improving the prior (the predictor that drives imagination), which is the quantity that determines dream quality.

DreamerV3 (Hafner et al., 2023) adds the changes that finally removed per-task tuning. Rewards and returns in different environments span wildly different magnitudes, from fractions to thousands, which destabilizes regression losses; V3 squashes them with the symlog transform, a signed logarithm that compresses large magnitudes while staying linear near zero and handling negatives:

The reward and value heads no longer regress a scalar at all. They predict a distribution over a fixed grid of symlog-spaced bins using a two-hot encoding (a target value falls between two grid points and is represented as the two-point interpolation that recovers it in expectation), trained with cross-entropy. This turns an unbounded regression into a stable classification problem that copes with heavy-tailed and sparse rewards without tuning.

V3 also reshapes the world-model loss into three explicitly weighted terms and clips each KL with free bits so the model never wastes effort driving an already-small KL toward zero (which would otherwise collapse the latent to the prior and starve reconstruction):

with $\beta_{pred} = 1$, $\beta_{dyn} = 1$, and $\beta_{rep} = 0.1$. The prediction loss $\mathcal{L}_{pred}$ collects the decoder, reward, and continuation log-likelihoods; the dynamics loss $\mathcal{L}_{dyn}$ (the KL-balanced prior-toward-posterior term) and the representation loss $\mathcal{L}_{rep}$ (posterior-toward-prior) carry over V2's KL balancing, now each floored at $1$ nat by the $\max(1, \cdot)$ free-bits clip. The down-weighting $\beta_{rep} = 0.1$ keeps the representation free to encode useful detail. Two further details complete the recipe: a 1% uniform mixture ("unimix") is blended into every categorical distribution so no class probability is ever exactly zero (which would give infinite KL and dead gradients), and returns are normalized by a running estimate of their percentile range rather than their variance, which is robust to the reward outliers that variance normalization mishandles. The payoff of this stack of small fixes is the headline result: one configuration, no per-task tuning, strong performance from Atari to continuous control to collecting diamonds in Minecraft from scratch.

9. TD-MPC2: Decoder-Free Latent Planning Advanced

Every model so far reconstructs pixels, and that decoder is expensive and arguably wasteful: to choose good actions you do not need to redraw the scene, only to predict reward and value. TD-MPC2 (Hansen et al., 2024) takes this to its conclusion with a decoder-free latent world model that never reconstructs an observation. An encoder maps the observation to a latent, $z = h_\theta(o)$; a latent dynamics model predicts the next latent, $z_{t+1} = d_\theta(z_t, a_t)$; a reward head $R_\theta$, a value head $Q_\theta$, and a policy prior $\pi_\theta$ read out of the latent. The crucial difference is what holds the latent together. With no reconstruction term, the latent is shaped purely by a self-consistency loss that forces the predicted next latent to match the encoder's actual encoding of the next observation:

The stop-gradient on the target $h_\theta(o_{t+1})$ prevents the trivial collapse where the encoder maps everything to a constant (which would make any prediction perfect). The full objective adds the reward loss and a temporal-difference value loss, both as two-hot cross-entropy as in DreamerV3, and crucially contains no reconstruction term. The latent is therefore free to discard everything about the observation that does not help predict reward and value, which is precisely the task-relevant compression a planner wants.

At decision time TD-MPC2 plans with MPPI (model-predictive path integral control), a sampling-based planner closely related to PlaNet's CEM but with a softer, reward-weighted update. It samples action trajectories (seeded partly from the policy prior $\pi_\theta$ to start from sensible behavior), rolls each through the latent dynamics, and scores each by the discounted sum of predicted rewards plus a terminal value bootstrap, $\sum_t \gamma^t \hat{r}_t + \gamma^H \hat{Q}$. Rather than keeping a hard top-K set of elites as CEM does, MPPI updates the sampling distribution by a reward-weighted softmax over all trajectories, so better trajectories pull the distribution proportionally to how good they are. The terminal $\gamma^H \hat{Q}$ is what lets a short planning horizon see past its own edge: the learned value summarizes everything beyond step $H$, so the planner need not imagine to the end of the episode.

10. Comparing the World-Model Family Intermediate

The six systems above are variations on one template, latent state plus learned transition, that differ along a few decisive axes: whether the stochastic latent exists and what form it takes, whether the policy gradient flows analytically through the model, how the prior-posterior KL is handled, whether the model reconstructs observations, and how an action is finally chosen. Reading the table row by row traces the field's trajectory from modular and hand-tuned toward unified, stable, and decoder-free.

Two trends jump out. The KL column moves from "none" through ever more carefully balanced and floored divergences and then back to "none" once the decoder is dropped, because the prior-posterior KL exists to keep a generative latent honest and a consistency-shaped latent does not need it. The reconstruction column stays "yes" for five systems and then turns "no," marking the decoder-free turn that Section 36.7 pushes further with the JEPA philosophy. The action-selection column tells the other story: planning (CEM, MPPI) and learned policies (Dreamer's actor-critic) are two answers to the same question, and TD-MPC2's seeding of MPPI from a learned policy prior shows the two answers converging.

11. Foundational Papers Intermediate

The six systems traced above are documented in a tight sequence of papers; the cards below give the primary sources, each with the single idea it contributed to the lineage.