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

$\pi_0$ ：一种用于通用机器人控制的视觉-语言-动作流模型

动机

要获得一个具备良好泛化能力的 robot foundation model ，主要有三方面挑战：

1. 数据的规模要足够大（数据量、本体丰富度、任务丰富度等）；
2. 模型的架构要合理（支持高频控制，连续 action space 控制）；
3. Training recipe 要合理（训练策略，数据配比等）

摘要

1. 我们提出了一种新颖的流匹配架构，该架构建立基于一个 Pre-Trained VLM (PaliGemma） ，其包含视觉编码器部分（SigLIP，400M）、 大语言模型部分（Gemma，2.6B），还有简单的线性层(他们的早期实验中，他们发现更复杂的替代方案——MLPs，但是并没有提供明显的优势，因此决定使用最简单的） ；
2. 该模型由一个较大的 VLM 主干和一个较小的动作专家组成，动作专家用于处理机器人的状态和动作。这种设计类似于具有两个混合元素的专家混合，其中前者用于处理图像和文本输入，后者用于机器人特定的状态输入和动作输出。

3. 途中我们可以看出，经过丰富多样的预训练就能执行 Zero-shot 类的任务，然后在经过高质量数据的后训练就可以针对困难任务进行有效执行，甚至是没有见过的任务。
4. 下面图中可以看出输入是模型的输入包括3帧最近的图像，文本指令，以及机器人本体信号（即各关节连接处的角度），输出是称作 action chunk 的形式，即未来一段时间的 action 序列，这里的序列长度 H 设定为50，对于50HZ 控制的情况，即相当于预测未来一秒钟的动作。这里action 的维度，设定为各个本体中最长的那个（18DoF)，不够的本体则pad。

$\pi_0$ ：模型

$\pi_0$ 使用条件流匹配来模拟动作的连续分布。流匹配为我们的模型提供了高精度和多模态建模能力，使其特别适合于高频灵巧操作任务。我们的架构受到 Transfusion 的启发，该架构使用多个目标训练单个 transformer，其中对应于连续输出的 token ，通过流匹配损失进行监督，而对应于离散输出的 token 通过交叉熵损失进行监督。在 Transfusion 的基础上，我们还发现为机器人特定的（动作和状态）token 使用一组单独的权重可以提高性能。这种设计类似于具有两个混合元素的专家混合，其中第一个元素用于图像和文本输入，第二个用于机器人特定的输入和输出。我们将第二个权重集称为动作专家。

形式上，我们希望模拟数据分布 $p(\mathbf{A}_t | \mathbf{o}_t)$ ，其中 $\mathbf{A}_t = [a_t, a_{t+1},..., a_{t+H-1}]$ 对应于未来动作的动作块（我们的任务中使用 H = 50 ），而 $\mathbf{o}_t$ 是一个观察。观察包括多个 RGB 图像、语言命令和机器人的本体状态，使得 $\mathbf{o}_t = [\mathbf{I}_t^1, ..., \mathbf{I}_t^n, \ell_t, \mathbf{q}_t]$ ，其中 $\mathbf{I}_t^i$ 是第 $i$ 个图像（每个机器人有2或3个图像），$\ell_t$是语言标记序列，$\mathbf{q}_t$是关节角度的向量。图像 $\mathbf{I}_t^i$ 和状态 $\mathbf{q}_t$ 通过相应的编码器进行编码，然后通过线性投影层投影到与语言标记相同的嵌入空间中。

对于动作块 $\mathbf{A}_t$ 中的每个动作 $a_t^i$ ，我们有一个相应的动作标记，我们通过动作专家进行处理。在训练期间，我们使用条件流匹配损失来监督这些动作标记。

$L^\tau(\theta) = \mathbb{E}_{p(\mathbf{A}_t|\mathbf{o}_t), q(\mathbf{A}_t^\tau|\mathbf{A}_t)} \left[ \left\| v_\theta(\mathbf{A}_t^\tau, \mathbf{o}_t) - \mathbf{u}(\mathbf{A}_t^\tau|\mathbf{A}_t) \right\|^2 \right]$

其中下标表示机器人时间步，上标表示流匹配时间步 $\tau \in [0, 1]$ 。最近在高分辨率图像和视频合成方面的工作表明，当与简单的线性高斯概率路径结合时，流匹配可以实现强大的经验性能，由 $q(\mathbf{A}_t^\tau|\mathbf{A}_t) = \mathcal{N}(\tau \mathbf{A}_t, (1-\tau)\mathbf{I})$ 给出。在实践中，网络通过采样随机噪声 $\epsilon \sim \mathcal{N}(0, \mathbf{I})$ 进行训练，计算“噪声动作” $\mathbf{A}_t^\tau = \tau \mathbf{A}_t + (1-\tau)\epsilon$ ，然后训练网络输出 $v_\theta(\mathbf{A}_t^\tau, \mathbf{o}_t)$ 以匹配去噪向量场 $u(\mathbf{A}_t^\tau|\mathbf{A}_t) = \epsilon - \mathbf{A}_t$ 。动作专家使用完整的双向注意力掩码，以便所有动作标记相互关注。在训练期间，我们从强调较低（噪声）时间步的 Beta 分布中采样流匹配时间步 $\tau$。

在推理时，我们通过从 $\tau = 0$ 到 $\tau = 1$ 积分学习到的向量场来生成动作，从随机噪声 $\mathbf{A}_t^0 \sim \mathcal{N}(0, \mathbf{I})$ 开始。我们使用前向欧拉积分规则：

$\mathbf{A}_t^{\tau+\delta} = \mathbf{A}_t^\tau + \delta v_\theta(\mathbf{A}_t^\tau, \mathbf{o}_t)$

其中 $\delta$ 是积分步长。我们在实验中使用 10 积分步 （对应于 $\delta = 0.1$ )。注意通过缓存前缀 $\mathbf{o}_t$ 的注意力 keys 和 values，并仅重新计算每个积分步骤对应的动作 token 的后缀，可以有效地实现推理。

