Что такое LIBERO и зачем его портировали на MuJoCo Warp
LIBERO: бенчмарк робоманипуляции из 130 задач, разбитых на пять наборов (spatial, object, goal, scene10, scene90). У каждой задачи есть человеческие демонстрации и языковая цель, поэтому бенчмарк стал стандартным тестбендом для имитационного обучения с языковыми условиями.
Обычно LIBERO работает на robosuite с CPU-версией MuJoCo: по одной среде за раз. Такой режим годится, чтобы проверить готовую политику, но не для быстрых экспериментов.
Почему CPU-версия тормозит обучение
Сбор данных упирается в арифметику. robosuite на CPU обрабатывает 38 шагов среды в секунду, и демонстрации для одного набора из 10 задач собирались 8-9 часов. Почти рабочий день ожидания, и только потом можно смотреть метрики и запускать следующую итерацию. При таком цикле перебор архитектур, гиперпараметров или наборов демонстраций становится дорогим: ошибка в конфиге стоит ещё один день.
Второе ограничение - последовательность процессов. Пока физика считается для одной сцены, масштабировать сбор можно только новыми копиями модели, и упирается это в CPU и память, а не в количество задач.
Что даёт MuJoCo Warp
MuJoCo Warp считает тысячи независимых миров одним батчем на GPU. Автор портировал на него весь LIBERO и получил на Radeon RX 9070 XT 17 137 шагов среды в секунду при 2 048 мирах. Генерация демонстраций для того же набора из 10 задач сжимается с 8-9 часов до минут.
Практическая ценность не в самой цифре, а в смене режима работы: вместо одного длинного запуска появляется цикл из сбора роллаутов, обучения, оценки и правки конфига. По словам автора порта, выигрыш именно в возможности итерировать, а не запускать сбор данных один раз.
Производительность на Radeon RX 9070 XT: цифры и сравнение с CPU
Замеры сделаны на одной конфигурации: Radeon RX 9070 XT, 16 ГБ, архитектура RDNA4, 2 048 миров. Цена карты около $700.
| Метрика | CPU (robosuite) | GPU (MuJoCo Warp, RX 9070 XT) |
|---|---|---|
| Шагов среды в секунду | 38 | 17 137 |
| Параллельных миров | 1 | 2 048 |
| BC-трансформер после 50 эпох | 50% | 42,5% |
| Демонстрации для набора из 10 задач | 8-9 часов | минуты |
По пропускной способности физики разница больше чем в 450 раз. Это агрегатная метрика: отдельный шаг симуляции в Warp не обязан быть быстрее шага на CPU, выигрыш набирается на батче.
Если сравниваете GPU-ускорение между разными решениями, заранее договоритесь, что именно считается эффективностью, иначе цифры плохо сопоставимы: об этом мы писали в разборе метрик эффективности GPU.
Почему BC-политика показывает 42,5% вместо 50%
BC-трансформер после 50 эпох обучения даёт 42,5% на Warp против 50% на CPU-версии robosuite. Разрыв 7,5 процентных пункта.
Причина не в одной физике. Behavioral cloning учится по изображениям, и политика чувствительна к визуальным условиям: если картинка из симулятора отличается от обучающих данных, ошибки накапливаются по ходу задачи. Часть разрыва пришлась именно на рендерер.
Разрыв не выглядит приговором: скорость сбора роллаутов позволяет быстрее проверять гипотезы и пытаться его закрыть.
Запуск Warp на AMD через HIP/ROCm: как это работает
Warp, фреймворк GPU-симуляции от NVIDIA, запустился на Radeon благодаря недавнему порту на AMD HIP/ROCm (в исходном посте упомянуты Tomas Thoresen и Strix Halo). Без этого порта ни MuJoCo Warp, ни LIBERO на AMD-карте не поехали бы.
Для AI-Manual тема знакомая: ROCm уже ускоряет инференс LLM в llama.cpp, и там набор компонентов похож. Разница в задаче: здесь на GPU считают не матрицы внимания, а физику контактов и столкновений. Правки ROCm-бэкенда в llama.cpp разбирали отдельно.
Компоненты стека: Warp, HIP, ROCm, JAX, PyTorch
- Warp компилирует HIP-ядра под gfx1201, целевую архитектуру RDNA4.
- JAX видит устройство как rocm:0.
- PyTorch видит cuda, работая поверх ROCm.
- Всё собирается из публичных исходников на ROCm.
Имя cuda в PyTorch часто сбивает с толку: код под CUDA запускается без правок, потому что ROCm-бэкенд сохраняет привычное название. На практике это значит, что часть пайплайна переносится на AMD без переписывания.
Возможные подводные камни при сборке
Порт на HIP/ROCm недавний, отсюда главный риск: стабильность и воспроизводимость могут отличаться от CUDA-пути. Автор не приводит пошаговой инструкции по сборке и не описывает конкретные ошибки компиляции, поэтому считать какой-то порядок действий гарантированным нельзя. Результаты получены на одной конфигурации (RX 9070 XT, 16 ГБ, RDNA4, 2 048 миров), и перенос на другую карту или версию драйвера стоит проверять отдельно.
Проблемы рендерера: почему BC-политика давала 0% и как это исправили
Первый прогон BC-политики на Warp дал 0%. Физика считалась, роллауты собирались, но политика не решала ни одной задачи. Дело оказалось в изображении.
Четыре правки рендерера подняли результат с 0% до 42,5%: вертикальный флип, константа теней 0.3 → 0.0, яркость ×1.15 и cube-map сэмплирование.
Вертикальный флип и константа теней
Вертикальный флип приводит изображение к правильной ориентации: без него политика видит сцену перевёрнутой относительно данных, на которых училась. Вторая правка меняет константу теней с 0.3 на 0.0, то есть убирает прежний вклад затенения в картинку. Автор перечисляет обе правки в общем списке четырёх исправлений, не разбивая отдельно их вклад в итоговый результат.
Яркость ×1.15 и cube-map сэмплирование
Яркость ×1.15 в одиночку дала 12,5 процентных пункта. Это самая крупная из четырёх правок, и по её эффекту видно, насколько сильно освещение влияло на предсказания: яркость в рендере Warp расходилась с той, на которой училась политика.
Cube-map сэмплирование потребовалось из-за текстуры дерева стола: она рендерилась плоско, тогда как на данных обучения выглядела объёмной.
Вывод для переноса BC-политик между симуляторами: расхождение в рендере может стоить дороже, чем кажется. При полностью рабочей физике результат оставался нулевым, пока картинка отличалась от обучающего распределения.
Как воспроизвести: сборка из исходников и готовые артефакты
Что уже готово в репозитории
Пример кода и чекпоинт BC-трансформера на 21,6 МБ выложены в репозитории автора, ссылка на него есть в исходном посте. Чекпоинт позволяет прогнать оценку без обучения с нуля и сверить свои цифры с заявленными 42,5%.
Пошаговая сборка на ROCm
Пошагового гайда в посте нет, и это стоит учитывать заранее: воспроизведение потребует самостоятельной настройки окружения. Из подтверждённого: сборка идёт из публичных исходников на ROCm, Warp компилирует HIP-ядра под gfx1201, JAX видит rocm:0, PyTorch видит cuda. Дальше придётся самому подобрать версию ROCm и драйвера, собрать Warp с поддержкой HIP и убедиться, что пример из репозитория отрабатывает на вашей карте.
Ограничения и что осталось незавершённым
Почему обучение пока офлайн
Warp собирает роллауты, а обучение политики идёт офлайн в PyTorch. Замкнутого цикла, где сбор данных и обновление весов происходят в одном процессе на GPU, пока нет. Для imitation learning это терпимо: демонстрации собирают отдельно, модель учат отдельно. Для обучения с подкреплением, где нужны миллионы шагов и постоянный обмен с политикой, отсутствие замкнутого цикла ограничивает выигрыш.
Перспективы и возможные улучшения
Автор планирует доделать сбор роллаутов внутри Warp. Если получится, обучение с подкреплением сможет идти целиком на GPU без обмена с PyTorch. Второе направление - рендерер: разрыв 42,5% против 50% частично связан с визуальными расхождениями, и его можно сокращать.
Стоит ли использовать AMD GPU для робоманипуляции: выводы
Кому подходит этот сетап
Сетап рассчитан на тех, кто готов настраивать окружение. Разработчикам и исследователям, которым нужна высокая пропускная способность симуляции и которые спокойно относятся к ROCm. Владельцам Radeon, которые хотят использовать карту для AI-задач, а не только для игр. Тем, кому важнее быстрый сбор роллаутов, чем максимальная точность на первой итерации.
Альтернативы на NVIDIA
На картах NVIDIA тот же стек работает через CUDA без портов: Warp, MuJoCo и robosuite рассчитаны на неё, поэтому меньше шансов упереться в сборку ядер. Плата за это - деньги: сопоставимая по производительности карта обойдётся дороже $700. Экосистема ROCm при этом развивается, AMD уже входит в партнёрские программы ML-инструментов, о чём мы писали в разборе партнёрства Hugging Face и AMD.
Короткий вывод: Radeon RX 9070 XT за $700 с 16 ГБ даёт рабочее GPU-ускорение LIBERO, но пока с ручной настройкой и разрывом в качестве BC-политики. Нужен минимальный порог входа - берите NVIDIA. Готовы собирать окружение и хотите сэкономить - AMD-путь уже проходим.