![](data:image/svg+xml;utf8,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="%23D95C41" stroke-width="1.8" stroke-linecap="round" stroke-linejoin="round"><line x1="12.1" y1="11.9" x2="18.9" y2="8.2" /><line x1="12.1" y1="12.1" x2="20.3" y2="12.9" /><line x1="12.2" y1="12.4" x2="16.6" y2="19.1" /><line x1="11.8" y1="12.4" x2="7.3" y2="19.2" /><line x1="11.9" y1="12.1" x2="3.7" y2="13.3" /><line x1="11.8" y1="11.7" x2="7.8" y2="4.4" /></svg>)
P03: Train a Dreamer Agent
Train a compact Dreamer agent with a world model and a latent Actor-Critic policy. This is a tutorial-scale demo: the objective is to show the Dreamer training loop, checkpoint wiring, and metric diagnostics, not to solve a hard control benchmark. There is no external gym dependency. A SyntheticEnv generates 64x64 RGB frames with a simple reward signal.
Prerequisite: P01 (vae_encoder.pt) and P02 (rssm.pt) if present. Otherwise the missing parts fall back to random initialization so the notebook still runs, but the trained agent is only meaningful with the pretrained checkpoints. This notebook saves the full agent to dreamer.pt for P05.
A noisy reward trace is acceptable here. The tutorial goal is a working world-model + policy pipeline, not a benchmark score.
Notebook source: p03_dreamer_agent.ipynb
%%bash
# Install dependencies for a fresh environment.
if command -v rocm-smi >/dev/null || [ -d /opt/rocm ]; then
pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.2
pip install matplotlib numpy
else
pip install torch torchvision matplotlib numpy
fi1. Setup
Define the shared environment, model dimensions, and training schedule. This notebook assembles four pieces from earlier lectures into the complete Dreamer pipeline: encode, predict forward with RSSM, evaluate imagined trajectories with an actor and critic, then execute in the real environment. One deliberate simplification up front, flagged here so the later code is not a surprise: the actor in this notebook is trained by imitating a hand-coded expert (Section 4's supervised_policy_update), not by backpropagating through imagined rewards as the lecture's latent Actor-Critic describes. Only the critic is trained on imagined returns here. This keeps the tutorial-scale agent stable to train in a few minutes; Section 3 still implements the full differentiable imagination machinery the real Dreamer actor update would use.
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from pathlib import Path
try:
from IPython import get_ipython
get_ipython().run_line_magic('matplotlib', 'inline')
except Exception:
pass
import matplotlib.pyplot as plt
from collections import deque
torch.manual_seed(42)
np.random.seed(42)
random.seed(42)
try:
import torch_xla.core.xla_model as xm
_XLA_AVAILABLE = True
except Exception:
xm = None
_XLA_AVAILABLE = False
def _resolve_device():
if _XLA_AVAILABLE:
return xm.xla_device()
if torch.cuda.is_available():
return torch.device('cuda')
return torch.device('cpu')
DEVICE = _resolve_device()
USE_TPU = DEVICE.type == 'xla'
USE_CUDA = DEVICE.type == 'cuda'
LOAD_DEVICE = torch.device('cpu') if USE_TPU else DEVICE
def optimizer_step(optimizer, scaler=None):
if USE_TPU:
xm.optimizer_step(optimizer)
elif scaler is not None:
scaler.step(optimizer)
scaler.update()
else:
optimizer.step()
print(f'Using device: {DEVICE}')
if USE_CUDA:
print(f'CUDA available: {torch.cuda.is_available()}')1.1 Hyperparameters
Keep the settings small so the full loop runs quickly. Three of these map directly onto lecture concepts: IMAGINE_H = 10 is the imagination horizon LAMBDA_RETURN = 0.95 is GAMMA = 0.99 is the discount factor
# Model dimensions
IMG_SIZE = 64
LATENT_DIM = 32 # VAE / stochastic latent z
HIDDEN_DIM = 128 # RSSM deterministic h
ACTION_DIM = 2 # binary action space
AC_HIDDEN = 128 # Actor / Critic hidden size
# Training schedule
EPISODE_LEN = 20 # steps per synthetic episode
N_ITERATIONS = 30 # outer training iterations
BATCH_SIZE = 4 # number of trajectories per world-model update
IMAGINE_H = 10 # imagination horizon
LAMBDA_RETURN = 0.95 # TD(lambda)
GAMMA = 0.99 # discount factor
# Learning rates
LR_WM = 3e-4
LR_AC = 3e-4
# Checkpoint paths (from earlier projects)
ENCODER_PATH = 'vae_encoder.pt'
RSSM_PATH = 'rssm.pt'
SAVE_PATH = 'dreamer.pt'1.2 VAE Encoder
Use the same encoder shape as P01, with one addition: VAEDecoder.forward here takes both z and h (torch.cat([z, h], dim=-1)), not just z as in P01 and P02. This matches the lecture's observation model h_t and the stochastic perception z_t, not on z_t alone as the simpler P01/P02 decoders did.
class VAEEncoder(nn.Module):
"""Encode 64x64 RGB frames into latent mean and logvar."""
def __init__(self, latent_dim=LATENT_DIM):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(3, 32, 4, stride=2, padding=1), # 32x32
nn.ReLU(),
nn.Conv2d(32, 64, 4, stride=2, padding=1), # 16x16
nn.ReLU(),
nn.Conv2d(64, 128, 4, stride=2, padding=1), # 8x8
nn.ReLU(),
nn.Conv2d(128, 256, 4, stride=2, padding=1),# 4x4
nn.ReLU(),
)
self.fc_mu = nn.Linear(256 * 4 * 4, latent_dim)
self.fc_logvar = nn.Linear(256 * 4 * 4, latent_dim)
def forward(self, x):
"""x: (B, 3, 64, 64) -> mu, logvar each (B, latent_dim)"""
h = self.conv(x).view(x.size(0), -1)
return self.fc_mu(h), self.fc_logvar(h)
def encode(self, x):
mu, logvar = self.forward(x)
std = (0.5 * logvar).exp()
eps = torch.randn_like(std)
return mu + eps * std, mu, logvar
class VAEDecoder(nn.Module):
"""Decoder: latent_dim -> 64x64 RGB reconstruction."""
def __init__(self, latent_dim=LATENT_DIM, hidden_dim=HIDDEN_DIM):
super().__init__()
self.fc = nn.Linear(latent_dim + hidden_dim, 256 * 4 * 4)
self.deconv = nn.Sequential(
nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(32, 3, 4, stride=2, padding=1),
nn.Sigmoid(),
)
def forward(self, z, h):
"""z: (B, latent_dim), h: (B, hidden_dim) -> (B, 3, 64, 64)"""
x = self.fc(torch.cat([z, h], dim=-1))
x = x.view(x.size(0), 256, 4, 4)
return self.deconv(x)1.3 RSSM
Reuse the latent dynamics interface from P02, restructured into three separate methods (prior, posterior, step) instead of P02's single forward loop, so this notebook can call each piece independently: step for advancing h, prior/posterior for sampling z. This split is what makes imagined_rollout below possible, since imagination needs prior alone (no observations) while training needs posterior (with observations), exactly the split the lecture's prior-vs-posterior callout describes.
One implementation detail worth flagging: prior and posterior use std = F.softplus(logvar) + 0.1 rather than P02's std = (0.5 * logvar).exp(). Both are valid ways to guarantee a positive standard deviation from an unconstrained network output; softplus(x) + 0.1 additionally enforces a minimum standard deviation of 0.1, which prevents the KL term from being driven to exactly zero by a near-deterministic posterior, a direct guard against the posterior collapse failure mode discussed in L04's Dreamer diagnostics.
class RSSM(nn.Module):
"""P02-compatible RSSM with a P03-friendly action interface."""
def __init__(self, latent_dim=LATENT_DIM, hidden_dim=HIDDEN_DIM, action_dim=ACTION_DIM):
super().__init__()
self.latent_dim = latent_dim
self.hidden_dim = hidden_dim
self.action_dim = action_dim
# P02 checkpoint expects a scalar action feature.
self.gru = nn.GRUCell(latent_dim + 1, hidden_dim)
# Prior: p(z_t | h_t)
self.prior_net = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.ELU(),
nn.Linear(hidden_dim, 2 * latent_dim),
)
# Posterior: q(z_t | h_t, e_t)
self.post_net = nn.Sequential(
nn.Linear(hidden_dim + latent_dim, hidden_dim),
nn.ELU(),
nn.Linear(hidden_dim, 2 * latent_dim),
)
self.recon = nn.Linear(latent_dim, latent_dim)
def _action_feature(self, action):
"""Convert scalar or one-hot actions to the P02 scalar action input."""
if action.dim() == 0:
action = action.view(1, 1)
elif action.dim() == 1:
action = action.unsqueeze(-1)
if action.shape[-1] > 1:
# Binary actions: keep the mass on action 1 so soft imagined
# actions remain differentiable during actor training.
action = action[..., 1:2]
return action.float()
def initial_state(self, batch_size):
h = torch.zeros(batch_size, self.hidden_dim, device=DEVICE)
z = torch.zeros(batch_size, self.latent_dim, device=DEVICE)
return h, z
def prior(self, h):
out = self.prior_net(h)
mu, logvar = out.chunk(2, dim=-1)
std = F.softplus(logvar) + 0.1
z = mu + std * torch.randn_like(std)
return z, mu, std
def posterior(self, h, enc_z):
out = self.post_net(torch.cat([h, enc_z], dim=-1))
mu, logvar = out.chunk(2, dim=-1)
std = F.softplus(logvar) + 0.1
z = mu + std * torch.randn_like(std)
return z, mu, std
def step(self, h, z, action_onehot):
"""Advance deterministic state. Returns new h."""
action_feat = self._action_feature(action_onehot)
inp = torch.cat([z, action_feat], dim=-1)
h_new = self.gru(inp, h)
return h_new1.4 Actor and Critic
Train both entirely in latent space, operating on (h, z) pairs rather than pixels, the same "no real observations needed" property that makes pure-imagination training possible. Actor.forward also accepts an optional bar_pos feature; this is a task-specific shortcut for the synthetic bar-balancing environment defined next, not part of the general Dreamer architecture, and defaults to zero when omitted (as it is during imagined_rollout below, where no real observation exists to compute it from). A RewardModel is defined alongside Actor and Critic: Dreamer needs to predict rewards inside imagined rollouts, since imagined states have no real environment to query for a reward signal, so a learned reward model is as essential a component as the dynamics model itself.
class Actor(nn.Module):
"""Map latent state (h, z) and a lightweight observation feature to action logits."""
def __init__(self, latent_dim=LATENT_DIM, hidden_dim=HIDDEN_DIM, action_dim=ACTION_DIM, ac_hidden=AC_HIDDEN, obs_feat_dim=1):
super().__init__()
inp = hidden_dim + latent_dim + obs_feat_dim
self.net = nn.Sequential(
nn.Linear(inp, ac_hidden),
nn.ELU(),
nn.Linear(ac_hidden, ac_hidden),
nn.ELU(),
nn.Linear(ac_hidden, action_dim),
)
def forward(self, h, z, bar_pos=None):
if bar_pos is None:
bar_pos = torch.zeros(h.shape[0], 1, device=h.device)
logits = self.net(torch.cat([h, z, bar_pos], dim=-1))
return logits
def sample(self, h, z, bar_pos=None):
logits = self.forward(h, z, bar_pos=bar_pos)
dist = torch.distributions.Categorical(logits=logits)
action = dist.sample()
return action, dist
class Critic(nn.Module):
"""Maps latent state (h, z) to scalar value estimate."""
def __init__(self, latent_dim=LATENT_DIM, hidden_dim=HIDDEN_DIM, ac_hidden=AC_HIDDEN):
super().__init__()
inp = hidden_dim + latent_dim
self.net = nn.Sequential(
nn.Linear(inp, ac_hidden),
nn.ELU(),
nn.Linear(ac_hidden, ac_hidden),
nn.ELU(),
nn.Linear(ac_hidden, 1),
)
def forward(self, h, z):
return self.net(torch.cat([h, z], dim=-1)).squeeze(-1)
class RewardModel(nn.Module):
"""Predict immediate reward from latent state and action."""
def __init__(self, latent_dim=LATENT_DIM, hidden_dim=HIDDEN_DIM, action_dim=ACTION_DIM, ac_hidden=AC_HIDDEN):
super().__init__()
inp = hidden_dim + latent_dim + action_dim
self.net = nn.Sequential(
nn.Linear(inp, ac_hidden),
nn.ELU(),
nn.Linear(ac_hidden, ac_hidden),
nn.ELU(),
nn.Linear(ac_hidden, 1),
)
def forward(self, h, z, a):
return self.net(torch.cat([h, z, a], dim=-1)).squeeze(-1)1.5 Synthetic Environment
SyntheticEnv is a minimal control task purpose-built for this notebook, not something introduced in the lectures: a red bar's horizontal position pos moves left or right by a fixed step depending on the binary action, and the reward is +1 whenever the action moved the bar closer to the center (abs(pos) < prev_abs) and -1 otherwise. This makes "keep the bar near the center" the optimal policy, simple enough that a working Dreamer loop can be demonstrated in a few minutes of training on CPU, while still requiring the agent to read position from pixels (via the encoder) and act on a multi-step signal (via the RSSM and actor) rather than solving the task with a one-line heuristic.
class SyntheticEnv:
"""Simple synthetic control environment with image observations."""
def __init__(self, episode_len=EPISODE_LEN, img_size=IMG_SIZE, seed=None):
self.episode_len = episode_len
self.img_size = img_size
self.rng = np.random.default_rng(seed)
self.pos = 0.0
self.step_count = 0
def _render(self):
img = np.zeros((self.img_size, self.img_size, 3), dtype=np.float32)
bar_x = int((self.pos + 1.0) / 2.0 * (self.img_size - 1))
bar_x = np.clip(bar_x, 0, self.img_size - 1)
img[:, max(0, bar_x - 2): bar_x + 3, 0] = 1.0 # red channel
# Add mild background noise to make the encoder non-trivial
img += self.rng.uniform(0, 0.05, img.shape).astype(np.float32)
return np.clip(img, 0, 1)
def reset(self):
self.pos = float(self.rng.uniform(-0.8, 0.8))
self.step_count = 0
return self._render()
def step(self, action):
"""action: int (0 or 1). Returns (obs, reward, done)."""
prev_abs = abs(self.pos)
delta = 0.1 if action == 1 else -0.1
self.pos = float(np.clip(self.pos + delta, -1.0, 1.0))
reward = 1.0 if abs(self.pos) < prev_abs else -1.0
self.step_count += 1
done = self.step_count >= self.episode_len
return self._render(), reward, done
# Quick sanity check
env = SyntheticEnv(seed=0)
obs = env.reset()
print(f'Observation shape: {obs.shape}, dtype: {obs.dtype}, range: [{obs.min():.2f}, {obs.max():.2f}]')
obs2, r, done = env.step(1)
print(f'After step: reward={r}, done={done}')1.6 Load or Initialize Models
Loads vae_encoder.pt (P01) and rssm.pt (P02) if present. Note the comment in _load_encoder_decoder_from_vae_checkpoint: only the encoder weights are reused from P01's checkpoint; P03's decoder is intentionally re-initialized and retrained from scratch, because P03's decoder architecture differs from P01's (it additionally conditions on h, as noted in Section 1.2 above), so P01's decoder weights would not be shape-compatible even if reused. If either checkpoint is missing, the corresponding module falls back to random initialization and the notebook still runs, but the resulting agent is a weaker demonstration since it is not building on the representations learned in P01/P02.
def obs_to_tensor(obs):
"""Convert HWC numpy float32 -> (1, 3, H, W) tensor."""
t = torch.from_numpy(obs).permute(2, 0, 1).unsqueeze(0)
return t.to(DEVICE)
encoder = VAEEncoder(latent_dim=LATENT_DIM).to(DEVICE)
decoder = VAEDecoder(latent_dim=LATENT_DIM, hidden_dim=HIDDEN_DIM).to(DEVICE)
rssm = RSSM(latent_dim=LATENT_DIM, hidden_dim=HIDDEN_DIM, action_dim=ACTION_DIM).to(DEVICE)
actor = Actor(latent_dim=LATENT_DIM, hidden_dim=HIDDEN_DIM, action_dim=ACTION_DIM, ac_hidden=AC_HIDDEN).to(DEVICE)
critic = Critic(latent_dim=LATENT_DIM, hidden_dim=HIDDEN_DIM, ac_hidden=AC_HIDDEN).to(DEVICE)
reward_model = RewardModel(latent_dim=LATENT_DIM, hidden_dim=HIDDEN_DIM, action_dim=ACTION_DIM, ac_hidden=AC_HIDDEN).to(DEVICE)
# Attempt to load weights from earlier projects.
def _load_encoder_decoder_from_vae_checkpoint(path):
ckpt = torch.load(path, map_location=DEVICE)
state = ckpt.get('model_state_dict', ckpt) if isinstance(ckpt, dict) else ckpt
if isinstance(ckpt, dict) and 'encoder' in ckpt:
enc_state = {k.replace('fc_log_var', 'fc_logvar'): v for k, v in ckpt['encoder'].items()}
encoder.load_state_dict(enc_state, strict=True)
return 'encoder only (decoder is intentionally re-initialized in P03)'
enc_state = {}
dec_state = {}
for key, value in state.items():
if key.startswith('encoder.'):
enc_key = key[len('encoder.'):].replace('fc_log_var', 'fc_logvar')
enc_state[enc_key] = value
elif key.startswith('decoder.'):
dec_state[key[len('decoder.'):]] = value
if not enc_state:
raise KeyError(f'Unrecognized VAE checkpoint format: {list(ckpt.keys())[:10] if isinstance(ckpt, dict) else type(ckpt)}')
encoder.load_state_dict(enc_state, strict=True)
return 'encoder only from model_state_dict (decoder is intentionally re-initialized in P03)'
vae_ckpt_candidates = [Path(ENCODER_PATH), Path('notebooks') / ENCODER_PATH]
vae_ckpt_path = next((p for p in vae_ckpt_candidates if p.exists()), None)
if vae_ckpt_path is not None:
try:
vae_ckpt_format = _load_encoder_decoder_from_vae_checkpoint(vae_ckpt_path)
print(f'Loaded encoder/decoder weights from {vae_ckpt_path} ({vae_ckpt_format})')
except Exception as e:
print(f'Could not load encoder/decoder from {vae_ckpt_path} ({e}), using random init')
else:
print('vae_encoder.pt not found; using random encoder init')
rssm_path = next((p for p in [Path(RSSM_PATH), Path('notebooks') / RSSM_PATH] if p.exists()), None)
if rssm_path is not None:
try:
state = torch.load(rssm_path, map_location=DEVICE)
if isinstance(state, dict) and 'rssm_state_dict' in state:
sd = state['rssm_state_dict']
rssm.load_state_dict(sd, strict=True)
print(
f"Loaded RSSM weights from {rssm_path} "
f"(hidden_dim={state.get('hidden_dim', HIDDEN_DIM)}, latent_dim={state.get('latent_dim', LATENT_DIM)}, action_dim={state.get('action_dim', 1)})"
)
elif isinstance(state, dict) and 'rssm' in state:
rssm.load_state_dict(state['rssm'], strict=False)
print(f'Loaded RSSM weights from {rssm_path} (legacy rssm key)')
else:
rssm.load_state_dict(state, strict=False)
print(f'Loaded RSSM weights from {rssm_path} (raw state_dict)')
except Exception as e:
print(f'Could not load RSSM ({e}), using random init')
else:
print('rssm.pt not found; using random RSSM init')
print('\nModel parameter counts:')
for name, m in [('encoder', encoder), ('decoder', decoder), ('rssm', rssm), ('actor', actor), ('critic', critic), ('reward_model', reward_model)]:
n = sum(p.numel() for p in m.parameters())
print(f' {name}: {n:,}')1.7 Replay Buffer and Optimizers
replay_buffer is a fixed-size FIFO queue (deque(maxlen=200)) of collected trajectories, the training data source for both the world-model update and the supervised policy update later. Note that opt_wm optimizes encoder + decoder + rssm + reward_model jointly as one parameter group, while opt_actor and opt_critic are separate optimizers: this separation matters because, as flagged in Section 1, the actor and critic are trained by different mechanisms (imitation vs. imagined returns) and should not share gradient updates.
# Replay buffer stores trajectory dictionaries.
replay_buffer = deque(maxlen=200)
# World model optimizer covers encoder + decoder + RSSM
wm_params = list(encoder.parameters()) + list(decoder.parameters()) + list(rssm.parameters()) + list(reward_model.parameters())
opt_wm = optim.Adam(wm_params, lr=LR_WM)
opt_actor = optim.Adam(actor.parameters(), lr=LR_AC)
opt_critic = optim.Adam(critic.parameters(), lr=LR_AC)
print('Optimizers initialized.')2. World Model Update
world_model_update is this notebook's version of the world-model training loop from P02, extended to also predict rewards. For each step in a trajectory: encode the current observation (encoder.encode), compute both the posterior z_post (conditioned on the real observation, via rssm.posterior) and the prior (conditioned on history alone, via rssm.prior), reconstruct the next observation from (z_post, h_next), and predict the immediate reward from (h, z_post, a_oh).
The kl computation is the same closed-form Gaussian-vs-Gaussian KL used in P02's RSSM, just written in terms of std directly ((post_std / prior_std).pow(2) + ...) rather than in terms of log_var as in P02: the two are algebraically equivalent, since log_var = 2 * log(std). h = h_next.detach() and z = z_post.detach() at the end of each step deliberately cut the backpropagation graph between consecutive timesteps (truncated backpropagation through time), trading a small amount of gradient accuracy for training stability and lower memory use over the 20-step episodes used here.
def action_to_onehot(action_int, action_dim=ACTION_DIM):
"""Scalar int -> (1, action_dim) one-hot tensor."""
oh = torch.zeros(1, action_dim, device=DEVICE)
oh[0, action_int] = 1.0
return oh
def world_model_update(batch):
"""Run one world-model ELBO update."""
encoder.train()
decoder.train()
rssm.train()
opt_wm.zero_grad()
total_recon = 0.0
total_kl = 0.0
total_reward = 0.0
count = 0
for traj in batch:
obs_list = traj['obs'] # list of T+1 numpy arrays (H,W,3)
act_list = traj['actions'] # list of T ints
T = len(act_list)
h, z = rssm.initial_state(1)
for t in range(T):
obs_t = obs_to_tensor(obs_list[t]) # (1,3,64,64)
obs_next = obs_to_tensor(obs_list[t + 1])
a_oh = action_to_onehot(act_list[t]) # (1, action_dim)
# Encode current observation
enc_z, enc_mu, enc_logvar = encoder.encode(obs_t)
# Posterior from encoder embedding
z_post, post_mu, post_std = rssm.posterior(h, enc_z)
# Prior from deterministic state
_, prior_mu, prior_std = rssm.prior(h)
# Reconstruct next observation
h_next = rssm.step(h, z_post.detach(), a_oh)
recon = decoder(z_post, h_next)
# Predict the actual transition reward from the latent state
reward_target = torch.tensor([traj['rewards'][t]], device=DEVICE, dtype=torch.float32)
reward_pred = reward_model(h, z_post, a_oh)
# Reconstruction loss (MSE per pixel, summed over spatial dims)
recon_loss = F.mse_loss(recon, obs_next, reduction='mean')
# KL divergence: posterior || prior (closed form, both Gaussian)
kl = 0.5 * (
(post_std / prior_std).pow(2)
+ ((post_mu - prior_mu) / prior_std).pow(2)
- 1
+ 2 * prior_std.log()
- 2 * post_std.log()
).sum(dim=-1).mean()
reward_loss = F.mse_loss(reward_pred, reward_target)
total_recon = total_recon + recon_loss
total_kl = total_kl + kl
total_reward = total_reward + reward_loss
count += 1
# Advance state (detach to avoid backprop through time across steps)
h = h_next.detach()
z = z_post.detach()
loss = (total_recon + total_kl + total_reward) / max(count, 1)
loss.backward()
nn.utils.clip_grad_norm_(wm_params, 100.0)
opt_wm.step()
return (total_recon / count).item(), (total_kl / count).item(), (total_reward / count).item()
print('world_model_update defined.')3. Behavior Learning (Imagination)
imagined_rollout implements the imagination step from the lecture's latent Actor-Critic training procedure: starting from a real (h, z) pair, repeatedly sample an action from the actor, advance the RSSM's prior (never the posterior, since no real observation exists once imagination begins), and predict a reward at each step, entirely without touching the real environment. The differentiable=True branch exists specifically to support the lecture-accurate actor training this notebook does not use by default: it replaces dist.sample() (non-differentiable, since sampling from a categorical distribution blocks gradient flow) with F.gumbel_softmax(..., hard=True), the standard technique for approximating a differentiable sample from a discrete distribution, which is what would let the actor's gradients flow back through the entire imagined trajectory as the lecture describes. behavior_update calls imagined_rollout in the non-differentiable mode (differentiable=False) and uses the result only to train the critic via lambda_returns, the same
lambda_returns implements the lecture's td = rewards[t] + gamma * values[t+1] is the one-step TD target, and the backward recursion G_next = (1 - lam) * td + lam * gamma * G_next builds the same "weighted average over all k-step returns" the lecture describes, computed efficiently backward through the horizon rather than as an explicit weighted sum.
def lambda_returns(rewards, values, gamma=GAMMA, lam=LAMBDA_RETURN):
"""Compute lambda-return targets for imagination."""
H = rewards.shape[0]
G = torch.zeros(H, device=DEVICE)
G_next = values[H]
for t in reversed(range(H)):
td = rewards[t] + gamma * values[t + 1]
G_next = (1 - lam) * td + lam * gamma * G_next
G[t] = G_next
return G
def set_requires_grad(module, flag):
for p in module.parameters():
p.requires_grad_(flag)
def imagined_rollout(start_h, start_z, horizon=IMAGINE_H, differentiable=False, tau=0.8):
"""Roll out latent trajectories in imagination."""
h, z = start_h, start_z
h_seq, z_seq, r_seq, ent_seq = [], [], [], []
for _ in range(horizon):
logits = actor(h, z)
dist = torch.distributions.Categorical(logits=logits)
ent_seq.append(dist.entropy().mean())
if differentiable:
a_oh = F.gumbel_softmax(logits, tau=tau, hard=True, dim=-1)
r_hat = reward_model(h, z, a_oh)
h_next = rssm.step(h, z, a_oh)
prior_out = rssm.prior_net(h_next)
prior_mu, _ = prior_out.chunk(2, dim=-1)
z_next = prior_mu
else:
a = dist.sample()
a_oh = F.one_hot(a, num_classes=ACTION_DIM).float()
with torch.no_grad():
r_hat = reward_model(h, z, a_oh)
h_next = rssm.step(h, z, a_oh)
prior_out = rssm.prior_net(h_next)
prior_mu, _ = prior_out.chunk(2, dim=-1)
z_next = prior_mu
h_seq.append(h)
z_seq.append(z)
r_seq.append(r_hat)
h, z = h_next, z_next
h_all = torch.stack(h_seq, dim=0)
z_all = torch.stack(z_seq, dim=0)
r_all = torch.stack(r_seq, dim=0).mean(dim=-1)
return h_all, z_all, r_all, ent_seq
def behavior_update(start_h, start_z, horizon=IMAGINE_H):
"""Train critic on imagined returns."""
critic.train()
reward_model.eval()
rssm.eval()
with torch.no_grad():
h_all, z_all, r_all, ent_seq = imagined_rollout(start_h.detach(), start_z.detach(), horizon=horizon, differentiable=False)
v_all = torch.zeros(horizon + 1, device=DEVICE)
for t in range(horizon):
v_all[t] = critic(h_all[t], z_all[t]).mean().detach()
v_all[horizon] = critic(h_all[-1], z_all[-1]).mean().detach()
G = lambda_returns(r_all, v_all)
opt_critic.zero_grad()
v_pred = torch.stack([critic(h_all[t], z_all[t]).mean() for t in range(horizon)])
critic_loss = F.mse_loss(v_pred, G.detach())
critic_loss.backward()
nn.utils.clip_grad_norm_(critic.parameters(), 100.0)
opt_critic.step()
return torch.stack(ent_seq).mean().item()
print('behavior_update defined.')4. Training Loop
Each of the N_ITERATIONS outer iterations runs four sub-steps: collect one real episode with the current actor (collect_episode), update the world model on a batch of replayed trajectories (world_model_update), update the critic on an imagined rollout from that episode's states (behavior_update), and update the actor by imitation (supervised_policy_update). That last step is the one flagged in Section 1: expert_action_from_obs is a hand-coded heuristic ("move toward center") baked directly into this cell, and the actor is trained with plain F.cross_entropy against the expert's action labels, standard supervised behavior cloning, not policy-gradient or backpropagation-through-imagination. This is a pragmatic substitution for a tutorial-scale demo: end-to-end differentiable actor training through imagined_rollout(differentiable=True) is implemented and available (Section 3), but is more sensitive to hyperparameters and slower to converge reliably within N_ITERATIONS = 30 iterations, so the default training loop uses the more stable imitation signal instead. If you want to see the lecture-accurate loss in action, replace the supervised_policy_update call with an actor loss computed from imagined_rollout(h, z, differentiable=True)'s returned r_seq, maximized via the critic exactly as the lecture describes.
With episode collection ready, initialize the metric history before the training loop starts writing results.
def collect_episode(env_seed=None, deterministic=False, epsilon=0.05):
"""Collect one episode with the current actor."""
env = SyntheticEnv(seed=env_seed)
obs = env.reset()
traj = {'obs': [obs], 'actions': [], 'rewards': []}
h, z = rssm.initial_state(1)
actor.eval()
encoder.eval()
rssm.eval()
total_reward = 0.0
done = False
with torch.no_grad():
while not done:
obs_t = obs_to_tensor(obs)
enc_z, _, _ = encoder.encode(obs_t)
z_post, _, _ = rssm.posterior(h, enc_z)
bar_pos = obs_to_bar_pos(obs)
logits = actor(h, z_post, bar_pos=bar_pos)
dist = torch.distributions.Categorical(logits=logits)
if deterministic:
a_int = int(torch.argmax(logits, dim=-1).item())
else:
if random.random() < epsilon:
a_int = random.randint(0, ACTION_DIM - 1)
else:
a_int = int(dist.sample().item())
obs_next, reward, done = env.step(a_int)
a_oh = action_to_onehot(a_int)
h = rssm.step(h, z_post, a_oh)
z = z_post
traj['obs'].append(obs_next)
traj['actions'].append(a_int)
traj['rewards'].append(reward)
obs = obs_next
total_reward += reward
return traj, total_reward
def obs_to_bar_pos(obs):
"""Estimate the bar position from the red channel as a normalized scalar."""
red_profile = obs[:, :, 0].mean(axis=0)
bar_x = int(np.argmax(red_profile))
denom = max(obs.shape[1] - 1, 1)
bar_pos = (2.0 * bar_x / denom) - 1.0
return torch.tensor([[bar_pos]], device=DEVICE, dtype=torch.float32)
def get_rssm_states_from_traj(traj):
"""Re-run RSSM on a trajectory to get posterior (h, z) pairs for imagination."""
encoder.eval()
rssm.eval()
h, z = rssm.initial_state(1)
h_list, z_list = [], []
with torch.no_grad():
for t, a_int in enumerate(traj['actions']):
obs_t = obs_to_tensor(traj['obs'][t])
enc_z, _, _ = encoder.encode(obs_t)
z_post, _, _ = rssm.posterior(h, enc_z)
h_list.append(h)
z_list.append(z_post)
a_oh = action_to_onehot(a_int)
h = rssm.step(h, z_post, a_oh)
return torch.cat(h_list, dim=0), torch.cat(z_list, dim=0) # (T, dim)
def expert_action_from_obs(obs):
"""Return the move-toward-center expert action for the synthetic bar task."""
red_profile = obs[:, :, 0].mean(axis=0)
bar_x = int(np.argmax(red_profile))
center = obs.shape[1] // 2
return 1 if bar_x < center else 0
def supervised_policy_update(batch):
"""Train actor to imitate the move-toward-center expert on replayed trajectories."""
actor.train()
losses = []
for traj in batch:
h_states, z_states = get_rssm_states_from_traj(traj)
targets = torch.tensor([expert_action_from_obs(obs) for obs in traj['obs'][:-1]], device=DEVICE, dtype=torch.long)
bar_pos = torch.cat([obs_to_bar_pos(obs) for obs in traj['obs'][:-1]], dim=0)
logits = actor(h_states, z_states, bar_pos=bar_pos)
losses.append(F.cross_entropy(logits, targets))
loss = torch.stack(losses).mean()
opt_actor.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(actor.parameters(), 100.0)
opt_actor.step()
return loss.item()
print('Collection utilities defined. Starting training...')With episode collection ready, initialize the metric history before the training loop starts writing results.
# Metrics history.
ep_rewards = []
recon_losses = []
kl_losses = []
reward_losses = []
policy_losses = []
actor_entropies = []
for iteration in range(N_ITERATIONS):
# --- Collect one episode ---
traj, ep_reward = collect_episode(env_seed=iteration, deterministic=False, epsilon=0.10)
replay_buffer.append(traj)
ep_rewards.append(ep_reward)
# --- World model update ---
buf_list = list(replay_buffer)
n_sample = min(BATCH_SIZE, len(buf_list))
batch = random.sample(buf_list, n_sample)
recon_l, kl_l, reward_l = world_model_update(batch)
recon_losses.append(recon_l)
kl_losses.append(kl_l)
reward_losses.append(reward_l)
# --- Critic update (imagination) ---
h_states, z_states = get_rssm_states_from_traj(traj)
entropy = behavior_update(h_states, z_states, horizon=IMAGINE_H)
actor_entropies.append(entropy)
# --- Actor update from a simple expert policy on replay ---
policy_l = supervised_policy_update(batch)
policy_losses.append(policy_l)
if (iteration + 1) % 10 == 0:
print(
f'Iter {iteration+1:3d} | '
f'ep_reward={ep_reward:+.1f} | '
f'recon={recon_l:.4f} | '
f'kl={kl_l:.4f} | '
f'reward={reward_l:.4f} | '
f'policy={policy_l:.4f} | '
f'actor_entropy={entropy:.4f}'
)
print('\nTraining complete.')Once the metrics are being recorded, turn them into a learning curve so you can watch the agent improve over time. ep_reward should trend upward over iterations if the imitation-trained actor is successfully learning the "move toward center" behavior; recon_losses and kl_losses are the same two ELBO components introduced in P01 and P02, now diagnosing whether the world model itself, not the policy, is degrading. Watch kl_losses in particular for the near-zero collapse pattern flagged in Section 1.3: a KL that drops to near zero within the first few iterations, while recon_losses stays high, is posterior collapse, not successful training.
fig, axes = plt.subplots(2, 2, figsize=(12, 8))
fig.suptitle('Dreamer Training Metrics', fontsize=14)
axes[0, 0].plot(ep_rewards, color='steelblue')
axes[0, 0].set_title('Episode Reward')
axes[0, 0].set_xlabel('Iteration')
axes[0, 0].set_ylabel('Total Reward')
axes[0, 0].axhline(0, color='gray', linestyle='--', alpha=0.5)
axes[0, 1].plot(recon_losses, color='tomato')
axes[0, 1].set_title('Reconstruction Loss (MSE)')
axes[0, 1].set_xlabel('Iteration')
axes[0, 1].set_ylabel('Loss')
axes[1, 0].plot(kl_losses, color='darkorange')
axes[1, 0].set_title('KL Divergence Loss')
axes[1, 0].set_xlabel('Iteration')
axes[1, 0].set_ylabel('KL')
axes[1, 1].plot(actor_entropies, color='mediumpurple')
axes[1, 1].set_title('Actor Entropy (Imagination)')
axes[1, 1].set_xlabel('Iteration')
axes[1, 1].set_ylabel('Entropy (nats)')
plt.tight_layout()
plt.show()5. Self-Evaluation Metrics
This section implements the exact diagnostic from L04's Dreamer worked example: compare rewards the world model imagines against rewards the real environment actually returns, for the same starting states. imagined_rollout_rewards is a read-only variant of imagined_rollout (Section 3) that additionally tracks the actor's entropy at each imagined step, since the lecture's diagnostic rule for a healthy Dreamer requires checking both reward correlation and entropy together, not reward correlation alone.
def imagined_rollout_rewards(start_h, start_z, horizon=10, deterministic=True):
"""Roll forward in imagination using the Actor and reward model."""
actor.eval()
reward_model.eval()
rssm.eval()
h, z = start_h.clone(), start_z.clone()
im_rewards = []
im_entropies = []
with torch.no_grad():
for _ in range(horizon):
logits = actor(h, z)
dist = torch.distributions.Categorical(logits=logits)
im_entropies.append(dist.entropy().mean().item())
if deterministic:
a = torch.argmax(logits, dim=-1)
else:
a = dist.sample()
a_oh = F.one_hot(a, num_classes=ACTION_DIM).float()
r_hat = reward_model(h, z, a_oh).mean().item()
im_rewards.append(r_hat)
h = rssm.step(h, z, a_oh)
prior_out = rssm.prior_net(h)
z, _ = prior_out.chunk(2, dim=-1)
return im_rewards, im_entropies
N_EVAL = 10
EVAL_H = EPISODE_LEN
real_reward_sums = []
imag_reward_sums = []
imag_entropies_ev = []
encoder.eval()
rssm.eval()
actor.eval()
reward_model.eval()
for ep_i in range(N_EVAL):
# Collect real episode
traj, ep_r = collect_episode(env_seed=1000 + ep_i, deterministic=True)
real_reward_sums.append(sum(traj['rewards']))
# Get RSSM state from first step as imagination seed
h0, z0 = get_rssm_states_from_traj(traj)
seed_h = h0[0:1] # single step
seed_z = z0[0:1]
# Imagined rewards and entropy
im_r, ents = imagined_rollout_rewards(seed_h, seed_z, horizon=EVAL_H, deterministic=True)
imag_reward_sums.append(sum(im_r))
imag_entropies_ev.append(float(np.mean(ents)))
print(f'Eval over {N_EVAL} episodes complete.')
print(f' Mean real reward (full episode): {np.mean(real_reward_sums):.3f}')
print(f' Mean imagined reward (predicted): {np.mean(imag_reward_sums):.3f}')
print(f' Mean imagined trajectory entropy: {np.mean(imag_entropies_ev):.4f}')After the imagination helper is defined, compare imagined reward against real return to check whether planning is aligned with the environment. N_EVAL = 10 fresh episodes are collected with deterministic=True (always take the actor's highest-logit action, no exploration noise), then for each episode, the RSSM's posterior state after the first real observation seeds an imagined rollout of the same length. Comparing real_reward_sums against imag_reward_sums for the same starting states is precisely the "starting from the same initial state" comparison the lecture recommends: if the two track each other, the world model's reward predictions are trustworthy inputs to planning; if they diverge, imagined rewards cannot be trusted to guide action selection, which is the model exploitation risk described in the Actor-Critic planning lecture.
# Pearson correlation between imagined and real reward sums
r_real = np.array(real_reward_sums)
r_imag = np.array(imag_reward_sums)
if r_real.std() > 1e-8 and r_imag.std() > 1e-8:
rho = np.corrcoef(r_real, r_imag)[0, 1]
else:
rho = 0.0
print(f'Reward correlation rho (predicted vs real, {EVAL_H}-step): {rho:.4f}')
print(f'Mean imagined trajectory entropy: {np.mean(imag_entropies_ev):.4f}')With the summary metric in place, lay out the diagnostic plots side by side for a fast sanity check. rho, computed in the previous cell with np.corrcoef, is exactly the Pearson correlation coefficient ρ ≥ 0.8 for a healthy world model; on this small tutorial-scale agent trained for only 30 iterations, expect a noisier value than that production target. The scatter plot's diagonal dashed line represents perfect agreement (imagined = real): points scattered far from that line, especially systematically above it, are the visual signature of the world model "lying" that the lecture warns about.
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
fig.suptitle('Dreamer Self-Evaluation', fontsize=13)
# Episode reward over training
axes[0].plot(ep_rewards, color='steelblue', label='train')
axes[0].set_title('Episode Reward (Training)')
axes[0].set_xlabel('Iteration')
axes[0].set_ylabel('Total Reward')
axes[0].legend()
# Reward correlation scatter
axes[1].scatter(r_real, r_imag, color='tomato', alpha=0.7)
axes[1].set_title(f'Reward Correlation (rho={rho:.3f})')
axes[1].set_xlabel(f'Real reward (full episode)')
axes[1].set_ylabel('Imagined reward (predicted)')
lims = [min(r_real.min(), r_imag.min()) - 0.5, max(r_real.max(), r_imag.max()) + 0.5]
axes[1].plot(lims, lims, 'k--', alpha=0.3, label='ideal')
axes[1].legend()
# Imagined trajectory entropy per eval episode
axes[2].bar(range(N_EVAL), imag_entropies_ev, color='mediumpurple', alpha=0.8)
axes[2].set_title('Imagined Trajectory Entropy')
axes[2].set_xlabel('Eval Episode')
axes[2].set_ylabel('Mean Entropy (nats)')
axes[2].axhline(np.mean(imag_entropies_ev), color='black', linestyle='--', label=f'mean={np.mean(imag_entropies_ev):.3f}')
axes[2].legend()
plt.tight_layout()
plt.show()6. Save Checkpoint
Saves every trained component (encoder, decoder, rssm, actor, critic, reward_model) plus the full metric history, not just the RSSM as P02 did, because P05's evaluation dashboard needs the complete agent, not only its dynamics core, to reproduce the reward-correlation and entropy diagnostics from Section 5 on demand.
checkpoint = {
'encoder': encoder.state_dict(),
'decoder': decoder.state_dict(),
'rssm': rssm.state_dict(),
'actor': actor.state_dict(),
'critic': critic.state_dict(),
'reward_model': reward_model.state_dict(),
'hyperparams': {
'latent_dim': LATENT_DIM,
'hidden_dim': HIDDEN_DIM,
'action_dim': ACTION_DIM,
'ac_hidden': AC_HIDDEN,
},
'metrics': {
'ep_rewards': ep_rewards,
'recon_losses': recon_losses,
'kl_losses': kl_losses,
'reward_losses': reward_losses,
'policy_losses': policy_losses,
'actor_entropies': actor_entropies,
'reward_corr_rho': float(rho),
'mean_imag_entropy': float(np.mean(imag_entropies_ev)),
'mean_eval_reward': float(np.mean(real_reward_sums)),
'mean_pred_reward': float(np.mean(imag_reward_sums)),
},
}
torch.save(checkpoint, SAVE_PATH)
print(f'Checkpoint saved to {SAVE_PATH}')
# Summary
print('\n--- Training Summary ---')
print(f' Final exploratory episode reward (last 10 avg): {np.mean(ep_rewards[-10:]):.2f}')
print(f' Final reconstruction loss: {recon_losses[-1]:.4f}')
print(f' Final KL loss: {kl_losses[-1]:.4f}')
print(f' Final reward loss: {reward_losses[-1]:.4f}')
print(f' Final policy loss: {policy_losses[-1]:.4f}')
print(f' Final actor entropy: {actor_entropies[-1]:.4f}')
print(f' Reward correlation rho: {rho:.4f}')
print(f' Mean imagined trajectory entropy: {np.mean(imag_entropies_ev):.4f}')
print(f' Mean eval reward: {np.mean(real_reward_sums):.2f}')