Architectures / Reinforcement Learning

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

From Hafner et al. (2023). Mastering Diverse Domains through World Models. This page is the graph, not the PDF: open it, edit it, verify it, export it.

Layers

Parameters

13.71M

Input

3 × 64 × 64

Output

3 × 64 × 64

Verifier

1 advisory

Every number on this page is computed from the graph by the same functions the app runs, not written by hand.

Open DreamerV3 World Model (RSSM) on the canvas Free, no account needed

When to pick it

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

Structure

21 layers. Output shapes are propagated from the input shape, batch dimension excluded.

	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

What the verifier says

The same 43 structural checks that run on every edit in the app, on this graph.

warn9 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.
deep-no-residual

info19 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.
deep-no-norm

The PyTorch it exports

Generated from the graph above. First 46 lines; the app exports the whole file, plus the training loop, the data contract and a deploy bundle.

# 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__()

For agents

This architecture is machine-readable end to end. An agent can list the set, fetch this graph, edit it, and have the edit verified before any GPU time is spent.

Markdown twin Index JSON This graph as JSON MCP server POST /api/v1/check