PlaNet:基于隐式 latent planning 的世界模型
一句话定位:PlaNet(Hafner et al., 2019)首次提出完整的 RSSM 架构,通过变分推断在隐状态空间中进行多步预测与规划,是 latent planning 路线的奠基之作。
前置依赖:变分推断(ELBO)、概率图模型、RSSM 状态空间模型基础概念。
一、前置知识与背景
1.1 问题定义:世界模型学习的目标
世界模型(World Model)的核心目标是:给定历史观测序列
学习这个世界模型有什么用?主要有两方面:
- 序列建模:压缩历史信息用于预测未来;
- 规划(Planning):在隐空间中进行决策优化,而无需在真实环境中交互。
1.2 早期方法对比
| 方法 | 代表工作 | 核心思想 | 局限性 |
|---|---|---|---|
| 无模型 RL | DQN、PPO | 直接学习策略,不建世界模型 | 样本效率极低 |
| World Models | Ha & Schmidhuber 2018 | VAE+MDN-RNN,但梦中学通过随机梦 rollout | 随机梦质量差、规划粒度粗 |
| SimPLe | Kaiser et al. 2019 | 基于 Worlds Models 的 PSL 方法 | 仅支持短 horizon(~10步) |
| PlaNet | Hafner et al. 2019 | RSSM + 变分 ELBO + latent imagination | 奠基性工作 |
1.3 核心挑战:预测与表征的双重需求
学习世界模型面临一个核心矛盾:
- 表征需求:需要从高维观测
中提取隐状态 ,这需要后验(posterior) 利用观测信息; - 预测需求:在规划时,我们没有未来观测,只能依赖先验(prior)
进行前向预测。
这两个分布必须保持一致,否则在 latent space 中的规划会崩溃——这就是 KL balancing 要解决的问题。
二、核心思想
2.1 三条关键设计原则
- 确定性路径(Deterministic Path):使用循环神经网络(RNN)的隐藏状态
作为不受随机性影响的记忆,解决长期依赖问题。相比完全随机的 SSTM,纯随机路径难以处理长期依赖。 - 随机路径(Stochastic Path):每个隐状态
是随机变量,捕捉观测中的不确定性和部分可观测性。 - 变分后验与先验的分离:训练时用后验
“校正”先验 ;推理时只用先验进行前向预测。
2.2 与 World Models 的本质区别
| 维度 | World Models (Ha 2018) | PlaNet (Hafner 2019) |
|---|---|---|
| 隐状态建模 | 随机采样 | RSSM:先验 + 后验分离 |
| 训练目标 | 重建 + MDN-RNN 预测下一帧 | 变分 ELBO,端到端优化 |
| 规划方式 | 随机梦 rollout(质量差) | latent imagination(在隐空间做 dreaming) |
| 多步预测 | 退化严重(compound error) | overshooting 技术 |
三、模型结构
3.1 RSSM 状态转移图
a_{t-1}
│
▼
┌───────────────┐
│ 动作输入 │
└───────┬───────┘
│
┌─────────────┼─────────────┐
▼ │ ▼
┌──────────┐ │ ┌──────────┐
│ h_{t-1} │ │ │ s_{t-1} │
│ (determ.)│ │ │(stoch.) │
└────┬─────┘ │ └────┬─────┘
│ │ │
▼ ▼ ▼
┌─────────────────────────────────────┐
│ 循环核心:RSSM 核心 │
│ h_t = f_h(h_{t-1}, s_{t-1}, a_{t-1})│
└─────────────────────────────────────┘
│ │ │
▼ │ ▼
┌──────────┐ │ ┌──────────┐
│ h_t │ │ │ s_t │
│(determ.) │ │ │(stoch.) │
└────┬─────┘ │ └────┬─────┘
│ │ │
▼ ▼ ▼
┌──────────┐ ┌──────────┐ ┌──────────┐
│ 先验 │ │ 后验 │ │ 奖励 │
│p(s_t|h_t)│ │q(s_t|h_t,│ │ r_t │
│ │ │ o_t)│ │ │
└──────────┘ └──────────┘ └──────────┘
│ │ │
└──────────────┴───────────┘
│
▼
┌──────────┐
│ 图像重建 │
│ p(o_t|s_t)│
└──────────┘
3.2 RSSM 状态转移公式
RSSM(Recurrent State Space Model)的状态转移由以下公式定义:
确定性路径(Deterministic Path):
其中
随机路径(Stochastic Path):
- 先验(Prior):预测下一步隐状态的分布,仅基于当前确定性状态
:
- 后验(Posterior):利用观测
校正隐状态分布(在训练时可用):
- 隐状态采样:
(训练时)或 (推理时)
观测模型(Observation Model):
奖励模型(Reward Model):
3.3 完整的概率图模型
完整的世界模型联合分布:
其中
四、数学推导:变分 ELBO
4.1 目标:最大化观测序列的对数似然下界
我们希望最大化
直接计算该似然是 intractable 的,因为需要对所有隐状态序列
引入变分分布
4.2 ELBO 推导(逐帧展开)
第一步:使用 Jensen 不等式
第二步:对联合分布进行因子分解
第三步:代入并整理
第四步:利用 KL 分解
注意到每一对相邻隐状态的 KL 可以提取出来:
综合整理,得到最终的 ELBO 目标函数:
其中
4.3 KL Balancing:平衡后验与先验
问题背景:
如果直接优化上述 ELBO,网络会学到一种”捷径”:让后验
解决方案:KL Balancing
PlaNet 引入 KL balancing 机制,将 KL 项分为两个部分:
但更常见的写法是将训练目标分离为两个阶段:
- 阶段一:冻结先验
,只优化后验 ,让后验去拟合观测信息; - 阶段二:冻结后验
,只优化先验 ,让先验学会生成与后验一致的分布。
论文中的实现方式是交替优化:每隔几步更新一次后验参数,再更新一次先验参数。
β-KL 权重:在损失函数中加入系数
4.4 训练算法总结
算法:PlaNet 世界模型训练
输入:交互数据 {(o_t, a_t, r_t)}_{t=1}^T
输出:世界模型参数 θ, φ
1. for each training batch do
2. for t = 1 to T do
3. h_t = f_h(h_{t-1}, s_{t-1}, a_{t-1}) // 确定性路径
4. q_\phi(s_t | h_t, o_t) = DiagNormal(μ_φ(h_t, o_t), σ_φ(h_t, o_t)) // 后验
5. s_t ~ q_\phi(s_t | h_t, o_t) // 采样
6. end for
7.
8. // 计算 ELBO
9. recon_loss = Σ_t [log p_θ(o_t | s_t) + log p_θ(r_t | s_t)]
10. KL_loss = Σ_t D_KL(q_φ(s_t | h_t, o_t) || p_θ(s_t | h_t))
11. loss = recon_loss - β * KL_loss
12.
13. // 更新参数(使用 KL balancing 或交替优化)
14. update θ, φ via gradient descent on loss
15. end for
五、Overshooting:多步 latent 预测
5.1 问题:复合误差(Compound Error)
当使用
随着步数
5.2 Overshooting 的核心思想
Overshooting 的关键是用后验信息”修正”多步预测。
具体地,对于
训练时,Overshooting 通过在 ELBO 中引入中间步骤的 KL 项来缓解这一问题:
其中
这相当于在每一步的前向预测中,都用后验信息校正先验,使得多步预测的分布不会过度偏离真实后验。
5.3 Overshooting 的物理直觉
无 Overshooting(k=0):
h_t → s_t → h_{t+1} → s_{t+1} → ... (误差累积)
有 Overshooting(k=1,2,...,K):
h_t → s_t → h_{t+1} → s_{t+1}
↗ (用 o_{t+1} 校正 h_{t+1} 和 s_{t+1})
→ h_{t+2} → s_{t+2}
↗ (用 o_{t+2} 校正 h_{t+2} 和 s_{t+2})
...
本质上是用观测信息定期修正 latent 空间中的预测轨迹,防止误差指数级增长。
六、Latent Imagination 与策略学习
6.1 Dreamer-style Actor-Critic
学习到世界模型后,PlaNet 在 latent space 中进行策略学习(Dreamer 风格的 actor-critic)。
核心思想:
- Imagined Trajectories(想象轨迹):从真实数据中取出初始隐状态,然后在 latent space 中使用当前策略
前向 rollout 多步。 - Value Function 学习:使用 TD/GAE 从想象轨迹中估计优势函数
。 - Actor 更新:通过策略梯度最大化期望累积回报。
6.2 策略模型
6.3 Critic(Value Function)
使用参数化近似:
6.4 想象 rollout 算法
算法:Latent Imagination Rollout
输入:初始后验分布 q_φ(s_1 | h_1, o_1),策略 π_ψ,世界模型 p_θ
输出:想象轨迹的回报
1. s_1 ~ q_φ(s_1 | h_1, o_1) // 从真实后验采样初始状态
2. for t = 1 to H do
3. a_t ~ π_ψ(a_t | s_t) // 用当前策略采样动作
4. h_{t+1} = f_h(h_t, s_t, a_t) // deterministic path
5. s_{t+1} ~ p_θ(s_{t+1} | h_{t+1}) // 用先验预测下一步隐状态
6. r_t ~ p_θ(r_t | s_t) // 预测奖励
7. end for
8. return imagined rewards r_1:H
注意:想象 rollout 不使用任何真实观测,完全在隐空间中进行。
6.5 Advantage Estimation(GAE)
使用 GAE(Generalized Advantage Estimation)估计优势函数:
其中
当
6.6 Actor Loss 推导
策略梯度目标(reinforcement learning world model style):
直觉:最大化优势大的时刻的动作概率。
6.7 Critic Loss 推导
使用 MSE 回归目标:
其中
七、完整训练流程
┌─────────────────────────────────────────────────────┐
│ PlaNet 完整训练流程 │
├─────────────────────────────────────────────────────┤
│ │
│ ┌─────────────┐ │
│ │ 真实环境 │ ← 收集交互数据 (o_t, a_t, r_t) │
│ └──────┬──────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────────────────────────────┐ │
│ │ 世界模型训练(ELBO + KL Balancing) │ │
│ │ • h_t = f_h(h_{t-1}, s_{t-1}, a_{t-1}) │ │
│ │ • q_φ(s_t | h_t, o_t) ← 后验 │ │
│ │ • p_θ(s_t | h_t) ← 先验 │ │
│ │ • ELBO = recon - β·KL │ │
│ │ • Overshooting: 多步 KL 正则 │ │
│ └───────────────┬───────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────────────────────────────┐ │
│ │ Latent Imagination Rollout │ │
│ │ • s_1 ~ q_φ(s_1 | h_1, o_1) │ │
│ │ • a_t ~ π_ψ(a_t | s_t) │ │
│ │ • s_{t+1} ~ p_θ(s_{t+1} | h_{t+1}) │ │
│ │ • r_t ~ p_θ(r_t | s_t) │ │
│ └───────────────┬───────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────────────────────────────┐ │
│ │ Actor-Critic 更新 │ │
│ │ • Critic: MSE loss on GAE │ │
│ │ • Actor: policy gradient with GAE │ │
│ └─────────────────────────────────────────────┘ │
│ │ │
│ └──────────────▶ 策略 π_ψ │
│ │
└─────────────────────────────────────────────────────┘
八、优点与局限
8.1 优点
| 优点 | 说明 |
|---|---|
| 隐空间规划 | 可以在低维隐空间中进行多步规划,计算效率远高于像素级规划 |
| 长期依赖 | 确定性路径 |
| Overshooting | 显著提升多步预测质量,12 步内的预测几乎不退化 |
| 端到端 | 观测模型、隐状态模型、奖励模型联合训练,协同优化 |
| 样本效率 | 相比 model-free 方法,样本效率提升显著(DeepMind Control Suite 上仅需 100 集) |
8.2 局限
| 局限 | 说明 |
|---|---|
| 后验延迟问题(Posterior Delay) | 后验 |
| 随机性处理 | 对纯随机性强的环境(如噪声很大的环境)建模困难 |
| 离散动作 | 策略输出为离散动作时,策略梯度在高维动作空间中梯度估计方差大 |
| Credit Assignment | 长 horizon 的信用分配问题仍然存在(虽然 GAE 有所缓解) |
九、与前后内容的衔接
9.1 前置内容(通往 PlaNet)
- World Models (Ha & Schmidhuber, 2018):提供整体框架(VAE + RNN),但没有变分推断和 KL balancing,梦 rollout 质量差。
- 信息论基础:KL 散度的非负性、ELBO 与变分推断的理论基础。
- 概率图模型:状态空间模型(SSM)、隐变量模型的因子图分解。
9.2 后续发展(从 PlaNet 出发)
PlaNet (2019)
│
├── Dreamer v1 (2020) ─── 完整实现 actor-critic,公开复现
│ │
│ ├── Dreamer v2 (2021) ─── 改进 RSSM、KL balancing
│ │ │
│ │ └── Dreamer v3 (2024) ─── 大型视觉控制
│ │
│ └── DayDreamer (2022) ─── 机器人 sim-to-real
│
├── SimPLe (2019) ─── 基于像素的 PSL baseline
│
├── GPS-1 (2020) ─── Guided Predictive Singing ...
│
└── 后来的工作:
├──深度隐式规划(2021)
├──深度RL的World Model综合(2022)
└── ...
十、可复现实现要点
10.1 关键超参数
| 超参数 | 值 | 说明 |
|---|---|---|
| 隐状态维度 | Mujoco 控制任务 | |
| 确定性隐藏状态维度 | GRU hidden | |
| 观测编码维度 | 公共 encoder 输出后直接拼接至 | - |
| KL 权重 | 1.0 或 0.1 | 论文中使用 0.1 |
| Overshooting 步数 | 12 | 论文中默认 |
| Imagined horizon | 15-30 | 策略更新的 rollout 步数 |
| GAE | 0.95 | - |
| Discount | 0.99 | - |
10.2 实现注意事项
-
重参数化技巧(Reparameterization Trick):
-
先验/后验使用相同的 sigmoid 函数但参数不共享:防止先验直接复制后验。
-
KL 项的计算:先计算两个正态分布的 KL,然后在 batch 上取平均:
-
策略梯度实现:使用 REINFORCE 或 PPO-style 的 clipping(Dreamer v2 开始使用)。
10.3 伪代码(核心部分)
# RSSM forward
def rssm_forward(h_prev, s_prev, a_prev, o_curr, encoder, prior_net, post_net):
h_curr = gru_cell(h_prev, s_prev, a_prev) # deterministic path
q_posterior = post_net(h_curr, encoder(o_curr)) # q(s_t | h_t, o_t)
p_prior = prior_net(h_curr) # p(s_t | h_t)
s_curr = q_posterior.sample() # training: use posterior
# s_curr = p_prior.sample() # inference: use prior
return h_curr, s_curr, q_posterior, p_prior
# ELBO loss
def compute_elbo(recon_loss, kl_loss, beta=1.0):
return recon_loss - beta * kl_loss
# Imagination rollout
def imagine_rollout(s_1, policy, world_model, horizon):
traj = []
s = s_1
for _ in range(horizon):
a = policy(s)
h_next, s_next, _, _ = rssm_forward(h, s, a, None, ...) # no obs
r = reward_net(s)
traj.append((s, a, r))
s = s_next
return traj十一、关键词
PlaNet, RSSM, 世界模型, 隐式规划, latent planning, 变分推断, ELBO, KL balancing, Overshooting, Dreamer, actor-critic, 隐空间规划, imagined trajectories, world model training, latent imagination, GAE, advantage estimation
章节摘要
PlaNet 是 latent planning 路线的奠基性工作,首次提出完整的 RSSM(Recurrent State Space Model) 架构。RSSM 通过两条并行路径处理时序信息:确定性路径(GRU 隐藏状态
在训练上,PlaNet 使用变分 ELBO 作为训练目标,其中 KL 项
Overshooting 是 PlaNet 的核心技术创新:通过在 ELBO 中引入多步 KL 项,强制让多步先验预测也能利用后验信息校正,从而显著缓解复合误差问题,使 latent space 中的多步预测(12 步内)几乎不退化。
在策略学习上,PlaNet 采用 Dreamer-style actor-critic:在隐空间中进行想象 rollout,使用 GAE 估计优势函数,通过策略梯度更新策略网络。这使得规划可以在低维隐空间高效进行,无需在真实环境中交互采样。
PlaNet 的局限性包括后验延迟问题(无法利用未来信息校正历史)、对强随机性环境建模困难,以及长 horizon 信用分配困难。这些问题在后续 Dreamer 系列工作中逐步得到改进。