DreamerV2 — Visual Experiments¶

Goal: See why discrete (categorical) latents and KL balancing make Dreamer's world model stronger and more stable.

What we'll explore¶

#	Experiment	Question answered
1	Multimodal latent	Why does a categorical latent beat a Gaussian on a branching observation?
2	Straight-through categorical	Can gradients train a discrete one-hot latent?
3	KL balancing	How does asymmetric KL keep the posterior informative?
4	Discrete state capacity	How much can a 32×32 categorical state represent?

---

Linked theory: research/02-dreamerv2.md

In [1]:

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
import warnings
warnings.filterwarnings('ignore')

torch.manual_seed(0)
np.random.seed(42)
plt.rcParams.update({
    'figure.dpi': 110,
    'font.size': 11,
    'axes.titlesize': 12,
    'axes.labelsize': 11,
    'figure.facecolor': 'white',
    'axes.facecolor': '#f8f9fa',
    'axes.grid': True,
    'grid.alpha': 0.35,
})
print('Setup complete.')

import numpy as np import matplotlib.pyplot as plt import torch import torch.nn.functional as F import warnings warnings.filterwarnings('ignore') torch.manual_seed(0) np.random.seed(42) plt.rcParams.update({ 'figure.dpi': 110, 'font.size': 11, 'axes.titlesize': 12, 'axes.labelsize': 11, 'figure.facecolor': 'white', 'axes.facecolor': '#f8f9fa', 'axes.grid': True, 'grid.alpha': 0.35, }) print('Setup complete.')

Setup complete.

---

Experiment 1: Multimodal latent¶

Question: From one state, the world may branch into distinct futures. A unimodal Gaussian latent must average them; a categorical latent can keep them separate. What does the difference look like?

The observation is one of two sprites (a left bar or a right bar) with equal probability. We compare the Gaussian posterior mean reconstruction against samples from a categorical latent.

In [2]:

np.random.seed(1)

S = 8
A = np.zeros((S, S)); A[:, 1] = 1.0      # bar on the left
B = np.zeros((S, S)); B[:, 6] = 1.0      # bar on the right
sprites = [A, B]

# Gaussian latent's best single prediction of a 50/50 branch is the pixel-wise mean
gaussian_mean = 0.5 * A + 0.5 * B
# Categorical latent samples one mode at a time
cat_samples = [sprites[np.random.randint(2)] for _ in range(3)]
print('A 50/50 branch -> Gaussian must average; categorical samples a crisp mode')

np.random.seed(1) S = 8 A = np.zeros((S, S)); A[:, 1] = 1.0 # bar on the left B = np.zeros((S, S)); B[:, 6] = 1.0 # bar on the right sprites = [A, B] # Gaussian latent's best single prediction of a 50/50 branch is the pixel-wise mean gaussian_mean = 0.5 * A + 0.5 * B # Categorical latent samples one mode at a time cat_samples = [sprites[np.random.randint(2)] for _ in range(3)] print('A 50/50 branch -> Gaussian must average; categorical samples a crisp mode')

A 50/50 branch -> Gaussian must average; categorical samples a crisp mode

In [3]:

fig, axes = plt.subplots(1, 5, figsize=(14, 3.2))
fig.suptitle('Categorical latent keeps modes crisp; Gaussian mean ghosts them', fontsize=13, fontweight='bold')

axes[0].imshow(A, cmap='magma', vmin=0, vmax=1); axes[0].set_title('true mode A')
axes[1].imshow(B, cmap='magma', vmin=0, vmax=1); axes[1].set_title('true mode B')
axes[2].imshow(gaussian_mean, cmap='magma', vmin=0, vmax=1)
axes[2].set_title('Gaussian mean\n(ghosted average)')
for ax, s in zip(axes[3:], cat_samples[:2]):
    ax.imshow(s, cmap='magma', vmin=0, vmax=1)
