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

发布于 2025 年 2 月 4 日

在 GitHub 上更新

168

我们已经将首批 机器人基础模型 移植到了 Hugging Face LeRobot！由 Physical Intelligence 开发的 π0 和 π0-FAST 现已在 LeRobot 代码库 中提供，为 Hugging Face 生态系统带来了通用机器人智能。如果你对视觉-语言-动作 (VLA) 模型与视觉-语言模型 (VLM) 的区别以及动作是如何表示的感到好奇，那就深入阅读这篇博文一探究竟吧！

在我们的代码库中探索模型集和模型的 PyTorch 版本：Huggingface Pi0 模型合集 | Huggingface Pi0+FAST 模型合集 | LeRobot 代码库

引言

罗伯特·海因莱因（Robert Heinlein）提出，一个全面发展的人应当能够处理广泛的任务——无论是智力上的还是体力上的——而不是局限于某一特定领域。将全面发展的人与机器智能进行类比：人工智能系统千差万别，但人类智能的卓越之处在于其多功能性——能够适应各种任务、环境和意外情况。尽管大型语言模型和视觉-语言模型 (LLM, VLM) 显示出巨大潜力，但它们缺乏与物理世界的互动。为了弥合这一差距，我们需要基于机器人数据训练的模型。通用机器人模型可以增强适应性，利用多样化的数据来提高泛化能力和鲁棒性。与在孤立任务上进行训练不同，在多样化的机器人数据上进行预训练——类似于 LLM 的做法——能够提升效率和性能。

开发通用机器人策略，或称机器人基础模型，面临三大关键挑战：

需要大规模研究 才能充分利用预训练的优势。
设计能够整合不同数据源同时捕捉复杂物理交互的模型架构。这方面的一个关键挑战是 跨形态训练 (cross-embodiment training)，即模型必须从具有不同配置、控制空间和动作表示的多种机器人类型中学习。现有方法通过以下方式解决这一问题：
- 结合来自不同机器人平台的多模态数据集 以增强泛化能力。
- 使用共享表示 来弥合不同机器人形态（如单臂、双臂和移动机械臂）之间的差距。
制定有效的训练方案，因为自然语言处理和视觉领域的最新进展在很大程度上依赖于精心的预训练和后训练策略。

在本文中，我们介绍由 Physical Intelligence 开发的原型模型和学习框架 π0 与 π0-FAST，旨在克服这些挑战。

🔍 π0 是什么？

论文 | Jax 代码

π0 (Pi-Zero) 是一个由 Physical Intelligence 团队开发的 视觉-语言-动作 (VLA) 模型，专为 通用机器人控制 设计。它建立在 大规模预训练 和 基于流匹配的动作生成 的基础上，使机器人能够在不同形态下执行 灵巧的操作任务。

π0 在来自 7 个机器人平台 和 68 个独特任务 的数据上进行训练，在复杂的真实世界任务（如 叠衣服、收拾餐桌、打包杂货、组装盒子和取物）中表现出强大的 零样本 和 微调性能。

与标准的机器人策略不同，π0 采用流匹配 来生成 流畅、实时的 50Hz 动作轨迹，使其在实际部署中 高效、精确且适应性强。流匹配曾用于连续归一化流，并提高了扩散模型中的生成质量。π0 使用的去噪过程以同样的方式工作，从一个随机噪声开始，逐步收敛到一个有意义的运动动作序列。

如何在 LeRobot 中使用 π0？

首先，你需要升级你的 lerobot 安装，它现在依赖于 transformers！只需在克隆 git 仓库后执行：

pip install -e ".[pi0]"

π0 模型是基础模型，就像 PaliGemma 一样，旨在适应各种框架、环境和场景输入。这里的基座模型可以直接使用，尤其是 π0。

在 π0 预训练模型上进行推理

python lerobot/scripts/eval.py \
--pretrained_policy.path=/path/to/pretrained/pi0

然而，由于这是从 jax 到 torch 的转换以及从特定环境的转换，性能有所下降。我们建议您像下面这样，将您自己的 π0 微调到您自己的环境中。

微调 π0 预训练模型

要使用来自 openpi 的 pi0_base 检查点微调 π0 模型，请运行以下命令：

python lerobot/scripts/train.py \
--policy.path=lerobot/pi0 \
--dataset.repo_id=danaaubakirova/koch_test

要在 PaliGemma 和 Expert Gemma 上微调 π0 神经网络（这些网络在 π0 微调之前使用了 VLM 默认参数进行预训练），请执行：

python lerobot/scripts/train.py \
--policy.type=pi0 \
--dataset.repo_id=danaaubakirova/koch_test

您也可以独立于 LeRobot 训练框架，使用以下代码来使用预训练的 π0 模型：

policy = Pi0Policy.from_pretrained("lerobot/pi0")

VLM 和 VLA 之间有什么区别？

