跳到主要内容

PlaNet:从像素学习潜变量动力学并规划

论文:Learning Latent Dynamics for Planning from Pixels
作者:Danijar Hafner、Timothy Lillicrap、Ian Fischer、Ruben Villegas、David Ha、Honglak Lee、James Davidson
会议:ICML 2019
原文:PMLR 论文页 · PDF · arXiv:1811.04551
对话来源:ChatGPT 分享对话

本文按分享对话的原始顺序整理,尽量保留对话中的解释、例子和追问:论文总览 → 为什么基于模型的规划更高效 → Algorithm 1 训练循环 → CEM 规划 → 用户对 Algorithm 1 的关键追问 → RSSM → 变分编码器与 ELBO。对话中重要的问题与内容以 标题标注,会出现在右侧目录 / 侧边栏中。

对话整理顺序

  1. 论文整体在解决什么问题:为什么不在像素空间直接预测,而是学 latent dynamics。
  2. 为什么基于模型的规划更省数据、更容易迁移(用户提问)。
  3. Algorithm 1 如何交替训练世界模型和采集真实数据。
  4. Algorithm 2 如何用 CEM 在潜空间内规划。
  5. 四个模型是否共同训练?未收敛时为什么还敢规划采数据(用户最关键追问)。
  6. RSSM 为什么组合确定性记忆和随机状态。
  7. filtering posterior、transition prior、变分编码器和 ELBO 如何共同工作。

一句话概括

PlaNet(Deep Planning Network)从图像交互数据中学习一个潜变量世界模型,然后在潜空间内用 CEM 搜索动作序列,并以 MPC 的方式每次只执行第一个动作、观察真实环境后重新规划。

图像历史 + 动作历史

filtering encoder

当前 latent belief (h_t, s_t)

RSSM 想象候选动作的未来 latent 轨迹

reward model 评价累计收益

CEM 筛选动作序列

只执行第一个动作,再观察并重规划

它不训练 policy network 或 value network,policy 本身就是在线规划算法。

⭐ 论文总览:为什么需要 Latent Dynamics

传统强化学习的问题

传统 RL 直接把观测映射到动作:

image observation

policy

action

问题是 sample efficiency 很低:每次都必须真实环境交互,机器人采集数据的成本极高。

Model-based RL 的想法是学习一个环境模型:

st,atst+1s_t, a_t \rightarrow s_{t+1}

然后先在模型里预测未来,再选择最好的动作,就像人开车时会先想象“向左打方向是否安全、加速是否会撞”。

为什么不能在像素空间直接预测

机器人输入不是状态而是图像,例如 224×224×3224\times224\times3。如果直接预测下一帧:

ItIt+1I_t \rightarrow I_{t+1}

非常困难,因为一张图像包含光照、纹理、背景、细节,而控制真正需要的是:

物体位置
速度
关节状态
空间关系

所以 PlaNet 学习 Latent Dynamics:

image

Encoder

latent state z

Dynamics model

future latent

核心思想一句话:

不要直接在像素空间预测未来,而是学习一个低维 latent space,在其中学习世界动态,并利用这个模型进行规划。

整体结构

          observation image
|
v
Encoder
|
v
latent state
|
+--------+--------+
| |
v v
Dynamics model Reward model
|
v
future latent states
|
v
Planning

论文使用 RSSM(Recurrent State Space Model) 作为世界模型,这是全篇最重要的贡献,后面单独展开。

为什么叫 Latent Dynamics

因为学习的目标不是:

image(t) → image(t+1)

而是:

image → encoder → latent z
z_t + action → dynamics → z_{t+1} → decoder → future image

也就是把 latent space 当作 imagination space。

Planning 部分:为什么不用普通 MPC

PlaNet 不用 policy network,而是用 MPC(Model Predictive Control):从当前状态随机采样候选动作序列,在模型中 rollout 未来、累计 reward,选最好的序列,只执行第一个动作,再观察真实环境重新规划。这叫 receding horizon control。

为什么不能直接用普通 MPC?因为真实状态未知。传统 MPC 需要 position / velocity / joint angle,而视觉输入只有 pixels。所以 PlaNet 先做:

pixels → latent state → MPC

训练过程概览

机器人随机探索收集 (ot,at,rt)(o_t,a_t,r_t),训练三个模型:

  1. Encoder:image → latent
  2. Dynamics(zt,at)zt+1(z_t,a_t)\rightarrow z_{t+1}
  3. Decoder:latent → image

损失主要是 reconstruction loss(保证 latent 保留视觉信息)和 KL loss(约束 prior 与 posterior 接近):

oto^t,DKL(q(ztht,ot)p(ztht))\|o_t-\hat{o}_t\|, \qquad D_{\mathrm{KL}}\bigl(q(z_t\mid h_t,o_t)\,\|\,p(z_t\mid h_t)\bigr)

