05. Pi0/Pi0.5: права доступа, короткое обучение и замкнутая оценка

В этой главе задаётся стандартный порядок работы с Pi0/Pi0.5: доступ к модели, загрузка, один шаг обратного прохода, короткое обучение, сохранение модели и замкнутая оценка по фиксированному протоколу. Представление действий, порядок строк данных и эксперименты с восстановлением описаны в записях экспериментов с Pi0/Pi0.5.

Результаты главы

• результат проверки доступа к Hugging Face;
• один прямой проход, обратный проход и обновление оптимизатора;
• модель короткого обучения и metrics.jsonl;
• results.jsonl, процент успеха и видео для фиксированных random seed;
• записи о размерности действий, частоте управления и статистике нормализации.

1. Права доступа

После запроса доступа на странице модели выполните:

hf auth login
hf auth whoami

hf download lerobot/pi0 \
  --include 'config.json' \
  --local-dir "$MODEL_ROOT/pi0-base"

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

2. Предварительная проверка окружения и данных

python - <<'PY'
import torch

print("torch:", torch.__version__)
print("hip:", torch.version.hip)
print("gpu:", torch.cuda.get_device_name(0))
assert torch.cuda.is_available()
assert torch.version.hip is not None
PY

Проверка данных должна охватывать как минимум:

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

3. Один шаг обучения

Скопируйте основную конфигурацию в короткую и измените только масштаб запуска и каталог вывода:

steps: 1
batch_size: 1
save_freq: 1
output_dir: ${RUN_ROOT}/pi0/smoke

После выполнения проверьте:

1. модель загружена на CPU и перенесена на GPU;
2. значение loss конечно;
3. обратный проход даёт ненулевые градиенты;
4. оптимизатор выполнил одно обновление;
5. каталог сохранения содержит конфигурацию модели и веса, которые можно загрузить снова.

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

4. Короткое и полное обучение

Рекомендуется три этапа:

Для каждого этапа используйте отдельный каталог:

$RUN_ROOT/pi0/<run_name>/
├── run_config.json
├── train.log
├── metrics.jsonl
└── checkpoints/

При возобновлении обучения загружайте train_config.json из каталога модели и состояние оптимизатора, а в журнале проверяйте, что число шагов продолжается с последней точки сохранения. Для вывода используйте новое имя запуска и сохраняйте предыдущую рабочую модель.

5. Представление действий

До обучения зафиксируйте контракт:

observation = images + robot_state + language
action = arm_or_eef + gripper
frequency = 20 Hz
chunk = N consecutive actions from the same episode

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

delta_position = target_tcp_position - current_tcp_position
delta_rotation = relative_rotation(current_tcp, target_tcp)

Если действия в концевом эффекторе получены из траекторий суставов, пересчитайте позу концевого эффектора для каждого кадра прямой кинематикой MuJoCo и проверьте, что воспроизведённая траектория совпадает с исходной демонстрацией.

6. Замкнутая оценка

После сохранения модели используйте пакетный замкнутый скрипт:

export POLICY_TYPE=pi0
export POLICY_PATH="$RUN_ROOT/pi0/<run_name>/checkpoints/<step>/pretrained_model"
export EVAL_SEEDS=1000,1001,1002,1003
export OUTPUT_ROOT="$RUN_ROOT/pi0/<run_name>/eval"
python code/run_closed_loop.py

Короткая оценка нужна, чтобы проверить направление действий и поведение контакта. Для итоговых результатов расширяйте набор random seed и записывайте:

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

7. Проверка визуальных и языковых параметров

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

groups = {
    "vision": [],
    "language": [],
    "action": [],
}

for name, param in model.named_parameters():
    if not param.requires_grad:
        continue
    key = "vision" if "vision" in name else "language" if "language" in name else "action"
    groups[key].append(param)

for key, params in groups.items():
    trainable = sum(p.numel() for p in params)
    with_grad = sum(p.numel() for p in params if p.grad is not None)
    print(key, trainable, with_grad)

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

8. Запись результатов

В сводке эксперимента обязательно укажите:

base model:
dataset:
episodes / frames:
observation keys:
state dimension:
action representation and dimension:
control frequency:
chunk size:
train steps / batch size / learning rate:
GPU / ROCm / PyTorch:
evaluation tasks and seeds:
physical success:
artifact directory:

Подробные сравнения для EEF-delta, порядка строк блока действий, данных восстановления и длительного обучения см. в записях экспериментов с Pi0/Pi0.5.