44.1 Latent Dynamics Models

1. Why Latent Dynamics?

Consider a robot arm performing chemical synthesis. Its camera produces 84x84 RGB images at each timestep, yielding an observation space of \(\mathbb{R}^{84 \times 84 \times 3}\). Predicting the next image pixel-by-pixel is a 21,168-dimensional regression problem, and most of those pixels (the static lab bench, the walls) carry no information about the dynamics. A latent dynamics model first encodes each observation into a compact representation \(z_t \in \mathbb{R}^d\) where \(d \ll 21{,}168\), then models transitions in this latent space: \(p(z_{t+1} | z_t, a_t)\). The savings are dramatic: both in the dimensionality of the prediction problem and in the amount of training data required.

A latent dynamics model is a neural network that predicts environment transitions in a compressed, learned coordinate system rather than in raw observation space. Low-dimensional latent spaces make prediction, planning, and uncertainty estimation tractable for environments whose raw observations (images, spectra, sensor arrays) are far too large to model directly. Three components implement this: an encoder maps each observation to a latent vector, a transition network predicts the next latent vector given the current one and an action, and a decoder maps latent vectors back to observations. Together, these form an end-to-end pipeline trained with reconstruction loss. Use a latent dynamics model when observations are high-dimensional and you need multi-step prediction or planning. Use a direct (observation-space) dynamics model only when the state is already compact (fewer than roughly 20 continuous dimensions) and you want to avoid learning an encoder and decoder.

Compression as Automated Variable Discovery

This compression does more than save compute. It forces the model to identify the variables that matter for prediction: arm position, liquid color, temperature. These are the same quantities a scientist would record in a lab notebook, and the latent dynamics model discovers them automatically from raw observations. In short: The latent space does not merely compress observations; it distills them into the variables that govern what happens next.

2. The Recurrent State-Space Model (RSSM)

Without the right latent structure, world models fail silently: they produce confident predictions that diverge from reality within a few steps, leading planners to choose actions that look optimal in imagination but perform catastrophically in the real environment. The RSSM architecture was designed to prevent exactly this failure.

Knowing that a compact latent space helps is only the first step; the critical design question is how to structure that space so it can faithfully capture both the predictable and the uncertain aspects of real-world dynamics.

The RSSM, introduced by Hafner et al. (2019) in the PlaNet paper and refined through DreamerV1, V2, and V3, is among the most successful latent dynamics architectures developed so far. Its key innovation is splitting the latent state into two components:

• Deterministic state \(h_t \in \mathbb{R}^{d_h}\): updated by a Gated Recurrent Unit (GRU) cell, this component carries forward information across time steps like a memory. It captures the predictable, deterministic aspects of the dynamics.
• Stochastic state \(s_t \in \mathbb{R}^{d_s}\): sampled from a learned distribution conditioned on \(h_t\), this component captures the inherent randomness in the environment. In DreamerV3, \(s_t\) uses a categorical distribution (32 classes across 32 dimensions) rather than a Gaussian, which substantially mitigates posterior collapse (where the KL regularization drives the learned posterior to match the prior so closely that the stochastic state carries no useful information).

The full RSSM consists of four components, each implemented as a neural network. Figure 44.1 illustrates how these components connect at a single timestep. Figure 44.1.1 illustrates RSSM architecture dataflow.

The sequence model is a GRU that updates the deterministic state. The prior predicts the stochastic state from the deterministic state alone (no observation), which is what we use during planning when we imagine trajectories without real observations. The posterior incorporates the actual observation \(o_t\) to produce a more accurate stochastic state during training. The decoder reconstructs the observation from the full latent state \((h_t, s_t)\), providing the reconstruction signal for learning.

The following implementation builds the core RSSM in PyTorch, with each component as a separate module composed together.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import OneHotCategorical


