QiuQiu

Dream To Control 精读

Learned world models summarize an agent’s experience…

例如机器人收集过:

(图像 ot, 动作 at, 奖励 rt)(\text{图像 }o_t,\ \text{动作 }a_t,\ \text{奖励 }r_t)

世界模型学习:

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

也就是:

在当前状态采取这个动作,接下来大概会进入什么状态?

st 是神经网络提取的低维 latent state。

  1. While learning world models from high-dimensional sensory inputs.

现在深度学习已经让我们可以直接从摄像头图像等高维输入中学习世界模型。

但问题是:

学会预测未来以后,怎么利用这些预测产生好的动作?

可以有很多方法:

Dreamer 主要创新的是后面两项。

    1. Dreamer solves long-horizon tasks from images purely by latent imagination

Dreamer 从图像得到当前 latent state,然后在 latent space 里模拟未来:

st→atst+1→at+1st+2→⋯s_t \xrightarrow{a_t} s_{t+1} \xrightarrow{a_{t+1}} s_{t+2} \xrightarrow{} \cdots

同时预测每一步奖励:

r^t,r^t+1,r^t+2,…\hat r_t,\hat r_{t+1},\hat r_{t+2},\ldots

这就是 latent imagination

注意,它不是生成未来图片,而是直接预测未来的 latent state,所以计算便宜很多。

“purely by latent imagination”也不是说完全不需要真实数据。准确意思是:

Dreamer 有三个关键部分:

Actor→World Model→Reward/Value\text{Actor} \rightarrow \text{World Model} \rightarrow \text{Reward/Value}

Actor 选择动作:

at=πϕ(st)a_t=\pi_\phi(s_t)

世界模型预测未来状态:

st+1=fθ(st,at)s_{t+1}=f_\theta(s_t,a_t)

reward/value model 判断这条未来轨迹有多好:

r^t,Vψ(st)\hat r_t,\quad V_\psi(s_t)

因为这些部分都是可微神经网络,所以可以使用链式法则:

∂未来回报∂ϕ=∂未来回报∂st+1∂st+1∂at∂at∂ϕ\frac{\partial \text{未来回报}}{\partial \phi} = \frac{\partial \text{未来回报}}{\partial s_{t+1}} \frac{\partial s_{t+1}}{\partial a_t} \frac{\partial a_t}{\partial \phi}

直观意思是:

如果把当前动作稍微向左改变一点,未来回报会提高还是下降?

然后梯度可以一路穿过想象轨迹,告诉 Actor 应该怎样改变参数。

这和真实环境不同:现实世界通常不能直接求导;但 Dreamer 学到的 world model 是神经网络,因此能够求导。

5. long-horizon 为什么成立?

世界模型不能无限准确地向未来预测,所以 Dreamer只想象有限的 HH 步:

r^t+γr^t+1+⋯+γH−1r^t+H−1\hat r_t+\gamma\hat r_{t+1}+\cdots+ \gamma^{H-1}\hat r_{t+H-1}

到轨迹结尾,再让 critic 估计更远的未来:

+γHVψ(st+H)+\gamma^H V_\psi(s_{t+H})

所以:

这就是它比只看有限 planning horizon 的方法更不容易“短视”的原因。

Ref:

https://www.youtube.com/watch?v=awyuuJoHawo

RSSM? 什么

RSSM 全称是 Recurrent State-Space Model

RSSM 是 Dreamer 的“内部世界状态跟踪器”:它把过去的图像和动作压缩成当前 latent state,并预测采取动作后未来状态怎样变化。

为什么需要它?

机器人每次只看到一张图片,但单张图片不足以描述完整状态。

例如,一张图片里有人站在前方,但看不出:

因此 Dreamer 需要把过去信息积累起来:

(o1,a1,o2,a2,…,ot)→st(o_1,a_1,o_2,a_2,\ldots,o_t) \rightarrow s_t

这个 sts_t 是对当前世界状态的内部估计。

RSSM 的两个部分

RSSM 的状态通常写成:

st=(ht,zt)s_t=(h_t,z_t)

1. Deterministic state hth_t

ht=f(ht−1,zt−1,at−1)h_t=f(h_{t-1},z_{t-1},a_{t-1})

它通常由 GRU 实现,负责保存历史信息,类似机器人的“记忆”。

例如可能隐式保存:

它是 deterministic 的:相同输入一定产生相同输出。

2. Stochastic state ztz_t

zt∼q(zt∣ht,ot)z_t\sim q(z_t\mid h_t,o_t)

它是从一个概率分布中采样的 latent variable,负责表示:

例如前方的人可能左转,也可能右转。纯确定性模型容易预测成“两种未来的平均值”;随机状态可以表达多个可能性。

RSSM 怎样工作?

在真实环境中,Dreamer看得到当前图片,因此使用 posterior:

q(zt∣ht,ot)q(z_t\mid h_t,o_t)

流程是:

ot→CNNeto_t\xrightarrow{\text{CNN}}e_t ht=f(ht−1,zt−1,at−1)h_t=f(h_{t-1},z_{t-1},a_{t-1}) zt∼q(zt∣ht,et)z_t\sim q(z_t\mid h_t,e_t)

最终:

st=[ht,zt]s_t=[h_t,z_t]

然后用这个状态重建图像、预测奖励:

o^t=D(ht,zt)\hat o_t=D(h_t,z_t) r^t=R(ht,zt)\hat r_t=R(h_t,z_t)

它怎样“想象未来”?

未来图片还没有发生,所以不能使用 posterior。Dreamer改用不看图片的 prior:

zt+1∼p(zt+1∣ht+1)z_{t+1}\sim p(z_{t+1}\mid h_{t+1})

然后循环预测:

(ht,zt)→at(ht+1,zt+1)→at+1(ht+2,zt+2)(h_t,z_t) \xrightarrow{a_t} (h_{t+1},z_{t+1}) \xrightarrow{a_{t+1}} (h_{t+2},z_{t+2})

这就是 latent imagination。

训练时通过 KL divergence 让 prior 接近看过真实图片的 posterior:

DKL[q(zt∣ht,ot)∥p(zt∣ht)]D_{\mathrm{KL}} \left[ q(z_t\mid h_t,o_t) \parallel p(z_t\mid h_t) \right]

意思是:

让“不看当前图片的预测结果”,尽量接近“看过当前图片以后得到的状态”。

因此未来没有图片时,RSSM也能根据历史和动作继续推演。

与普通RNN的区别

模型 状态 能表达不确定性 用途
RNN/GRU 只有确定性 hth_t 较弱 保存序列记忆
VAE 随机 latent zz 可以 压缩单张图片
RSSM 确定性 hth_t + 随机 ztz_t 可以 跟踪状态并预测未来

所以 RSSM 可以理解成:

RNN的长期记忆+VAE的随机latent state\boxed{\text{RNN的长期记忆}+\text{VAE的随机latent state}}

在 Dreamer 中,RSSM 是 world model 的核心;CNN只是把图片编码,真正负责“记住过去并预测未来”的是 RSSM。