视觉-语言模型 (VLM) 和视觉-语言-动作模型 (VLA) 共享一个共同的基础：transformer。然而，关键区别在于动作表示。VLM 处理和生成多模态表示（图像和文本），而 VLA 则通过加入动作和观察状态词元来扩展这一点。有了这些额外的词元，下一个挑战就是理解注意力是如何计算的。

机器人策略中的注意力机制

让我们扩展一下词汇，介绍一些关键术语：

状态词元 (State Token)

它是一个单一的词元，代表机器人的 当前环境状态（例如，关节角度、传感器数值或其他相关观测值）。
掩码规则允许该词元 关注前缀中的图像和文本，这意味着状态词元可以“看到”决策所需的所有视觉或文本线索。
它还以 三角方式关注之前的状态。如果使用多个状态词元，每个新的状态词元可以看到旧的词元，但反之不行。

动作词元 (Action Tokens)

代表 电机命令序列。
对除填充区域外的所有内容都具有 完全可见性。这意味着每个动作词元可以关注：
- 所有非填充的图像词元（整个场景），
- 所有非填充的文本词元（指令或描述），
- 状态词元（当前和之前的），
- 其他动作词元.

前缀词元 (Prefix Tokens)

代表 整个场景，并完全相互关注，类似于 PaliGemma。

核心思想

这些词元封装了：

机器人对环境的 内部表示 (状态)，
机器人发出的 命令或控制 (动作)，
时间或步骤索引的编码 (时间嵌入)。

它们被附加在前缀部分（图像+文本）之后，因此前缀作为上下文（例如，场景图像、语言指令如 “做个好机器人” 或 “移动方块”），而后缀则捕捉特定于策略的特征。

⚡ 迈向 π0 中更快的注意力机制

然而，在 π0 中高效地处理注意力机制也带来了一系列挑战。其注意力掩码的独特形状影响了注意力的计算方式——让我们深入了解细节！

处理二维注意力掩码

由此产生的 二维因果掩码 表现出强烈的 块稀疏性，但定义每个块的边界——尤其是在一批样本中——有点棘手。我们习惯于自回归建模中的三角结构因果掩码，但这并非如此。

正如您在下面的示例中所见：图像（第一个元素）有一些填充词元，代表空置的摄像头。然后，添加文本词元和状态词元。这个“前缀”部分形成了一个完全非因果的注意力，就像在 PaliGemma 中一样。然后，“后缀”（状态 + 动作/时间词元）具有一个因果块结构。简单的 eager 实现方式执行矩阵乘法并在整个输入上应用 softmax，这使其效率极低。

图 1：VLA 注意力掩码的可视化

我们能使用 FlashAttention2 吗？

FlashAttention2 提供了一个 varlen 接口，但 cu_seqlens（累积前缀长度）必须手动计算。它专为具有统一查询和键长度的 连续（或严格因果）注意力模式 而设计。
它 无法自然处理不规则的块掩码 或每个词元任意的“允许”位置，而这正是我们所需要的。
所以，虽然以一定的实现成本使用它是可能的，但我们决定转向……

在 PyTorch 中使用 FlexAttention

这看起来像是 FlexAttention 的活儿！它有一个纯 PyTorch 接口，我们探索了两种方案：

在我们因果掩码中注意力被关闭的位置 添加一个 score_mod。然而，即使是一个标量加法也会 显著降低 FlexAttention 的性能。这是因为在我们的案例中，score_mod 是在优化的 CUDA 内核之外添加的。
正确的选项是 对我们的因果掩码进行索引，并将得到的签名传递给 block mask 的创建。 这个 block mask 能有效地指示哪些地方需要计算注意力，哪些地方可以完全跳过。

# Example of indexing the causal mask and using mask_mod
causal_mask = generate_causal_mask(your_samples) # should be[batch, head, q_len, kv_len]
# Now we need to wrap this mask 
def precomputed_mask_factory(precomputed_mask: torch.Tensor) -> _mask_mod_signature:
    def mask_mod(b, h, q_idx, kv_idx):
        return precomputed_mask[b][h][q_idx][kv_idx]
    return mask_mod
flex_attention_output = flex_attention(query, key, value, mask_mod=mask_mod)

mask_mod = precomputed_mask_factory(causal_mask)
# create a block mask with that signature
block_mask = create_block_mask(
    mask_mod=mask_mod,
    # ...
)

# Call flex attention now!
attn_output, attention_weights = flex_attention(
    query,
    key,
    value,
    block_mask=block_mask,
)

目前的实现可以运行，但仍在进行中，目标是让它支持 torch.compile 并充分利用其优势！

如何有效表示“动作”？

既然我们知道动作只是可以被分词化的 n 维 向量，我们就可以探讨在视觉-语言-动作 (VLA) 模型中动作表示的挑战。动作的表示方式直接影响效率、泛化能力和执行保真度。

一种方法是 语义化动作表示，其中动作被描述为高级概念，如子任务或关键点。虽然这允许小样本和零样本学习，但它通常依赖于手动设计的底层控制器，限制了在不同机器人间的灵活性。相比之下，底层控制表示直接将动作映射到电机指令，能够实现精确运动，但训练过程不太稳定且难以扩展。