class RSSM(nn.Module):
    """Recurrent State-Space Model following the DreamerV3 architecture.

    Splits latent state into deterministic (h) and stochastic (s) parts.
    Stochastic state uses categorical distributions (num_classes x num_categoricals).
    """

    def __init__(
        self,
        obs_dim: int,
        action_dim: int,
        det_size: int = 256,
        num_categoricals: int = 32,
        num_classes: int = 32,
        hidden_size: int = 256,
    ):
        super().__init__()
        self.det_size = det_size
        self.num_categoricals = num_categoricals
        self.num_classes = num_classes
        self.stoch_size = num_categoricals * num_classes

        # Sequence model: GRU updates deterministic state
        self.gru_input_proj = nn.Linear(
            self.stoch_size + action_dim, hidden_size
        )
        self.gru_cell = nn.GRUCell(hidden_size, det_size)

        # Prior: predict stochastic state from deterministic alone
        self.prior_net = nn.Sequential(
            nn.Linear(det_size, hidden_size),
            nn.ELU(),
            nn.Linear(hidden_size, self.stoch_size),
        )

        # Posterior: predict stochastic state from deterministic + observation
        self.posterior_net = nn.Sequential(
            nn.Linear(det_size + obs_dim, hidden_size),
            nn.ELU(),
            nn.Linear(hidden_size, self.stoch_size),
        )

        # Observation decoder
        self.decoder = nn.Sequential(
            nn.Linear(det_size + self.stoch_size, hidden_size),
            nn.ELU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ELU(),
            nn.Linear(hidden_size, obs_dim),
        )

        # Reward predictor (for reinforcement learning tasks)
        self.reward_head = nn.Sequential(
            nn.Linear(det_size + self.stoch_size, hidden_size),
            nn.ELU(),
            nn.Linear(hidden_size, 1),
        )

    def initial_state(self, batch_size: int, device: torch.device):
        """Return zero-initialized deterministic and stochastic states."""
        h = torch.zeros(batch_size, self.det_size, device=device)
        s = torch.zeros(batch_size, self.stoch_size, device=device)
        return h, s

    def _categorical_straight_through(self, logits: torch.Tensor):
        """Sample from categorical with straight-through gradients."""
        shape = logits.shape[:-1]  # (batch, num_categoricals)
        logits_reshaped = logits.view(*shape, self.num_categoricals, self.num_classes)
        dist = OneHotCategorical(logits=logits_reshaped)
        sample = dist.sample()
        # Straight-through: forward uses sample, backward uses softmax
        probs = dist.probs
        sample_st = sample + probs - probs.detach()
        return sample_st.view(*shape[:-0], self.stoch_size), dist

    def sequence_step(self, h_prev, s_prev, action):
        """One step of the sequence model (GRU)."""
        gru_input = self.gru_input_proj(torch.cat([s_prev, action], dim=-1))
        h = self.gru_cell(F.elu(gru_input), h_prev)
        return h

    def prior(self, h):
        """Predict stochastic state distribution from deterministic state."""
        logits = self.prior_net(h)
        s, dist = self._categorical_straight_through(logits)
        return s, dist, logits

    def posterior(self, h, obs):
        """Predict stochastic state distribution from deterministic + observation."""
        logits = self.posterior_net(torch.cat([h, obs], dim=-1))
        s, dist = self._categorical_straight_through(logits)
        return s, dist, logits

    def decode(self, h, s):
        """Reconstruct observation from full latent state."""
        latent = torch.cat([h, s], dim=-1)
        obs_pred = self.decoder(latent)
        reward_pred = self.reward_head(latent)
        return obs_pred, reward_pred

    def observe(self, observations, actions):
        """Process a sequence of (observation, action) pairs.

        Args:
            observations: (batch, time, obs_dim)
            actions: (batch, time, action_dim)

        Returns:
            Dictionary of posterior states, prior states, predictions.
        """
        batch_size, seq_len = observations.shape[:2]
        device = observations.device
        h, s = self.initial_state(batch_size, device)

        prior_logits_list = []
        posterior_logits_list = []
        obs_preds = []
        reward_preds = []

        for t in range(seq_len):
            # Step 1: update deterministic state
            h = self.sequence_step(h, s, actions[:, t])

            # Step 2: compute prior and posterior
            _, _, prior_logits = self.prior(h)
            s, _, posterior_logits = self.posterior(h, observations[:, t])

            # Step 3: decode from posterior state (for training)
            obs_pred, reward_pred = self.decode(h, s)

            prior_logits_list.append(prior_logits)
            posterior_logits_list.append(posterior_logits)
            obs_preds.append(obs_pred)
            reward_preds.append(reward_pred)

        return {
            "prior_logits": torch.stack(prior_logits_list, dim=1),
            "posterior_logits": torch.stack(posterior_logits_list, dim=1),
            "obs_preds": torch.stack(obs_preds, dim=1),
            "reward_preds": torch.stack(reward_preds, dim=1),
        }

    def imagine(self, initial_h, initial_s, actions):
        """Roll out the model without observations (for planning).

        Args:
            initial_h: (batch, det_size) starting deterministic state
            initial_s: (batch, stoch_size) starting stochastic state
            actions: (batch, horizon, action_dim) planned actions

        Returns:
            Predicted observations, rewards, and states for each step.
        """
        h, s = initial_h, initial_s
        horizon = actions.shape[1]

        obs_preds = []
        reward_preds = []

        for t in range(horizon):
            h = self.sequence_step(h, s, actions[:, t])
            s, _, _ = self.prior(h)  # No observation: use prior
            obs_pred, reward_pred = self.decode(h, s)
            obs_preds.append(obs_pred)
            reward_preds.append(reward_pred)

        return {
            "obs_preds": torch.stack(obs_preds, dim=1),
            "reward_preds": torch.stack(reward_preds, dim=1),
        }