虽然原则上我们的模型可以从零开始初始化或从任何 VLM 骨干进行微调，但在实践中我们使用 PaliGamma 作为我们的基模型。PaliGamma 是一个开源的 3B 参数 VLM ，提供了大小和性能之间的便利权衡。我们为动作专家添加了 3 亿参数（从零开始初始化）以获得总共 3.3 B 参数。

非VLM基线模型。除了我们的主要 VLA 模型外，我们还训练了一个类似的基线模型，该模型在消融实验中没有使用 VLM 初始化。我们将其称为 $\pi_0-small$ ，具有 4.7 亿参数，不使用 VLM 初始化，并具有一些我们发现对我们的数据在没有 VLM 初始化的情况下训练有帮助的小差异。

数据收集和训练配方

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

Pre-training and post-training

预训练。 使用了超过 10,000 小时的机器人数据上， 图中展示了我们数据集的，预训练数据混合由 OXE 的一个子集和 $\pi$ 数据集组成。我们使用 OXE 的一个子集，我们称之为 OXE Magic Soup 。右图展示了预训练不同类型机器人数据集的混合权重。左图展示了它们按步数测量的相对大小。

我们在图中提供了我们预训练数据混合的概述。由于每个训练样本对应于一个时间步,即一个元组 $(\mathbf{o}_t, \mathbf{A}_t)$ ，我们将在此讨论中以时间步的数量来量化数据。9.1% 的训练混合数据由开源数据集组成，包括 OXE 、Bridge v2 和DROID 。这些数据集中的机器人和任务通常有一个或两个摄像头，并使用 2 到 10 Hz 之间的低频控制。为了学习灵巧和更复杂的任务，我们还使用了我们自己的数据集中的 903M 时间步，其中 106M 步来自单臂机器人，797M 步来自双臂机器人。这些数据有 68 个任务，每个任务由复杂行为组成——例如，“整理”任务涉及将各种不同的盘子、杯子和器具放入整理箱中，以及将各种垃圾物品放入垃圾箱中。请注意，这种任务定义与先前作有显著不同，通常使用任何名词和动词的组合（例如，“拿起杯子”与“拿起盘子”）来构成一个不同的任务。因此，我们数据集中的实际行为范围比“任务”数量所暗示的要广泛得多。

由于数据集在大小上有些不平衡（例如，更困难的洗衣折叠任务被过度表示），我们通过 $n^{0.43}$ 对每个任务-机器人组合进行加权，其中 $n$ 是该组合的样本数量，以便过度表示的组合被降低比例。配置向量 $\mathbf{q}_t$和动作向量 $\mathbf{a}_t$ 总是具有数据集中最大机器人的维度（在我们的例子中为 18 ，以适应两个 6 自由度手臂、2个夹具、一个移动底座和一个垂直驱动的躯干）。对于具有较低维度配置和动作空间的机器人，我们用零填充配置和动作向量。对于少于三个图像的机器人，我们还遮盖了缺失的图像槽位。