与 Pi3X / VGGT / PointMap 的定位差异

Pi3X / VGGTPlaNet
目标3D reconstructioncontrol
latentgeometry latentdynamics latent
预测3D structurefuture state
监督3D datainteraction data
输出point / cameraaction

它们并不冲突。老师说的“多预测光流、物体、动作,让预测互相辅助”,本质就是从静态几何理解走向动态 world model:

image

world representation

-------------------------------
| | |
3D geometry motion dynamics
(point map) optical flow prediction

planning / action

后续论文路线图

  • 第一阶段 World Model 基础:PlaNet(2018)→ Dreamer(2020)→ DreamerV2 / V3。
  • 第二阶段 机器人世界模型:MuJoCo World Models、RT-1 / RT-2、RoboCat。
  • 第三阶段 视觉 + 3D + 动作:GAIA-1(自动驾驶 world model)、Genie(interactive world model)、V-JEPA(video latent prediction)。

PlaNet 真正提出的思想:不要把视觉理解看成识别问题,而看成学习一个可以预测未来的世界模型。

问题设定与输入输出

像素观测通常不能唯一揭示环境状态,因此论文把问题视为 POMDP:

stp(stst1,at1),s_t\sim p(s_t\mid s_{t-1},a_{t-1}), otp(otst),rtp(rtst).(1)o_t\sim p(o_t\mid s_t), \qquad r_t\sim p(r_t\mid s_t). \tag{1}

其中:

  • sts_t:不可直接观测的环境状态(例如机械臂位置、物体位置、速度)。
  • oto_t:图像观测。
  • ata_t:连续动作向量。
  • rtr_t:标量奖励。

目标是实现 atp(atot,a<t)a_t\sim p(a_t\mid o_{\le t},a_{<t}) 并最大化期望累计奖励。

阶段输入输出
状态推断当前图像 oto_t、过去 latent 与动作posterior belief q(stot,a<t)q(s_t\mid o_{\le t},a_{<t})
世界模型训练轨迹片段 (ot,at,rt)(o_t,a_t,r_t)transition、observation、reward 与 encoder 参数
潜空间预测当前 belief、候选动作序列未来 latent state 与预测 reward
在线规划多条候选动作序列的预测收益当前执行动作 ata_t

⭐ 为什么 model-based planning 更节省数据?

论文原句 “Planning using learned models offers several benefits over model-free RL…” 怎么理解?

用户贴出的原文是:

Planning using learned models offers several benefits over model-free reinforcement learning. First, model-based planning can be more data efficient because it leverages a richer training signal and does not require propagating rewards through Bellman backups. Moreover, planning carries the promise of increasing performance just by increasing the computational budget for searching for actions, as shown by Silver et al. (2017). Finally, learned dynamics can be independent of any specific task and thus have the potential to transfer well to other tasks in the environment.

逐句拆解:

1. Model-free RL 是什么

直接学习 state → action 的映射:

Observation → Policy Network → Action

典型方法 DQN、PPO、SAC。它不知道“执行这个动作后未来会怎样”,只靠奖励慢慢调整。

2. Model-based RL 是什么

先学环境模型:

(st,at)st+1(s_t, a_t) \rightarrow s_{t+1}

再在模型里模拟未来、比较候选动作、选择最好的动作。真实交互很贵(机械臂抓取失败一次要几秒甚至几十秒),但在模型里模拟一万次几乎没有成本。

3. richer training signal(更丰富的训练信号)

Model-free 只知道“动作 → 最终奖励”,中间发生了什么不知道,只能靠 Bellman backup 把奖励一步步往回传播:

Q(s,a)=r+γmaxaQ(s,a)Q(s,a)=r+\gamma\max_a Q(s',a')

Model-based 直接知道“动作 → 下一状态 → 未来状态 → 奖励”的每一步,每一步都有监督,因此数据效率更高。

4. 不需要通过 Bellman backup 传播奖励

Model-free 中,第 10 步的动作为什么好,往往要等第 300 步的奖励一路传回来。Model-based 直接模拟“第 10 步动作 → 预测第 20 步 → 预测最终结果”,不用等待奖励传播。

5. 增加计算预算就能提升性能(Silver et al. 2017 = AlphaGo Zero)

规划的性能可以随着计算量增加而提升:尝试 100 个动作序列 vs 尝试 100 万个动作序列,后者可能找到更优动作。AlphaGo 的 MCTS 就是“搜索得越多,棋力越强”。这和现在 LLM 的 reasoning/search 类似:更多 inference compute 可能得到更好的结果。

6. dynamics 可以独立于任务、迁移到其他任务

