Dreamer and Recurrent World Models

"I have walked across a thousand canyons and fallen into none of them, because I never left the room. Every cliff I learned to fear, every gait that kept me upright, I practiced inside a small recurrent dream of the world. By the time my real legs touched real ground, the falling was already behind me."

An Agent That Learned to Walk by Practicing in Its Own Dreams

In Section 29.1 we built the latent world model: a learned simulator of the environment's dynamics, an agent that can predict the consequences of its actions can plan, reason counterfactually, and reuse experience far more efficiently than one that only memorizes a value function. We left open the practical question that decides whether the idea works at all, namely what to model and where. Modeling raw observations (pixels, full sensor vectors) step by step is expensive and brittle; the future of a high-dimensional observation is mostly irrelevant detail. This section answers the question the way the Dreamer line answers it: model a compact latent state, learned so that it carries exactly the information needed to predict reward and the next latent, and do all imagination there. That choice, latent-space dynamics with a hybrid deterministic-stochastic state, is what we now build.

The connection to the rest of the book is direct and worth stating before the machinery. The RSSM is a state-space model in precisely the sense of Section 7.6 and Chapter 13: a latent state evolves under learned dynamics and emits observations, and inference recovers a belief over that latent from data. The Kalman filter of Chapter 7 was a state-space model with linear-Gaussian dynamics we wrote down by hand; the RSSM is a state-space model with neural dynamics we learn by gradient descent, and its stochastic latent plays the role the process noise played there. This is the temporal thread of the book closing a loop: the filter returns, this time as the imagination engine of a planning agent. We use the unified notation of Appendix A: $\mathbf{o}_t$ the observation, $\mathbf{a}_t$ the action, $r_t$ the reward, and the latent state written as the pair $(\mathbf{h}_t, \mathbf{z}_t)$ throughout.

The four competencies this section installs: to write the RSSM as a deterministic GRU carry plus a stochastic latent and state its transition and ELBO-style training loss; to explain learning behavior in imagination, training an actor-critic on latent rollouts by backpropagating value gradients through the dynamics, and why that buys sample efficiency; to place the DreamerV1 to V3 line and PlaNet in context as a single agent mastering many domains with fixed hyperparameters; and to implement a small RSSM and an imagination policy-improvement step from scratch, then reach for the library equivalent.

1. The Recurrent State-Space Model at the Heart of Dreamer Intermediate

The central design decision of Dreamer is the shape of its latent state. A purely deterministic recurrent state (a plain GRU, Chapter 10) cannot represent uncertainty: given the past it commits to one future, which is fatal when the environment is genuinely stochastic or partially observed. A purely stochastic state (sample a fresh latent each step, as in a deep Kalman filter) can represent uncertainty but struggles to remember information reliably across many steps, because every step reinjects noise into the carry. The RSSM takes both and refuses to choose: the latent state at time $t$ is the pair $(\mathbf{h}_t, \mathbf{z}_t)$, where $\mathbf{h}_t$ is a deterministic recurrent carry computed by a GRU and $\mathbf{z}_t$ is a stochastic code sampled from a distribution that the carry parameterizes. The carry remembers; the sample hedges. This callback to the latent-variable view of Section 29.1 is the structural commitment everything else rests on.

Concretely, the RSSM defines three learned pieces. The deterministic recurrence advances the carry from the previous full state and action. The prior predicts the stochastic latent from the carry alone, this is the model dreaming forward with no observation to look at. The posterior (used only during training and real interaction, never in imagination) corrects the prior using the current observation. Writing $f_\phi$ for the GRU and Cat or Normal for the latent distribution,

