【具身智能】Cosmos Policy论文阅读随笔

简述

Cosmos Policy 构建于 Cosmos-Predict2（一种视频模型）之上，通过其海量互联网视频预训练的基础，使得模型能够处理机器人任务。

核心

潜在帧插入

所谓潜在帧插入，就是将额外的图像或机器人观测插入帧序列。

要搞清楚这到底是什么意思，我们必须先理解 Cosmos-Predict2 的原理。

前置：Cosmos-Predict2

Cosmos-Predict2 是一种视频生成模型，它吃一个初始图像和文本描述，生成后续视频。

采用扩散模型。

VAE

视频序列被 Wan2.1 编码为序列。

具体的，一个 \((1 + T) \times H \times W \times 3\) 的序列被压缩为 \((1+T') \times H' \times W' \times 16\) 的序列，\(T,H,W\)分别以4，8，8倍进行压缩。

对于\((1+T)\)，1 代表不进行时间压缩的初始图像。

扩散模型

扩散模型处理视频序列生成。

具体的，模型训练一个网络\(D_\theta\)，目标：

\[\mathcal{L}(D_\theta, \sigma) = \mathbb{E}_{\mathbf{x}_0, \mathbf{c}, \mathbf{n}} \left[ \| D_\theta(\mathbf{x}_0 + \mathbf{n}; \sigma, \mathbf{c}) - \mathbf{x}_0 \|_2^2 \right] \]

其中，\(\sigma\)为噪声强度，\(\mathbf{x}_0\)为干净图像序列，\(c\)为文本嵌入，\(\mathbf{n}\)为噪声。
（这是很平常的训练目标，本人并未在其中发现什么）

潜在帧插入

思考：如何让视频模型变为机器人任务模型（或者说世界动作模型）

在基础上，增加：

• 动作预测模型
• 世界模型
• 奖励预测模型

然而这是另一种模态的嵌入

它们采用一种“简单”的方法：
直接将多模态压缩后插进扩散模型处理序列中。

假设有11个潜帧：
1.空白占位符：一个保留位置。
2.机器人本体感知：如末端执行器姿态、关节角度。
3.腕部相机图像。
4.第一个第三人称相机图像。
5.第二个第三人称相机图像。
6.动作块：机器人要执行的动作序列。
7.未来机器人本体感知：预测的下一时刻本体状态。
8.未来腕部相机图像：预测的下一时刻腕部画面。
9.未来第一个第三人称相机图像。
10.未来第二个第三人称相机图像。
11.未来状态价值：预测的最终任务回报。
就是非常直接的，将其它模态的输入（或输出），插入其中了。
我来形象化的描述：😂
原本，扩散模型被要求处理1+T的视频序列，这里，T是很多个1帧的图像。
Cosmos Policy，将这个1帧的图像，变成了11小帧，每个小帧对应1个不同模态。

世界模型

接下来要从原模型中得到

• 策略：\(p(a, s', V(s') \mid s)\)，给定当前状态，生成动作、未来状态和价值。
• 世界模型：\(p(s', V(s') \mid s, a)\)，给定状态和动作，预测未来状态和价值。
• 价值函数模型：\(p(V(s') \mid s, a, s')\)，给定状态、动作和未来状态，评估价值。

将一步一步理清
（这里我将先讲微调和推理，后讲训练）

微调和推理

假设我们已经有了策略模型，世界模型，价值函数模型。

微调

先将多样的数据交由策略模型，生成rollout。
将这些由策略模型生成的rollout，90%交由世界模型，价值函数模型进行微调，10%交由策略模型微调。

这样，世界模型和价值函数模型能从更加多样的，且都是由策略模型展开的，rollout中进行学习预测，使得它们更加考虑策略模型的能力，且提升自己的探索性和预测力。

最后，微调出来的整个模型，将策略模型部分拆掉，装上微调之前的策略模型。
（其它两个部分本就是基于原策略模型的rollout进行微调的，这个微调后的策略模型实际上起到一个协同，保护另外两个部分的作用）

推理

推理模式：

• 并行解码：a，s‘，V’一同跑一遍扩散模型，其中通过交叉注意力等机制获取信息。
  （适用于实时动作生成任务）
• 自回归解码：s'先以a为条件，跑一遍扩散模型；然后V‘再以跑出来的最准确的s'为条件，跑一遍扩散模型，得到最准确的V’。
  （适用于规划任务）

规划机制：

• 1.从模型中采样N个动作，对于每个动作：
  • 通过世界模型3次，每次通过价值模型5次，生成15个预测；
  • 进行阈值归类后，进行类内平均得到最终预测。
• 2.选择最高预测的动作
• 3.N个动作由N块GPU并行加速

训练

基础训练数据以50%给策略模型，分别25%给世界模型和价值模型。

上文写（模型用）到很多处，以什么为条件生成什么，这在扩散模型中，使用掩码标记（加噪分类），进行条件和目标的归类。

总结

本文实际上只解析了论文中如何基于现有模型进行的transformation的方面。
只是想总结一下Cosmos Policy的最核心的架构吧。

然后官方地址在这👋