Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

The Dreamer Family › DreamerV2 and discrete latents

Michael BrenndoerferJuly 15, 20267 min read

World Models HandbookAll chapters

Part 8 of 12Decision-Centric Research Lineages

Chapter 45The Dreamer Family › DreamerV2 and discrete latents

Reading section 4 of 7 · DreamerV2 and discrete latents · 7 minSection 4 of 7 · 7 min

Section 4 of 7 · 7 min read

The Dreamer Family › DreamerV2 and discrete latentsOne section · Sections in this chapter4 / 7 ↓

Sections in this chapter7

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The original Dreamer established the imagined actor-critic recipe primarily on continuous-control tasks and also evaluated discrete actions and early termination on selected Atari and DMLab tasks. DreamerV2 scaled the approach to the full 55-game Atari benchmark and found that categorical stochastic states outperformed Gaussian states on 42 of 55 games in the reported comparison. The paper presents possible explanations as hypotheses, and the full Atari result also depends on KL balancing, actor-gradient choices, architecture, replay, and optimization. These results indicate the categorical change was central, but not sufficient on its own.

Why would discrete latents help so much? There are several converging reasons, and it is worth understanding each because the same reasoning recurs in later architectures.

The paper reports the categorical parameterization as a decisive empirical change. Its discussion offers hypotheses rather than a theorem: categorical variables may be easier to optimize and may impose a useful inductive bias. They do not guarantee sharper reconstructions or physically valid latent states.

Categorical and Gaussian KL terms have different geometry. For a $K$ -way posterior against a uniform prior, $\mathrm{KL}(q\|u)\leq\log K$ . This is not a general bound for $\mathrm{KL}(q\|p)$ : it can grow without bound as a learned prior assigns probability near zero where the posterior has mass. Gaussian KL is likewise unbounded.

Third, there is nominal information capacity. A $K$ -way categorical variable carries at most $\log_2 K$ bits. With 32 independent groups of 32 categories, DreamerV2's default stochastic state has maximum entropy of 160 bits, but its realized information content can be lower.

The categorical latent uses a straight-through estimator. It draws a one-hot sample, then adds the probabilities and subtracts detached probabilities, so the forward value stays discrete while the surrogate gradient follows the probabilities:

\text{soft} = \mathrm{softmax}(\text{logits}), \qquad \text{sample} \sim \mathrm{OneHotCategorical}(\text{soft}), \qquad z = \text{sample} + \text{soft} - \mathrm{sg}[\text{soft}],

where $\mathrm{sg}[\cdot]$ is stop-gradient. DreamerV2 draws a categorical one-hot sample, then uses sample + probs - stop(probs). This is a straight-through categorical estimator, not a deterministic argmax and not a Gumbel-softmax draw.

The categorical KL is computed for each group and summed. DreamerV2 also balances its gradients: one term fits the prior to a stopped posterior, while the other makes the posterior predictable by a stopped prior. With $\alpha=0.8$ , the reported objective gives greater weight to fitting the prior:

\mathcal{L}_\text{KL} = \alpha \, \mathrm{KL}\!\left( \mathrm{sg}[q] \,\|\, p \right) + (1 - \alpha) \, \mathrm{KL}\!\left( q \,\|\, \mathrm{sg}[p] \right),

with $\alpha$ typically around 0.8. The first term is the "dynamics" term: it trains the prior to match the posterior without letting gradients flow into the posterior. The second term is the "representation" term: it trains the posterior to be predictable by the prior, without letting gradients flow into the prior. Together, the weighted terms encourage prior-posterior agreement; they do not guarantee it or prevent either distribution from dominating during training. Without balancing, the posterior can drift away from the prior and imagination can lose grounding. The asymmetry in the weights reflects a deliberate priority. The dynamics term is weighted more heavily because a well-trained prior is what makes imagination work, while the representation term is kept lighter to reduce pressure to distort the posterior's ability to encode observations accurately.

The official DreamerV2 implementation initializes recurrent states to zero. Layer normalization appears in the implementation, but the paper's ablation describes its contribution as marginal. KL balancing and the categorical state are the documented central changes; the paper does not isolate mixed precision as an explanation for the headline result.

At 200 million Atari frames, the DreamerV2 paper reported a median human-normalized score of 2.15 across 55 games. Its Rainbow and IQN comparisons in that table use the same 200-million-frame interaction budget, so this result is not evidence of fewer interactions than those baselines. The notable scope claim is the single agent configuration across the suite.

Let's implement the categorical latent mechanism for our dot world. The conceptual pipeline stays the same, but categorical rollout, actor, and critic wiring would need compatible latent shapes and return signatures; the Gaussian helpers below are not drop-in compatible.

In[22]:

Code