机器人在厨房学到的是“物体如何移动、机械臂如何影响物体”这些世界规律,与具体任务无关。拿杯子、倒水、清理桌面共享同一套 dynamics,换任务时只需要重新学习 policy / planning objective,而 model-free 的抓取 policy 换任务往往要重新训练。

总结一句话:学习环境模型再在模型里规划,比直接学习策略更好——因为能利用每次交互产生更多监督信号(数据效率),能通过增加搜索计算提升性能,而且环境模型描述的是世界规律、不绑定具体任务,可以迁移。

⭐ Algorithm 1:训练世界模型与采集数据

PlaNet Algorithm 1:模型拟合与数据采集闭环

图源:Hafner et al., Learning Latent Dynamics for Planning from Pixels, Algorithm 1, p.2。

Algorithm 1 描述的是整个 model-based RL 的闭环训练过程,是典型的 Dyna-style algorithm:

真实环境 env → 采集 (o,a,r) → 数据集 D → 训练 Latent Dynamics Model
↑ ↓
└──────────── 执行新动作 ← imagination / planning ←──────┘

初始化

  1. Initialize dataset D with S random seed episodes:用随机策略采集 SS 个 episode,数据集 D={(ot,at,rt)}D=\{(o_t,a_t,r_t)\}。为什么先随机?因为模型一开始不知道世界规律,需要初始数据。
  2. Initialize model parameters θ\theta randomlyθ\theta 包括 encoder、RSSM dynamics、decoder、reward predictor 的全部参数。

外层循环 while not converged

循环内交替做“模型拟合”和“数据采集”两部分。

阶段 A:Model fitting

for update step s = 1...C:
从 D 采样 B 个长度为 L 的连续轨迹片段
计算 loss L(θ)
θ ← θ − α ∇_θ L(θ)

采样的是连续序列片段而不是单个 transition,因为 RSSM 是 recurrent model,需要历史 hth_t

损失由三部分构成:

L(θ)=Lobs+Lreward+LKL\mathcal{L}(\theta) = \mathcal{L}_{\mathrm{obs}} + \mathcal{L}_{\mathrm{reward}} + \mathcal{L}_{\mathrm{KL}}
  • observation reconstruction:oto^t\|o_t-\hat{o}_t\|
  • reward prediction:r^t\hat{r}_trtr_t 比较;
  • KL:DKL(q(ztht,ot)p(ztht))D_{\mathrm{KL}}(q(z_t\mid h_t,o_t)\,\|\,p(z_t\mid h_t))

阶段 B:Data collection

o_1 ← env.reset()
for time step t = 1...T/R:
推断当前 belief q(s_t | o_≤t, a_<t)
a_t ← planner(q(s_t), p) # CEM
a_t ← a_t + ε, ε ~ p(ε) # exploration noise
(r_t^k, o_{t+1}^k) ← env.step(a_t) 重复 R 次(action repeat)
r_t, o_{t+1} ← Σ_k r_t^k, o_{t+1}^R
D ← D ∪ {(o_t, a_t, r_t)}

infer belief:根据历史推断当前隐藏状态,因为真实状态看不到,只能得到 latent belief,类似 Kalman filter。

planner:输入当前 belief 和动力学模型 pp,用 CEM 在 latent 中想象未来、选择累计 reward 最大的动作序列。

exploration noise:避免模型只走熟悉路线、永远访问 planner 认为最优的狭窄区域,否则会 exploration failure。

action repeat:一个动作连续执行 RR 个环境帧,把期间奖励累加、保留最后一帧观测:

rt, ot+1k=1Rrtk, ot+1Rr_t,\ o_{t+1} \leftarrow \sum_{k=1}^{R} r_t^k,\ o_{t+1}^{R}

它降低控制频率,也让单次高层动作产生更明显的状态变化。

外层循环的含义

D0M1πplan,1D1M2πplan,2\mathcal D_0 \rightarrow M_1 \rightarrow \pi_{\mathrm{plan},1} \rightarrow \mathcal D_1 \rightarrow M_2 \rightarrow \pi_{\mathrm{plan},2} \rightarrow\cdots

MiM_i 是当前世界模型,πplan,i\pi_{\mathrm{plan},i} 是由当前模型和 CEM 临时产生的规划策略。

⭐ CEM 如何在 latent space 中规划?

PlaNet 补充材料中的 Algorithm 2:Latent planning with CEM

图源:Hafner et al., Supplementary Material, Algorithm 2, p.12。

为什么需要 planning

给定当前状态,三个候选动作不知道哪个最好,用动力学模型预测各自的未来和 reward:

a1 → future → reward = 5
a2 → future → reward = 10 ← 选它
a3 → future → reward = 2

为什么用 CEM 而不是枚举