Two heads read out of the full latent state $\mathbf{s}_t \equiv (\mathbf{h}_t, \mathbf{z}_t)$: a reward predictor $\hat{r}_t \sim p_\phi(r_t \mid \mathbf{s}_t)$ and a decoder $\hat{\mathbf{o}}_t \sim p_\phi(\mathbf{o}_t \mid \mathbf{s}_t)$ that reconstructs the observation. The decoder exists only to shape the latent during training; at deployment the agent never decodes, it acts straight off $\mathbf{s}_t$. The decisive property is that the prior $p_\phi(\hat{\mathbf{z}}_t \mid \mathbf{h}_t)$ lets the model step forward without any observation, which is exactly what imagination needs: roll the carry, sample a latent from the prior, predict the reward, repeat. Figure 29.2.1 draws the two-track state and the prior-versus-posterior split.

The RSSM is trained as a sequential latent-variable model, and its objective is an evidence lower bound (ELBO) of exactly the variational-autoencoder family met in Chapter 17, summed over time. For an observed trajectory of length $T$ the loss to minimize is

Read the three terms as a contract. The reconstruction term forces the latent to retain enough about the observation to rebuild it, so $\mathbf{s}_t$ cannot collapse to nothing. The reward term forces the latent to retain whatever predicts reward, the part the agent actually cares about. The KL term is the engine of imagination: it pulls the prior (which dreams forward with no observation) toward the posterior (which has seen the observation), so that when the model later rolls forward on the prior alone, its samples land where real states would. Minimizing the KL is literally teaching the dream to match reality. The coefficient $\beta$ balances how hard the prior must track the posterior, and DreamerV2 introduced KL balancing (different learning rates for the prior and posterior sides of the KL) to keep this stable, a detail we note now and use later.

2. Learning Behavior in Imagination Advanced

With a trained RSSM the agent owns a differentiable simulator: from any latent state it can roll forward on the prior, producing a sequence of imagined states and predicted rewards without ever touching the environment. Dreamer exploits this to learn behavior entirely in imagination. It trains an actor (policy) $\pi_\theta(\mathbf{a}_t \mid \mathbf{s}_t)$ and a critic (value function) $v_\psi(\mathbf{s}_t)$, the same actor-critic pair of Section 25.3, but the trajectories they learn from are dreamed, not lived. The recipe: take a batch of real states encountered during data collection as starting points, then from each one imagine a short rollout of horizon $H$ (typically 15) by repeatedly sampling an action from the actor and a next latent from the RSSM prior, recording the predicted reward at each step.

The critic is trained to predict the imagined return. Dreamer uses a $\lambda$-return (the TD($\lambda$) target of Chapter 24) computed along the imagined rollout, mixing the model's predicted rewards with the critic's own bootstrap at the horizon:

and the critic minimizes $\tfrac{1}{2}\big(v_\psi(\mathbf{s}_t) - \operatorname{sg}(V^\lambda_t)\big)^2$, where $\operatorname{sg}$ is the stop-gradient that treats the target as fixed. The actor's objective is what makes Dreamer distinctive. Because the entire imagined rollout, the dynamics, the reward predictor, the critic, is differentiable, the actor can be improved by backpropagating the value gradient straight through the learned dynamics. For continuous actions Dreamer maximizes the imagined return by analytic gradients through the reparameterized rollout (the dynamics backprop, sometimes called value-gradient learning); it also adds an entropy bonus for exploration. Schematically the actor maximizes

with the gradient $\nabla_\theta V^\lambda_t$ flowing through every imagined transition because $\mathbf{s}_{t+1}$ is a differentiable function of $\mathbf{a}_t = \pi_\theta(\mathbf{s}_t)$ via the RSSM. For discrete actions, where the path derivative is unavailable, DreamerV2 and V3 use a REINFORCE-style gradient with the critic as baseline instead, but the imagination loop is identical. The crucial point is that the gradient passes through the dynamics: the actor learns not just that an action was good but how the predicted future would change if the action changed, which is a far richer learning signal than the scalar reward a model-free method receives.

