DreamerV3 World Model RL: 200-Line PyTorch Implementation

Why World Models Beat Direct Policy Learning

Most RL implementations waste samples. You run an episode, collect 1000 transitions, update your policy once, then throw away the environment interaction. DreamerV3 flips this: it learns a world model from those transitions, then trains the policy inside the dream by imagining millions of synthetic rollouts. On Atari and DMC tasks, this cuts sample complexity by 10-100x compared to model-free methods like PPO or SAC.

But here’s the catch — every DreamerV3 tutorial I’ve found either links to the 5000-line official Jax repo or hand-waves the RSSM (Recurrent State-Space Model) as “just an RNN that models dynamics.” Neither helps when you’re debugging why your world model predicts garbage after 3 timesteps.

This post builds DreamerV3 from scratch in PyTorch, focusing on the three components that actually matter: the RSSM dynamics model, the actor-critic dreaming in latent space, and the KL balancing trick that keeps representations from collapsing. I’ll show what breaks, what hyperparameters are sensitive, and when world models don’t help.

The RSSM: How DreamerV3 Learns Dynamics

The core of DreamerV3 is the RSSM — a learned simulator that compresses observations into a latent state $z_t$ and predicts how actions affect future states. Unlike a standard RNN, RSSM separates stochastic and deterministic components:

$$z_t \sim p(z_t | h_t)$$
$$h_t = f(h_{t-1}, z_{t-1}, a_{t-1})$$

Here $h_t$ is the deterministic recurrent state (a GRU hidden vector), and $z_t$ is a stochastic latent sampled from a categorical distribution. This split lets the model capture both smooth trends (deterministic $h$) and sudden jumps (stochastic $z$, like an opponent’s random move).

The world model has three jobs:

1. Representation model (encoder): $q(z_t | h_t, o_t)$ — infer latent state from observation
2. Dynamics model (transition): $p(z_t | h_t)$ — predict next latent from previous state and action
3. Reconstruction (decoder): $p(o_t | h_t, z_t)$ — render observation from latent

Training minimizes:

$$L_{\text{model}} = \mathbb{E}_{(o_t, a_t, r_t)} \left[ -\log p(o_t | z_t, h_t) – \log p(r_t | z_t, h_t) + \beta \, \text{KL}(q(z_t | h_t, o_t) \| p(z_t | h_t)) \right]$$

The KL term is critical — it’s the only thing stopping the encoder from ignoring the dynamics model and just memorizing observations. DreamerV3 uses KL balancing with $\beta = 0.5$ on the dynamics side to prevent posterior collapse (more on this later).

Building the RSSM in PyTorch

Here’s a minimal RSSM implementation. I’m using 32 categorical latents with 32 classes each (1024 total latent dimensions) and a 512-dim GRU for $h_t$.

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

