Як використовувати NVIDIA Warp і MjWarp для прискорення робочих процесів симуляції та навчання в робототехніці

Класичний MuJoCo забезпечує швидку симуляцію роботів на CPU для розроблення, тестування та керування роботами, а також може розпаралелювати вибірку між ядрами CPU. Але зі зростанням навчальних навантажень питання змінюється: від того, наскільки швидко може працювати один світ, до того, скільки світів можуть працювати одночасно. Прискорення на GPU дає змогу просувати ці світи великими пакетами, зберігаючи дані симуляції та навчання поруч із пристроєм.

MuJoCo Warp (MJWarp), побудований на основі NVIDIA Warp, переносить сумісні моделі MuJoCo у цей масштабований на GPU режим. У цій статті ми перенесемо маніпулятор-наступник SO-101 зі звичного робочого процесу MuJoCo до 2 048 паралельних середовищ 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:

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 використовує ньютонівський розв’язувач обмежень 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, відтворена з CPU-симуляції MuJoCo. Завдання полягає в тому, щоб захопити червоний куб зі стороною 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 із супровідного репозиторію. Блокер публікації: перед публікацією цих інструкцій підтвердьте доступну 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)

Кожен масив пристрою має початковий вимір світу, тому стан хоста індексується як 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 виділяє буфери контактів та обмежень до виконання кроків. Перевищення цих місткостей робить відповідний rollout непридатним для перевірки або бенчмаркінгу, навіть якщо виконання триває з попередженням про переповнення, а не завершується помилкою. Збільште відповідне обмеження та повторно запустіть завдання. Більші буфери використовують більше пам’яті 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 повідомляє про контакти та обмеження, фактично використані сценою, і перериває rollout із ідентифікаторами світів-порушників, щойно будь-який світ переповнюється. Вважайте такі звіти помилками: збільште обмеження й повторіть запуск, перш ніж довіряти траєкторії або бенчмарку, а потім знову зменшуйте його щоразу, коли змінюються модель, геометрія зіткнень або завдання.

1. Масштабування до 2 048 світів

Після успішної перевірки еквівалентності для одного світу повторно виділіть ресурси в цільовому розмірі та продублюйте ініціалізований стан у пакет. Порівняно з етапом 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 до 2 048 незалежних станів 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
курс симуляції SO-101 із переходом до реального світу · навчальні курси з фізичного ШІ

Навчання поверх 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 SO-101 із переходом від симуляції до реального світу
• Newton наступна публікація: MJWarp як SolverMuJoCo та перенесення цього середовища

Інші матеріали цього автора