Part of World Models Handbook
Explains how Dreamer learns behaviors by latent imagination: RSSM dynamics, actor-critic training, and the V1, V2, and V3 design changes.
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 Dreamer Family
The PlaNet agent from the preceding chapter learned a compact latent dynamics model from pixels, then planned inside it with online trajectory optimization. At every environment step, PlaNet's cross-entropy method evaluates candidate action sequences, refits its sampling distribution to elite sequences, and executes the current-step mean action from the optimized distribution. More candidates and a longer horizon increase this decision-time cost. The search itself is transient, although the world model continues to learn from experience.
The Dreamer family, developed by Danijar Hafner and colleagues across a sequence of papers from 2019 to 2023, makes a decisive move that separates it from the PlaNet-style planners. Instead of using the learned model only as a forward simulator for planning, Dreamer uses it as a training ground for policies and value functions. The actor and critic are the learned parameters that carry the reusable skill. The model is still a latent dynamics model learned from sensory inputs, including pixels and, in DreamerV3, vector observations. But the agent does not search over action sequences at every step. It learns an actor (a policy) and a critic (a value function) by imagining trajectories inside the learned model, then acts by simply querying the actor. This is the essence of "learning behaviors by latent imagination". It is what makes Dreamer both fast at action time and capable of improving the behaviors it uses.
To see the computational shift, compare the controllers. PlaNet searches, evaluates, and compares action sequences before each action; Dreamer spends that control-learning compute during training and retains the result in actor parameters. PlaNet still trains its world model, so not all of its expense occurs at decision time. The distinction is between repeating a transient controller search and amortizing a rule for action selection.
This reframing matters more than it might first appear. PlaNet improves its learned model, but its controller remains an online optimizer. Dreamer also trains an actor and critic, so behavioral knowledge persists in parameters instead of being recomputed by a fresh action-sequence search. Both approaches can exploit errors in the model; an amortized policy does not automatically become conservative around uncertainty.
The family spans three major versions, and each one solves a problem the previous version exposed. DreamerV1 established the core recipe, DreamerV2 scaled it with categorical states and KL balancing, and DreamerV3 added cross-domain robustness techniques.
- DreamerV1 (Dream to Control; 2019 preprint, 2020 publication) established the core recipe: a recurrent state-space model (RSSM) for latent dynamics, imagined rollouts, and actor-critic learning through imagined dynamics. It reparameterized continuous actions and used straight-through gradients for discrete actions.
- DreamerV2 (Mastering Atari with Discrete World Models; 2021 publication) replaced Gaussian stochastic states with 32 groups of 32 categorical variables, introduced KL balancing, and used REINFORCE actor gradients on Atari while retaining dynamics gradients for continuous control.
- DreamerV3 appeared as a 2023 preprint and was published in Nature in 2025. It uses categorical states, split KL losses with free bits, one-percent uniform mixing, symlog vector inputs, symexp-spaced two-hot reward and value predictions, and percentile-normalized actor advantages. Image observations retain their convolutional path. Its fixed configuration was evaluated across a broad benchmark suite with benchmark-specific interfaces and budgets.
By the end of this chapter you should be able to explain why learning a policy in imagination differs from planning, how RSSM variants represent latent state, how their actor gradients differ, and what supports DreamerV3's cross-domain configuration robustness. We also build a compact Dreamer-style agent and compare its imagined training diagnostic with real-environment performance.
A note on how to read this chapter: the theory and the code are meant to be read together. We run a compact Dreamer-inspired agent end to end on a small "dot world," then define selected V2 and V3 components without training full variants. Try to keep the mapping between the math and the code in mind as you go, because the machinery of Dreamer is not exotic; it is a handful of well-understood components (a recurrent model, a variational latent, an actor, a critic) assembled in a careful way.
Recurrent state-space representations
Before we can imagine anything, we need a state representation that supports both perception (mapping observations to a belief) and prediction (rolling that belief forward under actions). Dreamer's representation is the recurrent state-space model, usually abbreviated RSSM. You may recognize the architecture if you worked through Part V, Chapter 1 on recurrent state-space models; here we focus on how Dreamer instantiates it and why each component is there.
The core problem is partial observability. From a single image, you cannot tell how fast an object is moving, what is behind it, or what happened a moment ago. A sufficient Markov state for control would include the physical configuration, velocities, and other variables that affect the future. In a POMDP, an exact belief is a distribution over environment states. An RSSM instead learns a compact history-dependent representation; its stochastic variables can encode ambiguity, but they are not automatically a calibrated Bayesian belief.
The RSSM splits the latent state into two pieces. The split is deliberate:
- A deterministic recurrent state , produced by a GRU (or similar recurrent cell) that carries long-range memory across time. This is the "hidden" part of the RSSM: it is a function of all past latents and actions, so it accumulates history without stochasticity.
- A stochastic latent state , sampled from a learned distribution. This gives the representation stochastic capacity and can encode multiple possibilities, but stochasticity alone does not establish calibrated uncertainty.
The deterministic state is the recurrent memory. The stochastic code is inferred by the posterior when an observation is available and predicted by the prior during imagination. Both variables participate in imagined rollouts. The stochastic code can represent ambiguity, but it need not correspond to calibrated uncertainty.
Writing this out, the RSSM has several learned components:
- Recurrent transition: , typically a GRU. It advances the deterministic memory by one step given the current latent and action.
- Prior (transition) distribution: , which predicts the successor latent from memory alone, without looking at the next observation. This is what you sample during imagination, when no observation is available.
- Posterior (representation) distribution: , which infers the successor latent from memory and the next observation. This is what you sample during training, when observations are available.
- Decoder (observation model): , which reconstructs the next observation from the successor latent. Reconstruction is the main signal that forces the latent to retain information about the world.
- Reward model: , which predicts the reward received after action from the successor latent. This is essential because the agent needs to imagine rewards to train the critic.
- Continue model: , which predicts whether the episode continues after action . It prevents rewards and value bootstraps from propagating across terminal boundaries.
The separation between prior and posterior is central to the design. During training, you have observations, so you can use the posterior to get an accurate latent. During imagination, we do not have future observations, so we must use the prior. If prior and posterior disagree badly, imagination drifts away from reality and the learned policy is trained on garbage. A central part of the RSSM loss is therefore the KL divergence between the posterior and the prior, which pulls the prior toward the posterior so that the two stay consistent. This is the same consistency pressure that the VAE world model in Part III, Chapter 3 relied on, but here it operates per timestep with a recurrent memory in the loop, so the prior must learn to predict a latent that is consistent with a whole history of observations, not only the current one.
Why have both parts? The recurrent state provides a persistent path for history, while the stochastic state gives the model additional representational capacity and a variational training objective. That design can represent ambiguity, but its samples should not be read as calibrated physical hypotheses without a separate calibration test.
A VAE whose per-step encoder sees only the current observation, with no recurrent state, temporal context, or transition prior, cannot represent long-range memory. An RSSM adds a recurrent path and a per-step KL that trains the prior to predict the posterior, allowing rollout without future observations.
The training objective for the representation combines reconstruction, reward prediction, continue prediction, and the KL term. Each term serves a distinct role: reconstruction forces the latent to retain information about the observation, reward and continue prediction give you the signals you need for imagination, and the KL keeps the prior and posterior aligned. A useful way to read the objective is as a tug-of-war between two pressures: the likelihood terms push the latent to be informative about the present, while the KL term pushes the prior to be a good predictor of the latent so that the model remains usable when observations are absent. The balance between these pressures is controlled by the weight , and choosing that weight is one of the more delicate parts of training a Dreamer-style agent.
where the expectation is over replay sequences and posterior latent samples, and scales the KL penalty. With Gaussian latents, reparameterization carries gradients through sampling. Reconstruction, reward, and continuation contribute negative log-likelihoods; the KL aligns the prior and posterior. In the toy code, binary-frame reconstruction and continuation BCE are Bernoulli negative log-likelihoods. Reward MSE equals the negative log-likelihood of a fixed-variance Gaussian with variance , up to an additive constant. DreamerV2 and V3 instead use categorical samples with straight-through surrogate gradients.
The index names a replay transition here: takes the model from to , while and are its reward and continuation outcome. The toy replay stores that tuple together, and model_loss predicts the next frame, reward, and continuation from the same successor latent. The later lambda-return uses the same transition indexing.
Two subtleties matter for closed-loop performance. First, the KL treatment is generation-specific: V1 uses a scaled KL term, V2 balances gradient flow between representation and dynamics, and V3 uses separately scaled dynamics and representation losses with free bits. Second, reconstruction can emphasize perceptually salient information that is not useful for control. That tension motivates the value-equivalent models discussed next.
Let's build the RSSM on a small synthetic task whose observation is a binary raster of a moving dot. We keep the example CPU-sized, seed every random generator, and make trajectory shapes explicit. Computation cells prepare the data; figure cells only render values computed earlier.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
torch.manual_seed(0)
np.random.seed(0)The world is a dot in a 2D box that we can nudge up, down, left, or right with an action. The observation is a directly rendered 16x16 binary image, which forces the model to compress pixels into state. We store transitions as tuples of (observation, action, reward, next observation, done).
class DotWorld:
"""A dot moves in a 2D box; action nudges it and we reward staying near center."""
def __init__(self, size=16.0):
self.size = size
# All randomness flows through the rng argument, so behavior is reproducible.
def reset(self, rng):
pos = rng.uniform(4.0, self.size - 4.0, size=2)
vel = np.zeros(2)
return pos, vel
def step(self, pos, vel, action, rng):
# action: 0 = push up, 1 = push down, 2 = push left, 3 = push right
thrust = {
0: np.array([0.0, -1.0]),
1: np.array([0.0, 1.0]),
2: np.array([-1.0, 0.0]),
3: np.array([1.0, 0.0]),
}[int(action)]
vel = 0.9 * vel + 0.5 * thrust
pos = pos + vel
# Wall reflection keeps the dot in the box
for d in range(2):
if pos[d] < 1.0:
pos[d] = 1.0
vel[d] = -0.5 * vel[d]
elif pos[d] > self.size - 1.0:
pos[d] = self.size - 1.0
vel[d] = -0.5 * vel[d]
center = np.array([self.size / 2, self.size / 2])
reward = -0.03 * np.linalg.norm(pos - center)
done = False
return pos, vel, reward, done
def render(self, pos, image_size=16):
img = np.zeros((image_size, image_size), dtype=np.float32)
x = int(
np.clip(pos[0] / self.size * (image_size - 1), 0, image_size - 1)
)
y = int(
np.clip(pos[1] / self.size * (image_size - 1), 0, image_size - 1)
)
img[y, x] = 1.0
return imgThis tiny environment is a stand-in for the real thing. Its simulator has a compact Markov state, with the dot position and velocity determining the next state, but the agent sees only a quantized position image, with velocity hidden. It must therefore infer motion from history. That is precisely the partial-observability structure the RSSM is designed to handle, making the dot world a fair testbed for the mechanism.
Now we collect a small dataset by running a random policy. In a full Dreamer setup, data collection is online. Here the fixed dataset is enough to exercise the RSSM and compare diagnostics; the experiment does not establish that its imagined images are plausible.
def collect_random_data(world, episodes=60, steps=50, image_size=16):
rng = np.random.default_rng(123)
obs, acts, rews, next_obs, dones = [], [], [], [], []
for _ in range(episodes):
pos, vel = world.reset(rng)
for _ in range(steps):
action = rng.integers(0, 4)
img = world.render(pos, image_size)
pos, vel, reward, done = world.step(pos, vel, action, rng)
next_img = world.render(pos, image_size)
obs.append(img)
acts.append(action)
rews.append(reward)
next_obs.append(next_img)
dones.append(done)
return (
np.stack(obs),
np.array(acts, dtype=np.int64),
np.array(rews, dtype=np.float32),
np.stack(next_obs),
np.array(dones, dtype=np.float32),
)
world = DotWorld()
obs, acts, rews, next_obs, dones = collect_random_data(world)Reward statistics: mean=-0.171, std=0.063 Action counts (0..3): [760 755 763 722]
The dot rarely sits near the center under random actions, so rewards are mostly negative. Actions are sampled uniformly, but momentum, walls, and resets make the resulting spatial occupancy nonuniform. The data provide broad exploratory coverage, not an even sample of the box.
Now we define the RSSM. We keep it small so training on CPU is quick. The GRU carries the deterministic state; a small MLP produces the prior and posterior parameters; another MLP decodes observations and predicts rewards.
class RSSM(nn.Module):
def __init__(
self,
action_dim=4,
deter_dim=64,
latent_dim=16,
hidden_dim=128,
obs_size=16,
):
super().__init__()
self.action_dim = action_dim
self.deter_dim = deter_dim
self.latent_dim = latent_dim
self.obs_size = obs_size
self.input_proj = nn.Linear(action_dim + 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, 2 * latent_dim),
)
self.post_net = nn.Sequential(
nn.Linear(deter_dim + obs_size * obs_size, hidden_dim),
nn.ELU(),
nn.Linear(hidden_dim, 2 * latent_dim),
)
self.decoder = nn.Sequential(
nn.Linear(deter_dim + latent_dim, hidden_dim),
nn.ELU(),
nn.Linear(hidden_dim, obs_size * obs_size),
)
self.reward_head = nn.Sequential(
nn.Linear(deter_dim + latent_dim, hidden_dim),
nn.ELU(),
nn.Linear(hidden_dim, 1),
)
self.continue_head = nn.Sequential(
nn.Linear(deter_dim + 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 forward_step(self, prev_action, prev_latent, deter, obs=None):
# prev_action: (B, action_dim) one-hot; obs: (B, obs_size*obs_size) flattened
x = torch.cat([prev_action, prev_latent], dim=-1)
x = torch.relu(self.input_proj(x))
deter = self.gru(x, deter)
prior_params = self.prior_net(deter)
prior_mean, prior_logvar = prior_params.chunk(2, dim=-1)
prior_logvar = torch.clamp(prior_logvar, -6.0, 2.0)
if obs is not None:
post_input = torch.cat([deter, obs], dim=-1)
post_params = self.post_net(post_input)
post_mean, post_logvar = post_params.chunk(2, dim=-1)
post_logvar = torch.clamp(post_logvar, -6.0, 2.0)
else:
post_mean, post_logvar = None, None
return deter, (prior_mean, prior_logvar), (post_mean, post_logvar)
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, contThe forward_step method is where the prior and posterior branches live side by side. When an observation is supplied (training time), the posterior branch reads it; when it is absent (imagination time), only the prior is available. Notice that the deterministic state update happens before the branch, so both distributions are conditioned on the same memory. This ordering matters: the recurrent state summarizes the past. Both the prior and posterior then interpret that summary given different amounts of current information.
We need a helper that samples a latent using the reparameterization trick and computes the KL divergence between two diagonal Gaussians in closed form. This is standard VAE machinery, but here the KL is between the posterior and the prior at each step rather than against a fixed standard normal.
def reparameterize(mean, logvar):
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
return mean + eps * std
def gaussian_kl(post_mean, post_logvar, prior_mean, prior_logvar):
# KL( N(post) || N(prior) ), diagonal covariance, summed over latent dims
var_ratio = torch.exp(post_logvar - prior_logvar)
diff = (post_mean - prior_mean) ** 2 / torch.exp(prior_logvar)
return 0.5 * torch.sum(
var_ratio + diff - 1.0 - (post_logvar - prior_logvar), dim=-1
)Now we write the sequence training step. We run the RSSM over a batch of sequences, using the posterior at each step because observations are available. The loss combines reconstruction (on the next observation, since we predict from the state after applying ), reward prediction, continue prediction, and the KL term. Next-observation reconstruction teaches the posterior representation to retain information about its observation; by itself it does not force dynamics learning because that posterior already receives the next observation. The action-conditioned prior and its KL agreement with the posterior provide the pressure to predict how latent state changes. That predictive pressure is essential for imagination to be meaningful.
def sequence_batch(
obs,
acts,
rews,
next_obs,
dones,
batch_size,
seq_len,
device,
episode_length=50,
):
if len(obs) % episode_length:
raise ValueError("transition count must contain complete episodes")
if seq_len > episode_length:
raise ValueError("seq_len must not exceed episode_length")
episodes = len(obs) // episode_length
episode_ids = np.random.randint(0, episodes, size=batch_size)
offsets = np.random.randint(
0, episode_length - seq_len + 1, size=batch_size
)
starts = episode_ids * episode_length + offsets
obs_seq = np.stack([obs[s : s + seq_len] for s in starts]).astype(
np.float32
)
act_seq = np.stack([acts[s : s + seq_len] for s in starts])
rew_seq = np.stack([rews[s : s + seq_len] for s in starts])
next_seq = np.stack([next_obs[s : s + seq_len] for s in starts]).astype(
np.float32
)
done_seq = np.stack([dones[s : s + seq_len] for s in starts])
obs_seq = torch.tensor(obs_seq, device=device).view(batch_size, seq_len, -1)
next_seq = torch.tensor(next_seq, device=device).view(
batch_size, seq_len, -1
)
act_seq = torch.tensor(act_seq, device=device)
rew_seq = torch.tensor(rew_seq, device=device)
done_seq = torch.tensor(done_seq, device=device)
return obs_seq, act_seq, rew_seq, next_seq, done_seq
def one_hot_action(a, action_dim, device):
return F.one_hot(a, num_classes=action_dim).float().to(device)
def model_loss(model, batch, beta=0.5):
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, _, (q_mean, q_logvar) = model.forward_step(
zero_action, zero_latent, deter, obs_seq[:, 0]
)
prev_latent = reparameterize(q_mean, q_logvar)
rec_loss = 0.0
rew_loss = 0.0
cont_loss = 0.0
kl_loss = 0.0
for t in range(T):
action = one_hot_action(act_seq[:, t], model.action_dim, device)
deter, (p_mean, p_logvar), (q_mean, q_logvar) = model.forward_step(
action, prev_latent, deter, next_seq[:, t]
)
latent = reparameterize(q_mean, q_logvar)
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_target = 1.0 - done_seq[:, t]
cont_loss = (
cont_loss
+ F.binary_cross_entropy(
cont_hat, cont_target, reduction="none"
).mean()
)
kl_loss = (
kl_loss + gaussian_kl(q_mean, q_logvar, p_mean, p_logvar).mean()
)
prev_latent = latent
rec_loss = rec_loss / T
rew_loss = rew_loss / T
cont_loss = cont_loss / T
kl_loss = kl_loss / T
total = rec_loss + rew_loss + cont_loss + beta * kl_loss
return total, dict(
rec=rec_loss.item(),
rew=rew_loss.item(),
cont=cont_loss.item(),
kl=kl_loss.item(),
)Let's train the world model. We keep it short so the chapter remains runnable, and the run measures reconstruction, KL, and held-out pixel diagnostics without visually validating imagined trajectories.
device = torch.device("cpu")
model = RSSM().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-3)
## Reserve the final eight complete episodes for evaluation.
train_end = 52 * 50
train_data = (
obs[:train_end],
acts[:train_end],
rews[:train_end],
next_obs[:train_end],
dones[:train_end],
)
loss_history = []
for step in range(600):
batch = sequence_batch(
*train_data,
batch_size=16,
seq_len=20,
device=device,
)
total, parts = model_loss(model, batch, beta=0.5)
optimizer.zero_grad()
total.backward()
nn.utils.clip_grad_norm_(model.parameters(), 100.0)
optimizer.step()
if step % 100 == 0:
loss_history.append((step, total.item(), parts["rec"], parts["kl"]))Total loss: 177.7027 -> 6.3306 Reconstruction BCE: 176.8755 -> 6.2999 KL: 0.1613 -> 0.0478
The printed diagnostics are observations, not proofs. Falling reconstruction BCE shows progress on the stated likelihood, while the KL reports prior-posterior agreement. A small KL can mean an easy-to-predict representation or posterior collapse; a finite nonzero KL alone does not distinguish them. Rollout evaluation below tests what the learned prior predicts.
DotWorld never terminates naturally, so every continuation target in this toy dataset is one. The head is present to preserve the Dreamer objective, but this example does not test terminal prediction.
Now we can imagine: starting from a real posterior state, roll forward using only the prior and a sequence of actions. We compute open-loop rollouts and, separately, posterior reconstructions conditioned on each target frame. These diagnostics distinguish recursive prediction from observation-conditioned reconstruction.
@torch.no_grad()
def imagine(model, obs_seed, start_action, horizon=15, action_dim=4):
# Encode the seed observation to get a posterior latent, then roll forward on the prior.
B = obs_seed.shape[0]
device = obs_seed.device
deter = model.initial_state(B, device)
prev_action = torch.zeros(B, action_dim, device=device)
prev_latent = torch.zeros(B, model.latent_dim, device=device)
deter, (_, _), (q_mean, q_logvar) = model.forward_step(
prev_action, prev_latent, deter, obs_seed.view(B, -1)
)
latent = q_mean # use the mean for a deterministic seed
images, rewards = [], []
action = torch.as_tensor(start_action, dtype=torch.long, device=device)
if action.ndim == 0:
action = action.expand(B)
elif action.shape != (B,):
raise ValueError(
f"start_action must be scalar or shape ({B},), got {tuple(action.shape)}"
)
for _ in range(horizon):
next_action = one_hot_action(action, action_dim, device)
deter, (p_mean, p_logvar), _ = model.forward_step(
prev_action=next_action, prev_latent=latent, deter=deter, obs=None
)
latent = p_mean # deterministic mean of the prior for a stable rollout
img, rew, _ = model.decode(deter, latent)
images.append(img.view(B, model.obs_size, model.obs_size))
rewards.append(rew)
prev_action = next_action
images = torch.stack(images, dim=1)
rewards = torch.stack(rewards, dim=1)
return images, rewards
seed = torch.tensor(next_obs[:4], dtype=torch.float32)
seed_actions = torch.tensor([2, 3, 0, 1]) # left, right, up, down
imagined_imgs, imagined_rews = imagine(model, seed, seed_actions, horizon=15)In this seeded run, the imagined_rews tensor computed above contains 60 predictions ranging from -0.1648 to -0.1514. All are negative, consistent with the environment's non-positive reward function; the hidden cell does not display them, and their sign alone does not establish reward calibration. The following hidden diagnostic compares one seed under constant-left and constant-right actions. It exercises action conditioning but, because it is neither displayed nor aggregated across seeds, does not establish general action sensitivity.
## Take one seed and compare repeated left versus repeated right actions.
single_seed = torch.tensor(next_obs[:1], dtype=torch.float32)
imgs_left, rews_left = imagine(model, single_seed, start_action=2, horizon=15)
imgs_right, rews_right = imagine(model, single_seed, start_action=3, horizon=15)Because this single-seed result is hidden, no conclusion about action sensitivity follows from it. The actor-critic section instead treats real-environment return as the transfer test.
Before we move on, compare posterior reconstruction with recursive open-loop rollout on the same held-out windows. The posterior diagnostic conditions on each target next frame and samples a latent; the open-loop diagnostic uses prior means after the seed. Their difference therefore combines observation conditioning with the sample-versus-mean choice.
@torch.no_grad()
def open_loop_rollout_error(model, obs_seq, act_seq, next_seq, steps=15):
# Seed with the first frame, then roll the prior forward with no further observations.
B = obs_seq.shape[0]
device = obs_seq.device
deter = model.initial_state(B, device)
prev_action = torch.zeros(B, model.action_dim, device=device)
prev_latent = torch.zeros(B, model.latent_dim, device=device)
deter, (_, _), (q_mean, _) = model.forward_step(
prev_action, prev_latent, deter, obs_seq[:, 0].view(B, -1)
)
latent = q_mean
open_loop_err = []
for t in range(steps):
action = one_hot_action(act_seq[:, t], model.action_dim, device)
deter, (p_mean, _), _ = model.forward_step(
prev_action=action, prev_latent=latent, deter=deter, obs=None
)
latent = p_mean
img_hat, _, _ = model.decode(deter, latent)
target = next_seq[:, t].view(B, -1)
open_loop_err.append(
F.mse_loss(img_hat, target, reduction="none").mean(dim=-1)
)
prev_action = action
open_loop_err = torch.stack(open_loop_err, dim=1)
return open_loop_err.mean(dim=0)
@torch.no_grad()
def posterior_reconstruction_error(model, obs_seq, act_seq, next_seq, steps=15):
# Infer the seed, then update with the action and resulting next observation.
B = obs_seq.shape[0]
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, _, (q_mean, q_logvar) = model.forward_step(
zero_action, zero_latent, deter, obs_seq[:, 0]
)
prev_latent = reparameterize(q_mean, q_logvar)
one_step_err = []
for t in range(steps):
action = one_hot_action(act_seq[:, t], model.action_dim, device)
deter, _, (q_mean, q_logvar) = model.forward_step(
action, prev_latent, deter, next_seq[:, t].view(B, -1)
)
latent = reparameterize(q_mean, q_logvar)
img_hat, _, _ = model.decode(deter, latent)
target = next_seq[:, t].view(B, -1)
one_step_err.append(
F.mse_loss(img_hat, target, reduction="none").mean(dim=-1)
)
prev_latent = latent
one_step_err = torch.stack(one_step_err, dim=1)
return one_step_err.mean(dim=0)
## Use one fixed within-episode window from each reserved episode.
hold_out_start = train_end
T_eval = 15
ho_obs_seq = torch.stack(
[
torch.tensor(
obs[hold_out_start + 50 * s : hold_out_start + 50 * s + T_eval],
dtype=torch.float32,
)
for s in range(8)
],
dim=0,
).view(8, T_eval, -1)
ho_next_seq = torch.stack(
[
torch.tensor(
next_obs[
hold_out_start + 50 * s : hold_out_start + 50 * s + T_eval
],
dtype=torch.float32,
)
for s in range(8)
],
dim=0,
).view(8, T_eval, -1)
ho_act_seq = torch.stack(
[
torch.tensor(
acts[hold_out_start + 50 * s : hold_out_start + 50 * s + T_eval],
dtype=torch.int64,
)
for s in range(8)
],
dim=0,
)
posterior_curve = posterior_reconstruction_error(
model, ho_obs_seq, ho_act_seq, ho_next_seq, steps=T_eval
).numpy()
open_loop_curve = open_loop_rollout_error(
model, ho_obs_seq, ho_act_seq, ho_next_seq, steps=T_eval
).numpy()
The curves must be read from the executed values. Their gap combines observation conditioning and different latent choices (posterior samples versus prior means). Every target contains one lit pixel, so an all-black prediction already has MSE ; the dashed zero line marks that trivial baseline, and negative percentages mean lower MSE. Both curves stay within about one percent of it. This small difference does not establish state tracking: even a predictor of average pixel occupancy could beat all-black without following the dot. Overlap alone is not diagnostic, since two accurate predictors could also overlap, and neither curve proves accurate long-horizon imagination.
A related evaluation pitfall appears in Part XI: Evaluation and Understanding: good one-step prediction need not yield useful multi-step imagined trajectories.
Latent imagination and actor-critic learning
We now have a model that predicts latent transitions, rewards, and continuations. Dreamer trains an actor and critic on imagined trajectories rather than unrolling the environment for each update. Depending on the generation and domain, the actor uses either analytic gradients through imagined dynamics or score-function gradients from returns while the model trajectory is stopped.
The actor update uses rewards and values predicted by the learned model rather than direct rewards from the environment. Those predictions are training targets, not ground truth. Real replay observations seed the imagined trajectories; after that, states and rewards come from the model. If the rollout leaves the data-supported region, the actor may optimize behavior that does not transfer.
The setup is as follows. At each iteration, you take a batch of real observations, encode them into latents, then roll out an imagined trajectory using the prior and the current policy for action selection. For each imagined step, we record the predicted reward, the predicted continuation probability, and the predicted value from the critic. The critic is trained against a continuation-weighted lambda-return that combines predicted rewards with critic bootstraps along the imagined trajectory and at its horizon. The actor is trained to maximize those imagined returns. This is essentially an actor-critic method, but the trajectory is entirely synthetic and lives in latent space.
Why use an actor instead of online planning? A policy amortizes decision-making across states and selects an action with one network evaluation rather than a fresh trajectory search. The actor-gradient estimator differs by generation: V1 backpropagates through imagined dynamics, reparameterizing continuous actions and using straight-through gradients for discrete actions; V2 uses REINFORCE on Atari and dynamics gradients on continuous control; V3 uses entropy-regularized REINFORCE and stops actor gradients through the world model.
An actor amortizes action selection across the states used in training. Unlike a fresh PlaNet-style CEM solve, it reuses learned parameters at the next decision. This does not guarantee smooth actions for nearby latents, and planners can also reuse computation through warm starts, learned priors, or value functions.
The toy implementation below uses a V3-style score-function update for discrete actions, but omits the cumulative imagination weights used in full Dreamer training. The world model generates states and rewards without gradients; the actor is then updated from the log-probabilities of exactly those sampled actions at exactly those decision states.
The toy actor loss maximizes expected advantage weighted by the log-probability of the chosen actions, with an entropy bonus to maintain exploration. Its unweighted surrogate is:
where:
- : the actor parameters
- : the critic parameters
- : the stop-gradient operator
- : the lambda-return, a weighted average of multi-step returns and bootstrapped critic predictions
- : the critic's value estimate at imagined state
- : the entropy coefficient
- : the imagination horizon
- : the imagined decision state before action
- : the entropy of the actor's action distribution at
The toy critic is trained to regress the lambda-return:
The lambda-return blends multi-step rewards with a value bootstrap at the imagination horizon. When it reduces to a one-step TD target. When , it includes rewards until a zero continuation; the horizon bootstrap is weighted by the product of intervening continuation probabilities:
where is the discount factor, is the predicted continuation probability, and the boundary is . A zero continuation truncates the return; a probability between zero and one softly downweights future terms.
For an episodic Dreamer objective, imagined steps also receive cumulative weights and . Without these weights, the actor, critic, and entropy losses would treat post-terminal imagined states as valid training data. Our toy code averages those losses across steps instead. DotWorld has no real terminal transitions and its continuation target is always one, so this example cannot test post-terminal masking. Even here, the code omits weighting and any downweighting from imperfect continuation predictions; it is not the published objective.
Why a lambda-return rather than the simple discounted return ? Because the imagination horizon is finite. An unbootstrapped recursion sets the unseen boundary , so its last included target is and earlier near-horizon targets omit all rewards beyond the rollout. They are truncated rather than generally zero, which can bias long-horizon value estimates. Bootstrapping with the critic's own prediction at the horizon mitigates this, and interpolates between low-variance (small , mostly bootstrapped) and low-bias (large , mostly Monte Carlo) returns. This is the same TD() idea that underlies much of modern deep RL.
Smaller relies more heavily on the critic; larger relies more heavily on accumulated model rewards. The familiar bias-variance description is conditional because model error and critic error both contribute. A value of 0.95 is an empirical setting, not a universal optimum.
The actor objective depends on both its action distribution and the action-conditioned latent trajectory. When every stochastic node on that path is reparameterizable, differentiating through the model gives a pathwise gradient of the learned-model objective; it can still be biased relative to real-environment return when the model is inaccurate and is often lower variance than REINFORCE. DreamerV2's categorical straight-through path is instead a biased surrogate estimator, and its Atari actor uses REINFORCE. DreamerV3 stops world-model gradients and uses a REINFORCE-style signal weighted by the critic's advantage.
Stopping the pathwise derivative does not prevent exploitation. A score-function update can still reinforce actions whose modeled transitions or rewards produce spuriously high returns. REINFORCE removes one gradient path and was reported to work across DreamerV3's domains; it does not make either learned head trustworthy outside the data distribution.
Now let's implement the actor and critic on top of our trained RSSM. We add a small MLP actor that takes the latent state and outputs a categorical action distribution, plus a critic that takes the latent state and outputs a scalar value.
class Actor(nn.Module):
def __init__(
self, latent_dim=16, deter_dim=64, action_dim=4, hidden_dim=128
):
super().__init__()
self.net = nn.Sequential(
nn.Linear(deter_dim + latent_dim, hidden_dim),
nn.ELU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ELU(),
nn.Linear(hidden_dim, action_dim),
)
def forward(self, deter, latent):
x = torch.cat([deter, latent], dim=-1)
return torch.distributions.Categorical(logits=self.net(x))
class Critic(nn.Module):
def __init__(self, latent_dim=16, deter_dim=64, hidden_dim=128):
super().__init__()
self.net = nn.Sequential(
nn.Linear(deter_dim + latent_dim, hidden_dim),
nn.ELU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ELU(),
nn.Linear(hidden_dim, 1),
)
def forward(self, deter, latent):
x = torch.cat([deter, latent], dim=-1)
return self.net(x).squeeze(-1)The rollout records each pre-action decision state, sampled action, predicted reward, and successor state. The world model is frozen, matching the stopped-dynamics estimator used by V3 and by this toy discrete-action example.
@torch.no_grad()
def encode_seed(model, obs_seed):
B = obs_seed.shape[0]
device = obs_seed.device
deter = model.initial_state(B, device)
prev_action = torch.zeros(B, model.action_dim, device=device)
prev_latent = torch.zeros(B, model.latent_dim, device=device)
deter, (_, _), (q_mean, _) = model.forward_step(
prev_action, prev_latent, deter, obs_seed.view(B, -1)
)
return deter, q_mean
@torch.no_grad()
def imagine_trajectory(model, actor, obs_seed, horizon=15, gamma=0.99):
deter, latent = encode_seed(model, obs_seed)
device = obs_seed.device
B = obs_seed.shape[0]
latents, deters, actions = [], [], []
rewards, continues = [], []
for _ in range(horizon):
deters.append(deter)
latents.append(latent)
dist = actor(deter, latent)
action = dist.sample()
action_onehot = F.one_hot(action, num_classes=model.action_dim).float()
deter, (p_mean, _), _ = model.forward_step(
prev_action=action_onehot, prev_latent=latent, deter=deter, obs=None
)
latent = p_mean
_, rew_hat, cont_hat = model.decode(deter, latent)
actions.append(action)
rewards.append(rew_hat)
continues.append(cont_hat)
return dict(
deters=torch.stack(deters, dim=1),
latents=torch.stack(latents, dim=1),
actions=torch.stack(actions, dim=1),
rewards=torch.stack(rewards, dim=1),
continues=torch.stack(continues, dim=1),
final_deter=deter,
final_latent=latent,
gamma=gamma,
)The imagine_trajectory function is the core of Dreamer. Notice that the model is asked for its prior at every step, never its posterior, because there is no new observation to condition on. The action is sampled from the current actor, the reward and continuation are read from the model's heads, and the latent is carried forward. The whole imagined episode is generated by the model playing both the world and the reward function, with the actor supplying the choices.
Now we compute the lambda-returns from the imagined rewards, continuation predictions, and critic values at each step. We use backward recursion, which is the cleanest way to compute and it avoids the recursion depth growing with the horizon.
def compute_lambda_returns(rewards, continues, values, gamma=0.99, lam=0.95):
# rewards and continues: (B, H); values: (B, H + 1)
B, H = rewards.shape
returns = torch.zeros_like(rewards)
next_return = values[:, H]
for t in reversed(range(H)):
next_value = values[:, t + 1]
bootstrap = (1.0 - lam) * next_value + lam * next_return
next_return = rewards[:, t] + gamma * continues[:, t] * bootstrap
returns[:, t] = next_return
return returns
def actor_critic_loss(
model,
actor,
critic,
obs_seed,
horizon=15,
gamma=0.99,
lam=0.95,
entropy_coef=1e-3,
):
traj = imagine_trajectory(
model, actor, obs_seed, horizon=horizon, gamma=gamma
)
deters = traj["deters"]
latents = traj["latents"]
rewards = traj["rewards"]
continues = traj["continues"]
B, H, D = deters.shape
flat_deter = deters.reshape(B * H, D)
flat_latent = latents.reshape(B * H, -1)
decision_values = critic(flat_deter, flat_latent).reshape(B, H)
final_value = critic(traj["final_deter"], traj["final_latent"]).unsqueeze(1)
values = torch.cat([decision_values, final_value], dim=1)
returns = compute_lambda_returns(
rewards, continues, values.detach(), gamma=gamma, lam=lam
)
advantage = (returns - decision_values).detach()
# log_probs must be recomputed FROM the sampled actions so gradients flow to the actor;
# the current log_probs come from @torch.no_grad() imagine_trajectory and carry no grad_fn.
dist_for_loss = actor(flat_deter, flat_latent)
log_probs = dist_for_loss.log_prob(traj["actions"].reshape(B * H)).reshape(
B, H
)
actor_loss = -(log_probs * advantage).mean()
critic_loss = F.mse_loss(decision_values, returns.detach())
dist = actor(flat_deter, flat_latent)
entropy = dist.entropy().mean()
total = actor_loss + critic_loss - entropy_coef * entropy
return total, dict(
actor=actor_loss.item(),
critic=critic_loss.item(),
entropy=entropy.item(),
return_mean=returns.mean().item(),
)The actor_loss line implements the toy's unweighted REINFORCE-with-advantage rule: the log-probability of the sampled action is multiplied by the advantage and used as a loss to maximize. The advantage is detached, so the actor is not tempted to move the value estimate itself. It can only change the action distribution. The critic_loss is a plain mean-squared error against the detached lambda-return. The entropy term is subtracted from the total loss because it is added to the objective, encouraging the actor to keep its action distribution spread out early in training. For an episodic task, the actor, critic, and entropy reductions would need the cumulative discount-and-continuation weights described above.
Now we train the actor and critic on imagined trajectories. Note that the model is frozen during this loop: only the actor and critic update. In a full Dreamer implementation, you would alternate between model updates and agent updates. You would also collect new real data between iterations. Here we keep the model fixed so we can isolate the actor-critic mechanism.
actor = Actor()
critic = Critic()
for parameter in model.parameters():
parameter.requires_grad_(False)
agent_optimizer = torch.optim.Adam(
list(actor.parameters()) + list(critic.parameters()), lr=3e-4
)
seed_pool = torch.tensor(next_obs[:256], dtype=torch.float32)
agent_history = []
for step in range(400):
idx = np.random.randint(0, len(seed_pool), size=16)
seeds = seed_pool[idx]
loss, parts = actor_critic_loss(model, actor, critic, seeds, horizon=15)
agent_optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(
list(actor.parameters()) + list(critic.parameters()), 100.0
)
agent_optimizer.step()
if step % 50 == 0:
agent_history.append(
(step, parts["actor"], parts["critic"], parts["return_mean"])
)Imagined return: -0.9208 -> -15.8317
Several train/evaluation differences matter here. seed_pool takes only the first 256 of 2,600 training transitions as isolated replay frames, so encode_seed starts RSSM memory at zero for each imagined trajectory; real-environment evaluation starts fresh episodes and carries recurrent memory through the steps of each episode. The actor samples actions during imagination, while evaluation below selects the most probable action. These seed-distribution, memory, and action-selection differences can affect real performance, so comparing the two return measurements does not isolate world-model exploitation.
The mean imagined lambda-return is a changing model-side diagnostic, not real environment return. Because its critic bootstrap changes during training, movement in this quantity cannot be attributed to the actor alone. In this run it decreases while the trained policy also performs worse than the random baseline in the real environment. These results do not establish an imagination-to-reality transfer failure or identify why the policy underperforms. Model exploitation, estimator noise, optimization failure, inadequate state learning, and the seed-state and action-selection mismatches just described remain possibilities. In a full system, poor real performance motivates collecting fresh data and updating the model.
Let's evaluate the trained actor greedily on the real environment, unrolling it step by step. Does the actor, trained on imagined rollouts seeded from real observations, keep the dot near the center in the real world?
@torch.no_grad()
def evaluate_actor(model, actor, world, episodes=20, steps=50, image_size=16):
rng = np.random.default_rng(999)
total_rewards = []
for _ in range(episodes):
pos, vel = world.reset(rng)
episode_reward = 0.0
deter = model.initial_state(1, torch.device("cpu"))
prev_action = torch.zeros(1, model.action_dim)
prev_latent = torch.zeros(1, model.latent_dim)
for _ in range(steps):
img = world.render(pos, image_size)
obs_t = torch.tensor(img, dtype=torch.float32).unsqueeze(0)
deter, _, (q_mean, _) = model.forward_step(
prev_action, prev_latent, deter, obs_t.view(1, -1)
)
latent = q_mean
dist = actor(deter, latent)
action = int(dist.probs.argmax(dim=-1).item())
pos, vel, reward, done = world.step(pos, vel, action, rng)
episode_reward += reward
prev_action = one_hot_action(
torch.tensor([action]), model.action_dim, torch.device("cpu")
)
prev_latent = latent
total_rewards.append(episode_reward)
return np.mean(total_rewards), np.std(total_rewards)
## Compute the random-policy baseline for comparison
rng = np.random.default_rng(777)
random_episode_rewards = []
for _ in range(20):
pos, vel = world.reset(rng)
r = 0.0
for _ in range(50):
action = rng.integers(0, 4)
pos, vel, reward, done = world.step(pos, vel, action, rng)
r += reward
random_episode_rewards.append(r)
random_mean = np.mean(random_episode_rewards)
random_std = np.std(random_episode_rewards)
trained_mean, trained_std = evaluate_actor(model, actor, world)
print(f"Random return: {random_mean:.4f} +/- {random_std:.4f}")
print(f"Trained return: {trained_mean:.4f} +/- {trained_std:.4f}")This comparison is diagnostic rather than guaranteed. Improvement would support transfer on this toy run; a worse return can reflect model exploitation, an inadequate model, estimator noise, optimization failure, or the seed-state and action-selection differences between training and evaluation. In practice, Dreamer alternates data collection, model learning, and actor-critic learning. The benchmark systems also include architectural, optimization, replay, and interface details beyond scaling networks and training time.
The most important takeaways from this section:
- Dreamer trains the actor and critic in imagination, using the learned model as a latent simulator.
- The actor's signal comes from the critic's advantage, not from real rewards at training time.
- The critic is trained with a lambda-return whose continuation factor truncates at zero and otherwise downweights future terms.
- This toy implementation and DreamerV3 stop actor gradients through the world model; V1 and V2 use the generation-specific estimators described above.
- The imagined return and the real return can diverge. Closing that gap is what the rest of the Dreamer pipeline is for.
DreamerV2 and discrete latents
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 -way posterior against a uniform prior, . This is not a general bound for : 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 -way categorical variable carries at most 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:
where 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 , the reported objective gives greater weight to fitting the prior:
with 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.
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, contThe _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.
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 -way prior is bounded by , 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.
DreamerV3 and cross-domain scaling
Earlier Dreamer versions used configurations developed for narrower benchmark suites. DreamerV3 evaluates one reported configuration across Atari, DMLab, robot control, and a modified Minecraft task. This establishes configuration robustness across the tested suites, not held-out-domain generalization.
One configuration reduces per-environment tuning within the evaluated scope. It does not by itself rule out suite-level engineering or establish generalization to unseen domains.
The problem with cross-domain robustness is scale mismatch. Different domains have:
- Very different reward scales. Atari games often use small discrete game rewards, continuous-control tasks accumulate dense rewards over an episode, and the modified Minecraft task has sparse rewards. The resulting return and value targets can occupy different numerical ranges under one optimizer configuration.
- Very different observation statistics. Pixel observations are in , but proprioceptive observations, robot states, and image statistics vary widely. A model that is not scale-invariant will be biased toward some domains.
- Very different horizons. Some tasks are solved in tens of steps; others (like Minecraft) require thousands of steps with sparse reward.
The following DreamerV3 transforms mainly address numerical scale and optimization. Long-horizon exploration and model accuracy remain separate challenges.
Symlog transformation. DreamerV3 applies symlog to vector observations as both encoder inputs and decoder reconstruction targets, and uses symexp-spaced reward and value bins:
where:
- : the input value (e.g., a reward or value target) in the original scale
- : the sign of , preserving the direction of the reward or value
- : a compressing transform of the magnitude, so large magnitudes grow logarithmically rather than linearly
This map is monotone, sign-preserving, logarithmically compressive, and unbounded. For magnitudes well above one, multiplying by 100 changes symlog by approximately ; that approximation fails near zero. Its inverse is symexp.
Two-hot prediction. In the published DreamerV3 formulation, the reward and critic heads predict categorical distributions over original-scale bins . A scalar target assigns weights to its two neighboring values in , providing linear interpolation within that finite support. Decoding takes the probability-weighted average of the original-scale bin locations, ; it is not symexp applied after an average in symlog space. Targets outside the support require clipping, so the representation is not precise at arbitrary unknown scales.
Return normalization. DreamerV3 divides the actor advantage by an exponentially tracked 95th-to-5th percentile range of returns, floored at one. It does not rescale rewards or critic targets this way. The entropy and KL loss scales are fixed configuration values, not adaptive coefficients.
Free bits. DreamerV3 clips each of its split KL losses from below at one nat. This removes additional KL pressure below the threshold; it does not guarantee a floor on mutual information or prevent the stochastic channel from becoming uninformative.
Together with one-percent uniform mixing, architectural choices, replay, and optimizer settings, these techniques supported one reported configuration across DM Control, DMLab, Atari, and a modified MineRL Diamond task. The final Nature paper reports that every training run discovered diamonds without human data or task-specific curricula. This is a milestone under that protocol, not evidence that long-horizon exploration is solved.
The following standalone helpers demonstrate symlog, symexp, and two-hot encoding. They are not integrated into the scalar critic above.
def symlog(x):
return torch.sign(x) * torch.log1p(torch.abs(x))
def symexp(x):
return torch.sign(x) * torch.expm1(torch.abs(x))
def value_bins(num_bins=41, symlog_limit=20.0, device=None):
symlog_bins = torch.linspace(
-symlog_limit, symlog_limit, num_bins, device=device
)
return symexp(symlog_bins)
def two_hot(target, num_bins=41, symlog_limit=20.0):
# Interpolate the original-scale target between symexp-spaced bins.
bins = value_bins(num_bins, symlog_limit, target.device)
clipped = torch.clamp(target, bins[0], bins[-1])
upper = torch.searchsorted(bins, clipped).clamp(1, num_bins - 1)
lower = upper - 1
lower_bin, upper_bin = bins[lower], bins[upper]
upper_weight = (clipped - lower_bin) / (upper_bin - lower_bin)
lower_weight = 1.0 - upper_weight
out = torch.zeros(target.shape[0], num_bins, device=target.device)
out.scatter_(1, lower.unsqueeze(1), lower_weight.unsqueeze(1))
out.scatter_add_(1, upper.unsqueeze(1), upper_weight.unsqueeze(1))
return out
def two_hot_value(critic_logits, num_bins=41, symlog_limit=20.0):
# Average original-scale bins; split signs to reduce cancellation error.
probs = F.softmax(critic_logits, dim=-1)
bins = value_bins(num_bins, symlog_limit, critic_logits.device)
weighted = probs * bins
midpoint = num_bins // 2
negative = weighted[..., :midpoint].flip(-1).sum(dim=-1)
if num_bins % 2:
center = weighted[..., midpoint]
positive = weighted[..., midpoint + 1 :].sum(dim=-1)
else:
center = torch.zeros_like(negative)
positive = weighted[..., midpoint:].sum(dim=-1)
return negative + center + positiveNow verify the symlog round trip and observe its logarithmic compression. This check does not imply scale invariance.
sample_rewards = torch.tensor([-100.0, -1.0, 0.0, 1.0, 100.0, 10000.0])
compressed = symlog(sample_rewards)
round_trip = symexp(compressed)Max round-trip error after symlog/symexp: 2.93e-03 Values from 1 to 10000 span four orders of magnitude before compression.
For inputs restricted to , symlog outputs lie near . Symlog itself remains unbounded. Two-hot prediction is also limited by its chosen finite bin support.
A central engineering detail in the published DreamerV3 formulation is the combination of symlog vector inputs, reward and value predictions over symexp-spaced original-scale bins, and percentile-normalized advantages. Its actor uses REINFORCE with an entropy term and stopped world-model gradients. The reported configuration does not use per-environment reward scaling or discount tuning.
We will not re-train DreamerV3 here. These transforms are individually compact, but the benchmark result belongs to their tested combination with the categorical RSSM, uniform mixing, split KL objective, architecture, optimizer, and replay design. Adding three transforms to V1 does not establish the same robustness.
Several limitations of DreamerV3 deserve attention, because a chapter that only reports the wins is not honest:
- Sample efficiency does not imply lower compute. RSSM training and imagination add work, but wall-clock comparisons with model-free methods depend on architectures, replay ratios, hardware, and implementations.
- Minecraft remains protocol-specific. Diamond discovery in the modified MineRL setting is a notable result without human data or curricula, not a general solution to sparse long-horizon exploration.
- Imagination rewards can diverge from real rewards. Advantage normalization and free bits address numerical scale and KL pressure; they do not by themselves prevent out-of-distribution model exploitation.
- One configuration does not mean assumption-free. The method still depends on the observation/action interfaces, replay distribution, training budget, and an environment regular enough to learn predictive latent dynamics.
- Interpretability is not guaranteed. Learned latent coordinates need not align with human concepts, although probes or visualizations can sometimes reveal structure. A decoder tests reconstructive information, not semantic transparency.
These limitations animate a great deal of subsequent work. Conservative and offline variants (discussed later in this part) address distribution shift. Value-equivalent and control-centric models (the next two chapters) address the perception-versus-control gap by making the latent directly optimized for decision quality rather than reconstruction. Tokenized and video-generator world models (in Part IX) revisit the representation question with larger-scale generative machinery. And evaluation frameworks (in Part XI) provide the tools for systematically measuring whether an imagined policy performs well in the environment, rather than only appearing self-consistently optimistic.
Now let's see whether the Dreamer-style policy we trained earlier performs better than a random policy in the real environment. This summary uses the imagined return curve already logged during training, but samples a fresh set of 30 episodes per policy for the real-return comparison.
## Precompute the trained-policy and random-policy returns before the plot cell, so figures only render data.
rng_eval = np.random.default_rng(2024)
random_returns = []
for _ in range(30):
pos, vel = world.reset(rng_eval)
r = 0.0
for _ in range(50):
action = rng_eval.integers(0, 4)
pos, vel, reward, done = world.step(pos, vel, action, rng_eval)
r += reward
random_returns.append(r)
trained_returns = []
with torch.no_grad():
for _ in range(30):
pos, vel = world.reset(rng_eval)
r = 0.0
deter = model.initial_state(1, torch.device("cpu"))
prev_action = torch.zeros(1, model.action_dim)
prev_latent = torch.zeros(1, model.latent_dim)
for _ in range(50):
img = world.render(pos, 16)
obs_t = torch.tensor(img, dtype=torch.float32).unsqueeze(0)
deter, _, (q_mean, _) = model.forward_step(
prev_action, prev_latent, deter, obs_t.view(1, -1)
)
latent = q_mean
dist = actor(deter, latent)
action = int(dist.probs.argmax(dim=-1).item())
pos, vel, reward, done = world.step(pos, vel, action, rng_eval)
r += reward
prev_action = one_hot_action(
torch.tensor([action]), model.action_dim, torch.device("cpu")
)
prev_latent = latent
trained_returns.append(r)
imagined_curve = [h[3] for h in agent_history]
Together, the plots show a changing imagined lambda-return diagnostic and measured real-environment return. Neither improves in this toy run: the imagined diagnostic decreases, and the trained greedy policy performs worse than the random baseline. These measurements use different horizons, seed states, and action-selection rules, so they cannot establish a transfer failure or identify why the policy underperforms.
PlaNet used an RSSM for online planning from pixels, building on earlier recurrent latent-state models. DreamerV1 retained PlaNet's RSSM-based world model and trained an actor and critic on imagined latent trajectories. DreamerV2 combined categorical states, KL balancing, and generation-specific actor gradients for the full Atari suite and continuous control. DreamerV3 evaluated a broader package (symlog vector encoder inputs and decoder targets, symexp two-hot reward and value heads, advantage normalization, free bits, uniform mixing, architecture, replay, and optimizer choices) with one reported configuration across several benchmark suites.
The next chapter on MuZero takes a different stance. MuZero learns latent dynamics that predict reward, value, and policy targets. When selecting actions in the real environment, its Monte Carlo tree search simulates transitions with that learned model rather than querying the environment for hypothetical outcomes.
Summary
The Dreamer family replaces per-step trajectory search with an actor and value function trained on latent imagination. Actor inference avoids a fresh online search. Published DreamerV3 experiments report one configuration across several benchmark suites, but additional imagination does not guarantee monotonic improvement or real-environment transfer.
The key ideas to take away:
- The RSSM combines recurrent memory with stochastic latent variables. It separates the prior used in imagination from the observation-conditioned posterior. Stochastic state can encode ambiguity but is not automatically calibrated uncertainty.
- Latent imagination rolls the model forward using the prior and current policy. V1 backpropagates through dynamics with reparameterized continuous actions or straight-through discrete actions; V2 selects between dynamics and REINFORCE estimators by domain, and V3 uses REINFORCE with stopped world-model gradients. The critic uses horizon-bootstrapped lambda-returns.
- DreamerV2 combines categorical latents, straight-through sampling, KL balancing, and domain-specific actor-gradient estimators. Its Atari result belongs to the complete system, not the latent change alone.
- DreamerV3 uses symlog vector encoder inputs and decoder targets, symexp-spaced two-hot reward and value predictions decoded as an original-scale weighted average, normalized actor advantages, free bits, uniform mixing, and further architecture and optimization choices; images retain their convolutional path. One reported configuration was evaluated across Atari, DMLab, control, and modified MineRL tasks.
- Limitations are real: compute comparisons are implementation-dependent, long-horizon sparse-reward exploration is not solved, imagination can diverge from reality, and semantic interpretability is not guaranteed.
The version history is qualitative rather than a set of commensurate capability scores:
| Version | Stochastic state | Actor-gradient estimator | Documented scope |
|---|---|---|---|
| DreamerV1 | Diagonal Gaussian | Dynamics/value gradients: reparameterized continuous actions, straight-through discrete actions | Image-based continuous control; selected Atari and DMLab tests |
| DreamerV2 | 32 groups of 32 categories | REINFORCE on Atari; dynamics gradients on continuous control | Atari and continuous control |
| DreamerV3 | Categorical with uniform mixing | Entropy-regularized REINFORCE; world-model gradients stopped | One reported configuration across multiple benchmark suites |
The most important skill to develop from this chapter is the ability to distinguish imagined performance from real performance. Dreamer's benchmark results show that latent imagination can train strong policies, while the toy experiment shows why transfer must be measured rather than inferred from the imagined objective.
Quiz
Ready to test your understanding? Take this quick quiz to reinforce what you've learned about the Dreamer family of world-model agents.
The Dreamer Family
Reference
Citation details
Cite or share this article.
Continue with the full handbook
This chapter is part of World Models Handbook. Use the handbook page to browse the complete table of contents and continue reading in sequence.
Explore World Models HandbookStay up to date
Get articles, book updates, and news delivered to your inbox.
No spam, unsubscribe anytime.
Join the community
Sign in to remove popups, track your reading progress, and join the discussion.
Comments
No comments yet. Be the first to share your thoughts!