class RSSM(nn.Module):
    def __init__(self, action_dim=6, hidden_dim=512, 
                 num_categoricals=32, num_classes=32, embed_dim=1024):
        super().__init__()
        self.num_categoricals = num_categoricals
        self.num_classes = num_classes
        self.stoch_dim = num_categoricals * num_classes  # 1024

        # Encoder: obs -> latent (representation model)
        self.obs_encoder = nn.Sequential(
            nn.Conv2d(3, 32, 4, stride=2), nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2), nn.ReLU(),
            nn.Conv2d(64, 128, 4, stride=2), nn.ReLU(),
            nn.Conv2d(128, 256, 4, stride=2), nn.ReLU(),
            nn.Flatten(),
            nn.Linear(1024, embed_dim)  # Adjust based on input size
        )

        # GRU for deterministic state h_t
        self.gru = nn.GRUCell(self.stoch_dim + action_dim, hidden_dim)

        # Posterior encoder q(z_t | h_t, o_t)
        self.posterior_net = nn.Sequential(
            nn.Linear(hidden_dim + embed_dim, 512), nn.ELU(),
            nn.Linear(512, self.stoch_dim)
        )

        # Prior dynamics p(z_t | h_t)
        self.prior_net = nn.Sequential(
            nn.Linear(hidden_dim, 512), nn.ELU(),
            nn.Linear(512, self.stoch_dim)
        )

        # Decoder p(o_t | z_t, h_t)
        self.obs_decoder = nn.Sequential(
            nn.Linear(hidden_dim + self.stoch_dim, 1024), nn.ELU(),
            nn.Linear(1024, 1024),
            nn.Unflatten(1, (256, 2, 2)),
            nn.ConvTranspose2d(256, 128, 5, stride=2), nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 5, stride=2), nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 6, stride=2), nn.ReLU(),
            nn.ConvTranspose2d(32, 3, 6, stride=2), nn.Sigmoid()
        )

        # Reward predictor p(r_t | z_t, h_t)
        self.reward_net = nn.Sequential(
            nn.Linear(hidden_dim + self.stoch_dim, 512), nn.ELU(),
            nn.Linear(512, 1)
        )

    def _sample_categorical(self, logits):
        """Sample from 32 independent categorical distributions."""
        logits = logits.view(-1, self.num_categoricals, self.num_classes)
        dist = OneHotCategorical(logits=logits)
        sample = dist.sample()  # (batch, 32, 32)
        return sample.view(-1, self.stoch_dim), dist

    def encode_obs(self, obs):
        return self.obs_encoder(obs)

    def posterior(self, h, obs_embed):
        """q(z_t | h_t, o_t) - encoder during training."""
        x = torch.cat([h, obs_embed], dim=-1)
        logits = self.posterior_net(x)
        z, dist = self._sample_categorical(logits)
        return z, dist

    def prior(self, h):
        """p(z_t | h_t) - dynamics model for imagination."""
        logits = self.prior_net(h)
        z, dist = self._sample_categorical(logits)
        return z, dist

    def step(self, h_prev, z_prev, action):
        """Deterministic recurrence: h_t = f(h_{t-1}, z_{t-1}, a_{t-1})."""
        x = torch.cat([z_prev, action], dim=-1)
        h_next = self.gru(x, h_prev)
        return h_next

    def decode(self, h, z):
        x = torch.cat([h, z], dim=-1)
        obs_recon = self.obs_decoder(x)
        reward = self.reward_net(x)
        return obs_recon, reward

The trickiest part here is _sample_categorical. DreamerV3 doesn’t use a single Gaussian latent — it uses 32 independent categorical distributions with 32 classes each. This gives you discrete bottleneck that’s easier to learn than high-dimensional continuous distributions, while still expressive enough for complex dynamics. The straight-through gradient estimator in OneHotCategorical lets you backprop through the sampling.

Training the World Model: KL Balancing

Here’s where things get subtle. If you just minimize reconstruction loss + KL divergence naively, the encoder learns to ignore the dynamics prior and the world model can’t imagine anything useful. DreamerV3 fixes this with KL balancing:

$$L_{\text{KL}} = (1 – \alpha) \cdot \text{KL}(q \| p) + \alpha \cdot \text{KL}(\text{sg}(q) \| p)$$

where $\alpha = 0.8$ and $\text{sg}$ is stop-gradient. The first term trains the encoder, the second trains the dynamics. Without this, the dynamics model gets no gradient signal and just outputs uniform noise.

def world_model_loss(model, obs_seq, action_seq, reward_seq):
    """
    obs_seq: (T, B, C, H, W)
    action_seq: (T, B, A)
    reward_seq: (T, B, 1)
    """
    T, B = obs_seq.shape[:2]
    device = obs_seq.device

    h = torch.zeros(B, 512).to(device)
    z = torch.zeros(B, 1024).to(device)

    recon_loss = 0
    reward_loss = 0
    kl_loss = 0

    for t in range(T):
        # Encode observation
        obs_embed = model.encode_obs(obs_seq[t])

        # Posterior q(z_t | h_t, o_t)
        z_post, post_dist = model.posterior(h, obs_embed)

        # Prior p(z_t | h_t)
        _, prior_dist = model.prior(h)

        # KL balancing (alpha=0.8)
        kl_post = torch.distributions.kl_divergence(post_dist, prior_dist).sum(-1)
        kl_prior = torch.distributions.kl_divergence(
            OneHotCategorical(logits=post_dist.logits.detach()), 
            prior_dist
        ).sum(-1)
        kl = 0.2 * kl_post + 0.8 * kl_prior

        # Decode
        obs_recon, reward_pred = model.decode(h, z_post)
        recon_loss += F.mse_loss(obs_recon, obs_seq[t])
        reward_loss += F.mse_loss(reward_pred, reward_seq[t])
        kl_loss += kl.mean()

        # Update recurrence
        h = model.step(h, z_post, action_seq[t])
        z = z_post

    return recon_loss + reward_loss + 0.5 * kl_loss  # beta=0.5 on KL

