状态空间模型 on WorldSense 技术笔记

从代码理解 RSSM（六）：默认配置、四条公式与代码↔数学↔语义对照

Mon, 24 Aug 2026 00:00:00 +0000

《从代码理解 RSSM》系列 · 第 6 篇 / 共 6 篇

系列目录（当前在第 6 篇，已加粗；上/下一篇见文末导航）：

（一）RSSM 的位置与 stochastic 状态

（二）先验/后验、straight-through 与 unimix

（三）_core()、deter=8192 与 Block GRU

（四）KL balancing、Free Nats 与最终 KL

（五）Imagine、Observe/Imagine 区别与 Reset 6. （六）默认配置、四条公式与对照表

二十五、默认配置：拆成"RSSM 架构"与"World-model 训练"两张表

为了避免把"RSSM 架构参数"和"world-model / agent 训练配置"混在一起，下面把配置拆成两张表。

RSSM architecture（RSSM 架构）

参数	默认值	含义
`deter`	8192	deterministic state 维度（会被 blocks 切分）
`stoch`	32	categorical variable 数量
`classes`	64	每个 categorical variable 的类别数
`blocks`	8	Block GRU 分组数量（8192 / 8 = 1024）
`hidden`	1024	RSSM 内部 MLP hidden dimension
`imglayers`	2	prior 网络层数
`obslayers`	1	posterior 网络层数
`dynlayers`	1	dynamics 网络层数

RSSM / world-model training（RSSM 与世界模型训练配置）

参数	默认值	含义	归属
`unimix`	0.01	categorical distribution 的均匀混合比例	RSSM 内部分布
`free_nats`	1.0	KL 的 loss floor（`max(KL, 1)`）	RSSM 内部
`dyn_scale`	1.0	dynamics KL 在最终 loss 中的权重	agent loss scale
`rep_scale`	0.1	representation KL 在最终 loss 中的权重	agent loss scale
`imag_length`	15	imagination rollout 的训练步数	agent 训练配置（非 RSSM 结构）

关键点：

从代码理解 RSSM（五）：Imagine、Observe/Imagine 区别、序列训练与 Reset

Sun, 23 Aug 2026 00:00:00 +0000

《从代码理解 RSSM》系列 · 第 5 篇 / 共 6 篇

系列目录（当前在第 5 篇，已加粗；上/下一篇见文末导航）：

（一）RSSM 的位置与 stochastic 状态

（二）先验/后验、straight-through 与 unimix

（三）_core()、deter=8192 与 Block GRU

（四）KL balancing、Free Nats 与最终 KL 5. （五）Imagine、Observe/Imagine 区别与 Reset

（六）默认配置、四条公式与对照表

二十、Imagine 阶段：没有 observation，RSSM 怎么跑？

这是 RSSM 最漂亮的地方。

训练真实轨迹时需要 observation 来计算 q(z_t|h_t,o_t)。但 imagination 阶段没有真实 observation。

直接使用 p(z_t|h_t)，也就是 prior：

z_t ~ p(z_t | h_t)

源码中的 imagine() 内部同样先走 _core()，再用 prior 产生 latent：

deter = self._core(
 carry['deter'],
 carry['stoch'],
 actemb
)

logit = self._prior(deter)

stoch = self._dist(logit).sample(...)

注意：

从代码理解 RSSM（四）：KL balancing、Free Nats 与最终 KL 组合

Sat, 22 Aug 2026 00:00:00 +0000

《从代码理解 RSSM》系列 · 第 4 篇 / 共 6 篇

系列目录（当前在第 4 篇，已加粗；上/下一篇见文末导航）：

（一）RSSM 的位置与 stochastic 状态

（二）先验/后验、straight-through 与 unimix

（三）_core()、deter=8192 与 Block GRU 4. （四）KL balancing、Free Nats 与最终 KL

（五）Imagine、Observe/Imagine 区别与 Reset

（六）默认配置、四条公式与对照表

十四、Observe 阶段完整数据流（Observe / Imagine 同一框架）

到这里，可以把整个 RSSM 的状态转移框架画出来，并且把 Observe 与 Imagine 放在同一个 _core() 转移框架里：

 ┌──────────────────────────┐
 │ RSSM transition │
 │ │
(h_{t-1},z_{t-1}) + a_{t-1}
 │
 ▼
 _core()
 │
 ▼
 h_t
 │
 ┌──────┴──────┐
 │ │
 ▼ ▼
 Prior Observation
 p(z_t | h_t) │
 │ ▼
 │ Posterior
 │ q(z_t|h_t,o_t)
 │ │
 │ ▼
 │ z_t
 └───────┬─────┘
 │
 ▼
 latent state
 │
 ▼
 next step

