Тонкая настройка модели Pi0: Адаптация к задачам роботов с пользовательскими наборами данных

Pi0 — это продвинутый модель Vision-Language-Action (VLA), которая может быстро адаптироваться к конкретным задачам роботов через few-shot fine-tuning.

Этот документ основан на фреймворке LeRobot и использует предварительно обученную модель Pi0 (lerobot/pi0) для тонкой настройки на пользовательских наборах данных роботов.

Предварительные требования

• Иметь набор данных в формате LeRobot (см. руководство по экспорту набора данных LeRobot в предыдущем разделе)
• Необходимость адаптировать модель Pi0 к конкретному оборудованию роботов, задачам или стратегиям управления
• Знакомство с процессом обучения PyTorch и экосистемой Hugging Face
• Базовый опыт обучения моделей глубокого обучения

Обзор тонкой настройки

Тонкая настройка (Fine-Tuning) — это процесс дальнейшего обучения предварительно обученной модели с использованием данных конкретной области, обычно с меньшим количеством итераций.

Модель Pi0 предварительно обучена на разнообразных задачах роботов, обладая общими возможностями визуального восприятия, понимания языка и генерации действий.

Через тонкую настройку Pi0 может достичь:

• Адаптация к окружению: Настройка на конкретные углы камеры, условия освещения и механические структуры
• Специализация задач: Оптимизация производительности для конкретных задач (например, захват объектов, классификация или размещение)
• Повышение точности: Значительное улучшение точности управления и коэффициента успеха в целевых задачах

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

Подготовка окружения

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

Убедитесь, что ваше окружение соответствует следующим требованиям:

• Python ≥ 3.8 (рекомендуется 3.10 или выше)
• GPU: Не менее 32 ГБ VRAM (Pi0 — крупномасштабная модель, рекомендуется NVIDIA V100 или GPU с более высокой производительностью)
• Память: Не менее 64 ГБ системной RAM
• Хранение: Достаточно дискового пространства для наборов данных и контрольных точек модели

Установка зависимостей

# Клонировать репозиторий LeRobot (основной)
git clone https://github.com/lerobot-ai/lerobot.git
cd lerobot/

# Установить фреймворк LeRobot (включая поддержку Pi0)
pip install -e ".[pi0]"

# Проверить установку
python -c "from lerobot.policies import Pi0Policy; print('Pi0 успешно установлен!')"

Зеркало (при необходимости): https://github.com/huggingface/lerobot

Подготовка вашего набора данных

Экспорт данных в формате LeRobot

Вы можете экспортировать аннотированные данные в набор данных формата LeRobot через страницу экспорта платформы IO Data, см. Экспорт набора данных LeRobot.

Важное замечание: Для обучения не нужно загружать данные в HuggingFace; используйте параметр --dataset.root= для указания данных локального каталога для тонкой настройки.

Пример структуры набора данных

Предположим, что экспортированные данные хранятся в ~/DualPiper_Pickup_Pen:

$ ls ~/DualPiper_Pickup_Pen
data  meta  videos

$ ls ~/DualPiper_Pickup_Pen/data
episode_000  episode_001  episode_002  ...

$ ls ~/DualPiper_Pickup_Pen/data/episode_000
observations  actions.npy  language_instruction.txt  metadata.json

Начало тонкой настройки

Базовая команда обучения

Используйте скрипт train.py для обучения:

# Установить политику распределения памяти CUDA (рекомендуется)
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# Запустить обучение
python3 -m lerobot.scripts.train \
  --dataset.root=~/DualPiper_Pickup_Pen \
  --dataset.repo_id=io-ai-data/DualPiper_Pickup_Pen \
  --policy.path=lerobot/pi0 \
  --policy.repo_id=lerobot/pi0 \
  --policy.device=cuda \
  --output_dir=./checkpoints/pi0_finetuned \
  --batch_size=1 \
  --policy.attention_implementation=flash_attention_2 \
  --training.learning_rate=1e-5 \
  --training.num_epochs=10

подсказка

Модель Pi0 имеет большое количество параметров, обучение на одном GPU может занять десятки часов. Например, тонкая настройка набора данных на 2 часа занимает около 50 часов на одном NVIDIA V100, в то время как использование 8 V100 может сократить до около 10 часов.

Детали параметров

Параметр	Значение	Рекомендуемое значение	Описание
--dataset.root	Локальный путь к набору данных	~/your_dataset	Указывает на каталог набора данных в формате LeRobot
--dataset.repo_id	ID набора данных Hugging Face	your-username/dataset	Для идентификации метаданных
--policy.path	Путь к предварительно обученной модели	lerobot/pi0	Указывает официальную предварительно обученную модель
--policy.repo_id	ID репозитория модели	lerobot/pi0	Репозиторий модели Hugging Face
--policy.device	Устройство обучения	cuda	Включить ускорение GPU
--output_dir	Каталог вывода	./checkpoints/pi0_finetuned	Путь сохранения тонко настроенной модели
--batch_size	Размер пакета	1	Настроить на основе VRAM
--policy.attention_implementation	Механизм внимания	flash_attention_2	Использовать эффективную реализацию для ускорения обучения
--training.learning_rate	Скорость обучения	1e-5	Использовать низкое значение при тонкой настройке, чтобы избежать катастрофического забывания
--training.num_epochs	Количество эпох обучения	10	Настроить на основе масштаба набора данных

