Сбор данных, обучение на ROCm и развёртывание в MuJoCo

Первые шесть задач посвящены проверке устройства, оценке физического успеха и диагностике моделей. Эта глава добавляет недостающее звено: от записи демонстраций с клавиатуры через набор данных LeRobot и обучение ACT, SmolVLA и pi_0 на AMD ROCm до замкнутой оценки checkpoint в MuJoCo.

Здесь «работает» означает выполнение пяти последовательных проверок:

1. окно MuJoCo принимает управление с клавиатуры и выполняет один реальный захват;
2. набор данных LeRobot содержит правильные изображения, состояние, действия и языковую инструкцию;
3. smoke-тест на 2 шага проходит загрузку данных, прямой и обратный проход и запись checkpoint;
4. полное обучение стабильно использует GPU AMD без OOM, NaN и аварийного завершения ядра;
5. checkpoint в замкнутом цикле MuJoCo с реально изменяемым начальным положением объекта проходит physical_success, и есть видео для проверки.

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

Путь изучения

При первом изучении начните с ACT и пройдите весь путь. Модель ACT компактна, обучается быстро, и если с данными или определением действий что-то не так, проблему легче найти. Когда замкнутый цикл ACT заработает, обучите SmolVLA на тех же языковых данных, а затем займитесь pi_0: правами на закрытую модель, загрузкой и контролем касания при малом объёме данных.

Два варианта старта

С нуля

Выполните 07, 08 и 11, чтобы получить первый результат замкнутого цикла для ACT. Затем запустите 09 и 10 и вернитесь к 11, чтобы переключить оценку на другую модель. Для этого пути не нужен заранее обученный checkpoint.

Есть данные или checkpoint

Если вы уже прошли вышестоящий проект «04 mujoco reproduction ACT, Pi0, SmolVLA», задайте переменные окружения ниже и переходите сразу к нужному Notebook обучения или развёртывания:

export PROJECT_ROOT=/path/to/every-embodied/mujoco-reproduction-act-pi0-smolvla
export DATA_ROOT=/path/to/large-disk/datasets/every_embodied
export MODEL_ROOT=/path/to/large-disk/checkpoints/every_embodied
export OUTPUT_ROOT=/path/to/large-disk/outputs/every_embodied

PROJECT_ROOT хранит лёгкий исходный код; данные, checkpoint, видео и кэш размещаются на диске достаточного объёма. Эти крупные файлы не следует коммитить в репозиторий учебника.

Границы сбора данных

Сбор данных не зависит от производителя GPU. Если совпадают MuJoCo, версия LeRobot, XML сцены, частота управления, схема состояния и действий и критерий успеха, можно собирать данные на настольной машине с GPU NVIDIA или только на CPU, а затем синхронизировать каталог набора с машиной для обучения на AMD.

Notebook 07 вносит пять исправлений в исходный интерактивный сбор данных:

• после сохранения каждой траектории запись отключается, поэтому процесс reset не попадает в следующий эпизод;
• инструкции для красной и синей кружки чередуются, что уменьшает перекос распределения задач на малом наборе;
• эпизод сохраняется только тогда, когда кружка действительно поднята, стоит вертикально, высота кружки и тарелки разумна и условие стабильно выполняется 5 шагов подряд;
• если тарелка сдвинута больше чем на 5 см, эпизод отклоняется, чтобы «кружка, вдавленная в тарелку», не засчитывалась как успешная демонстрация;
• каждая траектория одновременно записывается в MP4 с четырёх видов: Agent, Egocentric, Top и Side. Модель использует только первые две камеры, остальные два вида нужны для ручной проверки.

На устройстве AMD уже есть набор demo_data_language. По результатам измерений в нём 20 эпизодов, 2621 кадр, 20 Гц, состояние 6-мерное, действие 7-мерное. Notebook 07 показывает эту сводку и один реальный повтор с четырёх видов, поэтому структуру данных и итоговое видео можно увидеть ещё до начала собственного сбора.

Также проверьте генератор случайных чисел в mujoco_env/y_env2.py. В старом коде было np.random.seed(seed=0), поэтому любой переданный seed давал одно и то же положение объектов. Сейчас исправлено на np.random.seed(seed=seed). Старые данные подходят для smoke-теста на фиксированной сцене, но для проверки обобщения по положению их нужно пересобрать после исправления.

Если удалённый JupyterHub или Code Server не даёт окну MuJoCo получить фокус клавиатуры, не пытайтесь собирать данные в безоконной сессии. Используйте удалённый рабочий стол или соберите данные локально и синхронизируйте их.

Smoke-тест и полное обучение