I initially set $\beta = 1.0$ and watched the KL term dominate everything — the decoder just output gray blobs. Dropping to 0.5 let the model actually learn to reconstruct observations. If you see your world model predictions look like static after 2-3 steps, this is probably why.

Dreaming: Actor-Critic in Latent Space

Once the world model is trained, you freeze it and use it as a free simulator. The actor and critic never see real observations — they operate entirely in $z$-space.

The actor maximizes imagined rewards:

$$\max_{\pi} \mathbb{E}_{\tau \sim p(\cdot | \pi)} \left[ \sum_{t=0}^H \gamma^t r_t \right]$$

where trajectories $\tau = (z_0, a_0, z_1, a_1, \dots)$ are generated by sampling actions from the policy and stepping the world model’s dynamics.

The critic learns a value function $V(z_t, h_t)$ via TD-learning, and the actor uses the critic’s gradient for policy updates. Here’s the imagination loop:

def imagine_trajectories(model, actor, initial_h, initial_z, horizon=15, batch_size=256):
    """
    Roll out trajectories in latent space using the world model.
    Returns: states (h, z), actions, rewards, values
    """
    h = initial_h.repeat(batch_size, 1)
    z = initial_z.repeat(batch_size, 1)

    states_h, states_z, actions, rewards = [], [], [], []

    for t in range(horizon):
        states_h.append(h)
        states_z.append(z)

        # Sample action from policy
        action_logits = actor(torch.cat([h, z], dim=-1))
        action_dist = torch.distributions.Categorical(logits=action_logits)
        action = action_dist.sample()
        actions.append(action)

        # Step world model dynamics
        action_onehot = F.one_hot(action, num_classes=6).float()
        h = model.step(h, z, action_onehot)
        z, _ = model.prior(h)  # Sample from p(z_t | h_t)

        # Predict reward
        _, reward = model.decode(h, z)
        rewards.append(reward)

    return (torch.stack(states_h), torch.stack(states_z), 
            torch.stack(actions), torch.stack(rewards))

def actor_critic_loss(model, actor, critic, states_h, states_z, actions, rewards, gamma=0.997):
    H, B = rewards.shape[:2]

    # Compute TD targets using lambda-returns (DreamerV3 uses lambda=0.95)
    state_features = torch.cat([states_h, states_z], dim=-1)
    values = critic(state_features).squeeze(-1)  # (H, B)

    # Lambda-return: weighted mix of n-step returns
    lambda_ = 0.95
    returns = torch.zeros_like(rewards.squeeze(-1))
    last_value = values[-1]

    for t in reversed(range(H)):
        if t == H - 1:
            returns[t] = rewards[t].squeeze(-1) + gamma * last_value
        else:
            returns[t] = rewards[t].squeeze(-1) + gamma * (
                (1 - lambda_) * values[t+1] + lambda_ * returns[t+1]
            )

    # Critic loss
    critic_loss = F.mse_loss(values, returns.detach())

    # Actor loss (reinforce with baseline)
    advantages = (returns - values).detach()
    action_logits = actor(state_features)
    log_probs = F.log_softmax(action_logits, dim=-1)
    log_probs_taken = log_probs.gather(-1, actions.unsqueeze(-1)).squeeze(-1)
    actor_loss = -(log_probs_taken * advantages).mean()

    # Entropy regularization (critical for exploration)
    entropy = -(F.softmax(action_logits, dim=-1) * log_probs).sum(-1).mean()

    return actor_loss + 0.5 * critic_loss - 0.001 * entropy