Several design choices in Listing 44.1 deserve attention. The stochastic state uses a categorical distribution (32 categories each with 32 classes) rather than a Gaussian. DreamerV3 showed that categorical distributions substantially mitigate the posterior collapse problem that plagues Gaussian latent variables, where the KL term in the ELBO drives the posterior toward the prior, effectively ignoring the observations. The straight-through gradient estimator (a technique that uses the discrete sample in the forward pass but substitutes a continuous approximation in the backward pass, enabling gradient flow through discrete choices) lets us backpropagate through the discrete sampling operation by using the continuous softmax probabilities in the backward pass.

3. The ELBO Training Objective

How do we train an RSSM? The model defines a generative process: at each time step, it samples a latent state and produces an observation. The true posterior over latent states given observations is intractable, so we use variational inference. The objective is the ELBO, extended to sequential data.

For a sequence of observations \(o_{1:T}\) and actions \(a_{1:T}\), the ELBO is:

Each term has a clear interpretation:

• Reconstruction term: the posterior state \((h_t, s_t)\) must produce accurate observation predictions. This teaches the encoder to extract relevant features and the decoder to faithfully reconstruct them.
• KL term: the posterior (which sees the observation) must stay close to the prior (which does not). This ensures that the prior alone can make reasonable predictions, which is essential for planning when no observations are available. The coefficient \(\beta\) controls the balance; DreamerV3 uses a free-bits strategy where the KL is only penalized above a minimum threshold.

The following listing implements the ELBO loss function with the free-bits strategy from DreamerV3.