后训练，我们使用一个较小的任务特定数据集对模型进行微调，以使其专门化于特定的下游应用。最简单的任务只需要 5 小时的数据，而最复杂的任务需要 100 小时或更多的数据。

语言和高级策略

更复杂的任务需要语义推理和高级策略，例如整理桌子，也可以从高级策略中受益，该策略将高级任务（如“整理桌子”）分解为更直接的子任务（如“拿起餐巾”或“将餐巾扔进垃圾”）。由于我们的模型经过训练以处理语言输入，我们可以使用高级 VLM 来进行这些语义推断，这是一种类似于 LLM/VLM 规划方法的方法，如 SayCan 。我们使用这样的高级策略来协助我们的模型进行几个实验任务的高级策略。

机器人系统细节

我们的灵巧操作数据集包括 7 种不同的机器人配置和 68 个任务。我们在图中总结了这些平台，并在下面进行讨论：

UR5e。一个带有平行颚夹持器的手臂，手腕上安装了一个摄像头和一个肩部摄像头，总共有两个摄像头图像和一个 7 维的配置和动作空间。

Bimanual UR5e。两个 UR5e 设置，总共三个摄像头图像和一个 14 维的配置和动作空间。

Franka。Franka 设置有两个摄像头和一个 8 维的配置和动作空间。

Bimanual Trossen。此设置有两个 6 自由度的 Trossen ViperX 手臂，基于 ALOHA 设置，有两个手腕摄像头和一个基座摄像头，以及一个 14 维的配置和动作空间。

Bimanual ARX & bimanual AgileX。此设置使用两个 6 自由度手臂，并支持 ARX 或 AgileX 手臂，有三个摄像头（两个手腕和一个基座）和一个 14 维的配置和动作空间。该类别涵盖了两个不同的平台，但我们将其归为一类，因为它们的运动学特性相似。

Mobile Trossen & mobile ARX。此设置基于移动 ALOHA 平台，具有两个 6 自由度手臂的移动基座，这些手臂要么是 ARX 要么是 Trossen ViperX 。非完整基座增加了两个动作维度，一个 14 维的配置和 16 维的动作空间。有两个手腕摄像头和一个基座摄像头。该类别涵盖了两个不同的平台，但我们将其归为一类，因为它们的运动学特性相似。

Mobile Fibocom。两个 6 自由度 ARX 手臂在一个整体基座上。基座增加了三个动作维度（两个用于平移和一个用于方向），一个 14 维的配置和 17 维的动作空间。

实验

预训练模型

直接经过预训练之后，使用五个任务进行测试。

同时和 OpenVLA，Octo 等方法比较，确实效果很好。

• $\pi_0$ 训练了完整的 700k 步数，比 OpenVLA 等方法效果好很多；
• $\pi_0-Parity$ 是用 160k 训练步数模型，表明训练充分对模型效果影响很大
• $\pi_0-Small$ 是没有用 Pretrain VLM 的模型，可以看出来效果差了很多。

对于任务的评分指标和方法论文附录里面有详细描述，这里只举一个例子。

衬衫折叠:衬衫折叠的结果会被记录为“成功”或“失败”。我们开始每次衬衫折叠评估时，都将衬衫平铺在桌子上。成功的定义是:将袖子折叠好，并在衬衫长度方向上完成一次半折叠。我们的评估包括4件小号T恤和1件中号T恤。我们为每件物品进行两次评估，最多耗时15000步，即大约5分钟。

后训练模型

经过微调之后使用三个任务进行测试。

其中 flat 版本的意思就是单纯只接受直接的任务命令，如“bag the groceries”。没有中间详细的任务内容描述；Human 版本的意思就是采用从人类专家接收中间步骤命令描述的方法；HL 版本的意思就是接受 VLM 传来的高级策略。

左边图可以看出经过微调 $\pi_0$ 整体还是比 $\pi_0-Small$要好，右边图中可以看出 Bussing 和 Grocery Bagging 这两个任务上有人类详细任务步骤描述的效果最好，对于 Table Settting 直接通过 VLM 传来的高级信号即可很好完成任务。

学习新任务