动作空间是连续的(例如 at=[x,y,z,rotation]a_t=[x,y,z,\mathrm{rotation}]),不能像离散搜索那样枚举。规划长度 H=12H=12、每步 6 维时搜索空间是 6126^{12},巨大。所以用概率优化:不直接找动作序列,而是学习动作分布 q(at:t+H)q(a_{t:t+H})

CEM 的核心流程

CEM(Cross-Entropy Method)维护一个按时间分解的动作序列高斯分布:

q(at:t+H)=N(μt:t+H,σt:t+H2I)q(a_{t:t+H}) = \mathcal N(\mu_{t:t+H}, \sigma_{t:t+H}^2 I)

Algorithm 2 的输入:

参数含义默认值
HHplanning horizon(规划长度)12
IICEM iterations(迭代次数)10
JJcandidates per iteration(每轮候选数)1000
KKelite candidates(保留的优秀数)100

模型输入:当前 state belief q(stot,a<t)q(s_t\mid o_{\le t},a_{<t})、transition model p(stst1,at1)p(s_t\mid s_{t-1},a_{t-1})、reward model p(rtst)p(r_t\mid s_t)

每轮迭代:

  1. 初始化动作分布 q(at:t+H)N(0,I)q(a_{t:t+H})\leftarrow\mathcal N(0,I)
  2. 从当前分布采样 JJ 条候选动作序列 at:t+H(j)qa^{(j)}_{t:t+H}\sim q
  3. 从当前 belief 出发,用 transition model rollout 每条序列:
st:t+H+1(j)q(stot,a<t)τ=t+1t+H+1p(sτsτ1,aτ1(j))s_{t:t+H+1}^{(j)} \sim q(s_t\mid o_{\le t},a_{<t}) \prod_{\tau=t+1}^{t+H+1} p(s_\tau\mid s_{\tau-1},a_{\tau-1}^{(j)})
  1. 用 reward model 累加未来预测奖励:
R(j)=τ=t+1t+H+1E[p(rτsτ(j))]R^{(j)} = \sum_{\tau=t+1}^{t+H+1} \mathbb E[p(r_\tau\mid s_\tau^{(j)})]
  1. 按 reward 排序,取前 KK 条 elite sequences。
  2. 用 elite 的均值和方差重新拟合动作分布:
μ=1Ka(k),q(a)N(μ,σ2I)\mu=\frac{1}{K}\sum a^{(k)},\qquad q(a)\leftarrow\mathcal N(\mu,\sigma^2 I)
  1. 重复 II 次,最后返回当前时刻动作均值 μt\mu_t
一个具体数值例子:CEM 如何把随机动作逐渐集中到高收益区域?

假设只规划一个标量动作,初始 q(a)=N(0,1)q(a)=\mathcal N(0,1)。第一轮采到 6 个候选:

[1.2,0.4,0.1,0.5,0.8,1.4][-1.2,-0.4,0.1,0.5,0.8,1.4]

世界模型预测收益最高的两个是 0.50.50.80.8,elite 均值:

μ=0.5+0.82=0.65\mu=\frac{0.5+0.8}{2}=0.65

下一轮改从集中在 0.650.65 附近的高斯采样,更容易得到 0.6,0.7,0.750.6,0.7,0.75 等高收益动作。多轮后分布逐渐聚焦于模型认为的最优区域。

连续高维动作序列也是同样的过程,只是均值和方差具有“时间步 × 动作维度”的形状;每条候选序列产生一条独立的 imagined latent trajectory,而不是在候选之间共享未来状态。

为什么可以不生成图片

论文最后一句关键:planner can operate purely in latent space without generating images。很多模型需要 latent → decoder → image,但每次规划都生成图像非常慢。PlaNet 直接:

latent state → reward model → reward

image space 是 224×224×3224\times224\times3,latent space 只有几十维,搜索速度大幅提高。

CEM 与 MPC 的关系

CEM 搜索未来 H 步 → 只执行第一步 → 收到新图像更新 posterior → 重新从 N(0,I) 开始 CEM

这就是 receding horizon control / MPC。每一步重新初始化动作分布,减少陷入先前局部最优的风险。

⭐ 四个模型是否共同训练?未收敛时为什么还敢规划采数据?

用户追问:Algorithm 1 的第一个 for 就完成了一轮对 transition、observation、reward、encoder 四部分的联合训练吗?第二个 for 就是为了产数据吗?模型未收敛,怎么确保产的数据有利于下一次训练?

1. 第一个 for 循环联合训练整套 world model

是的,这个循环里联合训练 PlaNet 的整套世界模型,而不是先单独训练某一个模块。transition / dynamics、observation / decoder、reward model、以及 state inference / posterior model 都包含在同一个参数集合 θ\theta 中:

