π0：一个视觉-语言-动作的流模型用于通用机器人控制

24年10月来自具身智能创业公司 PI 的论文“π0: A Vision-Language-Action Flow Model for General Robot Control”。

机器人学习有望充分发挥灵活、通用和灵巧机器人系统的潜力，并解决人工智能领域一些最深层次的问题。然而，将机器人学习提升到有效现实世界系统所需的通用水平，在数据、泛化和鲁棒性方面面临重大障碍。本文讨论了通用机器人策略（即机器人基础模型）如何应对这些挑战，以及如何为复杂且高度灵巧的任务设计有效的通用机器人策略。本文提出一种基于预训练视觉-语言模型 (VLM) 的流匹配架构，继承互联网规模的语义知识。然后，在来自多个灵巧机器人平台（包括单臂机器人、双臂机器人和移动机械手）的大型多样化数据集上训练该模型。根据模型在按照直接提示执行任务、遵循来自人和高级 VLM 策略的语言指令、以及通过微调获取新技能等能力，对模型进行评估。结果体现在各种各样的任务，例如折叠衣物、清洁桌子和组装盒子。

如图所示：通用机器人策略使用预训练的视觉-语言模型 (VLM) 主干，以及具有各种灵巧操作任务的多样化跨具身数据集。添加单独的动作专家，其通过流程匹配产生连续动作，这样模型可适应机器人控制，从而实现精确流畅的操作技能。然后，该模型可直接用于基于一个提示执行任务，或根据高质量数据进行微调，实现复杂的多步任务，例如折叠多件衣物或组装盒子。

执行各种机器人行为的灵活通用模型具有巨大的实际意义，但它们也可能为当今机器人学习面临的一些最严峻挑战提供解决方案，例如数据的可用性、泛化和鲁棒性。在自然语言 [1] 和计算机视觉 [39] 中，在各种多任务数据上预训练的通用基础模型，往往优于狭义定制和专门的解决方案。例如，如果目标是识别照片中的鸟类，那么对许多不同的图像语言关联进行预训练，然后针对鸟类识别任务进行微调或提示，可能比仅对鸟类识别数据进行训练更为有利。同样，对于有效的专业机器人系统，首先对高度多样化的机器人数据进行预训练，然后针对所需任务进行微调或提示更为有效。这可以解决数据稀缺挑战，因为通用模型可以使用更多数据源（包括来自其他任务、其他机器人甚至非机器人来源的数据），并且可以解决鲁棒性和泛化挑战，因为多样化的数据展现出更大的观察和动作覆盖范围，提供了各种场景、修正和恢复行为，而这些可能在更狭窄的专业数据中恐怕不存在。因此，采用大规模预训练方法进行机器人学习有可能解决该领域的许多挑战，并使具有实用学习能力的机器人成为现实，同时进一步加深对人工智能最深层次问题的理解。

然而，开发这种通用的机器人策略——即机器人基础模型——涉及许多重大挑战。首先，任何此类研究都必须在非常大的模型规模上进行，因为大规模预训练的全部好处往往无法在较小的模型规模上体现出来 [54]。其次，它需要开发正确的模型架构，以有效地利用不同的数据源，同时能够表示与复杂物理场景交互所需的复杂和微妙的行为。第三，它需要正确的训练方法。这也许是最重要的因素，因为 NLP 和CV领域大模型的最新进展在很大程度上依赖于精心策划预训练数据和后训练数据的策略 [35]。

最近提出的视觉-语言-动作 (VLA) 模型使用针对机器人控制进行微调的预训练 VLM [7, 24, 55]。此类模型采用自回归离散化来表示动作，其方式类似于文本tokens。相比之下，本文模型采用一种新设计，可以微调 VLM 以通过流匹配 [32, 28]（扩散 [20, 46] 的变型）产生动作。这样能够处理高频的动作块 [57]（高达 50 Hz）和高度灵巧的任务，这个对之前的自回归 VLA [7] RT-2 构成了重大挑战。这类似于最近关于动作生成的扩散模型研究 [9, 60]。与这些研究相比，本文的模型使用预训练的 VLM 主干 [5]。该贡献根本上也是综合性的，重点关注机器人基础模型的框架，不仅包括模型架构本身，还包括预训练方法、预训练和后训练阶段，以及一系列现实世界的实验。

在机器人控制之外，已经提出许多将预训练语言模型与扩散相结合的模型 [40、41、14]，包括专门混合扩散和自回归大语言模型的模型 [19、29、59]。这类模型通常与图像生成有关，但动作生成模型建立在许多先前提出的概念之上。与Zhou [59] 的 Transfusion 一样，本文对单个序列元素应用扩散式（流匹配）损失来训练模型，代替仅解码器 Transformer 的标准交叉熵损失。与Liu [29] 的 Playground v3 一样，本文方法对与扩散相对应的tokens使用一组单独的权重。将这些概念整合到 VLA 模型中，该模型引入流匹配 VLA，为灵巧控制产生高频动作块。

