π₀-FAST (Pi0-FAST)

π₀-FAST 是一个用于通用机器人控制的视觉-语言-动作模型，使用自回归下一个令牌预测来建模连续机器人动作。

模型概述

π₀-FAST 将视觉-语言模型的强大功能与一种称为 FAST（频率空间动作序列标记化） 的新颖动作标记化方法相结合。这使得能够在使用标准基于分箱的离散化无法实现的高度灵巧任务上训练自回归 VLA，同时训练速度比基于扩散的方法（如 π₀）快 5 倍。

为什么选择 FAST？

机器人动作标记化的标准方法使用简单的每维度、每时间步分箱方案。虽然对于简单行为来说还可以，但对于需要精度和高频控制的复杂和灵巧技能，这会迅速崩溃。

FAST 通过使用信号处理技术压缩动作序列来解决这个问题，从而产生可以自回归预测的密集动作令牌序列——就像语言令牌一样。

FAST 标记化的工作原理

FAST 标记器通过以下步骤压缩动作序列：

1. 归一化：取形状为 (H, D) 的连续动作块，其中 H 是范围，D 是动作维度。使用支持的归一化方法之一进行归一化（建议使用分位数来处理异常值）。

2. 离散余弦变换（DCT）：对每个动作维度分别应用 DCT（通过 scipy）。DCT 是图像和音频编解码器（JPEG、MP3）中常用的压缩算法。

3. 量化：对每个动作维度进行舍入并删除不重要的系数，产生稀疏频率矩阵。

4. 展平：将矩阵展平为 1D 向量，低频分量在前。

5. 字节对编码（BPE）：训练 BPE 标记器将 DCT 系数压缩为密集动作令牌，通常实现比先前标记化方法压缩 10 倍。

这种方法可以通过训练它来预测这些 FAST 令牌，将任何现有的 VLM 转换为 VLA。

安装要求

1. 按照我们的安装指南安装 LeRobot。
2. 通过运行以下命令安装 π₀-FAST 依赖项：

pip install -e ".[pi]"

训练自定义 FAST 标记器

您有两个 FAST 标记器选项：

1. 使用预训练的标记器：lerobot/fast-action-tokenizer 标记器在 100 万+ 真实机器人动作序列上训练，可作为通用标记器使用。

2. 训练您自己的标记器：为了在您的特定数据集上获得最大性能，您可以在自己的数据上微调标记器。

训练您自己的标记器

lerobot-train-tokenizer \
    --repo_id "user/my-lerobot-dataset" \
    --action_horizon 10 \
    --encoded_dims "0:6" \
    --vocab_size 1024 \
    --scale 10.0 \
    --normalization_mode QUANTILES \
    --output_dir "./my_fast_tokenizer" \
    --push_to_hub \
    --hub_repo_id "username/my-action-tokenizer"

关键标记器参数

参数	描述	默认值
--repo_id	LeRobot 数据集仓库 ID	必需
--action_horizon	每个块中的未来动作数	10
--encoded_dims	要编码的逗号分隔维度范围（例如，"0:6,7:23"）	"0:6,7:23"
--vocab_size	BPE 词汇表大小	1024
--scale	用于量化的 DCT 缩放因子	10.0
--normalization_mode	归一化模式（MEAN_STD、MIN_MAX、QUANTILES、QUANTILE10、IDENTITY）	QUANTILES
--sample_fraction	每个回合采样的块的比例	0.1

使用

要在 LeRobot 中使用 π₀-FAST，请将策略类型指定为：

policy.type=pi0_fast

训练

对于训练 π₀-FAST，您可以使用 LeRobot 训练脚本：

lerobot-train \
    --dataset.repo_id=your_dataset \
    --policy.type=pi0_fast \
    --output_dir=./outputs/pi0fast_training \
    --job_name=pi0fast_training \
    --policy.pretrained_path=lerobot/pi0_fast_base \
    --policy.dtype=bfloat16 \
    --policy.gradient_checkpointing=true \
    --policy.chunk_size=10 \
    --policy.n_action_steps=10 \
    --policy.max_action_tokens=256 \
    --steps=100000 \
    --batch_size=4 \
    --policy.device=cuda

关键训练参数

