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

Класическият 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 ядра за придвижване на състоянията на симулацията върху NVIDIA GPU.

Това е втората статия от поредицата ни „Състояние на симулацията за Physical AI“. Първата статия очерта пейзажа на роботната симулация. Тук подготвяме и мащабираме средата за симулация; не обучаваме политика. Следващите части за Newton и Isaac Lab разглеждат следващите интеграционни слоеве.

Сглобяване на всичко

Кратък ориентир за избор:

Започнете с едно полезно Warp ядро

NVIDIA Warp е Python framework за писане на високопроизводителни ядра с 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; това не е път без копиране. За pipeline с PyTorch или JAX, разположен на устройството, използвайте адаптерите на Warp или споделяне, съвместимо с DLPack.
• Композиционни стартирания на ядра. Програмата може да стартира поредица от фокусирани ядра и да запише поддържаната CUDA работа в граф, за да намали повтарящите се разходи по изпращането. Записването на граф възпроизвежда стартиранията спрямо съществуващи буфери; то не слива произволни ядра.

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

Две допълнителни възможности на Warp заслужават внимание, въпреки че нито една не се използва в работния процес със SO-101 в тази статия. Ядрата на Warp са диференцируеми: wp.Tape записва извикванията на предните ядра, направени в неговия контекст, и възпроизвежда техните съпътстващи операции в обратен ред при извикване на backward(), поради което екипите изграждат диференцируема геометрия, CFD и персонализирана физика в Warp, включително CAE работни процеси за симулация и оптимизация на дизайна. Warp поддържа и детерминирано изпълнение, въведено в Warp 1.15: GPU атомарните операции по подразбиране зависят от планировчика, така че многократните стартирания на едно и също ядро могат леко да се различават, а режимите с включена детерминираност заменят част от производителността за възпроизводим ред в симулации, валидация и регресионни тестове. Това са възможности на Warp, а не гаранции за диференцируемост или детерминираност на целия rollout на 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 е реализация на физическия pipeline на MuJoCo в NVIDIA Warp, която разполага модела и пакет от независими състояния на NVIDIA GPU; едно извикване на 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 и приспиването, а не отделния физичен engine framework 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 mm и да се постави върху синия куб; същият робот и сцена се използват за валидация на 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 mm. Задачата използва този размер за праговете за успех. Основата на рамото е в началото на координатната система, обсегът му е по +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 rollout-а и преди качването на модела с mjw.put_model, така че двата backend-а да придвижват едно и също симулирано време:

mjm.opt.timestep = frame_dt / sim_substeps 

Без този ред всяко следващо измерване наследява несъответствието: сравненията за паритет, числата за пропускателна способност, цитирани като „симулирани секунди“, и всяка научена политика, чиято честота на действията вече не съответства на внедряването.

Проверете дали кубовете са успешно подредени. При кубове с размер 44 mm успехът се изразява в две измерими условия: хоризонтална грешка на центъра xy_err ≤ 0.015 m (измерена между центровете на кубовете) и вертикално разстояние 0.035 m ≤ dz ≤ 0.055 m между центровете на кубовете (един ръб на куба, с допуск за наместване). Оценявайте двете условия, след като кубовете се стабилизират; успешното приключване на процеса само по себе си не доказва успех на задачата.

Стартирайте 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, фиксирано към проверен commit, тъй като ресурсите на 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() синхронизират и копират данните към хоста при всяка подстъпка, така че това е път за валидация на задачата, а не benchmark за пропускателна способност. Той поддържа обратната кинематика, визуализацията и проверките на задачата на хоста. След копиране на qpos и qvel извикайте mujoco.mj_forward(mjm, mjd), за да обновите производните величини на хоста като mjd.xpos, преди да ги използвате за управление, визуализация или проверка на подреждането. Четенето на тези полета след цикъла не ги обновява автоматично. Врата 4 премахва тези копия към хоста при всяка стъпка от пътя за измерване на пропускателната способност.

1. Оразмеряване на капацитета за контакти и ограничения

MJWarp заделя буфери за контакти и ограничения преди стъпването. Превишаването на тези капацитети прави засегнатия rollout невалиден за проверка или benchmark, дори когато изпълнението продължава с предупреждение за препълване, вместо с изключение. Увеличете съответния лимит и стартирайте задачата отново. По-големите буфери използват повече 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-а с идентификаторите на засегнатите светове веднага щом някой свят препълни капацитета. Приемайте тези отчети като грешки: увеличете лимита и стартирайте отново, преди да се доверите на траекторията или benchmark-а, след което отново го затегнете, когато моделът, геометрията на сблъсъците или задачата се променят.

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 Graphs използват повторно модела и буферите с данни, записани тук. Обновявайте 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 обхожда размерите на пакета и отпечатва ms/step заедно с пропускателната способност и ускорението:

cd /tutorials/sim2real-blogs/notebooks/mujoco/part2
python solutions/so101_mjwarp_solution.py --headless-steps 600     # паритет, изисква 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: *Състоянието на симулацията за Physical AI: преглед* — Публикация 1 от тази поредица.
• NVIDIA Warp — GitHub · Документация · версия v1.15.0 (GPU детерминираност) · Ръководство за детерминирано изпълнение
• MuJoCo Warp — GitHub · Официална документация за MJWarp
• Изграждане на ускорен, диференцируем код за изчислителна физика за AI с NVIDIA Warp
• Представяне на програмирането, базирано на tiles, в Warp 1.5.0
• mjlab · arXiv:2601.22074
• MuJoCo Playground
• Курсът на NVIDIA за sim-to-real със SO-101
• Newton следваща публикация: MJWarp като SolverMuJoCo и пренасяне на тази среда

Още от този автор