axes[3].set_title('categorical\nsample 1'); axes[4].set_title('categorical\nsample 2')
for ax in axes:
    ax.set_xticks([]); ax.set_yticks([])

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 5, figsize=(14, 3.2)) fig.suptitle('Categorical latent keeps modes crisp; Gaussian mean ghosts them', fontsize=13, fontweight='bold') axes[0].imshow(A, cmap='magma', vmin=0, vmax=1); axes[0].set_title('true mode A') axes[1].imshow(B, cmap='magma', vmin=0, vmax=1); axes[1].set_title('true mode B') axes[2].imshow(gaussian_mean, cmap='magma', vmin=0, vmax=1) axes[2].set_title('Gaussian mean\n(ghosted average)') for ax, s in zip(axes[3:], cat_samples[:2]): ax.imshow(s, cmap='magma', vmin=0, vmax=1) axes[3].set_title('categorical\nsample 1'); axes[4].set_title('categorical\nsample 2') for ax in axes: ax.set_xticks([]); ax.set_yticks([]) plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The Gaussian-mean panel shows two faint bars at once — a blurry average that matches neither real outcome. The categorical samples each show one crisp bar, a genuine possible future.
2. Why it looks this way: A unimodal Gaussian centred between two modes places probability mass on the (empty) middle, so its best single reconstruction is the average. A categorical distribution puts mass on each discrete mode and samples one at a time.
3. Key takeaway: Discrete latents represent branching, multimodal dynamics faithfully — the main reason DreamerV2 switched from Gaussian to categorical stochastic states.

---

Experiment 2: Straight-through categorical¶

Question: A categorical sample is a non-differentiable one-hot. How does DreamerV2 still train it with gradients?

We learn the logits of a categorical latent so that its straight-through one-hot, decoded by a fixed linear map, reconstructs a target — gradients flow through the softmax probabilities.

In [4]:

torch.manual_seed(2)

K = 6                                    # number of classes
decoder = torch.randn(K, 5)              # fixed linear decoder: class -> 5-d output
target_class = 3
target = decoder[target_class].clone()   # we want the latent to pick class 3

logits = torch.zeros(K, requires_grad=True)
opt = torch.optim.Adam([logits], lr=0.2)
loss_hist, prob_hist = [], []
for it in range(150):
    opt.zero_grad()
    probs = F.softmax(logits, dim=0)
    hard = F.one_hot(torch.argmax(probs), K).float()
    onehot_st = hard + (probs - probs.detach())     # straight-through one-hot
    out = onehot_st @ decoder
    loss = ((out - target) ** 2).sum()
    loss.backward(); opt.step()
    loss_hist.append(loss.item()); prob_hist.append(probs.detach().numpy().copy())
prob_hist = np.array(prob_hist)
print(f'target class = {target_class}; final argmax = {int(np.argmax(prob_hist[-1]))}')

torch.manual_seed(2) K = 6 # number of classes decoder = torch.randn(K, 5) # fixed linear decoder: class -> 5-d output target_class = 3 target = decoder[target_class].clone() # we want the latent to pick class 3 logits = torch.zeros(K, requires_grad=True) opt = torch.optim.Adam([logits], lr=0.2) loss_hist, prob_hist = [], [] for it in range(150): opt.zero_grad() probs = F.softmax(logits, dim=0) hard = F.one_hot(torch.argmax(probs), K).float() onehot_st = hard + (probs - probs.detach()) # straight-through one-hot out = onehot_st @ decoder loss = ((out - target) ** 2).sum() loss.backward(); opt.step() loss_hist.append(loss.item()); prob_hist.append(probs.detach().numpy().copy()) prob_hist = np.array(prob_hist) print(f'target class = {target_class}; final argmax = {int(np.argmax(prob_hist[-1]))}')

target class = 3; final argmax = 3

In [5]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('Straight-through gradients train a discrete categorical latent', fontsize=13, fontweight='bold')

axes[0].plot(loss_hist, color='seagreen', lw=2)
axes[0].set_xlabel('iteration'); axes[0].set_ylabel('reconstruction loss')
axes[0].set_title('Loss falls as the latent snaps to the right class')
axes[0].set_yscale('log')