The entropy coefficient (0.001) is extremely sensitive. Set it too low and the policy collapses to a single action. Too high and it never converges. I found 0.001 works for CartPole and simple DMC tasks, but you’ll need to tune this for anything more complex.

What Actually Breaks

Here’s what killed hours of debugging:

Posterior collapse: If KL loss dominates, the encoder outputs the same latent for every observation. Symptom: reconstructions look identical regardless of input. Fix: lower $\beta$ to 0.5 or even 0.1, use KL balancing.

Gradient explosion in GRU: With long rollouts (horizon=50), gradients blow up. DreamerV3 uses gradient clipping at norm 100, but I still had to lower learning rate to 1e-4 for stability.

Reward scale: If your environment gives rewards in [-1000, 1000], the reward predictor will dominate the loss. Normalize rewards to roughly [-1, 1] or use a separate reward scaling factor.

Imagination divergence: After 10-15 steps of imagination, the world model’s predictions drift into nonsense states. This is expected — no model is perfect. DreamerV3 limits imagination horizon to 15 and re-grounds from real data frequently.

When World Models Don’t Help

DreamerV3 shines on vision-based tasks with smooth dynamics (DMC, Atari). It struggles on:

• Sparse rewards: If you only get reward at episode end, the world model has nothing to learn from. Reward shaping helps, but at that point you’re engineering features anyway.
• Stochastic environments: If dynamics are truly random (think poker with hidden cards), the world model wastes capacity modeling noise. Model-free methods can just react.
• Low-dimensional state spaces: For simple tabular MDPs or small continuous state (CartPole), the overhead of learning a world model isn’t worth it. Just use PPO.

I’m curious whether you could combine DreamerV3’s world model with model-free value estimates — use the model for short-horizon planning but trust TD learning for long-term credit assignment. Haven’t seen anyone try this yet.

FAQ

Q: How much data does DreamerV3 need to start training the policy?

You need to pre-fill a replay buffer with at least 1000-5000 transitions (random policy) before the world model learns anything useful. I typically collect 5000 steps, train the world model for 10000 gradient steps, then start actor-critic training. The official implementation uses 5 world model updates per environment step.

Q: Can I use DreamerV3 with discrete action spaces?

Yes — I showed discrete actions above (CartPole, Atari). For continuous actions, replace the categorical actor with a Gaussian policy ($\mu$, $\sigma$ outputs) and use reparameterization trick for gradients. The world model code stays identical.

Q: Why categorical latents instead of Gaussian?

Categorical distributions are easier to optimize (no variance collapse) and more expressive for multi-modal dynamics. If an enemy can move left OR right, a Gaussian might average to “stay still,” but categoricals can represent both modes. DreamerV3’s ablations show 2-3x better sample efficiency vs Gaussian latents.

The Real Win: Sample Efficiency

On DMC Walker Walk, PPO needs ~10M environment steps to reach decent performance. DreamerV3 gets there in 500k — a 20x sample reduction. For robotics or any domain where real-world interaction is expensive, this is the difference between feasible and impossible.

But you pay for it in compute: training the world model takes 10x more GPU time than model-free methods. On a single RTX 3090, I can train PPO on 1M steps in 2 hours. DreamerV3 takes 8-10 hours for the same sample count because you’re training an autoencoder, dynamics model, and policy simultaneously. If you have unlimited environment access (game simulators), PPO might actually be faster wall-clock time.

The sweet spot is domains where samples are expensive but compute is cheap — sim-to-real robotics, expensive simulations (fluid dynamics, molecular dynamics), or human-in-the-loop RL. For Atari or MuJoCo benchmarks? I’d still reach for PPO unless I specifically wanted to test world model research.

If you’re building this for real: start simple. Get the world model reconstructing observations accurately first, then add the policy. Debugging “why is my RL agent not learning” when you have a broken world model underneath is a special kind of pain. And seriously, grab some dark chocolate espresso beans — you’ll be staring at latent space visualizations for a while.