Перейти к основному содержимому

Дообучение моделей Pi0 и Pi0.5: официальный рабочий процесс на базе OpenPI

Pi0 и Pi0.5 — это модели класса Vision-Language-Action (VLA), разработанные компанией Physical Intelligence. Если вы планируете использовать данные, экспортированные из платформы EmbodyFlow, для дообучения этих моделей, данное руководство проведет вас через весь процесс на базе официального фреймворка OpenPI.

Почему стоит выбрать OpenPI, а не фреймворк LeRobot?

Хотя LeRobot поддерживает несколько популярных моделей, включая Pi0, для серии Pi0 мы настоятельно рекомендуем использовать официальный фреймворк для обучения OpenPI.
Он разработан на базе JAX, нативно поддерживает высокопроизводительное обучение на нескольких GPU и позволяет лучше раскрыть потенциал Pi0.

1. Подготовка: экспорт и размещение данных​

Для начала нам нужно конвертировать размеченные данные на платформе EmbodyFlow в формат, который распознает OpenPI.

Процесс экспорта​

  1. Выбор формата: на странице экспорта выберите стандартный формат LeRobot v2.1. export lerobot v2.1
  2. Локальная распаковка: скачайте сгенерированный файл .tar.gz и распакуйте его.
  3. Соблюдение структуры директорий: чтобы OpenPI мог без проблем найти данные, переместите их в локальный каталог кэша Hugging Face. Например: 
    # Подготовка каталога
    mkdir -p ~/.cache/huggingface/lerobot/local/mylerobot

    # Перемещение распакованных файлов (включая meta.json, data/ и т.д.)
    mv /path/to/extracted/data/* ~/.cache/huggingface/lerobot/local/mylerobot/

Справочник по соответствию полей (на примере Aloha с тремя камерами)​

В последующих конфигурациях вам нужно убедиться, что ключи в коде соответствуют полям в данных. По умолчанию мы рекомендуем использовать:

  • cam_high: вид сверху.
  • cam_left_wrist: вид с левого запястья.
  • cam_right_wrist: вид с правого запястья.
  • state: текущее состояние робота.
  • action: целевое действие (Внимание: в конфигурации OpenPI по умолчанию для ALOHA используется 14 измерений. Если размерность ваших данных отличается, обязательно обратитесь к разделу «Технические нюансы» ниже).

2. Суть: как выбрать подходящую конфигурацию обучения?​

Логика обучения в OpenPI сильно шаблонизирована. Выбор конфигурации — это, по сути, выбор шаблона политики, наиболее близкого к вашему роботу, с последующим его дообучением.

СценарийРекомендуемый путьНа что обратить внимание
Быстрая проверка / отладка цепочкиПуть симуляции (LIBERO / ALOHA Sim)Быстрое сопоставление входов/выходов, минимальные затраты.
Развертывание на реальном роботе (Aloha)ALOHA RealНеобходимо сопоставить ключи камер, размерность действий и логику захвата.
Однорукий / промышленный роботПример UR5Сначала решить вопрос совместимости интерфейса управления, затем — эффективности обучения.
Максимальная генерализацияНа базе данных DROIDИзучение стратегии нормализации (Norm Stats) DROID.

Проще говоря: если вы запускаете процесс впервые, используйте конфигурацию симуляции для отладки; если планируете работу с реальным роботом, выбирайте ALOHA Real и строго соблюдайте размерность state/action.

Технические нюансы: размерность вектора действий (14 vs 16)

Это «подводный камень», который очень легко пропустить. В политике ALOHA по умолчанию в OpenPI (aloha_policy.py) жестко прописана 14-мерная структура:

  • Структура по умолчанию: [6 суставов левой руки, 1 захват, 6 суставов правой руки, 1 захват] = 14 измерений.
  • Частая проблема: если вы используете 7-осевой манипулятор (например, [7, 1, 7, 1]), общая размерность составит 16. Если не изменить код, лишние измерения будут незаметно отсечены, и обученная модель не сможет управлять последними двумя суставами.

Рекомендации по изменению:

  1. Проверьте определение вашего вектора action.
  2. В файле aloha_policy.py замените все срезы :14 на вашу фактическую размерность (например, :16).
  3. Синхронно измените длину _joint_flip_mask, чтобы логика инверсии знаков соответствовала вашему оборудованию.

3. Написание конфигурации обучения​

Теперь нам нужно определить задачу дообучения в файле openpi/src/openpi/training/config.py.

# Пример: добавление пользовательской конфигурации для вашего робота
TrainConfig(
    name="pi0_aloha_mylerobot", 
    model=pi0_config.Pi0Config(),
    data=LeRobotAlohaDataConfig(
        repo_id="local/mylerobot",  # Указывает на каталог с данными, созданный ранее
        assets=AssetsConfig(
            assets_dir="/home/user/code/openpi/assets/pi0_aloha_mylerobot",
        ),
        default_prompt="fold the clothes", # Описание задачи, крайне важно
        repack_transforms=_transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "images": {
                            "cam_high": "observation.images.cam_high",
                            "cam_left_wrist": "observation.images.cam_left_wrist",
                            "cam_right_wrist": "observation.images.cam_right_wrist",
                        },
                        "state": "observation.state",
                        "actions": "action",
                    }
                )
            ]
        ),
    ),
    weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"),
    num_train_steps=20_000,
)

4. Запуск обучения: расчет статистики и выполнение​

Перед запуском обучения обязательно рассчитайте статистику нормализации. В противном случае диапазон значений, получаемых моделью, будет некорректным.

Шаг 1: Расчет Norm Stats​

uv run scripts/compute_norm_stats.py --config-name pi0_aloha_mylerobot

Шаг 2: Запуск дообучения​

Мы рекомендуем использовать режим JAX для максимальной производительности.

Режим с одним GPU:

export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9
CUDA_VISIBLE_DEVICES=0 uv run scripts/train.py pi0_aloha_mylerobot \
    --exp-name=my_first_experiment \
    --overwrite

Параллельный режим на нескольких GPU (FSDP):

uv run scripts/train.py pi0_aloha_mylerobot --exp-name=multi_gpu_run --fsdp-devices 4

5. Инференс и развертывание​

После завершения дообучения вы можете запустить сервер политики, чтобы робот начал выполнять задачу.

# Запуск сервера инференса, порт по умолчанию 8000
uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi0_aloha_mylerobot \
    --policy.dir=experiments/my_first_experiment/checkpoints/last

6. Поиск и устранение неисправностей (FAQ)​

  • В: Ошибка нехватки памяти (OOM)?
    • Уменьшите batch_size или проверьте правильность установки XLA_PYTHON_CLIENT_MEM_FRACTION. Использование FSDP на нескольких картах также помогает снизить нагрузку на видеопамять.
  • В: Движения модели странные или суставы вообще не двигаются?
    • Проверьте правильность сопоставления в RepackTransform.
    • Настоятельно рекомендуем вернуться к разделу о размерности 14 vs 16 и проверить, не отсекаются ли данные.
  • В: Loss при обучении вообще не снижается?
    • Проверьте точность описания в default_prompt.
    • Убедитесь, что файлы статистики, созданные compute_norm_stats, задействованы.

Дополнительные ресурсы​