Для ACT, SmolVLA и pi_0 (главы 08–10) создаются два файла YAML:

• конфигурация smoke запускает только 2 шага и нужна для проверки цепочки;
• полная конфигурация использует базовое число шагов обучения для модели и сохраняет промежуточные checkpoint.

В Notebook переменные RUN_SMOKE и RUN_FULL_TRAIN по умолчанию равны False. Сначала проверьте пути к данным, каталоги вывода и сгенерированную конфигурацию, а затем явно их включите. После smoke-теста проверьте:

• загрузку GPU и использование видеопамяти в rocm-smi выглядят разумно;
• значения loss и норма градиентов конечны;
• checkpoint записан в $MODEL_ROOT;
• на системном диске достаточно места;
• токен Hugging Face не записан в Notebook или журнал.

Для pi_0 нужны права на закрытую модель Hugging Face. Notebook проверяет только наличие HF_TOKEN и не выводит и не сохраняет его.

Главы 08–10 читают число шагов обучения, метрики запуска и текстовые индикаторы прогресса. Предварительные ячейки запускают короткий тест smoke или длинное обучение full и продолжают читать журнал в Notebook. Глава 11 использует checkpoint, полученный в текущем обучении, для замкнутой оценки, подсчитывает успешные попытки и генерирует видео.

Замкнутый цикл MuJoCo

Notebook 11 выполняет следующий цикл с частотой 20 Гц:

два изображения с камер + состояние суставов 6D + языковая инструкция
                    ↓
              ACT / SmolVLA / pi_0
                    ↓
          цель суставов 7D и команда захвата
                    ↓
                 MuJoCo
                    ↓
              следующее реальное наблюдение

По умолчанию оценка только показывает конфигурацию и не загружает большую модель автоматически. Убедитесь в POLICY_TYPE, MODEL_RUN_DIR, статистике данных и доступности экрана, затем задайте RUN_CLOSED_LOOP=True. Для небольшой панели можно взять 4 seed, но итоговый вывод требует не менее 20–30 held-out seed и отдельного отчёта для красной и синей кружки.

Выходной каталог содержит:

• results.jsonl: для каждого seed значения legacy_success, physical_success, подъём, xy и вертикальность;
• <policy>_seed<seed>.mp4: видео rollout для проверки реального поведения;
• сводку в терминале: число строгих успехов и общее число эпизодов.

Если legacy_success выше, чем physical_success, значит, есть толчки кружки, захват впустую, падение во время переноса или другие геометрические ложные срабатывания. В этом случае вернитесь к задачам 02–06 и продолжайте диагностику по observation, action, фазам контакта и видео.

Строгая проверка pi0 с входом strict-input

Notebook 12 фиксирует самое строгое из текущих сравнений, которое труднее всего понять неверно. В одной и той же фиксированной среде с seed 0 чистая модель raw pi0 даёт 0/12, а после добавления визуального и исторического GRU-блока она даёт 6/12. Этот составной блок читает только две камеры, язык, 6D-проприоцепцию робота и предыдущую команду 7D EEF/захват, которую стратегия выполнила сама. Он не читает координаты цели и тарелки, фазу GT, оракульное действие и текущие или будущие метки захвата GT. Однако это всё равно дополнительный обученный блок, поэтому его нельзя выдавать за raw pi0.

После исправления seed среды на действительно случайных средах 41–44 результат всего 1/4. Поэтому пересбор данных с несколькими положениями предметов важнее, чем вслепую добавлять шаги обучения pi0. Оценка в фиксированной среде отвечает на вопрос, стабильна ли выборка стратегии в одной и той же сцене, а оценка в случайной среде показывает, выполняется ли обобщение по положению. Эти два результата нужно обязательно сообщать в отдельных таблицах.

В материалах с четырьмя видами есть проблема входного протокола, которую легко пропустить. С полной инструкцией из обучения Pick up the blue mug and place it on the plate. вспомогательный блок активировался на 269 шагах и на 270 шаге строго успешно выполнил задачу. При замене на более короткую синонимичную инструкцию Place the blue mug on the plate. прототипный gate ни разу не перехватил управление, и попытка провалилась. Поэтому текущий вспомогательный блок на малом объёме данных ещё не доказывает обобщение на перефразированные инструкции. Для формальной оценки сначала зафиксируйте prompt, а следующий раунд добавьте аугментацию синонимичными инструкциями и проверку paraphrase gate.

Критерии завершения

После выполнения главы журнал эксперимента должен содержать как минимум:

Эти результаты можно сразу использовать для DAgger по ACT, взвешенной выборки SmolVLA и диагностики хвостовой фазы pi_0; базовую цепочку заново строить не нужно.