Understanding RSSM Through Code (6): Default Config, Four Formulas, and the Code↔Math↔Semantics Map

Understanding RSSM Through Code · 第 6 篇 / 共 6 篇

Series contents (you are on part 6, bolded; prev/next at the bottom):

1. (1) Where RSSM Sits & the Stochastic State
2. (2) Prior/Posterior, Straight-Through & unimix
3. (3) _core(), deter=8192 & Block GRU
4. (4) KL Balancing, Free Nats & Final KL
5. (5) Imagine, Observe vs. Imagine & Reset 6. (6) Default Config, Four Formulas & the Map

25. Default Config: Split into “RSSM Architecture” and “World-Model Training”

To avoid mixing “RSSM architecture parameters” with “world-model / agent training config,” the configuration is split into two tables below.

RSSM architecture

RSSM / world-model training

Key points:

• deter / stoch / classes / blocks are the RSSM architecture;
• unimix / free_nats are used inside the RSSM to construct distributions and KL;
• dyn_scale / rep_scale / imag_length are agent / world-model training config, not part of the RSSM structure itself.

These are default configurations, not fixed parameters for all DreamerV3 model scales. Different model scales (1M to 400M) change deter from 512 to 12288 and classes from 4 to 96.

8192 / 32 / 64 should be understood as a specific instance under DreamerV3’s default configuration, not RSSM’s fixed structure.

26. Compressing the Entire RSSM Into Four Formulas

① Deterministic transition

h_t = f(h_{t-1}, z_{t-1}, a_{t-1})    (_core(), no observation)

Historical state and action determine the deterministic state.

② Prior

p(z_t | h_t)

Depends only on deterministic state. Handles imagination.

③ Posterior

q(z_t | h_t, o_t)

Uses observation to correct the latent state. Handles state inference on real trajectories.

④ KL balancing

L_dyn = KL[sg(q(z_t|h_t,o_t)) || p(z_t|h_t)]
L_rep = KL[q(z_t|h_t,o_t) || sg(p(z_t|h_t))]
L_KL  = dyn_scale × L_dyn + rep_scale × L_rep

KL also has free_nats = max(KL, 1) applied before entering the final loss; dyn_scale / rep_scale come from the agent’s loss config, not from inside the RSSM class.

27. Code ↔ Math ↔ Semantics Mapping Table

To tie together the previous twenty-some sections, here is a table aligning “source symbol / math symbol / semantics”:

This table is the “index” of the whole article: any source symbol can be traced here to its mathematical meaning and role.

28. Re-Understanding RSSM From Code

Without looking at code, RSSM is easily understood as “a GRU + a Gaussian.”

But DreamerV3’s actual implementation is far richer:

                 ┌──────────────────────┐
                 │      RSSM            │
 action ────────►│    Block GRU         │
 z_{t-1} ───────►│        │             │
 h_{t-1} ───────►│       h_t            │
                 │      /   \            │
                 │  Prior   Posterior    │
                 │    │        ▲         │
                 │    │      token_t     │
                 │    ▼        ▼         │
                 │   p(z)     q(z)       │
                 │     └── KL ──┘        │
                 └──────────────────────┘

From the perspective of DreamerV3’s implementation, RSSM can be understood as simultaneously playing three functional roles:

Memory: deter preserves history through recurrent dynamics. State estimation: posterior corrects current latent state using observations. Prediction: prior predicts future latent states without observations.

Note: The above “RSSM plays three roles” is a functional generalization, not a strict definition of the RSSM. Precisely, it should be phrased as “from the … implementation perspective, it can be understood as …,” to avoid taking an implementation-level functional division as the mathematical definition of the RSSM.

During training, RSSM leverages real observations to learn latent dynamics; during imagination, it drops observations and relies solely on the prior to rollout forward in latent space.

This is the foundation enabling DreamerV3’s latent imagination.

29. Back to Source Code: Why Do These Details Matter?

From the paper alone, RSSM might be understood as RNN + latent distribution. But reading the code reveals extensive engineering:

Categorical latent + OneHot distribution + Unimix (p' = (1-ε)p + ε/K)
+ Block GRU (BlockLinear structured parameterization) + RMSNorm
+ Prior / Posterior + KL balancing (gradient routing: sg(q)||p and q||sg(p))
+ Stop Gradient + Free Nats (max(KL, 1)) + Scan (observe → nj.scan → _observe)

Individually none are complex. What matters is how they form a complete pipeline:

real observation → Encoder → observation token → RSSM → latent state
→ imagination → future latent states → reward / value / policy

From a code perspective, the real problem RSSM solves isn’t “how to predict the next observation.”

It’s:

What RSSM truly solves is not how to predict the next observation, but how to learn a latent state that can both be corrected by observations and, without observations, roll autonomously into the future on its own.

This is the core of DreamerV3’s RSSM.

Source Code References

• DreamerV3 GitHub Repository
• dreamerv3/rssm.py
• dreamerv3/configs.yaml