该领域的早期工作通常使用自监督或自主数据收集 [26, 22, 8]，为抓取 [18, 37] 或推动 [56] 等简单任务提供可处理的数据源，但没有更灵巧行为的复杂性。最近，已经收集许多用于机器人控制的高质量数据集，这些数据集可以实现广泛的泛化 [23, 10, 52, 33, 34, 43, 13, 6]，但通常用于由目标重定位和基本的家具操作（例如，打开抽屉）组成的简单任务 [31, 15]。更灵巧的任务已经在较小的规模下进行了研究，通常需要 10 秒或 100 秒的训练轨迹 [57]，相当于 10 小时或更少的时间。由于本文的目标之一是研究复杂而灵巧的行为，因此使用更大的数据集，其中包含大约 10,000 小时的演示，并辅以开源 OXE 数据集 [10]。在这个规模下，更复杂的预训练/后训练配方非常有效——类似于用于大语言模型的配方，预训练阶段为模型提供了广泛的知识基础，然后在后训练阶段使用更高质量的精选数据对其进行改进，以实现所需的行为。

这里任务的复杂性远远超出了以前的工作。虽然最近的研究已经展示了许多更复杂、更灵巧的行为，例如系鞋带 [58] 或烹饪虾 [17]，但对于结合身体灵活性和组合复杂性的行为，本文框架可以学习非常长的任务，有时长达数十分钟。例如，洗衣折叠任务要求机器人操纵可以从任何配置开始的各种衣物，并按顺序折叠多件衣物。桌子收拾任务需要辨别新目标的类别（垃圾或餐具）。一个单跨具身模型可以用作这些任务的基础模型。本文工作展示了端到端机器人最长的灵巧任务。

如图所示：π0 控制移动机械手折叠衣物。模型已预训练过来自 7 种不同机器人配置和 68 项任务的各种数据，然后可以直接提示或微调以执行复杂的下游任务，例如此衣物折叠策略，从烘干机中取出衣物，将其装入洗衣篮，将洗衣篮带到折叠桌上，然后折叠每件衣物。

如图所示：在提出的训练框架中，首先组装一个预训练混合数据集，该混合数据集由自己的灵巧操作数据集的加权组合组成，该数据集是在 7 种不同的机器人配置上针对 68 种不同的任务收集的，以及整个 OXE 数据集 [10]，其中包含来自 22 个机器人的数据。预训练阶段还使用各种语言标签，结合任务名称和片段注释（子轨迹的细粒度标签，通常长度约为 2 秒）。预训练阶段的目的，是训练一个表现出广泛功能和泛化能力的基础模型，但不一定专门针对任何一项任务的高性能。这个基础模型可以遵循语言命令并以基本的熟练程度执行各种任务。对于复杂而灵巧的任务，采用后训练程序，用高质量的精选数据使模型适应特定的下游任务。研究中分别使用少量到中等量数据进行高效的后训练，以及使用较大数据集进行高质量的后训练，完成诸如洗衣折叠和移动操作等复杂任务。

该模型基于 PaliGemma 视觉-语言模型 [5]，然后用数据混合对其进行进一步训练。为了将基础 PaliGemma VLM 变成 π0，添加使用流匹配 [32, 28] 生成连续动作分布的动作输出。注：用 PaliGemma 是为了方便，并且因为它的规模相对较小（这对于实时控制很有用），但本框架与任何基础预训练 VLM 兼容。

π0 模型主要由语言模型transformer主干组成。按照标准的后期融合 VLM 配方 [3、11、30]，图像编码器将机器人的图像观测嵌入到与语言tokens相同的嵌入空间中。进一步用机器人特定的输入和输出（即本体感受状态和机器人动作）增强此主干。π0 使用条件流匹配 [28、32] 来模拟动作的连续分布。流匹配为模型提供了高精度和多模态建模能力，使其特别适合高频灵巧任务。该架构灵感来自 Transfusion [59]，它使用多个目标训练单个transformer，其中对应于连续输出的 token 通过流匹配损失监督，对应于离散输出的 tokens 通过交叉熵损失监督。在 Transfusion 的基础上，对机器人特定（动作和状态）token 使用一组单独的权重可以提高性能。这种设计类似于混合专家 [45, 25, 12, 16]，其中有两个混合元素，第一个元素用于图像和文本输入，第二个元素用于机器人特定的输入和输出。第二组权重，称为动作专家。

