Dreamer and the Latent Space

Dreamer learns to act by rehearsing inside its own predictions: the world model is trained on real experience, and the policy is trained almost entirely on imagined experience in the model's latent space.

A world model on a 5x5 grid can be a table. On pixels it cannot, so the model has to compress first and predict second. This post follows that idea through two designs. The 2018 "World Models" paper splits the agent into a vision model, a memory model and a controller, each trained separately. Dreamer fuses representation and dynamics into one recurrent state-space model (RSSM) and trains an actor-critic on imagined latent trajectories using $\lambda$-returns. The code is PyTorch-style and minimal; it shows the shapes and the gradient paths, not a tuned implementation. The overview post motivates why this is worth doing; this one is about how.

V, M, C: the 2018 decomposition#

Ha and Schmidhuber's "World Models" (2018, arXiv:1803.10122) splits the agent into three parts, trained in sequence:

M predicts a mixture of Gaussians rather than a point, because the future is multimodal: a car might turn left or right, and the average of the two is a crash.

import math
import torch, torch.nn as nn, torch.nn.functional as F
 
class MDNRNN(nn.Module):
    def __init__(self, z_dim=32, a_dim=3, hidden=256, k=5):
        super().__init__()
        self.k, self.z_dim = k, z_dim
        self.lstm = nn.LSTMCell(z_dim + a_dim, hidden)
        self.head = nn.Linear(hidden, k * (1 + 2 * z_dim))   # logits, means, log-stds
 
    def forward(self, z, a, state):
        h, c = self.lstm(torch.cat([z, a], -1), state)
        out = self.head(h)
        logit, mu, log_std = out.split([self.k, self.k * self.z_dim, self.k * self.z_dim], -1)
        mu = mu.view(-1, self.k, self.z_dim); log_std = log_std.view(-1, self.k, self.z_dim)
        return logit, mu, log_std, (h, c)
 
def mdn_nll(logit, mu, log_std, z_next):
    """Negative log-likelihood of z_next under the mixture (diagonal Gaussians)."""
    z = z_next.unsqueeze(1)                                   # (B, 1, z_dim)
    log_n = -0.5 * (((z - mu) / log_std.exp()) ** 2 + 2 * log_std + math.log(2 * math.pi))
    return -torch.logsumexp(F.log_softmax(logit, -1) + log_n.sum(-1), -1).mean()
 
class Controller(nn.Module):
    def __init__(self, z_dim=32, hidden=256, a_dim=3):
        super().__init__()
        self.fc = nn.Linear(z_dim + hidden, a_dim)          # (32 + 256) * 3 + 3 = 867 params
    def forward(self, z, h):
        return torch.tanh(self.fc(torch.cat([z, h], -1)))

The controller is deliberately tiny, 867 parameters for the car-racing task, so that a gradient-free optimiser can search it directly. The paper also trains a controller entirely inside M's samples and transfers it to the real game, raising M's sampling temperature so the dream is harder than reality and the controller cannot exploit the model's quirks.

The design has two bottlenecks. First, V is trained only to reconstruct pixels, so it spends capacity on whatever is visually large rather than what matters for reward. Second, C is linear because evolution strategies do not scale to millions of parameters. Dreamer removes both.

The recurrent state-space model#

Dreamer (Hafner et al., 2019, "Dream to Control", arXiv:1912.01603) inherits its world model from PlaNet: a latent state with two parts.

• A deterministic state $h_t$, updated by a GRU. It carries memory across many steps without sampling noise washing it out.
• A stochastic state $z_t$, sampled at each step. It represents what cannot be predicted: which way the car turned, where the ball bounced.

A purely deterministic model cannot represent multiple futures; a purely stochastic one struggles to remember anything for long. The RSSM uses both:

class RSSM(nn.Module):
    def __init__(self, a_dim, embed=1024, deter=200, stoch=30, hidden=200):
        super().__init__()
        self.inp = nn.Linear(stoch + a_dim, hidden)
        self.gru = nn.GRUCell(hidden, deter)
        self.prior_net = nn.Sequential(nn.Linear(deter, hidden), nn.ELU(), nn.Linear(hidden, 2 * stoch))
        self.post_net = nn.Sequential(nn.Linear(deter + embed, hidden), nn.ELU(), nn.Linear(hidden, 2 * stoch))
 
    @staticmethod
    def dist(params):
        mean, std = params.chunk(2, -1)
        return torch.distributions.Normal(mean, F.softplus(std) + 0.1)
 
    def img_step(self, h, z, a):
        """Advance one step with no observation: the prior."""
        h = self.gru(F.elu(self.inp(torch.cat([z, a], -1))), h)
        prior = self.dist(self.prior_net(h))
        return h, prior.rsample(), prior
 
    def obs_step(self, h, z, a, e):
        """Advance one step and correct with the frame embedding e: the posterior."""
        h, _, prior = self.img_step(h, z, a)
        post = self.dist(self.post_net(torch.cat([h, e], -1)))
        return h, post.rsample(), post, prior

The prior and posterior share the same $h_t$. The only difference is whether the current frame is allowed in. That shared structure is what makes imagination possible: at training time the posterior sees frames and the prior learns to predict what the posterior will conclude; at imagination time the prior runs alone.

Training the world model#

The model is trained on sequences sampled from a replay buffer of real experience. The loss is a sequence ELBO: reconstruct the frame and the reward from $(h_t, z_t)$, and keep the posterior close to the prior.

