Тонкая настройка модели Pi0 с вашими данными

Pi0 — это мощная модель Vision-Language-Action (VLA), которая может быстро адаптироваться к вашим роботизированным задачам с помощью тонкой настройки на небольшом количестве данных.

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

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

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

Что такое тонкая настройка?

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

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

С помощью тонкой настройки вы можете заставить Pi0:

• Адаптироваться к вашей роботизированной среде: Адаптироваться к перспективе вашей камеры, световой среде и механической структуре
• Адаптироваться к конкретным задачам: Оптимизировать для ваших конкретных задач (таких как поднятие ручек, классификация, размещение)
• Улучшить точность управления: Значительно улучшить точность управления и успешность модели в ваших задачах

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

Настройка окружения

Системные требования

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

• Python ≥ 3.8 (рекомендуется 3.10 или 3.11)
• GPU: Рекомендуется не менее 32 ГБ видеопамяти (Pi0 — большая модель, рекомендуется V100 или более высокого уровня графические карты)
• Память: Рекомендуется не менее 64 ГБ ОЗУ
• Хранилище: Убедитесь, что у вас достаточно места на диске для наборов данных и моделей

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

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

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

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

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

Экспорт данных в формате 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 довольно большая, поэтому обучение на одной карте занимает как минимум десятки часов. Например, когда мы выполняли тонкую настройку 2 часов собранных данных, это заняло 50 часов на одной карте V100. 8 карт V100 ускорят процесс, но все равно потребуется почти 10 часов.

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

Параметр	Значение	Рекомендуемое значение	Описание
--dataset.root	Локальный путь к набору данных	~/your_dataset	Указывает на ваш каталог набора данных LeRobot
--dataset.repo_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	Использовать Flash Attention 2 для ускорения
--training.learning_rate	Скорость обучения	1e-5	Использовать меньшую скорость обучения для тонкой настройки
--training.num_epochs	Количество эпох обучения	10	Настроить в зависимости от объема данных

Ускоренное обучение на нескольких 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. Сетевое ускорение: Возможно, вам понадобится сетевое ускорение для загрузки модели Pi0 с HuggingFace, обратитесь к инструментам типа hg-mirror
2. Мониторинг VRAM: Пожалуйста, следите за использованием VRAM GPU во время обучения
3. Возобновление обучения: Если обучение прервано, вы можете продолжить с контрольных точек в --output_dir

Важные примечания: Совместимость модели Gemma

Модель Pi0 внутренне использует структуру модели Google PaLI-Gemma для языковой обработки. Поскольку это недавно добавленная модель Hugging Face:

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

• Transformers: Убедитесь, что вы используете последнюю версию transformers (рекомендуется ≥ 4.37.0)
• Проверка совместимости: Если вы столкнетесь с ошибками, связанными с embed_tokens, обновите transformers

Решения

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

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

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

Место сохранения модели

После завершения обучения модель будет сохранена в:

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

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

from lerobot.policies import Pi0Policy

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

# Выполнить вывод
action = policy.select_action(batch)

Анализ журналов обучения

Во время обучения будут генерироваться подробные журналы, включая:

• Кривые потерь: Изменения в потерях обучения и валидации
• Планирование скорости обучения: Изменения скорости обучения
• Информация о градиентах: Нормы градиентов и обрезка градиентов
• Использование памяти: VRAM GPU и использование системной памяти

Как оценить результаты?

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

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

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

Вы также можете оценить производительность Pi0 в ваших задачах следующими способами:

1. Количественная оценка

• Ошибка действия: Сравнить предсказанные действия с реальными действиями (такими как ошибка угла, ошибка позиции)
• Сравнение траекторий: Визуализировать предсказанные vs ground truth поведенческие траектории
• Метрики потерь: Записать тенденции изменения потерь на валидационном наборе

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

• Реальное развертывание: Проверить успешность задач на реальных роботах
• Наблюдение за поведением: Наблюдать, соответствуют ли действия, сгенерированные моделью, ожиданиям
• Обратная связь пользователей: Собирать обратную связь пользователей из фактического использования

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

import numpy as np
from lerobot.policies import Pi0Policy

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

# Функция оценки
def evaluate_model(policy, test_dataset):
    total_error = 0
    num_samples = 0
    
    for batch in test_dataset:
        predicted_action = policy.select_action(batch)
        true_action = batch['actions']
        
        # Вычислить ошибку действия
        error = np.mean(np.abs(predicted_action - true_action))
        total_error += error
        num_samples += 1
    
    return total_error / num_samples

# Выполнить оценку
avg_error = evaluate_model(policy, test_dataset)
print(f"Средняя ошибка действия: {avg_error:.4f}")

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

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

Часто задаваемые вопросы

Q: Модель сообщает, что embed_tokens не существует во время тонкой настройки?

A: Это происходит потому, что версия transformers слишком старая, просто обновите:

pip install -U transformers

Q: Набор данных не загружен в Hugging Face, можно ли использовать локальные данные?

A: Да, используйте параметр --dataset.root=your_local_path.

Q: Что делать, если VRAM недостаточно во время обучения?

A: Вы можете попробовать следующие методы:

# 1. Уменьшить размер пакета
--batch_size=1

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

Q: Сколько обычно занимает обучение?

A: Время обучения зависит от:

• Объема данных: Обычно 10+ часов
• Конфигурации оборудования: Производительность GPU имеет большое влияние
• Эпох обучения: Рекомендуется 5-20 эпох

Q: Как определить, сходится ли обучение?

A: Наблюдайте следующие показатели:

• Кривые потерь: Потери обучения и валидации стабилизируются
• Метрики валидации: Производительность на валидационном наборе больше не улучшается
• Переобучение: Остановите обучение, когда потери валидации начнут расти

Q: Как развернуть тонко настроенную модель?

A: Платформа IO Data уже поддерживает автоматическое развертывание Pi0 и SmolVLA и других распространенных моделей в области робототехники. Пожалуйста, проконсультируйтесь с нашими техническими экспертами для получения подробной информации.

Вы также можете обратиться к следующим шагам:

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

# 2. Подготовить входные данные
observation = {
    'images': camera_images,
    'language_instruction': "Поднимите красную ручку"
}

# 3. Сгенерировать команды действий
action = policy.select_action(observation)

Удачи в обучении!