高分辨率图像[14]和视频[38]合成方面的最新研究表明，当与简单的线性高斯（或最优传输）概率路径[28]相结合时，流匹配可以实现强大的经验性能。动作专家使用全双向注意掩码，以便所有动作 tokens 相互关注。在训练期间，从强调较低（噪声较大）时间步长的 beta 分布中抽取流匹配时间步长 τ。在推理时，通过从 τ = 0 到 τ = 1 积分学习的向量场从随机噪声开始来生成动作。通过缓存观测 ot 的注意机制 K-V，并仅重计算与每个集成步操作tokens 相对应的，可以有效地实现推理。

虽然原则上模型可以从头开始初始化或从任何 VLM 主干进行微调，但实际上用 PaliGemma [5] 作为基础模型。PaliGemma 是一个开源的 30 亿参数 VLM，可在大小和性能之间实现便捷的权衡。为动作专家添加 3 亿个参数（从头开始初始化），模型总共 33 亿个参数。

除了主要 VLA 模型外，还训练一个类似的基线模型，该模型未使用 VLM 初始化进行消融实验。这个模型，称之为 π0-small，有 470M 个参数，不使用 VLM 初始化，并且有许多小差异，有助于在没有 VLM 初始化的情况下对数据进行训练。

功能广泛的机器人基础模型不仅需要富有表现力和强大的架构，还需要正确的数据集，更重要的是，正确的训练方法。正如 LLM 训练通常分为预训练前和后训练阶段一样，对该模型采用多步训练程序。预训练阶段的目标是让模型接触各种任务，以便它能够获得广泛适用和一般的物理能力，而后训练阶段的目标是让模型能够熟练流畅地执行所需的下游任务。因此，预训练和后训练数据集的要求是不同的：预训练数据集应涵盖尽可能多的任务，并且在每个任务中应涵盖多种行为。后训练数据集应涵盖有助于有效执行任务的行为，这些行为应表现出一致且流畅的策略。直观地说，多样化（但质量较低）的预训练数据使模型能够从错误中恢复并处理高度多样化的情况，而这些情况可能不会发生在高质量的后训练数据中，而后训练数据可以教会模型很好地完成任务。

9.1% 的训练数据来自开源数据集，包括 OXE [10]、Bridge v2 [52] 和 DROID [23]。这些数据集中的机器人和任务通常有一两个摄像头，使用 2 到 10 Hz 之间的低频控制。然而，这些数据集涵盖了广泛的目标和环境。为了学习灵巧和更复杂的任务，还使用自己数据集中的 903M 个时间步长数据，其中 106M 个步长来自单臂机器人，797M 个步长来自双臂机器人。该数据有 68 个任务，每个任务都由复杂的行为组成 — — 例如，“收拾”任务包括将各种不同的碗碟、杯子和餐具放入收拾箱，将各种垃圾放入垃圾箱。请注意，此任务定义与以前的研究有显著不同，以前的研究通常使用名词和动词的任意组合（例如，“拿起杯子”与“拿起盘子”）来构成一项不同的任务。因此，数据集中的行为，实际范围比该“任务”数能暗示的要大得多。

由于数据集在大小上有些不平衡（例如，更困难的洗衣折叠任务，占比过高），用 n^0.43 来加权每个任务机器人组合，其中 n 是该组合的样本数量，这样占比过高的组合就会被降低权重。配置向量 qt 和动作向量 at 始终具有数据集中最大机器人的维数（在本文例子中为 18，即两个 6 自由度臂、2 个夹持器、一个移动基座和一个垂直驱动的躯干）。对于具有较低维数配置和动作空间的机器人，对配置和动作向量进行零填充。对于图像少于三幅的机器人，屏蔽缺失的图像槽。

在后训练阶段，用较小的特定于任务数据集对模型进行微调，以使其专门用于特定的下游应用。如前所述，对“任务”的定义相当广泛——例如，“搬运”任务需要操作各种不同的目标。不同的任务需要非常不同的数据集，最简单的任务只需要 5 个小时，而最复杂的任务则需要 100 个小时或更长时间的数据。

更复杂的任务需要语义推理和高级策略，例如收拾桌子，高级策略也可以发挥作用，将高级任务（例如“收拾桌子”）分解为更直接的子任务（例如“拿起餐巾纸”或“将餐巾纸扔进垃圾桶”）。由于模型经过训练可以处理语言输入，因此可以使用高级 VLM 进行语义推理，这种方法类似于 SayCan [2] 的 LLM/VLM 规划方法。 用这种高级策略来帮助模型制定高级策略，以完成几个实验任务。

灵巧操作数据集包括 7 种不同的机器人配置和 68 个任务，如图所示：