def world_model_loss(rssm, encoder, decoder, reward_head, obs, actions, rewards, free_nats=3.0):
    B, T = actions.shape[:2]
    h = torch.zeros(B, rssm.gru.hidden_size); z = torch.zeros(B, rssm.prior_net[-1].out_features // 2)
    e = encoder(obs.flatten(0, 1)).view(B, T, -1)
    recon, rew, kl, states = 0.0, 0.0, 0.0, []
    for t in range(T):
        h, z, post, prior = rssm.obs_step(h, z, actions[:, t - 1] if t else torch.zeros_like(actions[:, 0]), e[:, t])
        feat = torch.cat([h, z], -1); states.append(feat)
        recon = recon - decoder(feat).log_prob(obs[:, t]).mean()
        rew = rew - reward_head(feat).log_prob(rewards[:, t]).mean()
        kl_t = torch.distributions.kl_divergence(post, prior).sum(-1).mean()
        kl = kl + torch.clamp(kl_t, min=free_nats)            # no gradient below the floor
    return (recon + rew + kl) / T, torch.stack(states, 1).detach()

The KL term does two jobs. It regularises the posterior, and it trains the prior to predict the posterior, which is precisely the skill imagination needs. The free-nats floor stops the KL from collapsing the latent to nothing early in training, when matching a useless prior is the easiest way to lower the loss.

Learning behaviour in imagination#

Once the model is trained, behaviour learning never touches the environment. Start from posterior states computed on real sequences, then roll the prior forward with the actor choosing actions:

def imagine(rssm, actor, start_feats, horizon=15):
    """start_feats: (N, deter + stoch) posterior states from real data, detached."""
    deter = rssm.gru.hidden_size
    h, z = start_feats[:, :deter], start_feats[:, deter:]
    feats = []
    for _ in range(horizon):
        a = actor(torch.cat([h, z], -1)).rsample()            # reparameterised: gradients flow
        h, z, _ = rssm.img_step(h, z, a)
        feats.append(torch.cat([h, z], -1))
    return torch.stack(feats, 1)                               # (N, H, feat)

Every state in every training batch becomes a starting point, so one batch of 50 sequences of length 50 seeds 2500 imagined trajectories. The horizon is short, 15 steps in the paper, for the reason measured in the first post of this series: imagined trajectories drift away from reality, and drift compounds.

The $\lambda$-return#

A 15-step imagined trajectory sees only 15 rewards. To look further, bootstrap from a critic $v(s)$. The $\lambda$-return blends every $n$-step bootstrap into one target, computed backward:

With $\lambda = 0$ this is one-step TD; with $\lambda = 1$ it is the Monte Carlo return bootstrapped only at the horizon.

def lambda_return(reward, value, discount, lam=0.95):
    """reward, discount: (N, H); value: (N, H + 1). Returns (N, H) targets."""
    out, nxt = [], value[:, -1]
    for t in reversed(range(reward.shape[1])):
        nxt = reward[:, t] + discount[:, t] * ((1 - lam) * value[:, t + 1] + lam * nxt)
        out.append(nxt)
    return torch.stack(out[::-1], 1)

A numeric check with rewards [0, 0, 0, 1], a critic that says 0.5 everywhere (0 after the end), and $\gamma = 0.99$:

At $\lambda = 0$ the target trusts the critic's wrong 0.5 and ignores the reward three steps ahead. At $\lambda = 1$ it trusts only the imagined rewards. Dreamer uses $\lambda = 0.95$: mostly the rollout, with a little critic to reduce variance and to see past the horizon.

Actor and critic#

The critic regresses onto the $\lambda$-returns with the targets held fixed. The actor, in the original Dreamer, maximises them directly by backpropagating through the imagined dynamics: the reparameterised action, the GRU, the sampled latent, and the reward head are all differentiable.

def behaviour_step(rssm, actor, critic, reward_head, start_feats, opt_actor, opt_critic, gamma=0.99):
    frozen = [*rssm.parameters(), *reward_head.parameters()]
    for p in frozen: p.requires_grad_(False)                   # the model is fixed here
    feats = imagine(rssm, actor, start_feats)
    reward = reward_head(feats).mean
    value = critic(torch.cat([start_feats[:, None], feats], 1)).mean
    disc = torch.full_like(reward, gamma)
    targets = lambda_return(reward, value, disc)
    actor_loss = -targets.mean()                               # analytic value gradients
    opt_actor.zero_grad(); actor_loss.backward(); opt_actor.step()
    v = critic(torch.cat([start_feats[:, None], feats], 1).detach()[:, :-1]).mean
    critic_loss = 0.5 * ((v - targets.detach()) ** 2).mean()
    opt_critic.zero_grad(); critic_loss.backward(); opt_critic.step()
    for p in frozen: p.requires_grad_(True)

That gradient path is the conceptual leap over V-M-C. The controller is no longer searched blindly; it is trained by the model's own derivatives, so it can be a full neural network. The cost is that the actor now optimises the model's gradients, and wherever those are wrong it will push hard.

What DreamerV3 changes#

"Mastering Diverse Domains through World Models" (Hafner et al., 2023, arXiv:2301.04104) keeps the architecture and changes what makes it fragile across domains:

• Discrete latents. The stochastic state becomes a set of categorical variables (carried over from DreamerV2), trained with straight-through gradients.
• Symlog predictions. Rewards and values are predicted in $\operatorname{symlog}(x) = \operatorname{sign}(x)\ln(1 + |x|)$ space, so one set of hyperparameters handles rewards of very different scales.
• Balanced KL with free bits. The KL is split into a term that trains the prior and a smaller one that regularises the posterior, each clipped below one nat.
• Return normalisation. The actor's returns are scaled by the range between their 5th and 95th percentiles (never by less than 1) rather than by their standard deviation, so near-zero returns from sparse rewards are not amplified into noise.

The headline result is a single configuration across many benchmarks, including collecting diamonds in Minecraft from scratch without human data.

What to measure#

1. Open-loop prediction. Condition on 5 real frames, imagine 45, decode, and compare to the real frames. This is the most direct view of drift; the next post in the series, on evaluating a world model, turns it into numbers.
2. The prior-posterior KL over time. If it collapses to the free-nats floor, the latent carries nothing. If it never falls, the prior cannot predict.
3. Imagined versus real return for the same policy. A large positive gap means the actor is exploiting the model.
4. Sensitivity to horizon and $\lambda$. Sweep $H \in \{5, 15, 30\}$ and $\lambda \in \{0.8, 0.95, 1\}$. The best values tell you how far the model can be trusted.
5. Reward-head error on rare events. The decoder can be excellent while the reward head misses the one frame that matters.

For a contrast in how the model is used, the planning post replaces the learned actor with online search, and the foundation world models post scales the same predict-in-latent-space idea to interactive video.