Learning Latent Dynamics for Planning from Pixels (PlaNet)

一句话定位

PlaNet（Deep Planning Network）是纯模型式（model-based）强化学习 agent：只从像素学一个循环状态空间模型 RSSM（确定性 h_t + 随机 s_t 双通道 latent），然后在 latent 空间里用 CEM 做在线滚动规划选动作，不训练任何 policy/value 网络。它在 6 个 DeepMind Control 图像任务上以约 1000 个 episode 逼近甚至超过训练了 100,000 episode 的 model-free 强基线，平均省 ~200× 环境交互——RSSM 这个 backbone 正是后来 dreamer 系列世界模型的地基。

背景与定位

范式：从像素学 latent world model + 在 latent 空间规划（model-based RL / latent-space planning）。已知动力学下规划（AlphaGo、模拟机器人轨迹优化）非常成功，但未知环境要先从交互中把动力学学准到足以规划，长期是难题：模型误差、多步预测误差累积、无法表达多种可能未来、分布外过自信。

• 先前的 latent-规划工作（E2C, Watter 2015；RCE, Banijamali 2017）只能把图像嵌入 latent、学局部线性转移 + LQR，解 cartpole 平衡、2-link 臂等简单密奖励任务，且假设 Markov、难扩展。PlaNet 明确去掉 Markov 假设（走 POMDP），支持部分可观测、接触动力学、稀疏奖励。
• 直接前身是 world-models-ha-schmidhuber（Ha & Schmidhuber 2018）：PlaNet 复用了它的卷积/反卷积编解码网络，但关键区别是 PlaNet 不用 policy 网络——纯靠规划选动作，因此模型一改进就即时受益（世界模型是用 CMA-ES 在冻结的 latent 上进化一个 controller）。
• 相对纯视频预测（Oh 2015、Chiappa 2017、Denton & Fergus 2018）和 hybrid agent（在想象经验上训 policy：Kalweit 2017、Nagabandi 2017、Buckman 2018 STEVE），PlaNet 走的是”学模型 + 显式规划”这条更纯的路线。
• 发表于 ICML 2019（PMLR 97, pp. 2555–2565），arXiv v1 2018-11、v5 2019-06。作者 Danijar Hafner 等（Google Brain / DeepMind / U Toronto / U Michigan / Google Research）。

模型架构

Backbone = RSSM（Recurrent State-Space Model）。核心设计：把 latent 状态拆成确定性部分 h_t 和随机部分 s_t，两条路都必需（消融证明缺一不可，见评测）。

生成过程（式 4）：

• 确定性状态：h_t = f(h_{t-1}, s_{t-1}, a_{t-1})，f 用 GRU（200 units）——让模型能跨多步可靠地”记住”信息。
• 随机状态：s_t ~ p(s_t | h_t)，对角高斯，均值/方差由前馈网络给出。
• 观测模型：o_t ~ p(o_t | h_t, s_t)，高斯、均值由反卷积网络给出、单位协方差（等价于像素 MSE）。规划时不用观测模型（只提供训练信号）。
• 奖励模型：r_t ~ p(r_t | h_t, s_t)，标量高斯、均值由前馈网络给、单位方差。
• 编码器（近似后验/filtering）：q(s_t | h_t, o_t)，对角高斯，均值/方差由卷积网络 + 前馈网络给出。用 filtering 后验（只 condition 过去观测）因为最终要拿去规划。所有观测信息必须经过 s_t 的采样步，避免从输入到重建的确定性捷径。

对比图 2 的三种设计：(a) 纯确定性 RNN——无法表达多种未来、易被规划器钻模型空子；(b) 纯随机 SSM——难跨多步记忆；(c) RSSM = 两者拆分，既能记忆又能表达多模态未来。

Action-conditioning：转移显式条件于连续动作向量 a_{t-1}（喂进 GRU）。

关键配置数字：

• latent 分布：30 维对角高斯（随机 s_t），预测均值 + 标准差。
• 确定性通道：GRU 200 units。
• 其余所有函数：2 层全连接、每层 200 维、ReLU。
• 卷积/反卷积编解码：取自 world-models-ha-schmidhuber（Ha & Schmidhuber 2018）。
• 观测：第三人称相机图像 64×64×3，预处理把位深降到 5 bit。

动作选择 = 规划（无 policy/value 网络）：

• MPC，每步重新规划（replan every step），据新观测调整计划。
• CEM（交叉熵方法）（算法 2）：对动作序列维护时间相关对角高斯 belief Normal(μ, σ²I)，从零均值/单位方差起，反复采 J 条候选序列、在模型下评估、拿 top K 重拟合 belief，I 次迭代后返回当前步均值 μ_t。每步接新观测后 belief 重置为零均值单位方差以避局部最优。
• 评估一条候选序列：从当前 state belief 采一条轨迹、把沿途预测的均值奖励求和；因为是群体式优化器，每条序列只采单条轨迹，把算力花在”评更多不同序列”上。奖励是 latent 的函数 → 规划器纯在 latent 空间前推、不生成图像，可大批量快速评估。
• CEM 超参：horizon H=12, 迭代 I=10, 候选 J=1000, top K=100（即每步评估 10×1000 = 1 万条想象序列，与 blog 说的”每步比较 10,000 条想象动作序列”一致）。

