Как использовать NVIDIA Warp и MjWarp для ускорения рабочих процессов симуляции и обучения робототехники

Классический MuJoCo обеспечивает быстрое моделирование роботов на базе 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 отвечает за конфигурацию, выделение памяти и оркестрацию запуска.

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

Три ценностных предложения Warp:

Три свойства делают его полезным в робототехнике:

• Явное параллельное выполнение. wp.tid() идентифицирует точку, контакт, тело или мир, принадлежащий текущему логическому потоку.
• Явные массивы устройств. Массив находится на выбранном устройстве. Вызов .numpy() для массива CUDA синхронизирует и копирует его в память CPU; это не путь с нулевым копированием. Для конвейеров 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, которая размещает модель и пакет независимых состояний на GPU NVIDIA; один вызов mjw.step продвигает весь пакет.

Ценность MJWarp не обязательно заключается в более быстром шаге для одного мира. Это способность продвигать сотни или тысячи миров вместе, давая GPU достаточно параллельной работы для улучшения совокупной пропускной способности — общего количества шагов мира, выполненных в секунду. Это благоприятствует обучению с подкреплением и крупномасштабной выборке, где сбор опыта важнее, чем минимизация задержки одной среды.

Этот блог охватывает следующее:

• валидация одного мира MuJoCo,
• перенос его в MJWarp, формирование пакета,
• верификация и правильное измерение.

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

Тогда различие становится точным:

• Задержка — это время по настенным часам для одного шага моделирования.
• Совокупная пропускная способность — это общее количество шагов мира, выполненных за измеренную секунду по настенным часам.

Базовое использование: структуры, размеры пакетов и минимальный шаг

• Переход на основной 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. Конфигурация компактного решателя и многопроцессорная обработка выходят за рамки этого руководства; обратитесь к документации по настройке производительности MJWarp.

Для обучения политик на физике MJWarp:

• Isaac Lab через Newton
• mjlab (менеджер 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, отрендеренная с помощью симуляции MuJoCo CPU. Задача состоит в том, чтобы захватить красный куб размером 44 мм и поместить его на синий куб; те же робот и сцена используются для валидации MJWarp.

Для блока MJCF значения размера являются полуразмерами: size=”0.022 …” определяет куб с ребрами 44 мм. Задача использует этот размер для пороговых значений успеха. Основание манипулятора находится в начале координат, его досягаемость направлена вдоль +X, а кубы расположены вдоль Y.

В сопутствующем репозитории этот файл генерируется, а не пишется вручную: resolve_pick_place_scene() копирует манипулятор из Menagerie в .generated/, заполняет координаты стола и кубов из профиля робота и записывает scene_pick_place.xml. В руководстве используется профиль SO-101; опциональный вариант reBot описан ниже.

Загрузка. Компиляция и выполнение шагов — это обычный MuJoCo:

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

Согласуйте частоту симуляции и управления. При 50 кадрах управления в секунду и 10 физических подшагах на кадр используйте физический временной шаг 0,002 секунды. Установите его перед развертыванием CPU и перед загрузкой модели с помощью mjw.put_model, чтобы оба бэкенда продвигались на одно и то же симулированное время:

Без этой строки каждое последующее измерение унаследует несоответствие: сравнения паритета, показатели пропускной способности, указанные как «симулированные секунды», и любая обученная политика, чья частота действий больше не соответствует развертыванию.

Проверьте, успешно ли сложены кубы. При использовании кубов 44 мм успех становится двумя измеримыми условиями: горизонтальная ошибка центра xy_err ≤ 0,015 м (измеряется между центрами кубов) и вертикальное разделение 0,035 м ≤ dz ≤ 0,055 м между центрами кубов (одно ребро куба с допуском на оседание). Оценивайте оба условия после того, как кубы стабилизировались; успешный выход из процесса сам по себе не подтверждает успех задачи.

Запустите задачу CPU из сопутствующего репозитория. Блокировщик публикации: подтвердите доступный URL репозитория и закрепленные версии зависимостей и активов перед публикацией этих инструкций; заполнитель репозитория ниже не является исполняемым URL.

Запуск заканчивается выводом двух чисел выше (проверка стека: 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=…) вместо этого переносит всю инициализированную структуру за один вызов.

Цикл кадров — это цикл Gate 1 с внутренним шагом, перенаправленным на GPU и зеркально отраженным обратно:

Вызовы .numpy() синхронизируют и копируют данные на хост на каждом подшаге, поэтому это путь валидации задачи, а не тест пропускной способности. Он сохраняет обратную кинематику, просмотр и проверки задач на хосте. После копирования qpos и qvel вызовите mujoco.mj_forward(mjm, mjd), чтобы обновить производные величины хоста, такие как mjd.xpos, перед их использованием для управления, просмотра или проверки стека. Чтение этих полей после цикла не обновляет их автоматически. Gate 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 миров

Как только паритет одного мира пройден, перераспределите память при целевом размере и реплицируйте инициализированное состояние по всему пакету. Две вещи меняются по сравнению с Gate 2: nworld и тот факт, что ничего не пересекает шину PCIe за шаг.

np.tile дает каждому миру одинаковое начальное состояние, что является правильной базовой линией для измерения пропускной способности; рандомизация для каждого мира вместо этого записывала бы разные строки d.qpos на устройстве.

Графы CUDA повторно используют модель и буферы данных, захваченные здесь. Обновляйте d.ctrl на месте между повторами и захватывайте новый граф после замены буферов, изменения nworld или пересборки модели. Захват графа требует CUDA.

Рисунок 3. Масштабирование задачи SO-101 с одного мира CPU до 2048 независимых состояний GPU с использованием одной и той же совместимой модели. Один шаг MJWarp продвигает весь пакет. Эта концептуальная иллюстрация подчеркивает совокупную пропускную способность, измеряемую как количество шагов мира в секунду реального времени.

• Сначала проверка, затем измерение

Сначала проверка, затем измерение

Запуски GPU являются асинхронными, поэтому простой таймер измеряет скорость постановки задач в очередь Python, а не скорость их выполнения на GPU. Сначала выполните прогрев — первые запуски требуют компиляции ядра и выделения памяти — затем синхронизируйтесь непосредственно перед и после измеряемого участка:

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

Чтобы увидеть эту кривую на своем оборудовании, используйте scaling_study.py, который перебирает размер пакета и выводит мс/шаг вместе с пропускной способностью и ускорением:

Начало работы

Warp (уровень ядра) pip install warp-lang → python -m warp.examples.browse → docs ·

MJWarp (GPU MuJoCo) pip install mujoco-warp → mjwarp-viewer benchmarks/humanoid/humanoid.xml → · GitHub · Учебник в Colab

Контекст SO-101 SO-101 sim-to-real course · Пути обучения физическому ИИ

Обучение на базе MJWarp · MuJoCo Playground · Isaac Lab + Newton (предстоящие публикации)

Что дальше

В этой публикации были рассмотрены основы Warp → MJWarp: ядра GPU, пакетное пошаговое выполнение и сценарий SO-101 с использованием mjw.step.

Далее мы перенесем ту же среду MJCF в Newton, используя MuJoCo Warp в качестве решателя твердых тел (newton.solvers.SolverMuJoCo). Newton будет управлять моделью, состоянием, элементами управления и контактами, в то время как MJWarp будет работать в фоновом режиме.

Вы также увидите, что добавляет Newton: поддержку различных форматов ресурсов, сменные решатели, вспомогательные средства для датчиков/IK и путь интеграции с Isaac Lab.

Руководство по миграции продолжается с той же задачей SO-101 и ее дополнительным профилем reBot, объясняя изменения, необходимые для Newton, и отдельную интеграцию с Isaac Lab.

Если вы создаете что-то с помощью Warp или MJWarp, откройте issue в связанных репозиториях или найдите нас в Discord NVIDIA Omniverse.

Ссылки

• Блог 1: *Состояние симуляции для физического ИИ: обзор* — первая публикация этой серии.
• NVIDIA Warp — GitHub · · v1.15.0 release (GPU determinism) · Руководство по детерминированному выполнению
• MuJoCo Warp — GitHub · Официальная документация MJWarp
• MuJoCo Playground
• следующая публикация: MJWarp как SolverMuJoCo и перенос этой среды