for k in range(K):
    axes[1].plot(prob_hist[:, k], lw=2, label=f'class {k}' + (' (target)' if k == target_class else ''))
axes[1].set_xlabel('iteration'); axes[1].set_ylabel('class probability')
axes[1].set_title('Categorical probabilities concentrate on the target')
axes[1].legend(fontsize=8, ncol=2)

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13, 5)) fig.suptitle('Straight-through gradients train a discrete categorical latent', fontsize=13, fontweight='bold') axes[0].plot(loss_hist, color='seagreen', lw=2) axes[0].set_xlabel('iteration'); axes[0].set_ylabel('reconstruction loss') axes[0].set_title('Loss falls as the latent snaps to the right class') axes[0].set_yscale('log') for k in range(K): axes[1].plot(prob_hist[:, k], lw=2, label=f'class {k}' + (' (target)' if k == target_class else '')) axes[1].set_xlabel('iteration'); axes[1].set_ylabel('class probability') axes[1].set_title('Categorical probabilities concentrate on the target') axes[1].legend(fontsize=8, ncol=2) plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The loss drops to near zero while the probability of the target class (class 3) rises to one and the others fall away — the discrete latent learned to select the correct class.
2. Why it looks this way: The forward pass uses a hard one-hot, but the straight-through trick routes the gradient through the continuous softmax probabilities, so the logits receive a usable signal despite the non-differentiable argmax.
3. Key takeaway: Straight-through estimation — the same tool used for VQ tokens in Tier 9 — is what lets DreamerV2 backprop through its categorical world-model state.

---

Experiment 3: KL balancing¶

Question: Symmetric KL lets the model cheat by weakening the posterior (throwing away observation information) to reduce the KL term. How does KL balancing prevent that?

The posterior $q$ should match the data while a prior $p$ tracks it. We minimise data-fit + $\beta[\alpha\,\mathrm{KL}(\mathrm{sg}[q]\Vert p)+(1-\alpha)\mathrm{KL}(q\Vert\mathrm{sg}[p])]$ for a balanced ($\alpha{=}0.8$) and a posterior-heavy ($\alpha{=}0.2$) setting, and watch how well $q$ keeps fitting the data.

In [6]:

torch.manual_seed(3)

Kc = 8
data = torch.zeros(Kc); data[2] = 0.7; data[5] = 0.3   # informative target distribution
beta = 3.0


def kl(a_logits, b_logits, stop_a=False, stop_b=False):
    pa = F.softmax(a_logits.detach() if stop_a else a_logits, dim=0)
    logpa = F.log_softmax(a_logits.detach() if stop_a else a_logits, dim=0)
    logpb = F.log_softmax(b_logits.detach() if stop_b else b_logits, dim=0)
    return (pa * (logpa - logpb)).sum()


def train(alpha, steps=200):
    q = torch.zeros(Kc, requires_grad=True)         # posterior logits
    p = torch.randn(Kc, requires_grad=True)         # prior logits
    opt = torch.optim.Adam([q, p], lr=0.1)
    fit_hist = []
    for _ in range(steps):
        opt.zero_grad()
        logq = F.log_softmax(q, dim=0)
        data_fit = -(data * logq).sum()             # cross-entropy: q should match data
        kl_bal = alpha * kl(q, p, stop_a=True) + (1 - alpha) * kl(q, p, stop_b=True)
        (data_fit + beta * kl_bal).backward()
        opt.step()
        fit_hist.append(-(data * F.log_softmax(q, dim=0)).sum().item())   # cross-entropy q->data
    return np.array(fit_hist), F.softmax(q, 0).detach().numpy()


fit_bal, q_bal = train(0.8)
fit_heavy, q_heavy = train(0.2)
print(f'final data cross-entropy  balanced(α=0.8): {fit_bal[-1]:.3f}   posterior-heavy(α=0.2): {fit_heavy[-1]:.3f}')

