# π₀-FAST

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

## 模型概览

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

<img
  src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-pifast.png"
  alt="An overview of Pi0-FAST"
  width="85%"
/>

### 为什么选择 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. 按照我们的[安装指南](./installation)安装 LeRobot。
2. 运行以下命令安装 π₀-FAST 依赖项：

   ```bash
   pip install -e ".[pi]"
   ```

## 训练自定义 FAST 分词器

对于 FAST 分词器，您有两种选择：

1. **使用预训练的分词器**：`lerobot/fast-action-tokenizer` 分词器在超过 100 万个真实机器人动作序列上进行训练，可作为通用分词器使用。

2. **训练你自己的分词器**：为了在特定数据集上获得最佳性能，你可以在自己的数据上微调分词器。

### 训练你自己的分词器

```bash
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_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，请将策略类型指定为：

```python
policy.type=pi0_fast
```

## 训练

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

```bash
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 推理中广泛使用的优化技术。它通过缓存注意力机制中的键值对，避免在自回归解码过程中进行冗余计算。

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

### 推理示例

```python
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_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 缓存            | N/A                       | 支持                    |

## 复现 π₀Fast 的结果

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

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

- **π₀Fast LIBERO**：[lerobot/pi0fast-libero](https://huggingface.co/lerobot/pi0fast-libero)

使用以下训练命令：

```bash
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 实现来评估微调后的模型，运行以下命令：

```bash
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 | 平均  |
| ----------- | -------------- | ------------- | ----------- | --------- | -------- |
| **π₀-fast** | 70.0           | 100.0         | 100.0       | 60.0      | **82.5** |

完整的评估输出文件夹（包括视频）可在[此处](https://drive.google.com/drive/folders/1HXpwPTRm4hx6g1sF2P7OOqGG0TwPU7LQ?usp=sharing)获取

## 许可证

该模型遵循 **Apache 2.0 许可证**，与原始的 [OpenPI 仓库](https://github.com/Physical-Intelligence/openpi) 保持一致。

## 参考

- [FAST：高效的机器人动作分词](https://www.physicalintelligence.company/research/fast) - Physical Intelligence 博客
- [OpenPI 仓库](https://github.com/Physical-Intelligence/openpi) - 原始实现
- [Hugging Face 上的 FAST Tokenizer](https://huggingface.co/physical-intelligence/fast) - 预训练分词器