Как запустить тысячи симуляций роботов на GPU без тяжеловесного Omniverse

Project banner

Если вы когда-нибудь обучали роботов через Reinforcement Learning, то наверняка сталкивались с неприятным выбором. С одной стороны есть классический MuJoCo: быстрый, математически прозрачный, любимый исследователями по всему миру. Беда в том, что гонять его на процессоре пачками по несколько тысяч окружений раньше было слишком накладно. С другой стороны есть связка NVIDIA Isaac Sim и Isaac Lab. Она великолепно масштабируется на GPU и задает отличный стандарт архитектуры сред, но платить за это приходится сотнями гигабайт диска, тяжелым движком Omniverse и частыми конфликтами зависимостей.

Группа исследователей из Беркли (включая Питера Аббиля и Кевина Закку) выложила проект mjlab. Идея лежит на поверхности: взять удобную компонентную архитектуру Isaac Lab и перенести ее на быстрый движок MuJoCo Warp. В итоге получился легкий инструмент для обучения роботов на видеокартах без лишней инфраструктурной нагрузки.

Что находится под капотом

Главная деталь здесь — MuJoCo Warp, проект от инженеров Google DeepMind. Они переписали вычисления физики MuJoCo на базе фреймворка NVIDIA Warp. Физика теперь считается параллельно прямо на ядрах видеокарты, что дает возможность запускать тысячи копий одного и того же робота на одном чипе.

Создатели mjlab решили не изобретать собственный формат сред. Они взяли концепцию Manager-based API из Isaac Lab, разделив логику обучения на четыре независимых блока:

Менеджер наблюдений (Observation Manager), собирающий вектор состояний для нейросети
Менеджер наград (Reward Manager), вычисляющий функцию подкрепления
Менеджер команд (Command Manager), генерирующий целевые траектории или скорости
Менеджер событий (Event Manager), отвечающий за перезапуск упавших роботов и рандомизацию физики

Вы при этом не теряете доступ к нативным структурам данных MuJoCo. Здесь нет пяти слоев оберток, прячущих состояние симуляции. Если нужно достать матрицу контактов или приложить силу к конкретному сочленению напрямую, вы просто обращаетесь к стандартным полям движка.

Запуск в одну команду

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

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

uvx --from mjlab --refresh demo

Для тех, у кого нет локальной видеокарты с CUDA, авторы подготовили блокнот в Google Colab. Обучение требует NVIDIA GPU, однако запустить инференс и посмотреть поведение обученного агента можно и на процессоре macOS.

Если вы планируете менять код и собирать свои среды, проект клонируется стандартно:

git clone https://github.com/mujocolab/mjlab.git
cd mjlab
uv run demo

Для проверки качества кода разработчики предусмотрели короткие команды: make test-fast для быстрых тестов и make format для линтинга.

Практические сценарии

В репозитории сразу идут готовые примеры со сложными роботами вроде гуманоида Unitree G1.

Обучение удержанию скорости

Научить робота ходить с заданной скоростью по ровной поверхности можно одной командой, передав число параллельных миров:

uv run train Mjlab-Velocity-Flat-Unitree-G1 --env.scene.num-envs 4096

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

Когда мощностей одной карты не хватает, обучение масштабируется на несколько GPU флагом --gpu-ids:

uv run train Mjlab-Velocity-Flat-Unitree-G1 \
  --gpu-ids "[0, 1]" \
  --env.scene.num-envs 4096

Имитация движений по референсным данным

Вторая классическая задача — заставить робота повторять захваченные движения человека. mjlab поддерживает этот пайплайн из коробки:

uv run train Mjlab-Tracking-Flat-Unitree-G1 \
  --registry-name your-org/motions/motion-name \
  --env.scene.num-envs 4096

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

uv run play Mjlab-Tracking-Flat-Unitree-G1 --wandb-run-path your-org/mjlab/run-id

Проверка логики без нейросетей

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

uv run play Mjlab-Your-Task-Id --agent zero
uv run play Mjlab-Your-Task-Id --agent random

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

Ограничения и компромиссы

Проект вышел под открытой лицензией Apache 2.0, а часть утилит API заимствована из Isaac Lab под лицензией BSD-3.

Стоит трезво оценивать границы применимости. Движок MuJoCo Warp развивается быстро, но пока уступает Isaac Sim в сложных задачах компьютерного зрения, трассировке лучей и рендеринге многокамерных сетапов. Если вам нужна фотореалистичная картинка для сквозного обучения с камер (vision-to-action), Isaac Sim останется основным выбором.

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

Кому стоит попробовать

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

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