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

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

MuJoCo Warp (MJWarp), созданный на базе NVIDIA Warp, переносит совместимые модели MuJoCo в этот масштабируемый режим работы на GPU. В этой статье мы перенесём манипулятор SO-101 follower из привычного рабочего процесса MuJoCo в MJWarp, где можно запускать до 2048 параллельных окружений, и рассмотрим технологию и этапы проверки, делающие такой переход возможным.

Рисунок 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:

import numpy as np
import warp as wp

@wp.kernel
def integrate(
   positions: wp.array[wp.vec3],
   velocities: wp.array[wp.vec3],
   dt: float,
):
   i = wp.tid()
   velocities[i] += wp.vec3(0.0, 0.0, -9.81) * dt
   positions[i] += velocities[i] * dt

wp.init()
device = "cuda:0" if wp.is_cuda_available() else "cpu"
start = np.array([[0.0, 0.0, 0.5], [0.2, 0.0, 0.5]], dtype=np.float32)
positions = wp.array(start, dtype=wp.vec3, device=device)
velocities = wp.zeros_like(positions)

wp.launch(
   integrate,
   dim=len(start),
   inputs=[positions, velocities, 0.01],
   device=device,
)
wp.synchronize_device(device)
print(positions.numpy())

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

• Явная параллельная работа. wp.tid() определяет точку, контакт, тело или мир, принадлежащий текущему логическому потоку.
• Явные массивы устройств. Массив находится на выбранном устройстве. Вызов .numpy() для массива CUDA синхронизирует его и копирует в память CPU; это не путь с нулевым копированием. Для конвейера PyTorch или JAX, работающего на устройстве, используйте адаптеры Warp или совместное использование через DLPack.
• Комбинируемые запуски ядер. Программа может запускать последовательность специализированных ядер и записывать поддерживаемые операции CUDA в граф, чтобы уменьшить повторные накладные расходы на диспетчеризацию. Захват графа повторяет запуски с существующими буферами; произвольные ядра он не объединяет.

Дифференцируемость и детерминизм.

Стоит знать ещё о двух возможностях Warp, хотя в рабочем процессе SO-101 из этой статьи ни одна из них не используется. Ядра Warp являются дифференцируемыми: wp.Tape записывает запуски прямых ядер, выполненные внутри его контекста, и при вызове backward() воспроизводит их сопряжённые операции в обратном направлении. Именно поэтому команды создают в Warp дифференцируемую геометрию, CFD и пользовательскую физику, включая рабочие процессы 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 достаточно параллельной работы для повышения совокупной пропускной способности — общего числа завершённых шагов миров в секунду. Это особенно полезно для обучения с подкреплением и крупномасштабной выборки, где сбор опыта важнее минимизации задержки одного окружения.

В этой статье рассматриваются следующие этапы:

1. проверить один мир MuJoCo,
2. перенести его в MJWarp и сформировать пакет,
3. проверить его и корректно измерить производительность.

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

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

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

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

• Переход между основными API невелик:

Используйте mjw.make_data(), когда требуется состояние по умолчанию или новое состояние. Используйте mjw.put_data(), когда точное инициализированное состояние MuJoCo должно пересечь границу миграции.

Для выделения пакетных ресурсов необходимо определить следующие параметры (см. Размеры пакетов):

Настройка производительности

1. Захват графа CUDA: mjw.step состоит из множества запусков ядер; захватите их один раз и часто воспроизводите:

with wp.ScopedCapture() as capture:
     mjw.step(m, d)
wp.capture_launch(capture.graph)

2. Точно задавайте размеры nconmax / naconmax / njmax: объём памяти и вычислительная работа масштабируются вместе с ними. Настраивайте параметры с помощью mjwarp-testspeed: --measure_alloc и следите за переполнениями в mjwarp-viewer.

Дополнительные соображения по настройке. После определения размеров буферов контактов и ограничений проверьте ограничения на число итераций решателя, не меняя поведение задачи. Сетки и настройки CCD могут увеличить использование памяти; nccdmax / naccdmax могут уменьшить выделение буфера CCD, если измеренные значения числа контактов это позволяют. Компактный решатель MJWarp использует решатель ограничений Newton из MuJoCo и механизм засыпания, а не отдельный физический движок Newton. Настройка компактного решателя и нескольких GPU выходит за рамки этого руководства; обратитесь к документации 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

1. Создайте базовый вариант MuJoCo на CPU

Сцена. Пока здесь нет ничего специфичного для MJWarp: манипулятор SO-101, стол и два кубика для укладки, описанные в обычном MJCF.

Рисунок 2. Сцена SO-101 для захвата и перемещения, отрендеренная из симуляции MuJoCo на CPU. Задача состоит в том, чтобы захватить красный куб размером 44 мм и поставить его на синий куб; тот же робот и сцена используются для проверки MJWarp.

