Классический MuJoCo обеспечивает быструю симуляцию роботов на базе CPU для разработки, тестирования и управления роботами, а также позволяет распараллеливать выборку по ядрам CPU. Но по мере роста объемов задач обучения возникает вопрос не о том, насколько быстро может работать один мир, а о том, сколько миров может работать одновременно. Ускорение на GPU позволяет обрабатывать эти миры большими пакетами, удерживая данные симуляции и обучения непосредственно на устройстве.
MuJoCo Warp (MJWarp), созданный на базе NVIDIA Warp, переносит совместимые модели MuJoCo в режим масштабирования на GPU. В этой статье мы перенесем манипулятор-подражатель SO-101 из привычного рабочего процесса MuJoCo в среду, содержащую до 2048 параллельных окружений MJWarp, а также рассмотрим технологии и этапы проверки, которые делают этот переход возможным.
Рисунок 1. Как MJWarp связывает Python с симуляцией на GPU. MuJoCo загружает и компилирует модель MJCF; MJWarp реализует физику в NVIDIA Warp, которая компилирует ядра CUDA для продвижения состояний симуляции на GPU NVIDIA.
Это вторая статья в нашей серии «Состояние симуляции для физического искусственного интеллекта». В первой статье описывался ландшафт симуляции роботов. Здесь мы подготавливаем и масштабируем среду симуляции; мы не обучаем политику. В последующих публикациях, посвященных Newton и Isaac Lab, будут рассмотрены следующие уровни интеграции.
Подводим итоги
Краткий путь к решению:
Начните с одного полезного ядра Warp
NVIDIA Warp — это фреймворк на Python для написания высокопроизводительных ядер с ускорением на GPU. Warp позволяет разработчикам создавать статически типизированные ядра на Python и компилирует их для выполнения на CPU или CUDA. Первый запуск создает и кэширует нативный модуль; последующие запуски используют его повторно. Язык ядер представляет собой ориентированное на производительность подмножество Python, в то время как обычный Python отвечает за конфигурацию, выделение памяти и оркестрацию запускa.
Это небольшое ориентированное на робототехнику ядро обновляет положения точек под действием гравитации. Один логический поток обрабатывает одну точку, поэтому один и тот же код масштабируется от двух точек до миллионов без внедрения терминологии GPU в поток управления.
Три ценностных предложения Warp:
Три свойства делают этот инструмент полезным в робототехнике:
- Явное параллельное выполнение. wp.tid() определяет точку, контакт, тело или мир, принадлежащие текущему логическому потоку.
- Явные массивы устройства. Массив располагается на выбранном устройстве. Вызов .numpy() для массива CUDA синхронизирует и копирует его в память CPU; это не путь без копирования (zero-copy). Для конвейера PyTorch или JAX, работающего на устройстве, используйте вместо этого адаптеры фреймворка Warp или совместное использование, совместимое с DLPack.
- Композитные запускы ядер. Программа может запускать последовательность специализированных ядер и захватывать поддерживаемую работу CUDA в граф, чтобы снизить накладные расходы на повторяющиеся диспетчеризации. Захват графа воспроизводит запускы с существующими буферами; он не объединяет произвольные ядра.
Дифференцируемость и детерминизм.
Стоит знать о еще двух возможностях Warp, хотя ни одна из них не используется в рабочем процессе с SO-101 в этой статье. Ядра Warp дифференцируемы: wp.Tape записывает запускы ядер в прямом направлении, сделанные в своем контексте, и воспроизводит их сопряжения в обратном порядке при вызове backward(). Именно поэтому команды создают дифференцируемую геометрию, вычислительную гидродинамику (CFD) и пользовательскую физику в Warp, включая рабочие процессы CAE для симуляции и оптимизации дизайна. Warp также поддерживает детерминированное выполнение, появившееся в Warp 1.15: атомарные операции GPU по умолчанию зависят от планировщика, поэтому повторяющиеся запускы одного и того же ядра могут немного отличаться, а опциональные детерминированные режимы жертвуют частью производительности ради воспроизводимого порядка в симуляции, валидации и регрессионных тестах. Это возможности Warp, а не гарантии дифференцируемости или детерминизма для всего развертывания MJWarp. Подробности см. в документации Warp по дифференцируемости и детерминированному выполнению.
Попробуйте Warp: pip install warp-lang (≥ 1.15 для детерминизма на GPU), затем python -m warp.examples.browse или учебные ноутбуки.
Что такое MuJoCo Warp (MJWarp)?
Симулятор робота многократно вычисляет то, что происходит дальше: учитывая текущие положения суставов, скорости, управляющие воздействия и контакты, он продвигает сцену на один малый шаг по времени. В этой статье под миром понимается одна независимая копия этой сцены и ее состояния. Один мир может содержать манипулятор SO-101, тянущийся к кубу; другой может содержать тот же манипулятор, начинающий движение из немного другой позы.
MuJoCo и MJWarp могут запускать одну и ту же совместимую задачу и робота, но они организуют работу по-разному. MuJoCo естественным образом подходит для разработки и проверки одного или нескольких миров на CPU. MJWarp — это реализация физического конвейера MuJoCo на базе NVIDIA Warp, которая размещает модель и пакет независимых состояний на графических процессорах NVIDIA; один вызов mjw.step продвигает весь пакет вперед.
Ценность MJWarp заключается не обязательно в более быстром шаге для одного мира. Она заключается в возможности продвигать сотни или тысячи миров одновременно, предоставляя GPU достаточно параллельной работы для повышения совокупной пропускной способности — общего количества шагов мира, выполненных за секунду. Это благоприятствует обучению с подкреплением и крупномасштабной выборке, где сбор опыта важнее минимизации задержки одного окружения.
В этом блоге рассматривается следующее:
- проверка одного мира MuJoCo,
- перенос его в MJWarp, формирование пакета,
- проверка и корректное измерение.
Тонкая настройка решателя, представление якобиана и специализированные темы мульти-GPU или детерминизма не требуются для этой миграции и могут быть рассмотрены отдельно.
Тогда разница сформулирована точно:
- Задержка — это время по настенным часам для одного шага симуляции.
- Совокупная пропускная способность — это общее количество шагов мира, выполненных за измеренную секунду реального времени.
Базовое использование: структуры, размеры пакетов и минимальный шаг
- Переход на уровне основного API невелик:
Используйте mjw.make_data(), когда требуется состояние по умолчанию / новое состояние. Используйте mjw.put_data(), когда точное инициализированное состояние MuJoCo должно пересечь границу миграции.
Выделение пакетных ресурсов требует определения следующих параметров (см. раздел «Размеры пакетов»):
Тонкая настройка производительности
1. Захват графа CUDA: mjw.step — это множество запусков ядер; захватите один раз, воспроизводите часто:
2. Задавайте размеры nconmax / naconmax / njmax точно: память и объем работы масштабируются вместе с ними. Выполняйте настройку с помощью mjwarp-testspeed: --measure_alloc и следите за переполнениями в mjwarp-viewer.
Дополнительные соображения по настройке. После определения размеров буферов контактов и ограничений протестируйте пределы итераций решателя без изменения поведения задачи. Настройки мешей и CCD могут увеличить использование памяти; nccdmax / naccdmax могут уменьшить выделение памяти под буфер CCD, если это позволяют измеренные количества контактов. Компактный решатель MJWarp использует ньютоновский решатель ограничений MuJoCo и переход в спящий режим, а не отдельный фреймворк физического движка Newton. Конфигурация компактного решателя и мульти-GPU выходит за рамки этого руководства; обратитесь к документации по настройке производительности MJWarp.
Для обучения политик на физике MJWarp:
- Isaac Lab через Newton
- mjlab (manager API непосредственно поверх MJWarp + PyTorch)
- MuJoCo Playground через MJX (impl='warp')
Установка / тестирование: pip install mujoco-warp · mjwarp-viewer path/to/scene.xml · Учебник Colab
Рабочий процесс миграции сцены MuJoCo в MjWarp
- Создание базового уровня MuJoCo на CPU
Создание базового уровня MuJoCo на CPU
Сцена. Здесь пока нет ничего специфичного для MJWarp: рука SO-101, стол и два кубика для установки друг на друга, описанные в обычном формате MJCF.
Рисунок 2. Сцена захвата и размещения SO-101, отрендеренная из CPU-симуляции MuJoCo. Задача состоит в том, чтобы захватить красный кубик размером 44 мм и установить его на синий кубик; те же робот и сцена используются для валидации MJWarp.
Для коробочки MJCF значения size представляют собой полуразмеры: size=”0.022 …” определяет кубик с ребрами 44 мм. Задача использует этот размер для определения порогов успешности. Основание руки находится в начале координат, зона досягаемости простирается вдоль +X, а кубики расположены вдоль Y.
В сопутствующем репозитории этот файл генерируется, а не пишется вручную: resolve_pick_place_scene() копирует руку из Menagerie в .generated/, заполняет координаты стола и кубиков на основе профиля робота и записывает scene_pick_place.xml. В пошаговом руководстве используется профиль SO-101; опциональный вариант reBot описан ниже.
Загрузка. Компиляция и выполнение шагов происходят как в обычном MuJoCo:
Держите эту структуру в уме: вычислять управляющие воздействия один кадр за раз, выполнять шаги физики sim_substeps раз. В Шаге 2 изменяется только внутренний цикл, что и делает миграцию простой для проверки.
Согласуйте частоту симуляции и управления. При 50 управляющих кадрах в секунду и 10 физических подшагах на кадр используйте физический временной шаг 0,002 секунды. Установите его перед запуском на CPU и перед выгрузкой модели с помощью mjw.put_model, чтобы оба бэкенда продвигались на одинаковое симулированное время:
Без этой строки каждое последующее измерение наследует несоответствие: сравнения на идентичность (паритет), показатели пропускной способности, указываемые как «симулированные секунды», и любая обученная политика, чья частота действий больше не совпадает с развертыванием.
Проверьте, успешно ли сложены кубики. Для 44-мм кубиков успех определяется двумя измеряемыми условиями: горизонтальной погрешностью центра xy_err ≤ 0,015 м (измеряется между центрами кубиков) и вертикальным зазором 0,035 м ≤ dz ≤ 0,055 м между центрами кубиков (одна грань кубика с допуском на усадку). Оценивайте оба условия после того, как кубики улягутся; сам по себе успешный выход из процесса не доказывает успешность выполнения задачи.
Запустите задачу на CPU из сопутствующего репозитория. Блокирующий фактор для публикации: подтвердите доступный URL репозиторий, а также зафиксированные версии зависимостей и ассетов перед публикацией этих инструкций; заменитель URL репозитория ниже не является исполняемым.
Запуск завершается выводом двух приведенных выше чисел (stack check: xy_err=… dz=…), что является утверждением (ассертом), с которым сравнивается остальная часть статьи. so101_pick_place.py рядом с ним — это та же программа, но шаги физики оставлены в качестве упражнений.
Рука взята напрямую из MuJoCo Menagerie и зафиксирована на проверенном коммите, так как ассеты Menagerie меняются, поэтому рассматривайте сцену как шаблон. Опциональный вариант reBot. Сопутствующий код также предоставляет --robot rebot с отдельным профилем для компоновки сцены, захвата и пределов вместимости (nconmax=256, njmax=500). В этом руководстве используется SO-101. Проверьте ассет reBot и задачу отдельно, прежде чем сообщать их результаты.
- Валидация одномирового паритета MJWarp
Валидация одномирового паритета MJWarp
Сначала запустите один мир на GPU, оставив хост в контуре, чтобы вы могли наблюдать за той же задачей в том же визуализаторе и сравнить те же два числа. Загрузите модель, выделите пакетное состояние, инициализируйте его из инициализированного состояния хоста и выполните один прямой проход перед выполнением шага:
Каждый массив устройства имеет ведущее измерение мира, поэтому к состоянию хоста обращаются как mjd.qpos[None, :], размерность (1, nq) вместо (nq,). Масштабирование до тысяч миров позже меняет только это ведущее измерение, а не вызовы. mjw.put_model() также служит проверкой совместимости: он вызывает ошибку, если модель использует неподдерживаемые функции, вместо того чтобы молча отбрасывать их.
Явная инициализация трех полей является прозрачным вариантом и дает понять, что именно передается на устройство; вместо этого mjw.put_data(mjm, mjd, nworld=…) переносит всю инициализированную структуру за один вызов.
Цикл кадров тогда представляет собой цикл Шага 1 с перенаправлением его внутреннего шага на GPU и зеркальным отображением обратно:
Вызовы .numpy() синхронизируют и копируют данные на хост на каждом подшаге, поэтому это путь валидации задачи, а не бенчмарк пропускной способности. Он сохраняет обратную кинематику, визуализацию и проверки задач на хосте. После копирования qpos и qvel вызовите mujoco.mj_forward(mjm, mjd) для обновления производных величин хоста, таких как mjd.xpos, перед их использованием для управления, визуализации или проверки стека. Чтение этих полей после цикла не обновляет их автоматически. Шаг 4 удаляет эти копии на хосте для каждого шага из пути измерения пропускной способности.
- Определение емкости контактов и ограничений
Определение емкости контактов и ограничений
MJWarp выделяет буферы контактов и ограничений перед выполнением шагов. Превышение этой емкости делает соответствующий прогон недействительным для верификации или бенчмаркинга, даже если выполнение продолжается с предупреждением переполнения, а не с исключением. Увеличьте соответствующий предел и перезапустите задачу. Большие буферы используют больше памяти GPU, поэтому проверьте емкость для всей задачи перед окончательным сужением выделения памяти.
Установите лимиты контактов и ограничений для симулируемых робота и задачи. Профиль SO-101 использует nconmax=128 и njmax=300 в качестве начальной емкости. Убедитесь, что эти лимиты достаточны во время наиболее насыщенной контактами части задачи:
Рассчитывайте их под наиболее насыщенный контактами момент задачи: для захвата и размещения — момент, когда оба губных захвата и стол касаются кубика, а не когда рука зависает в свободном пространстве. Переполнение сообщается, а не вызывает ошибку: при значении Option.warn_overflow по умолчанию MJWarp выводит в терминал вашей программы или визуализатора бюджет, который нужно увеличить («narrowphase overflow - please increase nconmax to …»), и помечает затронутые миры в Data.overflow для чтения после шага. Только mjw.put_data вызывает ошибку напрямую, так как она может сравнить бюджеты с уже имеющимся у нее состоянием MuJoCo. mjwarp-testspeed --measure_alloc сообщает о контактах и ограничениях, которые сцена фактически потребила, и прерывает симуляцию с указанием идентификаторов проблемных миров, как только в каком-либо мире происходит переполнение. Относитесь к таким отчетам как к сбоям: увеличьте лимит и перезапустите симуляцию, прежде чем доверять траектории или бенчмарку, а затем снова уменьшайте лимит при каждом изменении модели, геометрии столкновений или задачи.
- Масштабирование до 2048 миров
Масштабирование до 2048 миров
Как только одномировый паритет пройдет успешно, перераспределите память под целевой размер и реплицируйте инициализированное состояние по всему батчу. По сравнению с Шагом 2 меняются две вещи: nworld и тот факт, что через шину PCIe на каждый шаг ничего не передается.
np.tile задает для каждого мира одинаковое начальное состояние, что является правильной отправной точкой для измерения пропускной способности; рандомизация для каждого мира вместо этого записывала бы разные строки d.qpos на устройстве.
CUDA Graphs повторно используют модель и буферы данных, захваченные здесь. Обновляйте d.ctrl на месте между перезапусками и захватывайте новый граф после замены буферов, изменения nworld или пересборки модели. Захват графа требует наличия CUDA.
Рисунок 3. Масштабирование задачи SO-101 с одного процессорного мира до 2048 независимых состояний GPU с использованием той же совместимой модели. Один шаг MJWarp продвигает весь батч. Эта концептуальная иллюстрация подчеркивает совокупную пропускную способность, измеряемую как шаги мира в секунду реального времени.
- Проверьте, затем измерьте
Проверьте, затем измерьте
Запуски GPU выполняются асинхронно, поэтому наивный таймер измеряет скорость постановки задач в очередь Python, а не скорость их выполнения на GPU. Сначала выполните прогрев — при первых запусках происходят компиляция ядра и выделение памяти — а затем синхронизируйте выполнение непосредственно перед и после измеряемой области:
Указывайте как совокупное количество шагов мира в секунду, так и миллисекунды на батчевый шаг вместе с размером батча. Используйте измеренную кривую, чтобы определить, где добавление миров увеличивает пропускную способность, а где ограничения памяти или вычислений снижают этот эффект. Результаты зависят от сцены, настроек симуляции и аппаратного обеспечения; сравнение задержки для одного мира не определяет пропускную способность батча.
Чтобы увидеть эту кривую на вашем собственном оборудовании, скрипт scaling_study.py перебирает размер батча и выводит мс/шаг наряду с пропускной способностью и ускорением:
Начало работы
Warp (слой ядер) pip install warp-lang → python -m warp.examples.browse → docs · GitHub
MJWarp (GPU MuJoCo) pip install mujoco-warp → mjwarp-viewer benchmarks/humanoid/humanoid.xml → docs · GitHub · Colab tutorial
Контекст SO-101 Курс SO-101 sim-to-real · Пути обучения физическому ИИ
Обучение поверх MJWarp mjlab · MuJoCo Playground · Isaac Lab + Newton (будущие публикации)
Что дальше
В этой статье рассматривались основы Warp → MJWarp: ядра GPU, батчевый шаг и сцена SO-101 с использованием mjw.step.
Затем мы перенесем ту же среду MJCF в Newton, используя MuJoCo Warp в качестве решателя твердых тел (newton.solvers.SolverMuJoCo). Newton будет управлять моделью, состоянием, управлениями и контактами, в то время как MJWarp работает «под капотом».
Вы также увидите, что добавляет Newton: многоформатные ассеты, заменяемые решатели, помощники сенсоров/ИК и путь Isaac Lab.
Руководство по миграции продолжается с той же задачей SO-101 и ее опциональным профилем reBot, объясняя изменения, требуемые Newton, и отдельную интеграцию с Isaac Lab.
Если вы создаете что-то с помощью Warp или MJWarp, откройте issue в связанных репозиториях или найдите нас в Discord NVIDIA Omniverse.
Ссылки
- Блог 1: *Состояние симуляции для физического ИИ: обзор* — Публикация 1 из этой серии.
- NVIDIA Warp — GitHub · Документация · Релиз v1.15.0 (детерминизм GPU) · Руководство по детерминированному выполнению
- MuJoCo Warp — GitHub · Официальная документация MJWarp
- Создание ускоренного дифференцируемого кода вычислительной физики для ИИ с помощью NVIDIA Warp
- Представляем плитное программирование (Tile-Based Programming) в Warp 1.5.0
- mjlab · arXiv:2601.22074
- MuJoCo Playground
- Курс NVIDIA SO-101 sim-to-real
- Следующая публикация Newton: MJWarp в качестве SolverMuJoCo и перенос этой среды