class CategoricalRSSM(nn.Module):
    def __init__(
        self,
        action_dim=4,
        deter_dim=64,
        num_groups=8,
        num_classes=8,
        hidden_dim=128,
        obs_size=16,
    ):
        super().__init__()
        self.action_dim = action_dim
        self.deter_dim = deter_dim
        self.num_groups = num_groups
        self.num_classes = num_classes
        self.latent_dim = num_groups * num_classes
        self.obs_size = obs_size

        self.input_proj = nn.Linear(action_dim + self.latent_dim, hidden_dim)
        self.gru = nn.GRUCell(hidden_dim, deter_dim)
        self.prior_net = nn.Sequential(
            nn.Linear(deter_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, num_groups * num_classes),
        )
        self.post_net = nn.Sequential(
            nn.Linear(deter_dim + obs_size * obs_size, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, num_groups * num_classes),
        )
        self.decoder = nn.Sequential(
            nn.Linear(deter_dim + self.latent_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, obs_size * obs_size),
        )
        self.reward_head = nn.Sequential(
            nn.Linear(deter_dim + self.latent_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, 1),
        )
        self.continue_head = nn.Sequential(
            nn.Linear(deter_dim + self.latent_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, 1),
        )

    def initial_state(self, batch_size, device):
        return torch.zeros(batch_size, self.deter_dim, device=device)

    def _logits_to_latent(self, logits):
        # logits: (B, num_groups * num_classes)
        B = logits.shape[0]
        logits = logits.view(B, self.num_groups, self.num_classes)
        soft = F.softmax(logits, dim=-1)
        sample = torch.distributions.OneHotCategorical(logits=logits).sample()
        straight_through = sample + soft - soft.detach()
        return straight_through.view(B, self.latent_dim)

    def forward_step(self, prev_action, prev_latent, deter, obs=None):
        x = torch.cat([prev_action, prev_latent], dim=-1)
        x = torch.relu(self.input_proj(x))
        deter = self.gru(x, deter)
        prior_logits = self.prior_net(deter).view(
            -1, self.num_groups, self.num_classes
        )
        if obs is not None:
            post_input = torch.cat([deter, obs], dim=-1)
            post_logits = self.post_net(post_input).view(
                -1, self.num_groups, self.num_classes
            )
        else:
            post_logits = None
        return deter, prior_logits, post_logits

    def decode(self, deter, latent):
        x = torch.cat([deter, latent], dim=-1)
        img = torch.sigmoid(self.decoder(x))
        reward = self.reward_head(x).squeeze(-1)
        cont = torch.sigmoid(self.continue_head(x)).squeeze(-1)
        return img, reward, cont

The _logits_to_latent method implements the paper's stochastic straight-through categorical sample. The forward value is one-hot; the backward surrogate follows the probabilities.

Two new helper functions: a categorical KL over groups, and KL balancing. Then a training step for the categorical model.

In[23]:

Code

def categorical_kl(post_logits, prior_logits):
    # post_logits, prior_logits: (B, num_groups, num_classes)
    post = F.softmax(post_logits, dim=-1)
    prior = F.softmax(prior_logits, dim=-1)
    log_post = F.log_softmax(post_logits, dim=-1)
    log_prior = F.log_softmax(prior_logits, dim=-1)
    # Sum categories and independent groups; caller averages the batch.
    kl = (post * (log_post - log_prior)).sum(dim=-1).sum(dim=-1)
    return kl  # shape (B,)


def balanced_kl_loss(post_logits, prior_logits, alpha=0.8):
    # alpha weights the "dynamics" term KL(sg[post] || prior); (1-alpha) weights "representation" KL(post || sg[prior])
    post_detached = post_logits.detach()
    prior_detached = prior_logits.detach()
    dynamics_kl = categorical_kl(post_detached, prior_logits)
    representation_kl = categorical_kl(post_logits, prior_detached)
    return alpha * dynamics_kl.mean() + (1.0 - alpha) * representation_kl.mean()


def categorical_model_loss(model, batch, beta=1.0):
    obs_seq, act_seq, rew_seq, next_seq, done_seq = batch
    B, T, _ = obs_seq.shape
    device = obs_seq.device

    deter = model.initial_state(B, device)
    zero_action = torch.zeros(B, model.action_dim, device=device)
    zero_latent = torch.zeros(B, model.latent_dim, device=device)
    deter, _, post_logits = model.forward_step(
        zero_action, zero_latent, deter, obs_seq[:, 0]
    )
    prev_latent = model._logits_to_latent(post_logits.view(B, -1))

    rec_loss = rew_loss = cont_loss = kl_loss = 0.0
    for t in range(T):
        action = one_hot_action(act_seq[:, t], model.action_dim, device)
        deter, prior_logits, post_logits = model.forward_step(
            action, prev_latent, deter, next_seq[:, t]
        )
        latent = model._logits_to_latent(post_logits.view(B, -1))
        img_hat, rew_hat, cont_hat = model.decode(deter, latent)
        rec_loss = (
            rec_loss
            + F.binary_cross_entropy(img_hat, next_seq[:, t], reduction="none")
            .sum(dim=-1)
            .mean()
        )
        rew_loss = (
            rew_loss
            + F.mse_loss(rew_hat, rew_seq[:, t], reduction="none").mean()
        )
        cont_loss = (
            cont_loss
            + F.binary_cross_entropy(
                cont_hat, 1.0 - done_seq[:, t], reduction="none"
            ).mean()
        )
        kl_loss = kl_loss + balanced_kl_loss(
            post_logits, prior_logits, alpha=0.8
        )
        prev_latent = latent

    total = (rec_loss + rew_loss + cont_loss) / T + beta * kl_loss / T
    return total, dict(rec=(rec_loss / T).item(), kl=(kl_loss / T).item())

We do not train or benchmark the categorical model here. The code illustrates stochastic sampling, tensor shapes, the summed group KL, and intended gradient routing, but comparative stability requires a matched multi-seed experiment and is not established by this example.

Three takeaways from the DreamerV2 section:

KL to a uniform $K$ -way prior is bounded by $\log K$ , but KL to a learned categorical prior is not generally bounded.
Straight-through estimation lets you train through discrete sampling without giving up end-to-end gradient flow.
KL balancing weights prior fitting more heavily while retaining a smaller posterior-regularization gradient; it encourages agreement but does not guarantee it.

Part IX: Foundation and World-Action Models is planned to cover tokenized game models and video generators.