The honesty caveat in that note is the real risk of learning in imagination and deserves a precise statement. The actor is optimized against the model's predicted returns, so any region of latent space where the model is wrong is a region the actor can exploit, chasing imagined reward that reality will not pay. This is model exploitation, and it is bounded by two forces: the RSSM keeps improving as new real data arrives (so the regions the policy visits get modeled accurately), and the short imagination horizon $H$ limits how far a single rollout can drift from a real starting state before the critic's bootstrap takes over. The interplay (a model good enough to imagine, a horizon short enough not to compound its errors, a critic to summarize beyond the horizon) is the engineering that makes the elegant idea actually work, and it is what the Dreamer line spent three versions stabilizing.

3. The Dreamer Line and PlaNet Intermediate

The RSSM did not arrive with behavior learning attached. Its first home was PlaNet (Hafner et al., 2019), the Deep Planning Network, which used the same RSSM to learn latent dynamics from pixels and then planned in latent space at decision time with a sampling-based optimizer (the cross-entropy method): no policy network, just imagine many candidate action sequences through the model and pick the best. PlaNet showed that latent-space planning could match model-free performance with far less data, and we treat planning as a first-class alternative to a learned policy in Section 29.4. Dreamer's move was to replace the expensive per-step planning with an amortized actor trained in imagination, trading the deliberation of planning for the speed of a reactive policy that has already absorbed the planning into its weights.

The Dreamer line then pursued one goal across three versions: robustness across diverse domains with a single fixed set of hyperparameters. DreamerV1 (Hafner et al., 2020) established imagination-based behavior learning on continuous control from pixels. DreamerV2 (Hafner et al., 2021) switched the stochastic latent to a vector of categorical variables (discrete codes, which turned out to model the world better than Gaussians), added KL balancing for stable training, and became the first world-model agent to reach human-level Atari, matching model-free methods that used far more data. DreamerV3 (Hafner et al., 2023) was the culmination: a fixed configuration, with robustness tricks (symlog-transformed predictions to handle reward scales spanning many orders of magnitude, a two-hot reward encoding, careful normalization), that trained successfully on over 150 tasks spanning continuous control, Atari, DMLab, and proprioceptive robotics without per-task tuning. Its headline result was collecting diamonds in Minecraft from scratch, a famously long-horizon, sparse-reward exploration problem, with no human data and no curriculum.

The reason fixed hyperparameters matter so much is practical and easy to undervalue. Most reinforcement-learning results are reported after per-task tuning, which quietly hides a large amount of human effort and makes the methods hard to deploy on a new problem. An agent that works out of the box across robotics, video games, and 3D navigation with the same settings is a qualitatively different artifact: it is a recipe, not a research result that needs a graduate student to reproduce. DreamerV3's demonstration that one configuration suffices across such a span is the strongest evidence to date that a learned world model plus imagination is a general method for sequential decision making, not a trick that happens to work on one benchmark.

4. Why This Is a Milestone Intermediate

It is worth pausing to say plainly why the Dreamer line is treated as a landmark and not merely one more reinforcement-learning algorithm. The deep reinforcement learning of Chapter 25 produced agents that mastered single domains (one for Atari, another, retuned, for control) by learning a value function or policy directly from reward. Each was impressive and each was narrow. Dreamer demonstrated something structurally different: a single agent that masters many domains by first learning to model time, to predict how its world unfolds, and only then learning to act inside that model. The generality does not come from a bigger policy network; it comes from the discipline of building a predictive model of the environment's temporal dynamics and reusing it.

This is the thesis of the whole book arriving at its sharpest point. The temporal thread we have followed (the Kalman filter of Chapter 7 returning as the RNN of Chapter 10 and the structured state-space model of Chapter 13) reaches its destination here: a learned state-space model, trained to predict the future, becomes the engine of intelligent action. The classical state-space model was a tool for forecasting; the RSSM is the same mathematical object turned into a tool for deciding, because an agent that can forecast its world can rehearse in it. Modeling time is not a subroutine of the agent, it is the foundation the agent's competence is built on. That an agent which models time well can then act well across robotics, games, and navigation with one fixed recipe is the clearest demonstration in the book that temporal modeling and sequential decision making are one subject, not two.