torch.manual_seed(3) Kc = 8 data = torch.zeros(Kc); data[2] = 0.7; data[5] = 0.3 # informative target distribution beta = 3.0 def kl(a_logits, b_logits, stop_a=False, stop_b=False): pa = F.softmax(a_logits.detach() if stop_a else a_logits, dim=0) logpa = F.log_softmax(a_logits.detach() if stop_a else a_logits, dim=0) logpb = F.log_softmax(b_logits.detach() if stop_b else b_logits, dim=0) return (pa * (logpa - logpb)).sum() def train(alpha, steps=200): q = torch.zeros(Kc, requires_grad=True) # posterior logits p = torch.randn(Kc, requires_grad=True) # prior logits opt = torch.optim.Adam([q, p], lr=0.1) fit_hist = [] for _ in range(steps): opt.zero_grad() logq = F.log_softmax(q, dim=0) data_fit = -(data * logq).sum() # cross-entropy: q should match data kl_bal = alpha * kl(q, p, stop_a=True) + (1 - alpha) * kl(q, p, stop_b=True) (data_fit + beta * kl_bal).backward() opt.step() fit_hist.append(-(data * F.log_softmax(q, dim=0)).sum().item()) # cross-entropy q->data return np.array(fit_hist), F.softmax(q, 0).detach().numpy() fit_bal, q_bal = train(0.8) fit_heavy, q_heavy = train(0.2) print(f'final data cross-entropy balanced(α=0.8): {fit_bal[-1]:.3f} posterior-heavy(α=0.2): {fit_heavy[-1]:.3f}')

final data cross-entropy  balanced(α=0.8): 0.615   posterior-heavy(α=0.2): 0.618

In [7]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('KL balancing keeps the posterior matched to the data', fontsize=13, fontweight='bold')

axes[0].plot(fit_bal, color='seagreen', lw=2, label='balanced α=0.8')
axes[0].plot(fit_heavy, color='crimson', lw=2, label='posterior-heavy α=0.2')
axes[0].set_xlabel('iteration'); axes[0].set_ylabel('posterior cross-entropy to data')
axes[0].set_title('Lower = posterior still fits the data')
axes[0].legend()

x = np.arange(Kc)
w = 0.35
axes[1].bar(x - w/2, data.numpy(), w, color='black', alpha=0.6, label='data target')
axes[1].bar(x + w/2, q_bal, w, color='seagreen', alpha=0.8, label='posterior (α=0.8)')
axes[1].bar(x + w/2, q_heavy, w, color='crimson', alpha=0.35, label='posterior (α=0.2)')
axes[1].set_xlabel('class'); axes[1].set_ylabel('probability')
axes[1].set_title('Balanced posterior keeps the data peaks')
axes[1].legend(fontsize=9)

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13, 5)) fig.suptitle('KL balancing keeps the posterior matched to the data', fontsize=13, fontweight='bold') axes[0].plot(fit_bal, color='seagreen', lw=2, label='balanced α=0.8') axes[0].plot(fit_heavy, color='crimson', lw=2, label='posterior-heavy α=0.2') axes[0].set_xlabel('iteration'); axes[0].set_ylabel('posterior cross-entropy to data') axes[0].set_title('Lower = posterior still fits the data') axes[0].legend() x = np.arange(Kc) w = 0.35 axes[1].bar(x - w/2, data.numpy(), w, color='black', alpha=0.6, label='data target') axes[1].bar(x + w/2, q_bal, w, color='seagreen', alpha=0.8, label='posterior (α=0.8)') axes[1].bar(x + w/2, q_heavy, w, color='crimson', alpha=0.35, label='posterior (α=0.2)') axes[1].set_xlabel('class'); axes[1].set_ylabel('probability') axes[1].set_title('Balanced posterior keeps the data peaks') axes[1].legend(fontsize=9) plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: With balanced KL (green) the posterior's cross-entropy to the data stays low and its bars match the data peaks at classes 2 and 5. With posterior-heavy KL (red) the posterior is dragged toward the vague prior, smearing its mass and fitting the data worse.
2. Why it looks this way: KL balancing weights the prior-update term more (α=0.8), so the prior chases the posterior while the posterior is barely pulled away from the data. A posterior-heavy weighting instead pulls the posterior toward the prior, destroying observation information.
3. Key takeaway: KL balancing lets the prior learn to predict without the posterior sacrificing the information it extracts from observations — a key stability fix in DreamerV2.