θθαθL(θ)\theta\leftarrow \theta-\alpha\nabla_\theta\mathcal L(\theta)

是一次端到端更新:

从 replay buffer D 取一段 trajectory

encoder / inference model q → latent states

RSSM transition p(s_t|s_{t-1},a_{t-1})

同时预测 image p(o|s) 和 reward p(r|s)

计算 ELBO / model loss

一起反向传播

不是“先训练 transition,再训练 observation,再训练 reward,再训练 q”,而是一个序列 batch 同时产生这些 loss,然后一起反向传播。

2. “C 次 update”不等于模型已收敛

for update step s=1...C 只表示这一轮交互前对已有数据做 C 次梯度更新,不是训练到收敛。真正的收敛发生在外层 while 循环上:

D0 → 训练 C 步 → θ1 → 用 θ1 采一段数据 → D1 = D0 + 新数据 → 再训练 → θ2 → ...

3. 第二个 for 基本就是产数据

对,t=1,,T/Rt=1,\dots,T/R 主要就是 data collection,但它同时做两件事:用当前模型真正执行任务、产生新的训练数据。

最关键的一点:模型生成的预测不会直接当成训练标签。真正加入 D\mathcal D 的仍然是真实环境产生的数据,不是模型 imagination 出来的数据。

4. 模型不准,规划出来的动作不也很烂吗?

完全正确。假设模型学到“猛踩油门 → reward 很高”,CEM 就会选猛踩油门,真实环境给出撞墙、reward = -100。但这个失败数据本身就是训练数据:下一轮模型必须解释 (st,猛踩油门)ot+1,rt=1(s_t,\text{猛踩油门})\rightarrow o_{t+1}, r_t=-1,于是逐渐修正。

错误规划真实环境纠错模型改善\boxed{\text{错误规划} \rightarrow \text{真实环境纠错} \rightarrow \text{模型改善}}

这实际上是一个 bootstrapping / self-improving loop:

随机模型 → 差的数据/失败 → 纠正模型 → 稍好的规划 → 更有意义的数据 → 更好的模型 → 更好的规划

5. 为什么不会陷入恶性循环

  • 初始不是零数据:先随机策略跑 SS 个 seed episodes,第一次规划时模型已经看过一批真实 transition。
  • exploration noise:即使 planner 认为 a=0.3a=0.3,实际执行可能是 0.27 / 0.35 / 0.22,扩大状态覆盖,避免 exploration failure。
  • 训练数据是真实数据:即使 s^t+1st+1\hat{s}_{t+1}\neq s_{t+1},最终存入 D\mathcal D 的仍是 ot+1real,rtrealo_{t+1}^{\mathrm{real}}, r_t^{\mathrm{real}},不会形成“模型预测错 → 把错误预测当 label → 越来越错”的循环。
  • planner 自动把数据采到相关区域:随机采样每个状态区域都有数据,但任务只需要 C → D → E;模型稍微学会后,CEM 逐渐探索 C/D/E 附近,数据从 random exploration 变成 task-directed exploration。
  • replay buffer / off-policy:训练时从整个 D\mathcal D 均匀采样,包含早期 random、旧 planner、当前 planner 的数据,单轮坏数据不会成为唯一训练来源。

6. 但 model exploitation 确实存在

Planner 会寻找模型认为 reward 高的地方,即使真实 reward 不高:

a=argmaxaRmodel(a),Rreal(a)Rmodel(a)a^{*}=\arg\max_a R_{\mathrm{model}}(a), \qquad R_{\mathrm{real}}(a^{*})\ll R_{\mathrm{model}}(a^{*})

PlaNet 能工作靠的是几个友好条件:连续控制任务动力学相对平滑;每次只执行规划序列中的第一个动作(plan long, execute short):

现在规划 H 步 a1..aH → 只执行 a1 → 得到真实 observation → 重新推断 state → 重新规划

因此模型误差不会无限开环积累;后续 PETS、MBPO、MOPO 等工作正是围绕 ensemble、uncertainty estimation 和 conservative model-based planning 研究这一问题。

⭐ RSSM:为什么同时需要确定性与随机性?

RNN、SSM 与 RSSM 的潜变量动力学设计对比

图源:Hafner et al., Figure 2, p.4。圆形是随机变量,方形是确定性变量;实线是生成过程,虚线是推断模型。

基本假设

我们看到的是图像 oto_t,但背后存在不可见真实状态 sts_t(机械臂位置、物体位置、速度……),形成 latent state-space model:

s1 → s2 → s3 → ...
↓ ↓ ↓
o1 o2 o3
r1 r2 r3

三个基本概率模型:

stp(stst1,at1),otp(otst),rtp(rtst)s_t\sim p(s_t\mid s_{t-1},a_{t-1}),\qquad o_t\sim p(o_t\mid s_t),\qquad r_t\sim p(r_t\mid s_t)