def compute_elbo_loss(
    rssm_output: dict,
    observations: torch.Tensor,
    rewards: torch.Tensor,
    free_bits: float = 1.0,
    kl_beta: float = 1.0,
    reward_weight: float = 1.0,
):
    """Compute the ELBO loss for RSSM training.

    Args:
        rssm_output: dictionary from RSSM.observe()
        observations: (batch, time, obs_dim) ground truth
        rewards: (batch, time, 1) ground truth rewards
        free_bits: minimum KL before penalty applies (nats)
        kl_beta: weight on KL term
        reward_weight: weight on reward prediction loss

    Returns:
        Total loss and a dictionary of component losses for logging.
    """
    # Reconstruction loss: MSE between predicted and true observations
    obs_loss = F.mse_loss(rssm_output["obs_preds"], observations)

    # Reward prediction loss
    rew_loss = F.mse_loss(
        rssm_output["reward_preds"], rewards
    )

    # KL divergence between posterior and prior (categorical distributions)
    prior_logits = rssm_output["prior_logits"]
    posterior_logits = rssm_output["posterior_logits"]

    # Reshape to (batch * time, num_categoricals, num_classes)
    batch, time = prior_logits.shape[:2]
    num_cat = 32  # num_categoricals
    num_cls = 32  # num_classes

    prior_reshaped = prior_logits.view(batch * time, num_cat, num_cls)
    posterior_reshaped = posterior_logits.view(batch * time, num_cat, num_cls)

    # Compute KL for each categorical dimension
    prior_probs = F.softmax(prior_reshaped, dim=-1)
    posterior_probs = F.softmax(posterior_reshaped, dim=-1)

    kl_per_categorical = (
        posterior_probs * (
            torch.log(posterior_probs + 1e-8)
            - torch.log(prior_probs + 1e-8)
        )
    ).sum(dim=-1)  # (batch*time, num_cat)

    # Free bits: only penalize KL above threshold per categorical
    kl_free = torch.clamp(kl_per_categorical, min=free_bits)
    kl_loss = kl_free.mean()

    # Total ELBO loss (negative because we maximize ELBO)
    total_loss = obs_loss + reward_weight * rew_loss + kl_beta * kl_loss

    return total_loss, {
        "obs_loss": obs_loss.item(),
        "reward_loss": rew_loss.item(),
        "kl_loss": kl_loss.item(),
        "total_loss": total_loss.item(),
    }

4. Transformer World Models

The ELBO gives us a principled way to train any latent dynamics model, but the RSSM is not the only architecture that can serve as the backbone; the same training objective applies equally well when the recurrent core is replaced by an attention mechanism.

The RSSM processes sequences one step at a time through its GRU cell. Transformer world models (Micheli et al., 2023; Robine et al., 2023) replace this recurrence with causal self-attention (an attention mechanism where each position can only attend to itself and earlier positions, preventing information from the future from leaking into predictions), processing entire trajectories in parallel. The key advantage is that attention can capture long-range dependencies without the information bottleneck of a fixed-size hidden state. The key disadvantage is that generation remains autoregressive (each predicted token is conditioned on all previously generated tokens, so outputs must be produced one at a time in sequence), so the wall-clock cost of rolling out \(H\) steps is \(O(H^2)\) rather than \(O(H)\).

The architecture tokenizes each timestep's state and action into a sequence of tokens, then applies a causal transformer (GPT-style) to predict the next timestep's tokens. IRIS (Micheli et al., 2023) tokenizes observations using a discrete autoencoder (Vector-Quantized VAE, or VQ-VAE), producing a sequence of discrete tokens per image. The transformer then models the joint sequence of observation tokens and action tokens.

class TransformerWorldModel(nn.Module):
    """Simplified transformer-based world model.

    Processes (state, action) sequences with causal self-attention.
    Suitable for low-dimensional state spaces (no image tokenization).
    """

    def __init__(
        self,
        obs_dim: int,
        action_dim: int,
        d_model: int = 256,
        nhead: int = 4,
        num_layers: int = 4,
        max_seq_len: int = 128,
    ):
        super().__init__()
        self.d_model = d_model

        # Project (obs, action) pair into transformer dimension
        self.input_proj = nn.Linear(obs_dim + action_dim, d_model)

        # Learned positional encoding
        self.pos_embedding = nn.Embedding(max_seq_len, d_model)

        # Causal transformer
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=4 * d_model,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.transformer = nn.TransformerEncoder(
            encoder_layer, num_layers=num_layers
        )

        # Output heads
        self.obs_head = nn.Linear(d_model, obs_dim)
        self.reward_head = nn.Linear(d_model, 1)

    def forward(self, observations, actions):
        """Predict next observations and rewards.

        Args:
            observations: (batch, time, obs_dim)
            actions: (batch, time, action_dim)

        Returns:
            obs_preds: (batch, time, obs_dim) next-obs predictions
            reward_preds: (batch, time, 1) reward predictions
        """
        batch, time = observations.shape[:2]
        x = self.input_proj(torch.cat([observations, actions], dim=-1))
        positions = torch.arange(time, device=x.device)
        x = x + self.pos_embedding(positions)

        # Causal mask: each position attends only to itself and earlier
        causal_mask = nn.Transformer.generate_square_subsequent_mask(
            time, device=x.device
        )
        x = self.transformer(x, mask=causal_mask)

        obs_preds = self.obs_head(x)
        reward_preds = self.reward_head(x)
        return obs_preds, reward_preds