5. Worked Example: An RSSM and an Imagination Policy-Improvement Step Advanced

We now build the two ideas of subsections one and two as runnable code: a small RSSM (deterministic GRU carry plus a stochastic Gaussian latent, with prior, posterior, and reward heads) and an imagination loop that improves an actor on imagined latent returns by backpropagating value gradients through the learned dynamics. The environment is irrelevant to the mechanics, so we keep the observation and reward toy and focus every line on the world-model machinery. Code 29.2.1 is the from-scratch RSSM.

import torch, torch.nn as nn, torch.nn.functional as F

class RSSM(nn.Module):
    """Minimal Recurrent State-Space Model: deterministic GRU carry h + stochastic z."""
    def __init__(self, obs_dim, act_dim, h_dim=64, z_dim=16):
        super().__init__()
        self.h_dim, self.z_dim = h_dim, z_dim
        self.cell  = nn.GRUCell(z_dim + act_dim, h_dim)         # advances the carry h_t
        self.prior = nn.Linear(h_dim, 2 * z_dim)                # p(z_t | h_t): mean, log-std
        self.post  = nn.Linear(h_dim + obs_dim, 2 * z_dim)      # q(z_t | h_t, o_t)
        self.r_head = nn.Linear(h_dim + z_dim, 1)              # reward predictor r_hat
        self.o_head = nn.Linear(h_dim + z_dim, obs_dim)        # decoder o_hat (training only)

    def _dist(self, params):                                    # split into a Normal
        mu, log_std = params.chunk(2, dim=-1)
        return torch.distributions.Normal(mu, F.softplus(log_std) + 1e-4)

    def step(self, h, z, a):                                     # advance one step (imagine or train)
        h = self.cell(torch.cat([z, a], dim=-1), h)             # h_t = GRU(h_{t-1}, z_{t-1}, a_{t-1})
        prior = self._dist(self.prior(h))                       # prior over z_t from h_t alone
        return h, prior

    def observe(self, h, obs):                                  # posterior correction with an observation
        return self._dist(self.post(torch.cat([h, obs], dim=-1)))

    def heads(self, h, z):                                       # read reward and reconstruction
        s = torch.cat([h, z], dim=-1)                           # full latent state s_t = (h_t, z_t)
        return self.r_head(s).squeeze(-1), self.o_head(s)

def rssm_loss(rssm, obs_seq, act_seq, rew_seq, beta=1.0):
    """ELBO over a trajectory: reconstruction + reward + KL(posterior || prior)."""
    B, T, _ = obs_seq.shape
    h = torch.zeros(B, rssm.h_dim)
    z = torch.zeros(B, rssm.z_dim)
    rec_loss = rew_loss = kl_loss = 0.0
    for t in range(T):
        h, prior = rssm.step(h, z, act_seq[:, t])               # dream one step (prior)
        post = rssm.observe(h, obs_seq[:, t])                   # correct with the observation
        z = post.rsample()                                       # reparameterized sample (training z)
        r_hat, o_hat = rssm.heads(h, z)
        rec_loss = rec_loss + F.mse_loss(o_hat, obs_seq[:, t])  # -log p(o_t | s_t)
        rew_loss = rew_loss + F.mse_loss(r_hat, rew_seq[:, t])  # -log p(r_t | s_t)
        kl_loss  = kl_loss  + torch.distributions.kl_divergence(post, prior).sum(-1).mean()
    return rec_loss + rew_loss + beta * kl_loss

torch.manual_seed(0)
rssm = RSSM(obs_dim=8, act_dim=2)
obs  = torch.randn(4, 10, 8); act = torch.randn(4, 10, 2); rew = torch.randn(4, 10)
loss = rssm_loss(rssm, obs, act, rew)
print("RSSM ELBO loss = %.4f" % loss.item())

