PlaNet 中的 RSSM 与 Overshooting：技术深度解析

一句话定位：深入解析 PlaNet 中 RSSM 的概率图模型结构与overshooting技术的数学动机与实现细节，是理解后续 Dreamer 系列工作的核心基础。

前置依赖：PlaNet 基础架构（1-PlaNet.md）、变分推断（ELBO）、KL 散度、概率图模型。

一、RSSM 的概率图模型视角

1.1 从有向图模型到 RSSM

标准状态空间模型（SSM） 通常定义为：

其中是隐状态，观测从条件生成。这种模型在语音识别（HMM）、机器人状态估计（Kalman Filter）等领域广泛应用。

RSSM 的核心创新：将每个时刻的状态分解为两部分：

确定性状态 （RNN 输出）：解决了纯随机模型（如标准 SSM）的长期依赖问题；
随机状态 ：捕捉部分可观测性（partial observability）和随机性。

1.2 RSSM 的因子图分解

RSSM 的联合分布按以下方式分解：

初 始 状 态 观 测 模 型 先 验 奖 励 模 型 确 定 性 转 移

写成更紧凑的形式：

1.3 隐状态后验的因子化

引入变分后验，使用均值场近似（mean-field approximation）进行因子分解：

实际上在实现中，由 RNN 确定性计算（给定），所以不需要近似，直接等于由和初始条件确定的 Dirac 分布：

因此后验只对进行近似：

二、状态转移的双路径架构

2.1 确定性路径的数学分析

转移方程：

其中表示拼接操作。

为什么需要确定性路径？

考虑纯随机模型（如标准 SSM）的前向传播：

设，在步后，的分布由个随机变量的非线性变换决定。根据中心极限定理，步后的分布趋向于高度混合/高方差，导致信息传递效率极低。

确定性路径提供了一条无损的信息传递通道：是由 确定性计算而来的，不引入额外随机性。这意味着 RNN 可以有效保留长期依赖，而不受复合随机性的影响。

2.2 随机路径的建模能力

先验与后验的分工：

	先验	后验
输入	仅（确定性状态）	+ （观测）
用途	推理/规划时的预测	训练时的校正
信息	无观测信息	包含当前观测的校正信息
性质	预测分布	真实后验的变分近似

先验的物理含义：在没有观测输入的情况下，包含了到时刻为止的全部历史信息。因此先验是对在已有历史信息条件下的下一隐状态的预测。

后验的物理含义：引入真实观测后，我们希望用来”校正”对的估计。后验就是这个校正后的分布。

2.3 两路径的信息流

信息流图（t 时刻）：

历史信息 ──► h_{t-1} ──► GRU ──► h_t ──► prior p(s_t|h_t) ──► 预测 s_t (规划用)
                  │                            │
                  s_{t-1}                      │ 用 h_t + o_t 校正
                  a_{t-1}                     ▼
                                            posterior q(s_t|h_t, o_t) ──► 校正 s_t (训练用)
                                               │
                                               ▼
观测 o_t ──────────────────────────────────────┘

三、变分后验的深度分析

3.1 后验网络结构

实现中，Encoder 将高维观测映射到一个低维向量，然后与拼接后分别通过两个 MLP 输出均值和对数方差。

3.2 先验网络结构

注意：先验网络和后验网络结构相似（都是输入，输出均值和对数方差），但参数不共享。这确保先验不能直接”复制”后验，必须从中独立预测分布。

3.3 重参数化技巧

为了能够通过随机变量进行反向传播，使用重参数化技巧：

其中表示元素-wise 乘法。这个技巧将随机性转移到上，使得可以通过和计算。

四、KL Balancing 的深度理解

4.1 KL 项的直觉

KL 散度衡量后验与先验的差异。

当此 KL = 0 时：后验 = 先验，意味着已经完美预测了，不需要观测的额外信息。

当此 KL > 0 时：后验与先验不一致，提供了关于的额外信息。

4.2 无 KL Balancing 时的问题

如果直接优化标准 ELBO：