---

Experiment 4: Discrete state capacity¶

Question: DreamerV2's stochastic state is 32 categorical variables of 32 classes each. How much can such a state represent?

We compute the number of distinct states and the information capacity in bits for different configurations of (variables × classes).

In [8]:

configs = [(8, 8), (16, 16), (32, 32), (32, 64)]
bits = [n * np.log2(c) for n, c in configs]
labels = [f'{n}×{c}' for n, c in configs]
for (n, c), b in zip(configs, bits):
    print(f'{n} vars × {c} classes  ->  {c**n:.2e} states  =  {b:.0f} bits')

configs = [(8, 8), (16, 16), (32, 32), (32, 64)] bits = [n * np.log2(c) for n, c in configs] labels = [f'{n}×{c}' for n, c in configs] for (n, c), b in zip(configs, bits): print(f'{n} vars × {c} classes -> {c**n:.2e} states = {b:.0f} bits')

8 vars × 8 classes  ->  1.68e+07 states  =  24 bits
16 vars × 16 classes  ->  1.84e+19 states  =  64 bits
32 vars × 32 classes  ->  1.46e+48 states  =  160 bits
32 vars × 64 classes  ->  6.28e+57 states  =  192 bits

In [9]:

fig, ax = plt.subplots(figsize=(8, 5))
fig.suptitle('A multi-categorical state packs huge capacity into a small vector', fontsize=13, fontweight='bold')

bars = ax.bar(labels, bits, color=['steelblue', 'steelblue', 'crimson', 'steelblue'])
ax.set_xlabel('configuration (variables × classes)'); ax.set_ylabel('capacity (bits)')
ax.set_title('DreamerV2 uses 32×32 (highlighted)')
for b, v in zip(bars, bits):
    ax.text(b.get_x() + b.get_width() / 2, v, f'{v:.0f} bits', ha='center', va='bottom', fontsize=9)

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(8, 5)) fig.suptitle('A multi-categorical state packs huge capacity into a small vector', fontsize=13, fontweight='bold') bars = ax.bar(labels, bits, color=['steelblue', 'steelblue', 'crimson', 'steelblue']) ax.set_xlabel('configuration (variables × classes)'); ax.set_ylabel('capacity (bits)') ax.set_title('DreamerV2 uses 32×32 (highlighted)') for b, v in zip(bars, bits): ax.text(b.get_x() + b.get_width() / 2, v, f'{v:.0f} bits', ha='center', va='bottom', fontsize=9) plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: Capacity grows with both the number of categorical variables and the classes per variable; the 32×32 configuration (red) carries 160 bits — about $10^{48}$ distinct states — in a compact one-hot representation.
2. Why it looks this way: Independent categoricals multiply: $n$ variables of $c$ classes give $c^n$ joint states, i.e. $n\log_2 c$ bits. The state is large in capacity yet sparse and discrete, ideal for a sequence model to predict.
3. Key takeaway: A multi-categorical state is both expressive and discrete — enough capacity to model rich dynamics while staying in the token-friendly regime that pairs naturally with straight-through training.

---

Summary¶

Experiment	Key insight
1. Multimodal latent	A categorical latent samples crisp modes; a Gaussian averages a branch into a ghost.
2. Straight-through categorical	Straight-through gradients train a discrete one-hot latent despite the argmax.
3. KL balancing	Weighting the prior-update term keeps the posterior informative instead of collapsing to the prior.
4. Discrete state capacity	32×32 categoricals carry ~160 bits ($10^{48}$ states) in a compact sparse vector.

Connection to the broader theory¶

DreamerV2 brings the discrete-latent ideas of Tier 9 into the world-model state, fixing DreamerV1's unimodal Gaussian with categoricals and stabilising training with KL balancing. It still reconstructs pixels and tunes hyperparameters per domain — the gap DreamerV3 closes with symlog, free-bits KL, and robust normalisation to win across hundreds of domains with one configuration.