大多数现有的 VLA 使用 离散动作分词，将连续动作转换为自回归生成的离散词元。最常用的方法——逐维度、逐时间步分箱——在处理高频控制任务时表现不佳，导致表示有损且训练效率低下。矢量量化 (VQ) 和时间序列压缩等替代方案有所帮助，但 VQ 对超参数敏感，使其在多样化的机器人设计中可靠性较低。

为了解决这些问题，频域动作序列分词 (FAST) 引入了一种新颖的时间序列压缩方法，使用离散余弦变换 (DCT)。FAST 减少了冗余，提高了效率，并增强了动作的保真度。

基于此，我们推出了 π0-FAST，这是 π0 的一个更快、自回归的版本，也已在 Lerobot 代码库中提供。它是 π0 的扩展，利用了这种新的分词器来更好地表示动作。

🚀 π0-FAST 是什么？

论文 | Jax 代码 | 我们在 Lerobot 中的实现

π0-FAST 是 π0 的一个 自回归版本，引入了 **FAST (频域动作序列分词)**——一种新的分词方案，可提高效率和性能。

π0-FAST 的主要优势：

与基于扩散的 VLA 相比，训练速度快 5 倍。
改进的动作表示，减少了动作序列中的冗余。
在未见过的环境和机器人形态上具有 更强的泛化能力。

🔗 π0-FAST 分词器 可在此处获取：FAST Tokenizer

🔗 预训练权重可在此处获取：Pi0+FAST

FAST 是如何工作的？

FAST 使用离散余弦变换 (DCT) 将连续的动作序列压缩为离散的词元。如图 2 所示，该过程首先对原始机器人动作进行归一化，将每个动作维度的第 1 和第 99 百分位数映射到 [-1, 1] 范围内。这种归一化用于确保不同机器人系统间的一致性，并提高对异常值的鲁棒性。

然后，每个动作维度使用 DCT 独立进行变换，将时域信号转换为频域。为了减少冗余，通过缩放和四舍五入操作去除不重要的系数，其中一个超参数平衡了压缩率和重建精度。得到的 DCT 系数矩阵通常是稀疏的，被展平成一个一维整数序列，首先跨维度交错低频分量以保留关键信息。

为了进一步压缩序列，应用了字节对编码 (BPE)。像往常一样，BPE 合并跨维度频繁出现的模式，同时保持固定大小的词汇表。

图 2：FAST 动作分词流程

由于所有操作都是可逆的，因此可以从词元中高效且无损地重建动作。该分词流程只有两个超参数：四舍五入前应用的缩放系数和 BPE 词汇表大小。这两个参数在不同数据集上都保持鲁棒性。

此外，一个名为 FAST+ 的通用版本 FAST 已经在一个包含单臂、双臂和移动操作机器人的一百万个动作序列上进行了训练，使其适用于各种机器人设置。FAST+ 作为 Hugging Face AutoProcessor 提供，允许用户仅用几行代码就能对动作序列进行分词。

为实现最佳压缩效果，输入动作在分词前应进行分位数归一化至 [-1, 1]。借助 AutoProcessor 模块，用户可以在自己的数据集上训练自定义的 FAST 分词器。

如何使用 FAST 分词器？

🔗 有关自定义动作分词器的使用和训练代码，请参见官方 FAST 代码库

FAST 已集成到 Hugging Face Transformers 中，可以轻松用于编码和解码机器人动作序列。

通用机器人智能的下一步是什么？

借助 π0 和 π0-FAST，我们向通用机器人智能迈出了重要一步，将可扩展、高效且多功能的视觉-语言-动作 (VLA) 模型引入 LeRobot。通过利用 FAST 分词，我们增强了动作表示，使机器人能够以更高的效率和适应性执行各种任务。这些进展为未来的多形态、实时机器人策略打开了大门，推动了机器人在现实世界中所能达到的极限。🚀

其他资源

参考文献

@book{heinlein2021time,
  title={Time enough for love},
  author={Heinlein, Robert A},
  year={2021},
  publisher={Penguin}
}

@article{black2024pi_0,
  title={$$\backslash$pi\_0 $: A Vision-Language-Action Flow Model for General Robot Control},
  author={Black, Kevin and Brown, Noah and Driess, Danny and Esmail, Adnan and Equi, Michael and Finn, Chelsea and Fusai, Niccolo and Groom, Lachy and Hausman, Karol and Ichter, Brian and others},
  journal={arXiv preprint arXiv:2410.24164},
  year={2024}
}

@article{pertsch2025fast,
  title={FAST: Efficient Action Tokenization for Vision-Language-Action Models},
  author={Pertsch, Karl and Stachowicz, Kyle and Ichter, Brian and Driess, Danny and Nair, Suraj and Vuong, Quan and Mees, Oier and Finn, Chelse|a and Levine, Sergey},
  journal={arXiv preprint arXiv:2501.09747},
  year={2025}
}

更多博客文章