为什么不用普通 RNN(图 2a)

Deterministic model 只有 hth_t

ht=f(ht1,at1)h_t=f(h_{t-1},a_{t-1})

给定 ht,ath_t,a_t,未来唯一。但现实有随机性:机器人推一下球,球可能向左、向右或停止。RNN 只能给出一个确定未来,论文说它 prevents the model from capturing multiple futures。更严重的是 planner 会利用这种错误:模型认为“动作 A → reward=100”,现实是“动作 A → reward=-10”,这就是 model exploitation。

为什么不用纯随机 SSM(图 2b)

状态完全随机:

stp(stst1,a)s_t\sim p(s_t\mid s_{t-1},a)

优点是可以表示多个未来。但长期信息容易丢:每一步都重新采样,过去 10 秒的“拿起钥匙 → 打开门 → 进入房间”这些历史难以保留。论文说 difficult to remember information over multiple time steps。

RSSM(图 2c):两者结合

状态拆成两部分:

st=(ht,zt)s_t=(h_t,z_t)
  • hth_t:确定性部分,负责 memory,类似 RNN hidden state,保存长期依赖。
  • ztz_t:随机部分,负责 uncertainty,表达多种未来。

推球例子:随机部分 ztz_t 表示“球可能向左/右/停止”的不确定性;确定性部分 hth_t 记住“看到球、碰到障碍、知道环境布局”这些历史背景。

RSSM 的具体过程

ht=f(ht1,zt1,at1)h_t=f(h_{t-1},z_{t-1},a_{t-1})

先更新确定性状态,然后:

p(ztht)(prior,不看图片,我猜现在状态)p(z_t\mid h_t) \quad(\text{prior,不看图片,我猜现在状态}) q(ztht,ot)(posterior,看到真实图片后修正)q(z_t\mid h_t,o_t) \quad(\text{posterior,看到真实图片后修正})

类似 Kalman filter:先预测,再观察,再纠正。

图中实线与虚线的含义

论文原句:solid lines denote generative process and dashed lines the inference model。实线是生成过程(state → observation),虚线是推断模型(observation → state)。学习时只有 oto_t,需要 inference model q(stot,a<t)q(s_t\mid o_{\le t},a_{<t}) 反推状态。

为什么 observation / transition / reward 都是 Gaussian

训练方便:如果 otN(μ,I)o_t\sim\mathcal N(\mu,I),log-likelihood 等价于 otμ2-\|o_t-\mu\|^2,即 MSE / L2。所以图像重建 loss 就是普通 L2。

三种模型总结

模型状态优点缺点
RNN只有 hth_t记忆强没有不确定未来
SSM只有 sts_t可以多个未来长期记忆差
RSSMht+zth_t+z_t记忆 + 随机未来复杂但效果最好

与 Pi3X / VGGT 的联系

Pi3X 学的是空间 latent state(image → 3D representation → point map / camera),PlaNet 学的是动态 latent state(image → latent → predict future)。未来机器人方向希望共享 latent world state:3D point + motion flow + dynamics prediction,一起支撑 planning / action。

⭐ Filtering posterior、prior 与 observation 的关系

为什么 posterior 要看当前图像,而 transition prior 不看?

transition prior:

p(stht)p(s_t\mid h_t)

只根据过去状态和动作形成的 hth_t 预测当前状态,回答“如果暂时没有新观测,我认为世界到了哪里”。

filtering posterior:

q(stht,ot)q(s_t\mid h_t,o_t)

额外看到真实图像,用来修正预测,回答“结合当前证据后,我认为世界在哪里”。

过去 latent + 上一步动作

prior prediction

当前图像 o_t
↓ correction
posterior belief

Filtering 在估计 tt 时只使用 oto_{\le t};smoothing 则可以利用未来 o>to_{>t} 回头修正过去。PlaNet 使用 filtering posterior,因为在线规划时未来观测尚不存在,训练与执行条件保持一致。

⭐ 变分编码器与 ELBO 训练目标

为什么需要 Variational Encoder

论文第一句话:Since the model is non-linear, we cannot directly compute the state posteriors。因为模型非线性,无法解析计算:

p(sto1:t,a1:t)p(s_t\mid o_{1:t},a_{1:t})

这是一个复杂的 inverse problem(只有 camera 图像,反推机械臂位置/物体位置/速度),所以采用 Variational Inference,学习近似 posterior q(st)q(s_t\mid\cdots) 代替 p(st)p(s_t\mid\cdots)

Encoder 的公式

q(s1:To1:T,a1:T)=t=1Tq(stst1,at1,ot)q(s_{1:T}\mid o_{1:T},a_{1:T}) = \prod_{t=1}^{T} q(s_t\mid s_{t-1},a_{t-1},o_t)