Ускоренное обучение на нескольких GPU

Если у вас несколько GPU, параллельное обучение может значительно ускорить процесс.

# Обучение на 8 GPU
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
torchrun --nproc_per_node=8 --master_port=29500 \
    -m lerobot.scripts.train \
    --dataset.root=/data/nfs2/export/lerobot/DualPiper_Pickup_Pen/ \
    --dataset.repo_id=io-ai-data/DualPiper_Pickup_Pen \
    --policy.path=lerobot/pi0 \
    --policy.repo_id=lerobot/pi0 \
    --policy.device=cuda \
    --output_dir=/home/puxk/DualPiper_Pickup_Pen_pi0_model \
    --batch_size=1 \
    --num_workers=4 \
    --policy.attention_implementation=flash_attention_2

Общий размер пакета — nproc_per_node × batch_size. С --batch_size=1, 8 карт равны общему размеру пакета 8; увеличьте, если памяти достаточно.

Для отладки рекомендуется сначала протестировать с --nproc_per_node=2, затем расширить до всех 8 карт.

Особые замечания

1. Доступ к сети: Загрузка моделей Hugging Face может потребовать прокси-инструментов, таких как hf-mirror
2. Мониторинг VRAM: Мониторить использование GPU в реальном времени во время обучения
3. Восстановление контрольной точки: Если обучение прервано, возобновить с последней контрольной точки в --output_dir

Модель Pi0 интегрирует Google PaliGemma для обработки языка. Как недавняя модель Hugging Face:

Совместимость версий

• Transformers: Требуется ≥ 4.37.0
• Проверка совместимости: Если возникает ошибка embed_tokens, обновите transformers

Решения

# Обновить до последней версии transformers
pip install -U transformers

# Или указать версию
pip install transformers>=4.37.0

Выбор политики (Pi0 vs SmolVLA)

• Для однотасковых, короткогоризонтных, высокочастотных задач с умеренными языковыми требованиями: политики в стиле Pi0/ACT часто более эффективны по вычислениям и отзывчивости.
• Для лучшей мультизадачной и кросс-сценарной генерализации либо при ориентации на одну GPU/потребительское железо: предпочтительнее SmolVLA; небольшая донастройка smolvla_base обычно даёт стабильные улучшения.
• При ограниченных ресурсах: включайте смешанную точность, накопление градиентов, снижайте разрешение входов и размер пакета.

Проверка результатов обучения

Сохранение модели

После обучения файлы модели находятся в:

./checkpoints/pi0_finetuned/
├── config.json
├── pytorch_model.bin
├── tokenizer.json
└── training_args.bin

Загрузка тонко настроенной модели

from lerobot import Pi0Policy  # Примечание: Путь импорта может быть lerobot.policies

# Загрузить модель
policy = Pi0Policy.from_pretrained("./checkpoints/pi0_finetuned")

# Пример вывода
action = policy.select_action(observation_batch)

Анализ логов

Логи обучения включают:

• Кривые потерь: Тренды потерь обучения/валидации
• Скорость обучения: Изменения планировщика
• Статистика градиентов: Нормы и информация о клиппинге
• Использование ресурсов: Использование памяти GPU/CPU

Оценка производительности

Оценка на платформе IO Data

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

Другие методы оценки

Количественные метрики

• Ошибка действий: Средняя абсолютная ошибка (MAE) или среднеквадратичная ошибка (MSE) между предсказанными действиями и ground truth
• Сходство траекторий: Сравнение траекторий с использованием динамического выравнивания времени (DTW) или расстояния Фреше
• Потери валидации: Мониторинг значений функции потерь на валидационном наборе

Качественная оценка

• Тестирование на реальном оборудовании: Развертывание на физических роботах и измерение коэффициента успеха задач
• Анализ поведения: Визуализация, разумны ли сгенерированные последовательности действий
• Оценка человеком: Сбор субъективных оценок экспертов на выходы модели

Пример скрипта оценки

import numpy as np
from lerobot import Pi0Policy
from torch.utils.data import DataLoader  # Предполагаем, что test_dataset — это DataLoader

# Загрузить модель
policy = Pi0Policy.from_pretrained("./checkpoints/pi0_finetuned")

# Функция оценки
def evaluate_model(policy, test_loader):
    total_mae = 0.0
    num_batches = 0
    
    for batch in test_loader:
        with torch.no_grad():
            predicted_action = policy.select_action(batch)
        true_action = batch['actions']
        
        # Вычислить MAE
        mae = np.mean(np.abs(predicted_action.cpu().numpy() - true_action.cpu().numpy()))
        total_mae += mae
        num_batches += 1
    
    return total_mae / num_batches

# Выполнить оценку (предполагаем, что test_loader подготовлен)
avg_mae = evaluate_model(policy, test_loader)
print(f"Средняя абсолютная ошибка (MAE): {avg_mae:.4f}")

Связанные ресурсы

• Домашняя страница проекта LeRobot
• Модель Pi0 на Hugging Face
• Оригинальная реализация Pi0 (JAX)
• Статья Pi0
• Документация LeRobot