网络可能学到的一种”捷径”是：

后验过强：变得非常 narrow（），几乎变成 Dirac delta ，使得重建项很大（因为已经”看”过了）；
先验退化：保持 broad（较大），但均值趋向于常数向量；
结果：KL 项变小（因为趋向于 Dirac，即使 broad，KL 也很小），但先验完全失去了预测能力。

这为什么是问题？因为规划时只有先验可用。退化的先验无法进行有效的 latent imagination。

4.3 KL Balancing 的解决机制

KL Balancing 通过交替优化迫使两个网络都学到有意义的东西：

阶段 1：冻结先验，更新后验

此时损失函数变为：

（ 固 定 ）

梯度主要流向（后验网络），学会从中提取信息。

阶段 2：冻结后验，更新先验

此时损失函数变为：

（ 固 定 ）

梯度主要流向（先验网络），学会生成与后验一致的分布。

交替比例：论文中通常每 1-2 次更新后验后更新一次先验，确保先验有足够时间去逼近后验。

4.4 KL Balancing 与 β-VAE 的区别

维度	β-VAE	KL Balancing
目标	解耦表示学习	保持先验预测能力
机制	全局加权 KL 项	交替优化
KL 权重	单一超参数	动态平衡
先验质量	不保证先验有意义	强制先验逼近后验

五、Overshooting 技术详解

5.1 复合误差的数学形式

设真实系统为：

如果我们用模型先验进行前向预测（无观测校正）：

定义步预测误差：

随着增大，通常会增大，因为每一步的预测误差会累积。

5.2 Overshooting 的数学动机

原始 ELBO（无 overshooting）：

带 Overshooting 的 ELBO：

注意：内层求和是从到 ，这意味着对于每个时间步，我们不仅考虑当前步的 KL，还考虑未来步的 KL。

5.3 Overshooting 的直观理解

k=0 项（当前步）：与标准 ELBO 相同，强制当前先验与后验一致。

k=1 项：对于时间步，还额外计算了步的 KL：

这意味着：即使我们在训练时用后验采样，但我们同时要求从预测的先验 也要与这个后验一致。

k=2,…,K 项：同理，要求多步先验都与对应后验一致。

5.4 为什么 Overshooting 能缓解复合误差？

关键洞察：先验只依赖于 。

如果我们直接用进行预测，相当于假设：

没有 overshooting 时，这个假设只在时被强制满足。有了 overshooting，这个假设在时也被强制满足。

因此， 学会了编码”如果有了正确的历史信息，隐状态应该是什么样”。这使得成为一个更好的”预测器”，即使在没有观测的情况下。

5.5 Overshooting 的实现细节

计算效率：朴素的 overshooting 需要计算每个对的 KL，时间复杂度。但实现时可以”借用”已经计算过的和，只需要额外计算项的 KL。

的选择：论文中是一个经验值。较大的能提升更长 horizon 的预测质量，但也会增加计算成本和过拟合风险。

与 KL balancing 的结合：Overshooting 的 KL 项也参与 KL balancing 的交替优化，确保先验在多步上都能逼近后验。

六、Latent Imagination 与 Planning

6.1 想象 rollout 的信息流

在 latent space 中进行 planning 时，信息流如下：

真实时刻 1 ────────────────────────────────────────►
                 │
                 │ (s_1 ~ q_φ(s_1|o_1))
                 ▼
           ┌───────────┐
           │  隐状态   │ ◄───── 仅在 t=1 使用真实后验
           │   s_1     │
           └─────┬─────┘
                 │
                 │ a_1 ~ π(a|s_1)
                 ▼
           ┌───────────┐      ┌───────────┐
           │   GRU     │ ───► │ h_2       │ ──► 先验 p(s_2|h_2)
           │ h_2 = f(h_1,s_1,a_1)        │
           └───────────┘      └───────────┘
                 │
                 ▼
           ┌───────────┐
           │   s_2     │ ◄───── 完全使用先验（无观测！）
           │(stoch.)   │        这就是 latent imagination
           └─────┬─────┘
                 │
                 │ a_2 ~ π(a|s_2)
                 ▼
                 ...

