π₀-FAST (Pi0-FAST)

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

Обзор модели

π₀-FAST сочетает возможности моделей Vision-Language с новым подходом к токенизации действий, названным FAST (Frequency-space Action Sequence Tokenization). Это позволяет обучать авто-регрессивные VLA на задачах с высокой степенью ловкости, которые невозможно решить с помощью стандартной дискретизации на основе бинов, при этом обучение проходит до 5 раз быстрее, чем в диффузионных подходах, таких как π₀.

Почему 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, следуя нашему Руководству по установке.

2. Установите зависимости π₀-FAST, выполнив:

   Скопировано

   pip install -e ".[pi]"

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

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

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

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"

Ключевые параметры токенизатора

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

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

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

Ключевые параметры обучения

Вывод

Кэширование KV для быстрого вывода

π₀-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, которые декодируются обратно в непрерывные действия.

Варианты конфигурации

Сравнение с π₀

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

Мы воспроизводим результаты π₀Fast на бенчмарке LIBERO, используя реализацию LeRobot. Мы берём базовую модель LeRobot PiFast lerobot/pi0fast-base и дообучаем её в течение дополнительных 40 тыс. шагов в формате bfloat16 с размером батча 256 на 8 GPU H100, используя набор данных 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:

Полная папка с результатами оценки, включая видео, доступна здесь

Лицензия

Эта модель распространяется по лицензии Apache 2.0 License, что соответствует исходному репозиторию OpenPI.

Ссылки

• FAST: эффективная токенизация действий робота — блог Physical Intelligence
• Репозиторий OpenPI — исходная реализация
• FAST Tokenizer на Hugging Face — предобученный токенизатор