The Dreamer Family › Recurrent state-space representations

Michael BrenndoerferJuly 15, 202620 min read
World Models HandbookAll chapters
Part 8 of 12Decision-Centric Research Lineages
Chapter 45The Dreamer Family › Recurrent state-space representations
Reading section 2 of 7 · Recurrent state-space representations · 20 minSection 2 of 7 · 20 min

Book contents

Section 2 of 7 · 20 min read

The Dreamer Family › Recurrent state-space representationsOne section · Sections in this chapter2 / 7 ↓
Reading tools Definitions & inline links

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

Article links

Make inline references clickable

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 hth_t, 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 ztz_t, 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 hth_t is the recurrent memory. The stochastic code ztz_t 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: ht+1=fϕ(ht,zt,at)h_{t+1} = f_\phi(h_t, z_t, a_t), typically a GRU. It advances the deterministic memory by one step given the current latent and action.
  • Prior (transition) distribution: pϕ(zt+1∣ht+1)p_\phi(z_{t+1} \mid h_{t+1}), 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: qϕ(zt+1∣ht+1,ot+1)q_\phi(z_{t+1} \mid h_{t+1}, o_{t+1}), 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): pϕ(ot+1∣ht+1,zt+1)p_\phi(o_{t+1} \mid h_{t+1}, z_{t+1}), 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: pϕ(rt∣ht+1,zt+1)p_\phi(r_t \mid h_{t+1}, z_{t+1}), which predicts the reward received after action ata_t from the successor latent. This is essential because the agent needs to imagine rewards to train the critic.
  • Continue model: pϕ(ct∣ht+1,zt+1)p_\phi(c_t \mid h_{t+1}, z_{t+1}), which predicts whether the episode continues after action ata_t. 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.

RSSM vs. a plain VAE sequence model

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 β\beta, and choosing that weight is one of the more delicate parts of training a Dreamer-style agent.

Lmodel(ϕ)=E[∑t(−log⁡pϕ(ot+1∣ht+1,zt+1)−log⁡pϕ(rt∣ht+1,zt+1)−log⁡pϕ(ct∣ht+1,zt+1)+β KL ⁣(qϕ(zt+1∣ht+1,ot+1) ∥ pϕ(zt+1∣ht+1)))],\mathcal{L}_\text{model}(\phi) = \mathbb{E}\left[ \sum_t \left( -\log p_\phi(o_{t+1} \mid h_{t+1}, z_{t+1}) - \log p_\phi(r_t \mid h_{t+1}, z_{t+1}) - \log p_\phi(c_t \mid h_{t+1}, z_{t+1}) + \beta \, \mathrm{KL}\!\left( q_\phi(z_{t+1} \mid h_{t+1}, o_{t+1}) \,\|\, p_\phi(z_{t+1} \mid h_{t+1}) \right) \right) \right],

where the expectation is over replay sequences and posterior latent samples, and β\beta 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 1/21/2, up to an additive constant. DreamerV2 and V3 instead use categorical samples with straight-through surrogate gradients.

The index tt names a replay transition here: ata_t takes the model from st=(ht,zt)s_t=(h_t,z_t) to st+1s_{t+1}, while rtr_t and ctc_t 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.

In[3]:
Code
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).

In[4]:
Code
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 img

This 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.

In[5]:
Code
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)
Out[6]:
Console
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.

In[7]:
Code
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, cont

The 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.

In[8]:
Code
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 ot+1o_{t+1} from the state after applying ata_t), 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.

In[9]:
Code
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.

In[10]:
Code
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"]))
Out[11]:
Console
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.

In[12]:
Code
@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.

In[13]:
Code
## 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.

In[14]:
Code
@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()
Out[15]:
Visualization
Line chart of posterior reconstruction and open-loop pixel MSE as percent difference from the dashed all-black baseline over fifteen held-out steps; both curves stay within about one percent of zero.
Mean pixel MSE on eight held-out episode windows, plotted as a percentage difference from the 1/256 all-black baseline (zero); negative values are lower error than that baseline. Posterior reconstruction conditions on each target frame and samples a latent; open-loop rollout uses prior means after the seed, so the comparison includes both conditioning and sampling differences.

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 1/2561/256; 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.