# π₀-FAST (Pi0-FAST)

π₀-FAST — это **модель Vision-Language-Action для общего управления роботами**, которая использует авторегрессионное предсказание следующего токена для моделирования непрерывных действий робота.

## Обзор модели

π₀-FAST сочетает возможности моделей Vision-Language с новым подходом к токенизации действий, названным **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 сжимает последовательности действий с помощью следующих шагов:

1. **Normalize**: возьмите непрерывный блок действий формы `(H, D)`, где `H` — горизонт, а `D` — размерность действия. Нормализуйте с помощью одного из поддерживаемых методов нормализации (рекомендуется Quantiles для обработки выбросов).

2. **Дискретное косинусное преобразование (DCT)**: применяйте DCT (через scipy) к каждому измерению действия отдельно. DCT — это алгоритм сжатия, часто используемый в кодеках изображений и аудио (JPEG, MP3).

3. **Квантование**: округление и удаление незначимых коэффициентов для каждого измерения действия, что приводит к разреженной частотной матрице.

4. **Flatten**: Развернуть матрицу в одномерный вектор, начиная с низкочастотных компонент.

5. **Кодирование байтовых пар (BPE)**: обучите токенизатор BPE для сжатия коэффициентов DCT в плотные токены действий, как правило, обеспечивая **10-кратное сжатие** по сравнению с предыдущими подходами к токенизации.

Этот подход позволяет превратить **любую существующую VLM** в VLA, обучив её предсказывать токены FAST.

## Требования к установке

1. Установите LeRobot, следуя нашему [Руководству по установке](./installation).
2. Установите зависимости π₀-FAST, выполнив:

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

## Обучение собственного токенизатора FAST

Для токенизатора FAST у вас есть два варианта:

1. **Используйте предобученный токенизатор**: токенизатор `lerobot/fast-action-tokenizer` обучен на более чем 1 млн реальных последовательностей действий роботов и работает как универсальный токенизатор.

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                                                     | Обязательный     |
| `--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`        |

## Использование

Чтобы использовать π₀-FAST в LeRobot, укажите тип политики как:

```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 для быстрого вывода

π₀-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                      |
| --------------------- | ------------------------- | ---------------------------- |
| Представление действия | Flow Matching (диффузия) | Авто регрессионные токены (FAST) |
| Скорость обучения        | 1x                        | **В 5 раз быстрее**                |
| Сноровка             | Высокий                      | Высокий                         |
| Метод инференса      | Итеративное удаление шума       | Авто регрессионное декодирование      |
| Кэширование KV            | N/A                       | Поддерживается                    |

## Воспроизведение результатов π₀Fast

Мы воспроизводим результаты π₀Fast на бенчмарке LIBERO, используя реализацию LeRobot. Мы берём базовую модель LeRobot PiFast [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) и дообучаем её в течение дополнительных 40 тыс. шагов в формате bfloat16 с размером батча 256 на 8 GPU H100, используя [набор данных HuggingFace LIBERO](https://huggingface.co/datasets/HuggingFaceVLA/libero).

Дообученная модель доступна здесь:

- **π₀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 Spatial | LIBERO Object | 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 License**, что соответствует исходному репозиторию [OpenPI](https://github.com/Physical-Intelligence/openpi).

## Ссылки

- [FAST: эффективная токенизация действий робота](https://www.physicalintelligence.company/research/fast) — блог Physical Intelligence
- [Репозиторий OpenPI](https://github.com/Physical-Intelligence/openpi) — исходная реализация
- [FAST Tokenizer на Hugging Face](https://huggingface.co/physical-intelligence/fast) — предобученный токенизатор