Understanding RSSM Through Code · 第 2 篇 / 共 6 篇
Series contents (you are on part 2, bolded; prev/next at the bottom):
4. Translating the Source Code into Mathematical Formulas (Watch the Time Index)
DreamerV3’s recurrence relation:
h_t = f(h_{t-1}, z_{t-1}, a_{t-1})
Then compute separately:
p(z_t | h_t) [prior, no observation]
q(z_t | h_t, o_t) [posterior, with observation]
Where p is the prior, q is the posterior, o_t is the current observation, and h_t is the deterministic state.
A time-index note that is critical but where beginners often get stuck:
In
h_t = f(h_{t-1}, z_{t-1}, a_{t-1}), theh_talready fuses the latent/action history up tot-1, but it does not contain the currento_t. The current observation only enters the posterior branch to inferz_t; it never enters the deterministic transition.
In other words, beginners often ask: “Since h_t is the current state, why doesn’t it contain o_t?” The answer: h_t here is the “current deterministic memory derived from history (h_{t-1}, z_{t-1}, a_{t-1})”; o_t is a separate stream of information at the same timestep, used only by the posterior to correct z_t, and never fed into _core(). Understanding this index is the key to reading Dreamer’s RSSM.
The entire RSSM can be understood as:
┌─────────────────────────────────────────────────┐
│ (h_{t-1}, z_{t-1}, a_{t-1}) → h_t → │
│ ┌──────────┐ │
│ │ p(z_t|h_t)│ │
│ │ q(z_t|h_t,│ │
│ │ o_t) │ │
│ └──────────┘ │
└─────────────────────────────────────────────────┘
This is also the key to understanding imagination and KL loss later.
5. Why Two Distributions — Prior and Posterior?
These two distributions actually solve two different problems.
1. Posterior: After seeing the observation, what do I think the current state is?
q(z_t | h_t, o_t)
It can see both historical information h_t and the current observation o_t, so it has more information.
During training on real trajectories, we use the posterior to get the stochastic state:
logit = posterior(deter, token)
stoch = sample(logit)
The posterior can be understood as:
“After seeing the real world, the estimate of the current latent state.”
2. Prior: If there were no observation, what do I predict the state would be?
p(z_t | h_t)
It only sees h_t, so it doesn’t know the current real observation.
It expresses:
“Based only on historical state and action, what I predict the next latent state will be.”
This is exactly the capability needed during imagination.
6. How Is Categorical Latent Sampled? (Straight-Through)
Both posterior and prior ultimately output logits:
logits
[B, 32, 64]
Each [64] corresponds to a categorical distribution for the i-th categorical variable (64 class probabilities).
In the source code:
def _dist(self, logits):
out = embodied.jax.outs.OneHot(
logits,
self.unimix
)
out = embodied.jax.outs.Agg(
out, 1, jnp.sum
)
return out
This is not a Gaussian distribution, but a one-hot categorical distribution.
A precise description of the sampling is needed here, to avoid a common misunderstanding:
z_tis NOT “first sample an integer[B, 32], then convert to one-hot.”
More accurately:
- Each stochastic variable corresponds to a 64-class categorical distribution;
- The sampled stochastic state participates in downstream networks in the forward pass as a one-hot categorical representation (a one-hot-like representation);
- The implementation uses a straight-through estimator so that discrete sampling can participate in backpropagation — i.e., the forward pass uses the discrete one-hot sample, while the backward pass approximates the gradient back to the network parameters that produced the logits.
So in the forward computation z_t is a one-hot-like tensor of shape [B, 32, 64]. It “looks like one-hot,” but its gradient path is straight-through, not the “round to integer then one-hot” operation that would completely cut the gradient.
logits
│
▼
categorical distribution (each [64] a 64-class dist)
│
▼
sample ── forward: one-hot-like representation
│ backward: straight-through estimator
▼
stoch.shape = [B, 32, 64]
7. What Does unimix=0.01 Do?
DreamerV3’s categorical distribution has another easily overlooked parameter:
unimix: 0.01
Its purpose is to keep a small portion of uniform distribution in the categorical distribution.
If the network is already very confident:
class 7: 0.9999
other: 0.0001
The distribution becomes very sharp.
unimix mixes it slightly with a uniform distribution, ensuring every class retains some probability. Written as a formula:
U_i = 1 / K (K = number of classes, here K = 64)
p'_i = (1 - ε) × p_i + ε × U_i
= (1 - ε) × p_i + ε / 64
Where ε = 0.01.
So it adds a small probability floor ε/K to every single class, rather than just “blunting the distribution overall.”
Its role has two layers:
- Improving categorical latent training stability, preventing the distribution from becoming too sharp too early;
- Providing numerical stability for categorical KL / entropy computations (avoiding log(0) and division-by-zero caused by zero probabilities).
8. Why Does DreamerV3 Make the Stochastic State So Many Categorical Variables?
This is an important design in DreamerV3’s latent representation.
Default: stoch=32, classes=64
Not z ∈ R^32, but:
z = [
categorical(64),
categorical(64),
...
categorical(64)
]
× 32
An important benefit is forming a very rich discrete combinatorial space.
From the perspective of the combinatorial space, 32 variables with 64 classes each correspond to:
64^32
discrete combinations.
Mathematically, this factorized categorical distribution is the product of independent per-variable distributions:
p(z_t | h_t) = Π_{i=1}^{32} p(z_t^i | h_t)
That is, the joint distribution is the product of 32 independent categorical distributions.
An important qualification is needed:
64^32describes the size of the combinatorial space that this factorized categorical distribution “covers.” It does NOT mean the model explicitly enumerates64^32states. The model merely parameterizes this space using32 × 64logits (a factorized categorical distribution); it does not maintain a lookup table or state set of size64^32.
So a more precise statement is: “this structure provides combinatorial representational capacity on the order of 64^32,” not “the model uses a giant discrete state space with 64^32 explicit states.” This is why DreamerV3 can express complex environment states in a relatively compact latent space.
Comments