<mujoco model="so101_pick_place">
  <include file="so101.xml"/>

  <worldbody>
    <light pos="0.3 0 1.5" dir="0 0 -1" directional="true"/>
    <geom name="floor" type="plane" size="0 0 0.05"/>

    <geom name="table" type="box" pos="0.35 -0.04 0.012"
          size="0.16 0.26 0.012" rgba="0.32 0.32 0.32 1"
          friction="1 0.005 0.0005" condim="3"/>

    <body name="red_cube" pos="0.33 -0.13 0.046">
      <freejoint name="red_cube_joint"/>
      <geom type="box" size="0.022 0.022 0.022" mass="0.08"
            rgba="0.85 0.05 0.04 1" friction="1.2 0.005 0.0005" condim="3"/>
    </body>

    <body name="blue_cube" pos="0.33 0.06 0.046">
      <freejoint name="blue_cube_joint"/>
      <geom type="box" size="0.022 0.022 0.022" mass="0.08"
            rgba="0.05 0.20 0.90 1" friction="1.2 0.005 0.0005" condim="3"/>
    </body>
  </worldbody>
</mujoco>

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

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

Загрузка. Компиляция и выполнение шагов в MuJoCo являются стандартными:

import mujoco

mjm = mujoco.MjModel.from_xml_path("scene_pick_place.xml")
mjd = mujoco.MjData(mjm)

fps = 50 
sim_substeps = 10 
frame_dt = 1.0 / fps
mjm.opt.timestep = frame_dt / sim_substeps

controller = PickPlaceController(spec=spec) 

for _ in range(600): 
    ctrl = controller.step(mjm, mjd, frame_dt)
    for _ in range(sim_substeps):
        mjd.ctrl[: mjm.nu] = ctrl
        mujoco.mj_step(mjm, mjd)

Запомните эту структуру: вычислять управляющие воздействия один раз за кадр, а физику продвигать sim_substeps раз. На этапе 2 меняется только внутренний цикл, что упрощает проверку миграции.

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

mjm.opt.timestep = frame_dt / sim_substeps 

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

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

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

git clone https://github.com/NVIDIA/accelerated-computing-hub.git blogs
cd blogs/tutorials/sim2real-blogs/notebooks/mujoco
uv venv --python 3.12 && source .venv/bin/activate
uv pip install -r requirements.txt

cd /tutorials/sim2real-blogs/notebooks/mujoco
python solutions/so101_pick_place_solution.py --headless-steps 600 --debug

В конце запуска печатаются два указанных выше числа (проверка укладки: xy_err=… dz=…), что является утверждением, с которым сравниваются остальные части статьи. Файл so101_pick_place.py рядом с ним содержит ту же программу, но оставляет шаги физики в качестве упражнений.

Манипулятор взят непосредственно из MuJoCo Menagerie и зафиксирован на заведомо рабочем коммите, поскольку ресурсы Menagerie меняются; рассматривайте сцену как шаблон. Необязательный вариант reBot. Сопутствующий код также предоставляет параметр --robot rebot с отдельным профилем для компоновки сцены, захвата и ограничений ёмкости (nconmax=256, njmax=500). В этом руководстве используется SO-101. Перед публикацией результатов отдельно проверьте ресурс reBot и задачу.

1. Проверьте эквивалентность MJWarp для одного мира

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

wp.init()
import mujoco_warp as mjw

device = wp.get_device()

m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=1, nconmax=spec.nconmax, njmax=spec.njmax)

wp.copy(d.qpos, wp.array(mjd.qpos[None, :], dtype=wp.float32, device=device))
wp.copy(d.qvel, wp.array(mjd.qvel[None, :], dtype=wp.float32, device=device))
wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device))
mjw.forward(m, d)

Каждый массив устройства имеет ведущую размерность world, поэтому состояние хоста индексируется как mjd.qpos[None, :], с формой (1, nq), а не (nq,). При масштабировании до тысяч миров позднее изменится только эта ведущая размерность, а не вызовы. mjw.put_model() также служит проверкой совместимости: он выдаёт ошибку, если модель использует неподдерживаемые возможности, вместо того чтобы молча отбрасывать их.

Явная инициализация трёх полей — прозрачный вариант, показывающий, что именно пересекает границу устройства; mjw.put_data(mjm, mjd, nworld=…) вместо этого переносит всю инициализированную структуру одним вызовом.

Затем цикл кадров представляет собой цикл этапа 1, где внутренний шаг перенаправлен на GPU и скопирован обратно:

def simulate_frame() -> None:
    ctrl = controller.step(mjm, mjd, frame_dt)
    for _ in range(sim_substeps):
        mjd.ctrl[: mjm.nu] = ctrl
        wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device))
        mjw.step(m, d)
        mjd.qpos[:] = d.qpos.numpy()[0]
        mjd.qvel[:] = d.qvel.numpy()[0]
    mujoco.mj_forward(mjm, mjd)

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

1. Определите ёмкость контактов и ограничений

MJWarp выделяет буферы контактов и ограничений до начала шагов. Превышение этих ёмкостей делает соответствующий прогон непригодным для проверки или бенчмаркинга, даже если выполнение продолжается с предупреждением о переполнении, а не с исключением. Увеличьте соответствующий предел и повторно запустите задачу. Большие буферы используют больше памяти GPU, поэтому проверяйте ёмкость на всём протяжении задачи, прежде чем уменьшать выделение.

Задайте пределы контактов и ограничений для моделируемых робота и задачи. Профиль SO-101 использует nconmax=128 и njmax=300 в качестве начальных значений ёмкости. Убедитесь, что этих пределов достаточно в наиболее насыщенной контактами части задачи:

d = mjw.make_data(mjm, nworld=nworld, nconmax=spec.nconmax, njmax=spec.njmax)

Определяйте их по наиболее насыщенному контактами моменту задачи: для захвата и перемещения это момент, когда обе челюсти и стол касаются кубика, а не момент, когда манипулятор зависает в свободном пространстве. О переполнении сообщается, а не выдаётся исключение: при значении Option.warn_overflow по умолчанию MJWarp печатает в терминал, где запущен скрипт или визуализатор, бюджет, который нужно увеличить («narrowphase overflow - please increase nconmax to …»), и помечает затронутые миры в Data.overflow, чтобы вы могли прочитать это после шага. Только mjw.put_data сразу выдаёт ошибку, поскольку может сравнить бюджеты с уже имеющимся у него состоянием MuJoCo. mjwarp-testspeed --measure_alloc сообщает о фактически использованных сценой контактах и ограничениях и прерывает прогон с идентификаторами проблемных миров, как только какой-либо мир переполняется. Считайте такие сообщения ошибками: увеличьте предел и повторите запуск, прежде чем доверять траектории или бенчмарку, а затем снова уменьшите его, если изменились модель, геометрия столкновений или задача.

1. Масштабируйте до 2048 миров

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

nworld = 2_048
d = mjw.make_data(mjm, nworld=nworld, nconmax=spec.nconmax, njmax=spec.njmax)

wp.copy(d.qpos, wp.array(np.tile(mjd.qpos, (nworld, 1)), dtype=wp.float32, device=device))
wp.copy(d.qvel, wp.array(np.tile(mjd.qvel, (nworld, 1)), dtype=wp.float32, device=device))
wp.copy(d.ctrl, wp.array(np.tile(mjd.ctrl, (nworld, 1)), dtype=wp.float32, device=device))
mjw.forward(m, d)

with wp.ScopedCapture() as capture:
    mjw.step(m, d)
step_graph = capture.graph

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

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

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

1. Сначала проверьте, затем измеряйте

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

import time
for _ in range(10): 
    wp.capture_launch(step_graph)
wp.synchronize()

t0 = time.perf_counter()
for _ in range(200):
    wp.capture_launch(step_graph)
wp.synchronize() 
elapsed = time.perf_counter() - t0
total = 200 * nworld
print(f"{total / elapsed:,.0f} world-steps/second")

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

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

cd /tutorials/sim2real-blogs/notebooks/mujoco/part2
python solutions/so101_mjwarp_solution.py --headless-steps 600     # parity, needs CUDA
python scaling_study.py --worlds 1 64 1024 2048 8192 --steps 100

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

Warp (уровень ядер)
pip install warp-lang → python -m warp.examples.browse → документация · GitHub

MJWarp (MuJoCo на GPU)
pip install mujoco-warp → mjwarp-viewer benchmarks/humanoid/humanoid.xml → документация · GitHub · учебник в Colab

Контекст SO-101
курс sim-to-real для SO-101 · учебные курсы по Physical AI

Обучение поверх MJWarp
mjlab · 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 (детерминизм GPU) · руководство по детерминированному выполнению
• MuJoCo Warp — GitHub · официальная документация MJWarp
• Создание ускоренного дифференцируемого кода вычислительной физики для ИИ с помощью NVIDIA Warp
• Представляем тайловое программирование в Warp 1.5.0
• mjlab · arXiv:2601.22074
• MuJoCo Playground
• курс NVIDIA sim-to-real для SO-101
• Newton следующая публикация: MJWarp как SolverMuJoCo и перенос этого окружения

Другие материалы этого автора