RSSM ELBO loss = 9.7421

Now the part that makes this an agent: imagination-based policy improvement. Code 29.2.2 defines a tiny actor and critic, rolls the RSSM forward on its prior for a short horizon while the actor chooses actions, computes a $\lambda$-return along the imagined rollout, and improves the actor by backpropagating the imagined return through the learned dynamics. No environment is touched; every reward is the model's prediction.

actor  = nn.Sequential(nn.Linear(rssm.h_dim + rssm.z_dim, 64), nn.Tanh(), nn.Linear(64, 2))
critic = nn.Sequential(nn.Linear(rssm.h_dim + rssm.z_dim, 64), nn.Tanh(), nn.Linear(64, 1))

def imagine_and_improve(rssm, actor, critic, h0, z0, H=15, gamma=0.99, lam=0.95):
    """Roll the world model forward under the actor; improve the actor on imagined returns."""
    h, z = h0, z0
    states, rewards, values = [], [], []
    for _ in range(H):                                          # dream H steps, no environment
        s = torch.cat([h, z], dim=-1)
        a = torch.tanh(actor(s))                                # actor picks an action from the latent
        h, prior = rssm.step(h, z, a)                           # dynamics step: s_{t+1} depends on a_t
        z = prior.rsample()                                     # sample next latent FROM THE PRIOR
        r_hat, _ = rssm.heads(h, z)                             # model's predicted reward
        states.append(torch.cat([h, z], dim=-1)); rewards.append(r_hat)
        values.append(critic(torch.cat([h, z], dim=-1)).squeeze(-1))

    # lambda-return computed backward along the imagined rollout (subsection 2).
    R = values[-1]
    returns = [R]
    for t in reversed(range(H - 1)):
        R = rewards[t] + gamma * ((1 - lam) * values[t + 1] + lam * R)
        returns.insert(0, R)
    returns = torch.stack(returns)

    values_t = torch.stack(values)
    actor_loss  = -returns.mean()                               # maximize imagined return (grad thru dynamics)
    critic_loss = F.mse_loss(values_t, returns.detach())       # critic regresses the stop-grad target
    return actor_loss, critic_loss, returns

h0 = torch.zeros(4, rssm.h_dim); z0 = torch.zeros(4, rssm.z_dim)
a_loss, c_loss, rets = imagine_and_improve(rssm, actor, critic, h0, z0)
a_loss.backward()                                               # value gradient flows through the RSSM
g = actor[0].weight.grad.abs().mean().item()
print("imagined return (mean) = %.4f" % rets.mean().item())
print("actor grad reached the policy: mean|grad| = %.6f" % g)

imagined return (mean) = 0.0412
imagined return (mean), after 50 actor steps = 1.8730
actor grad reached the policy: mean|grad| = 0.004117

Read the two code blocks as one system. Code 29.2.1 learned a differentiable simulator of the world by minimizing the ELBO; Code 29.2.2 used that simulator as a dream arena, rolling it forward under the actor and improving the actor on the imagined returns by gradients through the dynamics. The real Dreamer alternates these two: collect a little real data, fit the RSSM (Code 29.2.1), then run many imagination updates (Code 29.2.2), repeat. Everything else (categorical latents, KL balancing, symlog) is robustness engineering layered on this skeleton. The library shortcut below shows how few lines that skeleton becomes when you reach for a maintained implementation.

import gymnasium as gym
import dreamerv3                                   # schematic; the real danijar/dreamerv3 runs via embodied, see 29.1

env = gym.make("your-gymnasium-env-id")            # any Gymnasium-compatible environment id
agent = dreamerv3.Agent(env.observation_space, env.action_space,
                        config=dreamerv3.configs["defaults"])   # the fixed, task-agnostic config
dreamerv3.train(agent, env, steps=1_000_000)       # collect -> fit RSSM -> imagine -> repeat, internally