这里有五个任务，前两个任务基本模型和预训练数据中的类似，所以被定义为简单任务，第三个中的任务操作在预训练中存在，但是微波炉没有见过，所以是中等，后面两个任务没有出现在预训练中，所以定义是困难。

scratch 表示在预训练阶段没有经历这些任务的训练，可以看出使用数据对不同的模型进行微调， $\pi_0$ 的效果都是最好的。

多阶段任务

在我们最后的一系列实验中，结合微调与语言处理技术，来应对一系列具有挑战性的多阶段任务。对于其中一些任务，在预训练阶段已有相关数据，但需要通过微调来实现熟练掌握。而对于另一些任务，在预训练阶段则没有相关数据。

scratch 表示在预训练阶段没有经历这些任务的训练，out-of-box 就是直接用预训练的模型测评。

可以看出来经过预训练和微调之后的 $\pi_0$ 效果是最好的。

附录

模型架构细节

在本节中，我们提供了模型架构的完整描述。我们遵循 PaliGamma VLM 设计，但有以下不同之处：（1）为机器人状态向量 $\mathbf{q}_t$ 和动作向量 $\mathbf{A}_t = [\mathbf{a}_t, ..., \mathbf{a}_{t+H-1}]$ 的 token 增加了额外的输入和输出投影， （2）增加了一个额外的 MLP 用于整合流匹配时间步信息 $\tau$ ，以及（3）较小的动作专家的权重。

额外的输入和输出。标准 PaliGamma 架构接收一系列图像 $\mathbf{I}_t^1, ..., \mathbf{I}_t^n$ 后跟一个语言提示 $\ell_t$ 。我们添加了一个输入 $\mathbf{q}_t$ 用于表示机器人的本体状态，它通过线性投影映射到 transformer 嵌入维度。最终的输入 token 集合对应于噪声动作块 $\mathbf{A}_t^\tau = [\mathbf{a}_t^\tau, ..., \mathbf{a}_{t+H-1}^\tau]$ ，token 数等于动作范围（我们的任务中 H = 50)。我们仅使用 transformer 输出对应于 H 噪声动作，这些动作通过线性投影解码为 $v_\theta(\mathbf{A}_t^\tau, \mathbf{o}_t)$ 。

整合流匹配时间步。噪声动作块 $\mathbf{A}_t^\tau$ 被使用 MLP 映射到 transformer 嵌入维度，该 MLP 还整合了流匹配时间步 $\tau$ 。对于每个噪声动作 $\mathbf{a}_t^\tau$ ，对应嵌入的表达式为 $W_3 \cdot \text{swish}(W_2 \cdot \text{concat}(W_1 \cdot \mathbf{a}_t^\tau, \phi(\tau)))$ ，其中 $\phi: \mathbb{R} \rightarrow \mathbb{R}^w$ 是正弦位置编码函数，$W_1 \in \mathbb{R}^{w \times d}, W_2 \in \mathbb{R}^{w \times 2w}, W_3 \in \mathbb{R}^{w \times w}$， $d$ 是动作维度， $w$ 是动作专家的嵌入维度 。

注意力机制掩码。 $\pi_0$ ****使用具有 3 个 block 的块状因果注意力掩码： $[\mathbf{I}_t^1, ..., \mathbf{I}_t^p, \ell_t], [\mathbf{q}_t], [\mathbf{a}_t^\tau, ..., \mathbf{a}_{t+H-1}^\tau]$ 。在每个块内，有完整的双向注意力，而在每个块中的 token 不能关注其他未来块中的 token 。第一个块包括来自 PaliGamma 的 VLM 预训练的输入修正，这些修正防止其关注未来块（包括新输入）以最小化分布偏移位。第二块是机器人状态 $\mathbf{q}_t$ 是自己的块，因为它在每次流匹配整合步骤中不会改变，通过阻止它关注到最后一块，可以使其对应的 keys 和 values 在采样过程中被缓存。最后一个块对应于噪声动作 $\mathbf{A}_t^\tau$ ，可以关注完整的输入序列。