When should you choose a transformer world model over an RSSM? Empirical evidence (circa 2023) shows that transformer world models such as IRIS match or exceed RSSM-based agents on the Atari 100k benchmark, a standardized evaluation using only 100,000 environment interactions. However, RSSMs generate faster because the GRU state update costs \(O(1)\) per step, while the transformer must attend to all previous steps. For planning applications that roll out thousands of trajectories (Section 44.3), this difference matters. (A planner evaluating 1,000 candidate trajectories of length 50 executes roughly 50,000 GRU steps with an RSSM, but over 1.25 million attention operations with a transformer.) As of 2025, DreamerV3 remains broadly competitive with transformer world models on most published benchmarks, though results vary by domain. Hybrid architectures that combine RSSM-style latent states with transformer sequence models are an active area of development.

from dreamer_torch import DreamerV3

agent = DreamerV3(
    obs_space={"image": (3, 64, 64)},
    act_space={"action": 6},
    config=DreamerV3.defaults.update(
        rssm={"deter": 4096, "units": 1024, "stoch": 32, "classes": 32},
        encoder={"mlp_keys": "", "cnn_keys": "image"},
    ),
)
# agent.train(replay_buffer) handles ELBO optimization internally

5. RSSM vs. Transformer: Architectural Trade-offs

The choice between RSSM and transformer world models involves several trade-offs that depend on the specific discovery application:

For scientific discovery applications where uncertainty quantification matters (and in most experimental sciences it does), the RSSM's explicit stochastic state provides a natural advantage. The prior-posterior divergence gives a built-in signal for epistemic uncertainty (uncertainty about the model's own knowledge, as opposed to aleatoric uncertainty from inherent randomness in the environment): when the prior and posterior disagree strongly, the model is uncertain about the dynamics. Transformers require additional machinery (ensembles, conformal prediction) to achieve comparable uncertainty estimates.

6. Training an RSSM: The Complete Loop

With the architectural trade-offs understood, we can now turn from model design to the practical mechanics of fitting an RSSM to data collected from a real or simulated environment.

To train an RSSM world model, we need trajectories collected from the environment (or simulator). Each trajectory is a sequence of (observation, action, reward) tuples. The training loop alternates between collecting data using the current policy and updating the world model on batches sampled from a replay buffer (a data structure that stores past experience tuples and serves random mini-batches for training, decoupling data collection from gradient updates).

import gymnasium as gym
import numpy as np
from collections import deque


def collect_trajectories(env, policy, num_episodes=10):
    """Collect trajectories from a Gymnasium environment.

    Args:
        env: Gymnasium environment
        policy: callable(obs) -> action
        num_episodes: number of episodes to collect

    Returns:
        List of trajectory dicts with obs, actions, rewards arrays.
    """
    trajectories = []
    for _ in range(num_episodes):
        obs_list, act_list, rew_list = [], [], []
        obs, _ = env.reset()
        done = False
        while not done:
            action = policy(obs)
            next_obs, reward, terminated, truncated, _ = env.step(action)
            obs_list.append(obs)
            act_list.append(action)
            rew_list.append(reward)
            obs = next_obs
            done = terminated or truncated
        trajectories.append({
            "observations": np.array(obs_list, dtype=np.float32),
            "actions": np.array(act_list, dtype=np.float32),
            "rewards": np.array(rew_list, dtype=np.float32),
        })
    return trajectories


