# DreamerV3 World Model (RSSM)

> Hafner 2023 - the agent learns a model of the environment and plans inside it. A deterministic recurrent path runs beside a categorical stochastic latent.

Pick when environment steps are expensive and imagined rollouts are cheaper than real ones. The state width is deterDim + stochDim x stochClasses, not deterDim.

- Category: Reinforcement Learning
- Paper: Hafner et al. (2023). Mastering Diverse Domains through World Models (https://arxiv.org/abs/2301.04104)
- Layers: 21
- Parameters: 13.71M
- Input shape (batchless): 3 × 64 × 64
- Output shape: 3 × 64 × 64
- Verifier verdict: warn
- Graph JSON: https://neurarch.com/templates/dreamer-v3/model.json
- Open on the canvas: https://neurarch.com/?template=dreamer-v3

## Structure

| # | Layer | Type | Parameters | Output shape |
| --- | --- | --- | --- | --- |
| 1 | Observation (64x64 RGB) | Input | shape=[3, 64, 64] | 3 × 64 × 64 |
| 2 | Encoder Conv 32 | Conv2D | outChannels=32, inChannels=3, kernelSize=4 | 32 × 32 × 32 |
| 3 | SiLU | Swish |  | 32 × 32 × 32 |
| 4 | Encoder Conv 64 | Conv2D | outChannels=64, inChannels=32, kernelSize=4 | 64 × 16 × 16 |
| 5 | SiLU | Swish |  | 64 × 16 × 16 |
| 6 | Encoder Conv 128 | Conv2D | outChannels=128, inChannels=64, kernelSize=4 | 128 × 8 × 8 |
| 7 | SiLU | Swish |  | 128 × 8 × 8 |
| 8 | Encoder Conv 256 | Conv2D | outChannels=256, inChannels=128, kernelSize=4 | 256 × 4 × 4 |
| 9 | SiLU | Swish |  | 256 × 4 × 4 |
| 10 | flatten | Flatten |  | 4096 |
| 11 | RSSM (deter 512 + 32x32 stoch) | RSSM (Dreamer) |  | 1536 |
| 12 | Decoder Projection | Linear | outFeatures=4096, inFeatures=1536 | 4096 |
| 13 | to [256, 4, 4] | Reshape | shape=[256, 4, 4] | 256 × 4 × 4 |
| 14 | Decoder Deconv 128 | TransposeConv2D | outChannels=128, kernelSize=4, stride=2 | 128 × 8 × 8 |
| 15 | SiLU | Swish |  | 128 × 8 × 8 |
| 16 | Decoder Deconv 64 | TransposeConv2D | outChannels=64, kernelSize=4, stride=2 | 64 × 16 × 16 |
| 17 | SiLU | Swish |  | 64 × 16 × 16 |
| 18 | Decoder Deconv 32 | TransposeConv2D | outChannels=32, kernelSize=4, stride=2 | 32 × 32 × 32 |
| 19 | SiLU | Swish |  | 32 × 32 × 32 |
| 20 | Decoder Deconv 3 | TransposeConv2D | outChannels=3, kernelSize=4, stride=2 | 3 × 64 × 64 |
| 21 | reconstructed frame | Output |  | 3 × 64 × 64 |

## Verifier findings

- **warn** `deep-no-residual`: 9 conv/linear layers detected but no residual (Add/Skip) layers. Networks deeper than 8 layers are highly prone to vanishing gradients without skip connections. Fix: Add Residual or Add layers every 2-4 layers (ResNet-style). For transformers, use the built-in TransformerBlock which includes residuals.
- **info** `deep-no-norm`: 19 layers with no BatchNorm, LayerNorm, or GroupNorm. Without normalization, activations can explode or vanish across layers, causing slow or unstable training. Fix: Add BatchNorm after Conv2d (CV tasks), LayerNorm after attention/FFN (NLP/LLM), or GroupNorm for small batch sizes.

## Exported PyTorch (first 46 lines)

```python
# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class RSSM(nn.Module):
    """Dreamer's recurrent state-space model: a deterministic recurrent path
    beside a categorical stochastic latent. Everything downstream reads their
    CONCATENATION, so the state width is deter_dim + stoch_dim * stoch_classes,
    which is the number a plain GRU would get wrong."""

    def __init__(self, embed_dim: int, deter_dim: int = 512, stoch_dim: int = 32,
                 stoch_classes: int = 32, hidden_dim: int = 512):
        super().__init__()
        self.deter_dim, self.stoch_dim, self.stoch_classes = deter_dim, stoch_dim, stoch_classes
        flat = stoch_dim * stoch_classes
        self.cell = nn.GRUCell(flat, deter_dim)
        self.prior = nn.Sequential(
            nn.Linear(deter_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, flat))
        self.post = nn.Sequential(
            nn.Linear(deter_dim + embed_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, flat))

    def forward(self, embed: torch.Tensor, state=None) -> torch.Tensor:
        b = embed.size(0)
        flat = self.stoch_dim * self.stoch_classes
        if state is None:
            deter = embed.new_zeros(b, self.deter_dim)
            stoch = embed.new_zeros(b, flat)
        else:
            deter, stoch = state
        deter = self.cell(stoch, deter)
        # self.prior(deter) is the imagination branch: it predicts the same
        # latent WITHOUT an observation, and the KL between the two is the
        # world model's training signal. Observed steps use the posterior.
        logits = self.post(torch.cat([deter, embed], dim=-1))
        stoch = logits.view(b, self.stoch_dim, self.stoch_classes).softmax(-1).reshape(b, flat)
        return torch.cat([deter, stoch], dim=-1)


class DreamerV3WorldModelRSSM(nn.Module):
    def __init__(self):
        super().__init__()
```

## Machine access

- Every architecture: https://neurarch.com/a/index.json
- Verify a graph of your own: `POST https://www.neurarch.com/api/v1/check` (see https://neurarch.com/developer.html)
- MCP server, so an agent edits the graph with the checks in the loop: https://neurarch.com/docs/mcp.md