数据

任务与规模（都是”边规划边收集”的在线 RL，无外部数据集）：

• 6 个 DeepMind Control Suite 图像连续控制任务：cartpole swingup（固定相机、小车会出画 → 要跨帧记忆、需长规划 horizon）、reacher easy（稀疏奖励）、cheetah run（多关节、地面接触）、finger spin（手指与物体接触动力学）、cup catch（只在接住球时给稀疏奖励）、walker walk（先站起再走、难预测的地面碰撞）。观测只有 64×64×3 第三人称图像。
• 数据采集流程：从 S=5 个随机动作 seed episode 起，每 C=100 个更新步用当前（半训练）模型规划采 1 个新 episode 加入数据集；采集时加高斯探索噪声 ε ~ Normal(0, 0.3)。
• 动作重复 R（唯一按任务不同的超参，缩短有效规划 horizon、给模型更清晰信号）：cartpole R=8、reacher R=4、cheetah R=4、finger R=2、cup R=4、walker R=2。
• 训练用的序列 chunk：批 B=50 条、每条长 L=50，从数据集均匀随机抽。
• PlaNet 解各任务总量级 ~1000 episode（对照 model-free 基线的 100,000）。
• 除动作重复外，所有任务共用同一套超参；还做了”一个 agent 学全部 6 任务”（不告诉它当前是哪个任务，靠图像自己推断）。

训练方法

目标函数 = 变分下界（ELBO）（式 3）：对数据对数似然的变分界 = 重建项（观测 + 奖励的对数似然，即 latent 解码的负 MSE）− KL 复杂度项 KL[q(s_t|o≤t,a<t) ‖ p(s_t|s_{t-1},a_{t-1})]。用编码器给近似后验、单个重参数化样本估外层期望、梯度上升优化。

Latent overshooting（本文提出的多步目标）：

• 动机：标准 ELBO 里转移函数 p(s_t|s_{t-1}) 只被一步预测的 KL 训练（梯度只穿一次转移），但规划需要准的多步预测；有限容量/受限分布族下，只训一步预测最优 ≠ 多步预测最优。
• 做法（式 5–7）：把标准界推广到训练所有距离 1≤d≤D 的多步预测，全部在 latent 空间做（不必再解码额外图像，故很快），用 β-VAE 式权重 {β_d}。对 d>1 的多步预测停掉后验梯度（让多步先验往有信息的后验靠、而非反向）。
• 理论：借数据处理不等式论证多步预测分布 p_d(o_{1:T}) 在数据集期望下也是一步分布 p(o_{1:T}) 的下界。
• 重要结论：最终用 RSSM 的 agent 并不需要 latent overshooting（Appendix D 显示它对 RSSM 略有损、但能大幅提升 DRNN 等较弱模型）。即最终 PlaNet 只用标准 ELBO。

关键超参：Adam，学习率 1e-3，ε=1e-4，梯度裁剪范数 1000。不对 KL 相对重建项缩放，而是给模型 3 free nats（KL 低于此值就裁掉不罚）。作者点名最重要的超参是：动作重复 R、KL 权重 β、学习率。

Infra（训练 / 推理工程）

• 训练硬件：单张 Nvidia V100 GPU，每任务 10–20 小时（视任务而定）。
• 框架：TensorFlow Probability（Dillon 2017）；实现依赖 tensorflow-gpu 1.13.1 + tensorflow_probability 0.6.0 + dm_control。
• 尽管 episode 数比 A3C/D4PG 少约两个数量级，总计算时间与之相当甚至更少（省的是环境交互，不是算力）。
• 并行/精度：未披露（2018 单卡实验，未涉及分布式或混合精度）。
• 推理/控制：MPC 每步在线规划，每步评估 1 万条 latent 想象序列；因规划纯在 latent（不解码图像）故可大批量快评。具体控制频率 Hz / 单步 latency / 边缘硬件未披露。

评测 benchmark

主表（Table 1）：6 任务最终性能（5 seeds × 10 trajectories 均值），baseline 数字取自 DeepMind Control Suite（Tassa 2018）。

方法	模态	Episodes	Cartpole Swingup	Reacher Easy	Cheetah Run	Finger Spin	Cup Catch	Walker Walk
A3C	proprioceptive	100,000	558	285	214	129	105	311
D4PG	pixels	100,000	862	967	524	985	980	968
PlaNet (ours)	pixels	1,000	821	832	662	700	930	951
CEM + 真实模拟器	simulator state	0	850	964	656	825	993	994