左边是整个时间序列的 latent state 分布;右边表示每一步独立递推:当前状态由上一个 latent state、上一个 action、当前 observation 决定,类似 Kalman filter。

为什么 q 里面需要 o_t

transition model p(stst1,at1)p(s_t\mid s_{t-1},a_{t-1}) 没有 oto_t,因为它是“只知道过去,预测未来”;encoder q(stst1,at1,ot)q(s_t\mid s_{t-1},a_{t-1},o_t) 加入当前图片,因为要“根据真实观察修正状态”。两者关系:

Prediction: p(s_t | s_{t-1},a) → prior
Correction: q(s_t | s_{t-1},a,o) → posterior

这就是 Bayesian filtering。

为什么叫 filtering posterior

  • Filtering:只用过去和当前,q(stot,a<t)q(s_t\mid o_{\le t},a_{<t}),实时估计当前状态。
  • Smoothing:使用未来信息,q(sto1:T)q(s_t\mid o_{1:T}),视频结束后回头重新估计第 10 帧。

例子:视频 t=1..5t=1..5,估计 t=3t=3 的状态——filtering 只看 1,2,3;smoothing 看 1,2,3,4,5。

PlaNet 用 filtering,因为最终用途是 planning:机器人运行时未来图像不存在,只能 oto_{\le t},训练也保持一致。

Training Objective:ELBO

论文目标是最大化:

lnp(o1:Ta1:T)\ln p(o_{1:T}\mid a_{1:T})

直接最大化困难,改用变分下界(ELBO):

lnp(o1:Ta1:T)t=1T(Eq[lnp(otst)]Eq[DKL(q(stot,a<t)p(stst1,at1))])(3)\ln p(o_{1:T}\mid a_{1:T}) \ge \sum_{t=1}^{T} \left( \mathbb E_q[\ln p(o_t\mid s_t)] - \mathbb E_q\left[ D_{\mathrm{KL}}\bigl( q(s_t\mid o_{\le t},a_{<t}) \|p(s_t\mid s_{t-1},a_{t-1}) \bigr) \right] \right) \tag{3}

reward likelihood 的形式与 observation 项类似,论文为简洁在式 (3) 中省略。

第一项:reconstruction

Eq[lnp(otst)]\mathbb E_q[\ln p(o_t\mid s_t)]

要求 latent state 能重建观察:image → encoder → sts_t → decoder → image reconstruction。如果重建不好,说明 latent 没有包含视觉信息。论文的 observation model 是单位协方差高斯,因此最大化 log-likelihood 在常数意义下等价于最小化像素 MSE。

第二项:complexity / KL

DKL(qposteriorpprior)D_{\mathrm{KL}}\bigl(q_{\mathrm{posterior}}\,\|\,p_{\mathrm{prior}}\bigr)

posterior 看到了真实图片,prior 只能根据动力学预测。KL 要求 qpq\approx p:看到图片后的状态不要偏离动力学预测状态太多。

为什么需要:如果没有 KL,encoder 可以作弊——每帧图像都 memorized 成任意 latent code,decoder 再还原,reconstruction 很好但 dynamics 完全学不到。KL 强迫 latent 必须符合“过去状态 + 动作”的演化规律。

Sequential VAE

普通 VAE:image → zz → image。PlaNet:video sequence → s1,s2,s_1,s_2,\dots → future prediction。增加了 temporal dynamics 和 action conditioning,可以理解为“一个带动作控制的 VAE + RNN dynamics”。

与 Dreamer 的关系

PlaNet 训练 q(st)q(s_t) 后做 CEM planning;Dreamer 同样用 RSSM,但利用 latent imagination 直接训练 policy。RSSM 是 Dreamer / DreamerV2 / DreamerV3 的核心。

Latent overshooting

标准 ELBO 的 KL 主要直接训练一步 transition。即使一步预测误差不大,多步 rollout 仍可能累积偏差;而规划恰恰依赖多步 reward prediction。

标准变分目标、observation overshooting 与 latent overshooting

图源:Hafner et al., Figure 3, p.5。波浪箭头表示 latent KL,指向灰色观测的箭头表示 reconstruction likelihood。

三种方案:

  1. Standard variational bound:每一步 posterior 重建观测,并用一步 prior 对齐 posterior。
  2. Observation overshooting:从更早状态 rollout 多步,并把每个未来 latent 解码成图像监督;准确但在图像域很贵。
  3. Latent overshooting:只在潜空间把多步 prior 对齐对应 posterior,无需反复生成图像。

其本质是要求:

