简述

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

核心

潜在帧插入

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

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

前置:Cosmos-Predict2

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

采用扩散模型。

VAE

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

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

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

扩散模型

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

具体的,模型训练一个网络 D θ D_\theta Dθ,目标:
L ( D θ , σ ) = E x 0 , c , n [ ∥ D θ ( x 0 + n ; σ , c ) − x 0 ∥ 2 2 ] \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] L(Dθ,σ)=Ex0,c,n[Dθ(x0+n;σ,c)x022]

其中, σ \sigma σ为噪声强度, x 0 \mathbf{x}_0 x0为干净图像序列, c c c为文本嵌入, n \mathbf{n} 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 ′ ) ∣ s ) p(a, s', V(s') \mid s) p(a,s,V(s)s),给定当前状态,生成动作、未来状态和价值。
  • 世界模型: p ( s ′ , V ( s ′ ) ∣ s , a ) p(s', V(s') \mid s, a) p(s,V(s)s,a),给定状态和动作,预测未来状态和价值。
  • 价值函数模型: p ( V ( s ′ ) ∣ s , a , s ′ ) p(V(s') \mid s, a, s') p(V(s)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的最核心的架构吧。

然后官方地址在这👋

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