数据效率增益（PlaNet 相对 D4PG，达到 PlaNet 最终性能所需 episode 之比，Table 1 末行）：cartpole 250×、reacher 40×、cheetah 500+×、finger 300×、cup 100×、walker 90×，平均 ~200×（论文 contribution 与 discussion 口径；blog 更保守地写”平均省 5000% = 50×”）。

• 100 episode 内 PlaNet 在所有任务上超过（从 proprio 状态训了 100,000 episode 的）A3C；500 episode 后除 finger 外逼近 D4PG（从像素训 100,000 episode）。
• cheetah run 上最终超过 D4PG，相对提升 26%；PlaNet 662 vs D4PG 524。

消融：

• 模型设计（Fig 4）：RSSM vs 纯确定性 GRU vs 纯随机 SSM——两条通道都关键。确定性部分负责跨多步记忆；随机部分更关键——没有它 agent 根本学不会（任务从 agent 视角因初始状态部分可观测而随机；噪声也给规划目标加安全裕度、令动作序列更鲁棒）。
• Agent 设计（Fig 5）：在线采集 vs 随机采集——在线对所有任务有帮助，对 cartpole/finger/walker 是必需；CEM 迭代优化 vs random shooting（每步取 1000 条里最好一条不迭代）——CEM 在所有任务上更好。
• Latent overshooting（Appendix D, Fig 8）：大幅提升 DRNN（双 RNN 编解码 + 中间随机状态序列），但略降 RSSM。
• 激活函数（Appendix E, Fig 9）：平滑 ELU 帮到纯随机/纯确定性模型，RSSM 对激活选择鲁棒。
• 规划参数（Appendix J, Fig 12，cheetah + 真实模拟器）：horizon=6 不够、约 8 步最佳、太长反而伤（搜索空间变大，132→837 分区间）；评更多序列、迭代更多、重拟合到更少的 top proposals 都更好。
• 多任务（Appendix C, Fig 6–7）：单 agent 学全 6 任务、全部学会，仅中等减速（每个任务的学习都比单独训练的 agent 略慢，但最终都解出）。
• 开环视频预测（Appendix H）：cheetah 上给 5 帧 context 后能逐像素准确预测 50 步未来。状态诊断（Appendix I）表明真实位置/速度/奖励大多能从 latent 线性预测出、且外推超过所用规划 horizon。

创新点与影响

贡献：

1. latent 空间规划：证明从像素学动力学 + 在 compact latent 里高效规划，能解比先前 learned-model 规划难得多的连续控制任务（接触、部分可观测、稀疏奖励），并超过 model-free A3C、逼近/局部超过 D4PG，平均省 ~200× 交互。
2. RSSM：确定性 + 随机双通道 latent 转移模型，实验证明二者对高性能规划都必需。这是本文最持久的遗产——dreamer（Dreamer v1/v2/v3）直接把 RSSM 当 backbone，只是把”CEM 规划”换成”在 latent imagination 里学 actor-critic”。
3. latent overshooting：把标准变分界推广到 latent 空间多步预测的通用正则，兼容任意 latent 序列模型（虽然最终 RSSM 未采用）。

改变了什么：是”用 learned model 规划在图像任务上首次超过 model-free”的早期范例之一，把 model-based RL 从低维状态推到高维像素，奠定了 latent world model + 规划/想象这条主线。

论文自陈局限 / 未来方向：

• 用固定动作重复而非学到的时间抽象（可用分层模型改进）。
• 无 value function 估计 planning horizon 之外的回报（加 value 可提最终性能）。
• 基于梯度的规划可能比 CEM 更高效。
• 无重建的表征学习（不靠像素重建）有望解决视觉多样性更高的任务（这正是后续 Dreamer→DreamerV2/V3、以及 contrastive/JEPA 系世界模型的方向）。
• 多任务控制靠共享动力学模型只是起点。

原始链接

• arXiv 摘要页：https://arxiv.org/abs/1811.04551
• arXiv PDF：https://arxiv.org/pdf/1811.04551
• 官方博客（Google Research，原 ai.googleblog.com）：https://research.google/blog/introducing-planet-a-deep-planning-network-for-reinforcement-learning/
• GitHub 开源实现：https://github.com/google-research/planet
• 在线论文（HTML）：https://planetrl.github.io/
• 作者项目页：https://danijar.com/project/planet/

一手源存档（sources/）

• planet—blog — Google Research 官方博客快照（sources/world-model/2018/planet—blog.md）
• planet—github-readme — GitHub README 快照（sources/world-model/2018/planet—github-readme.md）
• arXiv 原文 PDF（1811.04551，全文已读；PDF 不入 git，见上方 arXiv 链接）