想象 horizon H 内，全部使用先验 p(s_{t+1}|h_{t+1}) 预测

6.2 先验质量的直接影响

Latent imagination 的质量完全取决于先验的质量。

如果先验退化（偏离真实后验），则：

想象的轨迹会与真实环境轨迹分叉（divergence）；
在想象的轨迹上学习的策略迁移到真实环境时会失败；
Value function 估计会有偏差。

这正是 Overshooting 和 KL Balancing 如此重要的原因：它们确保先验在多步上都能保持与后验一致的分布，从而保证 latent imagination 的质量。

6.3 信用分配与 GAE

问题：在想象的 horizon 内，每一步的奖励都由世界模型预测。但模型预测的与真实有偏差，导致策略梯度估计不准确。

GAE 的缓解作用：GAE 通过指数加权平均的 TD 误差来平衡偏差与方差：

当时，更多依赖近期奖励（低偏差但可能高方差）；当时，等价于 Monte Carlo 回报（低方差但高偏差）。

七、与 Dreamer 系列的演进关系

7.1 PlaNet → Dreamer v1 的改进

改进点	PlaNet	Dreamer v1
训练框架	分阶段训练	联合训练 actor-critic 与世界模型
KL Balancing	显式交替优化	使用固定的和 target prior 技巧
后验网络	标准 MLP	更深的网络
Overshooting	核心创新	保留
RSSM	核心架构	保留并改进

7.2 关键架构继承

PlaNet (2019)
    │
    ├── 确定性路径 h_t = GRU(h_{t-1}, s_{t-1}, a_{t-1})
    │       └── 保留到 Dreamer v3
    │
    ├── 随机路径 s_t ~ p(s_t|h_t) / q(s_t|h_t, o_t)
    │       └── 保留到 Dreamer v3，但分布从 DiagNormal 扩展
    │
    ├── 先验/后验分离训练
    │       └── 演化为 Dreamer 的 target prior 技巧
    │
    └── Overshooting (K=12)
            └── 大部分后续工作保留或调整 K 值

八、核心公式速查

8.1 RSSM 状态转移

确 定 性 路 径

先 验

后 验

8.2 重参数化

8.3 KL 散度（对角正态）

8.4 ELBO with Overshooting

8.5 GAE

九、关键词

RSSM, Deterministic Path, Stochastic Path, 先验 Prior, 后验 Posterior, 重参数化 Reparameterization, Overshooting, KL Balancing, 复合误差 Compound Error, Latent Imagination, Dreamer, ELBO, GAE, Advantage Estimation, 隐状态后验, 均值场近似 Mean-field Approximation, 状态空间模型 SSM

章节摘要

本章节深入解析了 PlaNet 中 RSSM（Recurrent State Space Model） 的概率图模型结构与 Overshooting 技术的数学动机与实现细节。

RSSM 的核心创新在于双路径架构：确定性路径（GRU 隐藏状态）负责长期记忆的保持，随机路径（隐状态）负责建模部分可观测性和随机性。这两条路径的分离使得模型既能处理长期依赖，又保留了概率建模的灵活性。

KL Balancing 通过交替优化迫使先验学会逼近后验，防止后验过强导致先验退化。这保证了在 latent imagination（纯先验前向预测）时的质量。

Overshooting 是解决多步预测中复合误差的核心技术。通过在 ELBO 中引入步的 KL 项，强制先验在多步上都与后验保持一致。这使得确定性状态学会了编码”在给定历史信息条件下，隐状态应该如何预测”，从而在无观测条件下也能进行可靠的多步预测。

这两个技术（KL Balancing + Overshooting）共同保证了 PlaNet 的 latent imagination 质量，使得在低维隐空间中的多步规划成为可能。后续 Dreamer 系列工作在此基础上进行了改进，但核心思想（RSSM + Overshooting）保持不变。

Huajishe Notes

Explorer

2-PlaNet中的RSSM与overshooting