在旁边注明两种用法，其实这就是全文最重要的两句话：

从代码理解 RSSM（三）：确定性转移 _core()、deter=8192 与 Block GRU

Fri, 21 Aug 2026 00:00:00 +0000

《从代码理解 RSSM》系列 · 第 3 篇 / 共 6 篇

系列目录（当前在第 3 篇，已加粗；上/下一篇见文末导航）：

（一）RSSM 的位置与 stochastic 状态

（二）先验/后验、straight-through 与 unimix 3. （三）_core()、deter=8192 与 Block GRU

（四）KL balancing、Free Nats 与最终 KL

（五）Imagine、Observe/Imagine 区别与 Reset

（六）默认配置、四条公式与对照表

九、真正的 deterministic transition：`_core()`

这是整个 rssm.py 中最值得读的部分。

源码并不是简单调用：

nn.GRUCell(...)

而是自己构造了一个 block-wise GRU。

核心输入包括：

deter
stoch
action

首先：

stoch = stoch.reshape((stoch.shape[0], -1))

也就是：

[B, 32, 64]
 │
 ▼
[B, 2048]

然后把 deter、stoch、action 各自通过一个输入映射（注意：这里的"分别 Linear 映射再 concat"是概念示意图，真正的核心在于 BlockLinear 的结构化参数化，见下一节）：

deter ──► Linear ──┐
 │
stoch ──► Linear ──┼──► concat ──► Block GRU
 │
action ─► Linear ──┘

注意：

从代码理解 RSSM（二）：先验/后验、straight-through 采样与 unimix

Thu, 20 Aug 2026 00:00:00 +0000

《从代码理解 RSSM》系列 · 第 2 篇 / 共 6 篇

系列目录（当前在第 2 篇，已加粗；上/下一篇见文末导航）：

（一）RSSM 的位置与 stochastic 状态 2. （二）先验/后验、straight-through 与 unimix

（三）_core()、deter=8192 与 Block GRU

（四）KL balancing、Free Nats 与最终 KL

（五）Imagine、Observe/Imagine 区别与 Reset

（六）默认配置、四条公式与对照表

四、把源码翻译成数学公式（注意时间索引）

DreamerV3 的递推关系可以写成：

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

然后分别计算：

p(z_t | h_t) [prior，不读 observation]
q(z_t | h_t, o_t) [posterior，读 observation]

其中：

p：prior；
q：posterior；
o_t：当前 observation；
h_t：deterministic state。

一个非常重要、但初学者很容易卡住的时间索引说明：

h_t = f(h_{t-1}, z_{t-1}, a_{t-1}) 中的 h_t 已经融合了截至 t-1 的 latent / action 历史，但它并不包含当前 o_t。当前 observation 只在 posterior 分支中用于推断 z_t，而不进入 deterministic transition。

从代码理解 RSSM（一）：RSSM 在 DreamerV3 中的位置与 stochastic 状态

Wed, 19 Aug 2026 00:00:00 +0000

《从代码理解 RSSM》系列 · 第 1 篇 / 共 6 篇

系列目录（当前在第 1 篇，已加粗；上/下一篇见文末导航）： 1. （一）RSSM 的位置与 stochastic 状态 2. （二）先验/后验、straight-through 与 unimix 3. （三）_core()、deter=8192 与 Block GRU 4. （四）KL balancing、Free Nats 与最终 KL 5. （五）Imagine、Observe/Imagine 区别与 Reset 6. （六）默认配置、四条公式与对照表

源码阅读提示

本文有意不从"标准 RSSM 伪代码"开始，而是按照 rssm.py 的实际执行路径展开。阅读时可以重点关注下面几个函数：
observe() → sequence-level wrapper，沿时间维调用 _observe()
_observe() → 单个时间步的状态推断（真实轨迹）
_core() → 单个时间步的 deterministic dynamics
imagine() → 无 observation 的 latent rollout
这里最容易混淆的一点是：observe() 并不是一个"处理单步 transition"的函数。它只是 sequence-level 的封装，内部通过 Ninjax 的 nj.scan(...) 沿着时间维度反复调用 _observe()；真正处理单个时间步的是 _observe() 与 _core()。把 observe() / _observe() / _core() / imagine() 这一层调用关系理清，基本就能理解 DreamerV3 RSSM 的主体。