Latent state-space and recurrent world models (DreamerV3)

"They asked me to imagine a hundred futures before moving a muscle. So I compressed the world into forty numbers and dreamed inside them instead."

A Rehearsing AI Agent

The whole chapter revolves around one question: can a machine hold a useful model of the physical world in its head? This section answers it with the most battle-tested recipe available, the recurrent latent state-space model, and dissects the three problems it solves: what the latent state looks like, how the model is trained from sensor sequences, and how imagination inside that state replaces expensive interaction with the real world.

The RSSM: a hybrid deterministic-stochastic state

What the model maintains is a latent state split into two parts: a deterministic recurrent vector \(h_t\) carried by a GRU, and a stochastic vector \(z_t\) drawn from a distribution that \(h_t\) parameterizes. Why the split matters is subtle and is the central design insight of the RSSM. A purely deterministic recurrence cannot represent uncertainty (a fork in the road, an occluded object, sensor noise), so its long rollouts collapse to a single blurry average future. A purely stochastic chain, by contrast, forgets: it has no reliable channel to carry information many steps forward. Combining them gives you memory that persists (through \(h_t\)) and genuine multi-modal uncertainty that can branch (through \(z_t\)). In DreamerV3 the stochastic \(z_t\) is a vector of categorical variables (for example 32 classes across 32 groups) rather than a Gaussian, which sidesteps the posterior-collapse and gradient pathologies that plague Gaussian latents on high-dimensional observations.

How the pieces connect defines the model. At each step the recurrent core advances the deterministic state, and two distributions over \(z_t\) are produced: a prior that predicts the next stochastic state from memory and action alone, and a posterior that additionally sees the current observation \(x_t\):

The prior \(p_\phi\) is the dynamics predictor, the part that lets the model dream forward with no sensor input. The posterior \(q_\phi\) is the encoder that corrects the belief when a real observation arrives, exactly the predict-then-update rhythm of the Kalman family in Chapter 9, but with both the transition and the observation model learned by neural networks. Attached to the full model state \(s_t = (h_t, z_t)\) are lightweight prediction heads: a decoder that reconstructs \(x_t\), a reward head, and a continue head that predicts episode termination.

Training: reconstruction, KL balancing, and free bits

The world model is trained by maximizing an evidence lower bound (ELBO) on observed sensor sequences, the same variational objective you met in Chapter 4. The loss has a reconstruction term (decode \(x_t\), reward, and continuation from \(s_t\)) and a dynamics term that pulls prior and posterior together:

How DreamerV3 makes this stable across wildly different sensor scales is a set of small but decisive tricks. KL balancing uses different learning rates for the two directions of the KL, so the prior is pulled hard toward the (slowly moving) posterior while the posterior is only lightly regularized, which stops the representation from collapsing. Free bits clip the KL term below a floor (around one nat) so the model does not waste capacity over-compressing already-simple states. Crucially, DreamerV3 wraps every regression target (rewards, reconstructions of unbounded signals) in a symlog transform, \(\operatorname{symlog}(x) = \operatorname{sign}(x)\ln(1+|x|)\), and predicts scalars with a two-hot discretized distribution. Why this earns its place: symlog squashes signals that range from millivolts to kilonewtons into a common scale, and it is exactly this normalization, plus percentile-based return scaling in the actor-critic, that lets one hyperparameter set survive across domains. That is the property that makes the model practical for a new sensor without re-tuning.

Imagination: learning behavior without touching the world

What the trained world model unlocks is a second learner that never sees a raw sensor. Dreamer trains an actor and a critic purely on trajectories that the world model dreams: starting from real encoded states drawn from the replay buffer, it rolls the prior forward under the current policy, producing imagined states, rewards, and values for, say, 15 steps. Why this is transformative for sensing systems is data efficiency and safety. Real interaction with a physical plant, a patient, or a vehicle is slow, expensive, and sometimes hazardous; imagined interaction is a few matrix multiplies. The actor is optimized to maximize an imagined \(\lambda\)-return, and because the entire rollout lives in a differentiable latent space, gradients can flow straight through the learned dynamics. The world model turns a scarce resource (real experience) into an abundant one (dreamed experience), which is why Dreamer reaches strong performance from far fewer environment steps than model-free baselines.

The relationship to other temporal models in this book is worth naming. The RSSM is a learned, nonlinear, stochastic state-space model; the deterministic linear-time state-space models of Chapter 16 (S4, Mamba) share the "carry a compact state, update it in linear time" philosophy but drop the explicit stochastic branch and the generative decoder that make imagination possible. Choose an RSSM when you need to generate and plan; choose a Mamba-style backbone when you only need to discriminate or forecast a fixed target.

import torch, torch.nn as nn

class RSSMStep(nn.Module):
    """One recurrent state-space update: predict prior, then correct with obs."""
    def __init__(self, deter=512, stoch=32, classes=32, act_dim=6, embed=1024):
        super().__init__()
        self.stoch, self.classes = stoch, classes
        self.gru   = nn.GRUCell(stoch * classes + act_dim, deter)
        self.prior = nn.Linear(deter, stoch * classes)            # dynamics predictor
        self.post  = nn.Linear(deter + embed, stoch * classes)    # encoder posterior

    def _dist(self, logits):                                      # categorical latent
        logits = logits.view(-1, self.stoch, self.classes)
        return torch.distributions.OneHotCategoricalStraightThrough(logits=logits)

    def forward(self, h, z, action, embed):
        h = self.gru(torch.cat([z.flatten(1), action], -1), h)   # advance memory
        prior = self._dist(self.prior(h))                        # imagine next z
        post  = self._dist(self.post(torch.cat([h, embed], -1))) # correct with sensor
        z = post.rsample().flatten(1)                            # reparameterized sample
        kl = torch.distributions.kl_divergence(post, prior).sum(-1)
        return h, z, kl                                          # kl drives prior->post

When to reach for a recurrent latent world model

When this design is the right tool: your system is partially observable (any single sensor frame underdetermines the true state, so you need a belief that integrates history), you want to plan or rehearse rather than merely classify, and real-world interaction is costly enough that learning from imagined rollouts pays off. Wearables predicting the next few seconds of gait, industrial controllers rehearsing a valve sequence, and robots anticipating occluded obstacles all fit. When it is overkill: if you only need a fixed forecast horizon on a clean signal with abundant labels, a temporal convolutional or state-space forecaster from Chapter 14 is simpler, cheaper, and easier to certify. And as always, evaluate on leakage-safe splits: because the model trains on logged episodes, an episode that straddles the train/test boundary silently inflates every reported number.

What's Next

In Section 53.3, we question the one assumption DreamerV3 never drops: that the world model must reconstruct its observations. Joint-embedding predictive architectures (I-JEPA and V-JEPA 2) predict in an abstract representation space instead of pixel space, throwing away the decoder to spend capacity on what matters for prediction rather than on rendering every noisy detail of a sensor frame.