参数	描述	默认值
--policy.gradient_checkpointing=true	在训练期间显著减少内存使用	false
--policy.dtype=bfloat16	使用混合精度训练以提高效率	float32
--policy.chunk_size	要预测的动作步数（动作范围）	50
--policy.n_action_steps	要执行的动作步数	50
--policy.max_action_tokens	每个动作块的最大 FAST 令牌数	256
--policy.action_tokenizer_name	要使用的 FAST 标记器	lerobot/fast-action-tokenizer
--policy.compile_model=true	启用 torch.compile 以加快训练速度	false

推理

用于快速推理的 KV 缓存

π₀-FAST 支持 KV 缓存，这是 LLM 推理中广泛使用的优化。这会缓存注意力机制中的键值对，避免在自回归解码期间进行冗余计算。

# KV 缓存默认启用
policy.use_kv_cache=true

推理示例

from lerobot.policies.pi0_fast import PI0FastPolicy, PI0FastConfig

# 加载策略
policy = PI0FastPolicy.from_pretrained("your-model-path")

# 在推理期间
actions = policy.predict_action_chunk(batch)

模型架构

π₀-FAST 使用基于 PaliGemma 的架构：

• 视觉编码器：用于图像理解的 SigLIP 视觉塔
• 语言模型：Gemma 2B，用于处理语言指令和预测动作令牌

该模型将图像、文本指令和机器人状态作为输入，并输出离散的 FAST 令牌，这些令牌被解码回连续动作。

配置选项

参数	描述	默认值
paligemma_variant	VLM 主干变体（gemma_300m、gemma_2b）	gemma_2b
max_state_dim	最大状态向量维度（填充）	32
max_action_dim	最大动作向量维度（填充）	32
temperature	采样温度（0.0 表示贪婪）	0.0
max_decoding_steps	最大解码步数	256
use_kv_cache	启用 KV 缓存以加快推理速度	true

与 π₀ 的比较

特性	π₀	π₀-FAST
动作表示	流匹配（扩散）	自回归令牌（FAST）
训练速度	1x	快 5 倍
灵巧性	高	高
推理方法	迭代去噪	自回归解码
KV 缓存	不适用	支持

重现 π₀Fast 结果

我们使用 LeRobot 实现在 LIBERO 基准测试上重现 π₀Fast 的结果。我们采用 LeRobot PiFast 基础模型 lerobot/pi0fast-base，并在 8 个 H100 GPU 上使用批量大小为 256 的 bfloat16 进行额外 40k 步的微调，使用 HuggingFace LIBERO 数据集。

微调后的模型可以在这里找到：

• π₀Fast LIBERO：lerobot/pi0fast-libero

使用以下训练命令：

lerobot-train \
  --dataset.repo_id=lerobot/libero \
  --output_dir=outputs/libero_pi0fast \
  --job_name=libero_pi0fast \
  --policy.path=lerobot/pi0fast_base \
  --policy.dtype=bfloat16 \
  --steps=100000 \
  --save_freq=20000 \
  --batch_size=4 \
  --policy.device=cuda \
  --policy.scheduler_warmup_steps=4000 \
  --policy.scheduler_decay_steps=100000 \
  --policy.scheduler_decay_lr=1e-5 \
  --policy.gradient_checkpointing=true \
  --policy.chunk_size=10 \
  --policy.n_action_steps=10 \
  --policy.max_action_tokens=256 \
  --policy.empty_cameras=1 \

然后我们使用 LeRobot LIBERO 实现评估微调后的模型，运行以下命令：

tasks="libero_object,libero_spatial,libero_goal,libero_10"
lerobot-eval \
  --policy.path=lerobot/pi0fast-libero \
  --policy.max_action_tokens=256 \
  --env.type=libero \
  --policy.gradient_checkpointing=false \
  --env.task=${tasks} \
  --eval.batch_size=1 \
  --eval.n_episodes=1 \
  --rename_map='{"observation.images.image":"observation.images.base_0_rgb","observation.images.image2":"observation.images.left_wrist_0_rgb"}'

注意： 我们设置 n_action_steps=10，类似于原始 OpenPI 实现。

结果

我们在 LIBERO 基准测试上获得以下结果：

模型	LIBERO Spatial	LIBERO Object	LIBERO Goal	LIBERO 10	平均
π₀-fast	70.0	100.0	100.0	60.0	82.5

完整的评估输出文件夹（包括视频）可在此处获得

许可证

此模型遵循 Apache 2.0 许可证，与原始 OpenPI 仓库一致。

参考

• FAST：高效的机器人动作标记化 - Physical Intelligence 博客
• OpenPI 仓库 - 原始实现
• Hugging Face 上的 FAST 标记器 - 预训练标记器