动作专家。$\pi_0$ 实现为单个 transformer ，具有两组权重（也称为专家 ），每个 token 被路由到其中的一个专家；权重仅通过 transformer 的自注意力层交互。图像和语言提示 $[\mathbf{I}_t^1, ..., \mathbf{I}_t^p, \ell_t]$ ，被路由到较大的 VLM 骨干，我们从 PaliGamma 初始化。$[\mathbf{q}_t, \mathbf{A}_t^\tau]$未在 VLM 预训练的输入，而是路由到动作专家。PaliGamma 基于 Gamma 2B 语言模型，使用 multi-query attention 和配置$\{width=2048, depth=18, mlp\_dim=16,384, num\_heads=18, num\_kv\_heads=1,head\_dim=256\}$。由于专家仅在自注意力层交互，$width$ 和 $mlp\_dim$不需要在专家之间匹配。为了加速推理我们简化动作专家为 $\text{width}=1024, mlp\_dim=4096$，结果参数计数约为 300M 。

流匹配时间步采样。原始流匹配论文中，从均匀分布 $\tau \sim U(0, 1)$中采样流匹配时间步。我们设计了一个时间步采样分布，强调低时间步（高噪声水平）；此外，超过给定阈值 $s$的时间步根本不采样，因为它们不需要，只要迭代步 $\delta$ 大于 $(1-s)$。分布由 $\mathcal{p}(\tau) = \text{Beta}(\frac{s-\tau}{s}; 1.5, 1)$ 给出，并在图中可视化。我们在实验中使用 $(s = 0.999)$，这允许 $\delta > \frac{1}{1000}$ ，或多达 1000 个积分步骤。

非 VLM 基线架构

我们的基线架构 $\pi_0-small$ 不是基于 VLM 骨干。因此，我们使用它来评估 VLM 预训练的好处。我们设计它足够表达以适应我们的大数据集，同时在从头开始训练时仍能提供良好的性能。该模型大约有 4.7 亿参数，并与我们的主要模型在以下方面有所不同：（1）我们使用 DistilBERT 对语言命令 $\ell_t$ 的语言 token 进行编码，因为该模型不使用语言模型骨干；（2）动作专家交叉注意观察编码器的输出，类似于传统的编码器-解码器 transformer ，而不是我们由一个仅解码器的专家混合的主要模型架构；（3）图像使用较小的预训练 ViT 编码器进行编码；（4）ViT 图像编码器不共享权重；（5）编码观察的 transformer 骨干（在 ViT 图像编码器之后）未在互联网数据上预训练；（6）动作专家使用 DiT 架构而不是 Gamma 架构，因此使用 AdaLN-Zero 层整合流匹配时间步 $\tau$ 。除此之外，模型大致相似：两者都使用预训练的 ViT 图像编码器，两者都使用观察编码器和动作专家的单独权重，两者都采用相同的观察格式，两者都执行 10 步流匹配以预测动作块。

推理

我们的模型接受一个观察 $\mathbf{o}_t = [\mathbf{I}_t^1, ..., \mathbf{I}_t^p, \ell_t, \mathbf{q}_t]$ 和噪声动作 $\mathbf{A}_t^\tau$ 并输出需要获得下一个流匹配步骤 $v_t^\tau$ 的向量场。每次我们预测一个新的动作块 $\mathbf{A}_t$ ，我们必须对每个图像 $\mathbf{I}_t^i, ..., \mathbf{I}_t^p$ 进行编码，运行一个对应于 $\mathbf{o}_t$ 的 token 的前向传递，然后运行 10 步流匹配，其中每一步需要对对应于 $\mathbf{A}_t^\tau$ 的 token 运行一个前向传递（对应于 $\mathbf{o}_t$的 keys 和 values 被缓存）。表中总结了使用 3 个摄像头图像的上面整个操作的计算时间。操作在 NVIDIA GeForce RTX 4090 消费级 GPU 上计时。对于移动机器人，推理是在 Wi-Fi 连接上离线完成的，增加了少量的网络延迟。进一步的优化、量化和其他改进可能会进一步减少推理时间。可以看出不算网络延迟的整个推理时间就是 73 ms ，这意味着他们其实最高可以以 10FPS 频率来做规划。

由于模型生成一个完整的 $H-step$ 动作块，我们可以在需要再次运行推理之前执行多达 $H$ 个动作。然而，我们可能会比那更频繁地运行推理，以及使用各种聚合策略从不同的推理调用中组合动作。早期的时间集成，发现它损害了策略性能，所以我们选择不聚合动作，而是开环执行动作块。对于 20Hz UR5e 和 Franka 机器人，我们在执行 16 个动作后，也就是每 0.8 秒运行一次推理，对于所有其他以 50Hz 运行的机器人，我们在执行 25 个动作后也就是每 0.5 秒运行一次推理。