本文提出一种兼具可交互性和可扩展性的世界模型架构 iVideoGPT,基于视频的压缩标记化和自回归 Transformer,支持在百万条人类和机器人操作轨迹上预训练,并适配到动作条件视频预测、视觉规划和基于模型的强化学习等多种控制相关任务。
作者
:吴佳龙*,尹绍沣*,冯宁亚,和煦,李栋,郝建业,龙明盛
链接
:https://arxiv.org/pdf/2405.14369
主页
:https://thuml.github.io/iVideoGPT
代码 & 预训练模型
:https://github.com/thuml/iVideoGPT
近年来,以Sora为代表的视频生成模型,能够生成逼真的长视频,产生了广泛的应用。这引发了人们对于构建世界模型的设想:
视频生成模型是世界模型吗?
我们的回答是
还不是
。
-
从
任务
的视角,
世界模型需要具备可交互性
。它使得智能体能够在虚拟世界中根据当前观测产生候选动作,推演动作产生的后果,并根据想象的观测再次执行动作,如此往复,逐步找到解决问题的途径,相比于在真实世界中尝试更加高效和安全。
-
从
数据
的视角,
世界模型需要具备可扩展性
。它需要能够建模足够复杂的真实世界转移,而不仅仅局限在视觉简单的游戏上;更进一步地,我们希望世界模型能够从来自互联网的超大规模数据中学习世界演变的通用知识,从而能够泛化或快速迁移到特定任务场景。
-
然而,从
模型
的视角,在本文的工作开展时,主流的模型并没有同时满足交互性和扩展性的高要求:基于循环网络的世界模型缺乏高可扩展性,而视频生成模型则缺乏高可交互性(展开讨论见下文)。因此,本文的核心研究问题是:
如何基于可扩展的视频生成模型的进展来开发交互式的视觉世界模型?
世界模型是由智能体学习用于模拟环境的内部模型。环境通常被建模为部分可观测的马尔可夫决策过程 (POMDP)
,其中状态转移概率定义为
,且智能体往往仅能获得状态的不完全观测
。
尽管世界模型作为一种通用形式可以学习多种类型的数据,
视频是一种与任务无关、广泛可获取且包含丰富知识的模态
。因此,我们专注于学习视觉世界模型(即
),并将其形式化为
交互式视频预测
问题,这是一个智能体产生动作和世界模型推演转移的交替过程:
基于循环架构的世界模型
,如 Dreamer,天然支持如上
逐帧的可交互性
,但循环神经网络的可扩展性不足,难以建模复杂的现实世界数据。
视频生成模型
,如 Stable Video Diffusion 和 VideoGPT 等,虽然发展了更加可扩展的架构,但其往往具有时间维度上非因果的信息融合模块,导致它们仅能提供
轨迹级的交互性
:
即,仅允许在视频生成开始前输入文本/动作条件,缺乏在模拟过程中进行中间动作控制的能力,且通常生成固定长度的视频。
iVideoGPT是一种通用且高效的世界模型架构。它将视觉观测(通过
压缩标记化
)、动作和奖励等多模态信号整合为一个标记(token)序列,并通过
自回归 Transformer
进行下一个标记预测(next-token prediction)提供交互式体验。
3.1 视频的压缩标记化
Transformer 更加擅长处理离散标记序列。不同于使用每帧独立的图像 tokenizer 导致序列过长,或使用 3D tokenzier 进行时空压缩时牺牲交互性,我们提出了一种具有
双编码解码器
的条件 VQGAN
来进行视频的压缩标记化。
具体地,初始帧
包含丰富的上下文信息,被每帧独立地离散化为
个标记:
相比之下,
由于初始帧和未来帧之间存在时空冗余,仅需编码关键的动态信息
,例如移动物体的位置和姿态等。为此,采用了条件编码器和解码器,仅需更少数量的
个标记:
其中条件机制通过多尺度特征图之间的交叉注意力实现(与前序工作
ContextWM
相同)。
在本文中,
且
,带来了
每帧 16 倍的标记数量压缩
。这样的压缩标记化带来了两个主要益处:
-
更短的标记序列,可加快
生成速度
,从而加速基于模型的规划和强化学习的 rollout;
-
更容易保持视频内容的
时间一致性
,专注于关键动态信息的建模。
3.2 基于自回归 Transformer 的交互式预测
经标记化后,视频被展平为一个标记序列:
其中插入了特殊的槽标记(slot token)
,用于指示帧之间的边界,并便于融入额外的低维模态(见下)。一个 GPT 式的自回归 Transformer 被训练用于生成下一个标记,从而实现未来的逐帧预测。在本文中,我们采用了 GPT-2 尺寸的 LLaMA 架构。
采用下一个token预测的序列建模范式为我们带来了指定
不同任务、输入、输出的高度灵活性
。
(1) 序列建模的灵活性:多模态
iVideoGPT 可以灵活地整合不同模态的输入输出。在本文中,我们通过线性映射将
动作
信息添加到槽标记的嵌入中;对于
奖励
预测,我们在每帧观测的最后一个标记的隐藏状态后接一个线性预测头并使用 MSE 损失训练。
实现完全的多模态离散标记序列也是可行的,可以利用相关工作中已经发展了的动作、奖励量化(quantization)等技术。
(2) 序列建模的灵活性:多任务
为了初步展示兼容不同任务的灵活性,我们训练了一种 iVideoGPT 的变体,用于
基于目标的视频预测
(goal-conditioned video prediction):
。这可以
简单通过重新排列帧序列实现,而保持模型架构和训练过程不变
:
基于 GPT 的架构具有良好的可扩展性,这使我们可以在大规模数据上进行有效的预训练。
世界模型的
视频预训练范式
(action-free video pre-training)以视频预测作为自监督预训练任务。我们基于这一通用任务预训练 iVideoGPT,优化交叉熵损失来预测未来帧的标记:
其中
是需要预测的帧的第一个标记的下标。
本文中,我们专门为机械臂操作领域预训练了 iVideoGPT,使用
来自 Open X-Embodiment (OXE) 机械臂操作和 Something-Something v2 人手操作的混合数据集,总计 140 万条轨迹
。不同机械臂和人手具有高度异构的动作空间,但可以被轻松统一在视频预训练框架下。
预训练的 iVideoGPT 可以作为基础模型,微调到多种下游任务,包括基于动作的视频预测、基于视觉的模型预测控制和强化学习等。
5.1 视频预测
我们展示了 iVideoGPT 在不同数据集和设置上的视频预测结果。在大规模 Open X-Embodiment 数据上,基于
无动作视频预测
训练的 iVideoGPT 能够生成与真实轨迹不同,清晰、自然、多样化的机械臂操作轨迹,这说明它理解了与环境交互的一般规律。
基于目标的视频预测
训练的 iVideoGPT 则能够更加准确地预测到达目标的轨迹。
与之相比,基于循环网络架构的 DreamerV3-XL 模型(2 亿参数,与 iVideoGPT 相当)则缺乏在真实世界数据上进行大规模训练的能力,其在 OXE 上的预训练未能准确建模自然的机械臂动力学转移,仅能产生质量低且模糊的预测。
更多视频示例见:https://thuml.github.io/iVideoGPT/
iVideoGPT 可以自然扩展到基于动作的预测和高分辨率,并且不会有明显的时间不一致和闪烁伪影的问题。在下游数据集 BAIR robot pushing 和 RoboNet 上,iVideoGPT 取得了与最先进的模型相当的性能,并在架构上同时实现了交互性和可扩展性。