def train_world_model(
    rssm,
    trajectories,
    num_epochs=100,
    batch_size=16,
    seq_len=50,
    lr=3e-4,
):
    """Train RSSM on collected trajectories.

    Args:
        rssm: RSSM model
        trajectories: list of trajectory dicts
        num_epochs: training epochs
        batch_size: batch size
        seq_len: subsequence length for training
        lr: learning rate

    Returns:
        List of loss dictionaries per epoch.
    """
    optimizer = torch.optim.Adam(rssm.parameters(), lr=lr)
    history = []

    for epoch in range(num_epochs):
        epoch_losses = []

        # Sample random subsequences from trajectories
        for _ in range(max(1, len(trajectories) // batch_size)):
            obs_batch, act_batch, rew_batch = [], [], []
            for _ in range(batch_size):
                traj = trajectories[np.random.randint(len(trajectories))]
                max_start = len(traj["observations"]) - seq_len
                if max_start <= 0:
                    start = 0
                    sl = len(traj["observations"])
                else:
                    start = np.random.randint(max_start)
                    sl = seq_len
                obs_batch.append(traj["observations"][start:start + sl])
                act_batch.append(traj["actions"][start:start + sl])
                rew_batch.append(traj["rewards"][start:start + sl])

            # Pad to same length and convert to tensors
            actual_len = min(len(o) for o in obs_batch)
            obs_t = torch.tensor(
                np.stack([o[:actual_len] for o in obs_batch])
            )
            act_t = torch.tensor(
                np.stack([a[:actual_len] for a in act_batch])
            )
            rew_t = torch.tensor(
                np.stack([r[:actual_len] for r in rew_batch])
            ).unsqueeze(-1)

            # Ensure actions are float and correct shape
            if act_t.dim() == 2:
                act_t = act_t.unsqueeze(-1).float()

            # Forward pass
            output = rssm.observe(obs_t, act_t)
            loss, loss_dict = compute_elbo_loss(output, obs_t, rew_t)

            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(rssm.parameters(), 100.0)
            optimizer.step()

            epoch_losses.append(loss_dict)

        avg_loss = {
            k: np.mean([d[k] for d in epoch_losses])
            for k in epoch_losses[0]
        }
        history.append(avg_loss)

        if (epoch + 1) % 20 == 0:
            print(
                f"Epoch {epoch+1}: total={avg_loss['total_loss']:.4f} "
                f"obs={avg_loss['obs_loss']:.4f} "
                f"kl={avg_loss['kl_loss']:.4f}"
            )

    return history

7. Connecting to the Discovery Workbench

In the Discovery Workbench architecture (Chapter 6), the world model serves as a simulation backend that slots into the experiment planning pipeline. When the Workbench's optimization module (Chapter 45) or experiment designer (Chapter 46) needs to evaluate a candidate action sequence, it queries the world model rather than the real environment. The RSSM's imagine() method is the Workbench API for this: pass an initial state and a sequence of proposed actions, receive predicted observations and rewards. The Workbench tracks the model's prediction uncertainty (via prior-posterior KL divergence) and automatically triggers real experiments when the world model's confidence drops below a configurable threshold.

Exercises

1. (Conceptual) Explain why the RSSM needs both a prior and a posterior distribution over the stochastic state. What would go wrong if we trained only with the posterior and then tried to use the model for planning?
2. (Coding) Modify the RSSM to use Gaussian stochastic states instead of categorical ones. Train both versions on CartPole trajectories and compare the KL divergence curves. Do you observe posterior collapse with the Gaussian version?
3. (Analysis) The ELBO objective has a reconstruction term and a KL term. Plot how the balance between these two terms (controlled by \(\beta\)) affects one-step and five-step prediction accuracy. What value of \(\beta\) gives the best five-step predictions? Is it the same as the value that gives the best one-step predictions?
4. (Research) DreamerV3 uses 32 categorical dimensions with 32 classes each, giving a stochastic state space of size \(32^{32} \approx 10^{48}\). Is this capacity necessary for simple environments like CartPole? Experiment with smaller stochastic states (e.g., 8 categoricals with 8 classes) and measure the effect on prediction quality.