p(ststd,atd:t1)q(stot,a<t),d=1,,Dp(s_t\mid s_{t-d},a_{t-d:t-1}) \approx q(s_t\mid o_{\le t},a_{<t}), \qquad d=1,\ldots,D

注意:Latent overshooting 是论文提出的重要泛化方法,但最终 RSSM agent 并不依赖它——附录报告它能明显帮助某些其他动力学模型,却使 RSSM 性能略降。不能把 PlaNet 的成功简单归因于 overshooting。

训练配置

官方补充材料给出的主要配置:

  • 图像缩小为 64×64×364\times64\times3,并降低到 5-bit 色深。
  • RSSM 确定性路径:200-unit GRU。
  • 随机 latent:30 维对角高斯。
  • 其他函数:两层 200-unit ReLU MLP。
  • batch size B=50B=50,sequence chunk length L=50L=50
  • Adam 学习率 10310^{-3},gradient clipping norm 1000。
  • KL 使用 3 free nats。
  • 5 个随机 seed episodes;每 100 次模型更新收集一个新 episode。
  • 探索噪声:ϵN(0,0.3)\epsilon\sim\mathcal N(0,0.3)

实验结论

实验覆盖 Cartpole、Reacher、Cheetah、Finger、Cup 和 Walker,观测均为像素图像,包含部分可观测、接触动力学和稀疏奖励任务。

主要结论:

  • RSSM 通常优于纯 deterministic GRU 和纯 stochastic SSM,支持“确定性记忆 + 随机状态”的设计。
  • PlaNet 用 1000 个 episodes 达到接近或部分超过强 model-free baseline 的最终表现。
  • 论文按各任务估算,相对 D4PG 达到相同性能所需数据少约 40 到 500 倍以上;摘要概括为平均约 200 倍更少环境交互。
  • CEM 优于 random shooting;用 planner 引导数据采集优于始终随机采集。
  • 学得的 latent 中包含可用于预测 simulator 位置、速度和 reward 的信息。
注意

这些效率数字来自特定 DeepMind Control Suite 设置、baseline 曲线估算和 2019 年方法,不能直接外推到真实机器人或与今天的算法横向比较。

与 Dreamer 的关系

PlaNet 与 Dreamer 共享 RSSM 世界模型,但动作选择方式不同:

PlaNetDreamer
每个环境步在线运行 CEM在 latent imagination 中训练 actor/critic
没有 policy/value network有可快速执行的 policy 与 value
增加搜索预算可改善动作部署动作选择更便宜
在线规划计算较重训练行为模型更复杂

可以概括为:

PlaNet = RSSM + online planning
Dreamer = RSSM + latent imagination actor-critic

与三维视觉研究的联系

VGGT、Pi3 一类模型主要学习静态空间表示:

imagescamera / depth / point map\text{images}\rightarrow\text{camera / depth / point map}

PlaNet 学习的是可预测、动作条件化的时间状态:

(images,actions)latent dynamicsplanning(\text{images},\text{actions}) \rightarrow \text{latent dynamics} \rightarrow \text{planning}

二者并非互相替代。面向机器人的世界表示可以同时包含:

shared visual representation
├── 3D geometry / camera / depth
├── optical flow / object motion
├── action-conditioned dynamics
└── reward / task prediction

老师说的“不要只做三维重建,多预测一些东西,让预测之间互相辅助”,从理论上就是让 latent 不仅满足 p(os)p(o\mid s),还满足 p(depths)p(\mathrm{depth}\mid s)p(flows)p(\mathrm{flow}\mid s)p(actions)p(\mathrm{action}\mid s)——从 representation learning 走向 world model learning。

局限

  • 像素重建会迫使 latent 表达与任务无关的视觉细节。
  • 世界模型在分布外区域可能被 CEM 利用(model exploitation)。
  • 没有显式 ensemble uncertainty 或 conservative objective。
  • CEM 每个真实时间步评价大量动作序列,在线计算昂贵。
  • 固定短 horizon 可能忽略很远的奖励。
  • 任务 reward 仍需要提供,整套 agent 不是完全 task-agnostic。
  • 结果主要来自模拟环境,真实机器人中的视觉变化、安全和数据代价更难。

最终理解

PlaNet 的核心不是简单的“图像经过 encoder 再做 MPC”,而是把三个目标绑定在同一个 latent 上:

能解释当前观测+能按动作预测未来+能支持奖励规划\boxed{ \text{能解释当前观测} +\text{能按动作预测未来} +\text{能支持奖励规划} }

ELBO 让 posterior 与 dynamics prior 对齐,RSSM 同时保存确定性历史和随机不确定性,CEM 在 latent 中搜索动作,MPC 用真实新观测不断校正模型误差。整套方法的价值在于:让视觉表示从“描述当前画面”变成“可用于预测与决策的世界状态”。