LeRobot documentation

π₀-FAST

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

π₀-FAST

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

模型概览

π₀-FAST 将视觉语言模型的能力与一种名为 FAST(Frequency-space Action Sequence Tokenization) 的新型动作分词方法相结合。这使得在标准基于分箱的离散化方法无法处理的高灵巧性任务上训练自回归 VLA 成为可能,同时训练速度比 π₀ 等基于扩散的方法快至 5 倍。

An overview of Pi0-FAST

为什么选择 FAST?

机器人动作离散化的标准方法通常采用简单的逐维度、逐时间步分箱方案。虽然这种方法对于简单的行为尚可接受,但在需要高精度和高频控制的复杂灵巧技能中,它会迅速失效。

FAST 通过利用信号处理技术对动作序列进行压缩来解决这一问题,从而生成密集的动作标记序列,使其能够像语言标记一样进行自回归预测。

FAST 分词的工作原理

FAST 分词器通过以下步骤压缩动作序列:

  1. 归一化:获取形状为 (H, D) 的连续动作块,其中 H 为时间范围,D 为动作维度。使用受支持的归一化方法之一进行归一化(推荐 Quantiles 以处理异常值)。

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

  3. 量化:对每个动作维度进行舍入并移除不重要的系数,从而生成稀疏频率矩阵。

  4. Flatten:将矩阵展平为一维向量,低频分量在前。

  5. 字节对编码(BPE):训练一个 BPE 分词器,将 DCT 系数压缩为稠密的操作 token,通常相比先前的分词方法可实现 10 倍压缩。

该方法可以通过训练模型预测这些 FAST token,将任何现有的 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 true \
    --hub_repo_id "username/my-action-tokenizer"

关键分词器参数

参数描述默认
--repo_idLeRobot 数据集仓库 ID必需
--action_horizon每个块中未来动作的数量10
--encoded_dims要编码的逗号分隔的维度范围(例如,"0:6,7:23")"0:6,7:23"
--vocab_sizeBPE 词汇表大小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/pi0fast-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
--policy.optimizer_fused=true使用融合 AdamW;恢复时忽略None

推理

用于快速推理的 KV-Caching

π₀-FAST 支持 KV 缓存,这是一种在 LLM 推理中广泛使用的优化技术。它通过缓存注意力机制中的键值对,避免在自回归解码过程中进行冗余计算。

# KV-caching is enabled by default
policy.use_kv_cache=true

推理示例

from lerobot.policies.pi0_fast import PI0FastPolicy, PI0FastConfig

# Load the policy
policy = PI0FastPolicy.from_pretrained("your-model-path")

# During inference
actions = policy.predict_action_chunk(batch)

模型架构

π₀-FAST 采用基于 PaliGemma 的架构:

  • 视觉编码器:用于图像理解的 SigLIP 视觉塔
  • 语言模型:Gemma 2B,用于处理语言指令并预测动作标记

模型以图像、文本指令和机器人状态作为输入,输出离散的 FAST 标记,并将其解码回连续动作。

配置选项

参数描述默认
paligemma_variantVLM 骨干变体(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 缓存N/A支持

复现 π₀Fast 的结果

我们使用 LeRobot 实现复现了 π₀Fast 在 LIBERO 基准测试上的结果。我们采用 LeRobot PiFast 基础模型 lerobot/pi0fast-base,在 bfloat16 精度下,使用 HuggingFace LIBERO 数据集,在 8 块 H100 GPU 上以 256 的批大小额外微调 40k 步。

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

使用以下训练命令:

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 空间LIBERO 对象LIBERO 目标LIBERO 10平均
π₀-fast70.0100.0100.060.082.5

完整的评估输出文件夹(包括视频)可在此处获取

许可证

该模型遵循 Apache 2.0 许可证,与原始的 OpenPI 仓库 保持一致。

参考

Update on GitHub