π0——用于通用机器人控制的流匹配VLA模型:一套框架控制7种机械臂,实现丝滑迁移与平替
前言
在此文之前,我花了一天半,详细解读了清华这个机器人扩散大模型RDT,包括其每一个附录,并在上文中预告说:下一篇是一个3B的机器人大模型打通7种不同品牌的机械臂,这几个工作宣告机器人真正进入大模型时代
故,本文来了
第一部分 整体理解π0
1.1 π0——用于通用机器人控制的流匹配VLA模型
1.1.1 背景
在自然语言[1]和计算机视觉[39]领域,预先在多任务数据上训练的通用基础模型往往比那些专门定制的模型表现更佳
例如,如果目标是识别照片中的鸟类,那么与其仅仅在鸟类识别数据上进行训练,不如先在许多不同的图像-语言关联上进行预训练,然后针对鸟类识别任务进行微调或提示
同样地,作者会发现
- 对于有效的专业机器人系统,先在高度多样化的机器人数据上进行预训练,然后针对所需任务进行微调或提示更为有效。这可以解决数据稀缺性挑战,因为通用模型可以获得更多的数据来源——包括其他任务、其他机器人,甚至非机器人来源的数据
- 同时也可以解决鲁棒性和泛化性挑战,因为多样化的数据展示了更广泛的观察和行动覆盖,提供了更多样的场景、修正和恢复行为,这些可能在更狭窄的专业数据中不存在
然而,开发这样的通用机器人策略——即机器人基础模型——涉及许多重大挑战
- 首先,任何此类研究都必须在非常大规模上进行,因为大规模预训练的全部好处通常在较小规模上无法显现[54]
- 其次,需要开发合适的模型架构,这些架构能够有效利用多样化的数据源,同时能够表现出与复杂物理场景交互所需的复杂和微妙行为
- 第三,需要正确的训练策略,包括最近在NLP和计算机视觉领域的大模型进展在很大程度上依赖于精细的策略来策划预训练和后训练数据[35]
1.1.2 预训练的视觉-语言模型VLM主干 + 动作专家通过流匹配输出动作
24年10月底,来自的研究者们提出了一个原型模型和学习框架,称之为π0,其对应的论文为《π0: A Vision-Language-Action Flow Model for General Robot Control》
下图图1中展示了该模型和系统
为了整合多样的数据源,作者
- 首先利用一个预训练的视觉-语言模型(VLM)来导入互联网规模的经验。通过基于VLM构建我们的模型,继承了语言和视觉-语言模型的通用知识、语义推理和问题解决能力
- 然后,进一步训练他们的模型以整合机器人动作,使其成为一个视觉-语言-动作(VLA)模型。为了能够利用多种不同的机器人数据源,作者采用跨化身训练[10],将多种机器人类型的数据合并到同一个模型中。这些不同的机器人类型具有不同的配置空间和动作表示,包括单臂和双臂系统,以及移动操纵器
其实这点和上一篇文章介绍的RDT是同样的问题,详见此文的.. - 此外,为了能够执行高度灵巧和复杂的物理任务,作者使用带有流匹配「[32,28],即扩散[20,46]的变体」的动作分块架构[57-Learning fine-grained bimanual manipulation with
low-cost hardware,即动作分块算法ACT,详见此文:]来表示复杂的连续动作分布[28,32]
相当于通过流匹配微调VLM以生成动作
这使得我们的模型能够以高达50 Hz的频率控制机器人进行如折叠衣物(见上图图1)这样的灵巧任务
且为了将流匹配与VLM结合,他们使用了一种新颖的动作专家,它通过流基输出增强了标准VLM
总之,虽然许多模型被提议结合预训练语言模型与扩散[40,41,14],包括专门将扩散与自回归大型语言模型混合的模型[19,29,59]
这些模型通常关注图像生成,但作者的动作生成模型与Zhou等[59]类似,其通过应用于单个序列元素的扩散风格(流匹配)损失训练他们的模型,以替代仅用于解码器的transformers的标准交叉熵损失。与Liu等[29]类似,为对应于扩散的标记使用了一套独立的权重。将这些概念融入VLA模型中,作者引入了据他们所知的第一个流匹配VLA,用于生成高频动作块以实现灵巧控制
此外,为了灵活且稳健地执行复杂任务,我们需要正确的训练策略,具体而言,模型首先在一个非常大且多样化的语料库上进行预训练,然后在更狭窄且更精心策划的数据上进行微调,以引导出所需的行为模式
1.2 模型架构
作者首先组装了一个预训练混合体,该混合体由我们自己的灵巧操作数据集(第V-C节)与整个OXE数据集[10]的加权组合组成,该数据集是在7种不同的机器人配置下为68个不同的任务收集的,而OXE数据集包含来自22个机器人的数据
- 预训练阶段(第V-A节)还使用了多样的语言标签,结合了任务名称和段注释(用于子轨迹的细粒度标签,通常长度约为2秒)
预训练阶段的目的是训练一个基础模型,该模型展示了广泛的能力和泛化性,但不一定专注于在任何一个任务上达到高性能。这个基础模型可以遵循语言指令并以基础水平执行多种任务 - 对于复杂和灵巧的任务,随后采用后训练程序(第V-A节),使用高质量的精心策划数据将模型调整为特定的下游任务。他们研究了使用少量到中等量数据的高效后训练,以及使用较大数据集进行复杂任务(如折叠衣物和移动操作)的高质量后训练
如下图图3所示
- π0模型主要由一个语言模型transformer骨干组成。遵循标准的后期融合视觉语言模型(VLM)方法[3,11,30],图像编码器将机器人的图像观测嵌入到与语言token相同的嵌入空间中。且进一步通过特定于机器人学的输入和输出进行增强——即本体感觉状态和机器人动作
- π0使用条件流匹配[28,32]来建模动作的连续分布。流匹配为他们的模型提供了高精度和多模态建模能力,使其特别适合高频灵巧任务
该架构灵感来自Transfusion [59],它通过多重目标训练单一transformer,使用tokens。对应于通过流匹配损失监督的连续输出和通过交叉熵损失监督的离散输出的token
在Transfusion的基础上,他们还发现,为机器人特定的(动作和状态)标记使用一组单独的权重可以提高性能
这种设计类似于专家混合[45,25,12,16],其中有两个混合元素
- 第一个元素用于图像和文本输入
- 第二个用于机器人特定的输入和输出
他们将第二组权重称为动作专家
正式地,他们希望对数据分布建模,其中
- 对应于未来动作的动作块(作者在任务中使用H=50)
- 而是一个观察。观察由多个RGB图像、语言命令和机器人的本体状态组成,因此
其中是第i个图像(每个机器人有2或3个图像)
是语言token序列
是关节角度向量
图像和状态通过相应的编码器进行编码,然后通过线性投影层投影到与语言token相同的嵌入空间中
对于动作块中的每个动作,我们有一个相应的动作token,通过动作专家来处理。在训练过程中,使用条件流匹配损失[28,32]监督这些动作token
其中下标表示机器人时间步,上标表示流匹配时间步,
- 最近在高分辨率图像[14]和视频[38]合成方面的研究表明,当与简单的线性高斯(或最优传输)概率路径[28]结合时,流匹配可以实现强大的经验性能,其由下述表达式给出
- 在实践中,网络通过采样随机噪声进行训练,计算“噪声动作”
然后训练网络输出以匹配去噪向量场 - 动作专家使用全双向注意力掩码,以便所有动作token彼此关注「The action expert uses a full bidirectional attention mask, so that all action tokens attendto each other」
在训练期间,从一个强调较低(更嘈杂)时间步的贝塔分布中采样流匹配时间步τ。更多细节请参见附录B「During training, we sample the flow matchingtimestep τ from a beta distribution that emphasizes lower(noisier) timesteps. See Appendix B for more details.」
在推理时,作者通过从到积分学习到的向量场来生成动作,从随机噪声开始,且作者使用前向欧拉积分规则「At inference time, we generate actions by integrating thelearned vector field from τ = 0 to τ = 1, starting with randomnoise A0t ∼N(0, I). We use the forward Euler integrationrule:」
其中是积分步长
在他们的实验中,使用了10个积分步骤(对应于)。注意,可以通过缓存前缀的注意力键和值,并且仅对每个积分步骤重新计算与动作token对应的后缀,从而高效地实现推理
他们在附录D中提供了有关推理过程的更多细节,包括模型每个部分的推理时间。虽然原则上他们的模型可以从头初始化或从任何VLM骨干微调,但实际上我们使用PaliGemma [5-PaliGemma: A versatile 3B VLM for transfer]作为他们的基础模型
PaliGemma是一个开源的30亿参数VLM,提供了大小和性能之间的便利折衷。作者为动作专家添加了3亿参数(从头初始化),总共达到33亿参数,且作者在附录B中提供了模型架构的完整描述
非VLM基线模型,除了他们的主要VLA模型外,他们还训练了一个类似的基线模型,该模型在消融实验中没有使用VLM初始化
作者称之为π0-small的这个模型拥有470M参数,没有使用VLM初始化,并且在没有使用VLM初始化的数据训练中,作者发现了一些有助于训练的小差异,这些差异在附录C中进行了总结。该模型用于我们的比较中,以评估合并VLM相关的好处
1.3 数据收集和训练方案
就像大型语言模型(LLM)的训练通常分为预训练和后训练阶段一样,作者对他们的模型也采用多阶段训练程序
- 预训练阶段的目标是让模型接触到各种各样的任务,以便它能够获得广泛适用和一般的物理能力,而后训练阶段的目标是使模型能够熟练和流畅地执行所需的下游任务
- 因此,预训练和后训练数据集的要求是不同的:
预训练数据集应涵盖尽可能多的任务,并在每个任务中涵盖多样化的行为
后训练数据集则应涵盖有助于有效任务执行的行为,这些行为应表现出一致且流畅的策略。
直观地说,多样化(但质量较低)的预训练数据允许模型从错误中恢复并处理高度变化的情况,这些情况可能在高质量的后训练数据中不会出现,而后训练数据教会模型良好地执行任务
1.3.1 预训练和后训练
在下图图4中提供了作者预训练混合的概述「数据集的概览:预训练混合物由OXE[10]的一个子集和π数据集组成。他们使用OXE的一个子集,称之为OXE魔法汤[24]。右图展示了预训练混合物中不同数据集的权重。左图展示了通过步数衡量的相对大小」,注意,由于每个训练示例对应一个时间步——即一个元组,——在本次讨论中,将以时间步来量化数据
- 训练混合中有9.1%由开源数据集组成,包括OXE [10]、Bridgev2 [52] 和 DROID [23]。这些数据集中的机器人和任务通常配备一到两个摄像头,并使用低频控制,频率在2到10 Hz之间。然而,这些数据集涵盖了广泛的物体和环境
- 为了学习灵巧且更复杂的任务,作者还使用了来自作者自有数据集的903M时间步长的数据,其中106M步来自单臂机器人,797M步来自双臂机器人
这些数据涵盖了68个任务,每个任务由复杂的行为组成——例如,“清理”任务涉及将各种不同的盘子、杯子和餐具放入清理箱,以及将各种垃圾物品放入垃圾桶
请注意,这种任务的定义与之前的工作显著不同,之前的工作通常使用任何名词和动词的组合(例如,“拾起杯子”与“拾起盘子”)来构成一个独立的任务
因此,作者数据集中行为的实际范围比这些“任务”数量所暗示的要广泛得多。他们将在第V-C节中更详细地讨论我们数据集中的具体机器人和任务
此外,由于数据集在大小上有些不平衡(例如,更困难的洗衣折叠任务被过度代表),作者对每个任务-机器人组合赋予权重,其中是该组合的样本数量,以便对过度代表的组合进行降权
- 配置向量和动作向量始终具有数据集中最大机器人的维度大小(在作者的案例中为18,以适应两个6-DoF机械臂、两个夹持器、一个移动底座和一个垂直驱动的躯干)
- 对于配置和动作空间维度较低的机器人,对配置和动作向量进行零填充。对于少于三张图像的机器人,还会屏蔽掉缺失的图像槽
- 在训练后阶段,使用一个较小的任务特定数据集对模型进行微调,以使其专门化用于特定的下游应用。如前所述,对“任务”的定义相当广泛——例如,“收拾”任务需要操作多种不同的物体。不同的任务需要非常不同的数据集,最简单的任务只需5小时,而最复杂的任务需要100小时或更多的数据
1.3.2 语言和高级策略
更复杂的任务需要语义推理和高级策略,例如清理桌子,也可以通过将高级任务(如“清理桌子”)分解为更直接的子任务(如“拿起餐巾”或“将餐巾扔进垃圾桶”)的高级策略来受益
由于作者的模型经过训练可以处理语言输入,故可以使用高级视觉语言模型(VLM)来进行这些语义推理,这种方法类似于LLM/VLM的规划方法,例如SayCan [2]。通过使用这种高级策略来辅助他们的模型在多个实验任务中制定高级策略,正如将在第六节中讨论的那样
1.3.3 可操作的机器人系统:涉及7种机械臂和68个任务
作者灵巧操作数据集包括7种不同的机器人配置和68个任务,如下图图5所示
- UR5e
一个配备平行爪夹持器的手臂,带有一个腕部安装和一个肩上摄像头,总共提供两张相机图像和一个7维的配置和动作空间 - 双臂UR5e
两个UR5e设置,总共提供三张相机图像和一个14维的配置和动作空间 - Franka
Franka设置有两个摄像头和一个8维的配置和动作空间 - 双臂Trossen
此设置有两个基于ALOHA设置[4,57]的6自由度Trossen ViperX手臂,配有两个腕部摄像头和一个基座摄像头,以及一个14维的配置和动作空间 - 双臂ARX和双臂AgileX
此设置使用两个6自由度手臂,支持ARX或AgileX手臂,配有三个摄像头(两个腕部和一个基座),以及一个14维的配置和动作空间。此类别涵盖两个不同的平台,但由于它们的运动学特性相似,我们将它们归为一类 - 移动Trossen和移动ARX
此设置基于移动ALOHA[57]平台,具有两个安装在移动基座上的6自由度手臂,可以是ARX手臂或Trossen ViperX手臂。非全向基座增加了两个动作维度,形成一个14维的配置和16维的动作空间。有两个腕部摄像头和一个基座摄像头。此类别涵盖两个不同的平台,但由于它们的运动学特性相似,我们将它们归为一类 - 移动Fibocom
两个安装在全向基座上的6自由度ARX手臂。基座增加三个动作维度(两个用于平移,一个用于方向),形成一个14维的配置和17维的动作空间。在图4中总结了我们数据集中每个机器人的比例
1.4 实验效果及验证
为了做一系列验证对比,作者
- 作者与OpenVLA[24]进行比较,这是一个最初在OXE数据集[10]上训练的7B参数VLA模型。作者在完整混合物上训练OpenVLA。这对于OpenVLA来说是一个非常困难的混合物,因为它不支持动作分块或高频控制
- 此外,还与Octo [50]进行比较,这是一个较小的93M参数模型。虽然Octo不是VLA,但它确实使用扩散过程生成动作,为作者的流匹配VLA提供了一个有价值的比较点
作者还在与他们模型相同的混合物上训练Octo。由于时间限制,作者无法为OpenVLA和Octo训练与他们完整模型相同的周期数。因此,作者还与“计算平价”版本的模型进行比较,该版本仅训练160k步(而我们的主模型训练了700k步),这相当于或低于为基线提供的步骤数量(OpenVLA为160k,Octo为320k) - 还包括一个仅在UR5e数据上微调的OpenVLA模型版本,不进行跨化身训练,希望在UR5e任务上提供更强的基线
- 最后,作者还包括与第四节描述的π0-small模型的比较,该模型可以视为作者模型的缩小版,没有进行VLM预训练。评估指标使用在每个任务和方法的10个回合中平均的标准化得分,其中一次回合在完全成功时得分为1.0,部分成功时得分为小数
例如,搬运的得分是正确放置在适当容器中的物体的比例。我们在附录E中描述了评分标准
结果如图7所示,π0在所有零样本任务上取得了迄今为止最好的结果,在衬衫折叠和较简单的搬运任务上接近完美的成功率,并且相对于所有基线有很大改善。经过160k步训练的“parity”版本的π0仍然优于所有基线,甚至π0-small也优于OpenVLA和Octo。OpenVLA在这些任务上表现不佳,因为其自回归离散化架构不支持动作块。仅限UR5e的OpenVLA模型表现较好,但仍远低于π0的性能。Octo确实支持动作块,但其表示能力相对有限
此比较说明了结合大型、具有表现力的架构与通过流匹配或扩散来建模复杂分布的能力的重要性。此外,与π0-small的比较说明了结合VLM预训练的重要性。不幸的是,最后的比较很难做到公平:π0-small使用的参数更少,但较大的模型在没有预训练的情况下难以使用
总体而言,这些实验表明π0提供了一个强大的预训练模型,能够有效地执行各种机器人任务,性能远优于先前的模型
// 待更