Як NVIDIA переносить симуляцію робочу руку з одного CPU-ядра на 2048 GPU-світів?

NVIDIA Warp — це Python-фреймворк, що компілює кернели для виконання на GPU. MuJoCo Warp базується на ньому, щоб одночасно запускати тисячі незалежних фізичних сцен MuJoCo на відеокартах NVIDIA.

Стек і зміна парадигми

Класичний MuJoCo працює швидко на CPU-ядрах. Він розпаралелює вибірку по цих ядрах. Але підсилюване навчання та масове семплінг потребують іншого: не меншої латентності для одного світу, а вищої сумарної пропускної здатності по багатьох. Сумарна пропускна здатність — це загальна кількість кроків світу за секунду реального часу. MJWarp досягає цього, розміщуючи сумісну модель MuJoCo та пакет незалежних станів на GPU. Один виклик mjw.step просуває весь пакет.

Тестовим випадком слугує рука-підписник SO-101. Сцена містить руку, стіл. До того ж є два куби по 44 мм для складання. Завдання: схопити червоний куб і покласти його на синій. Та сама MJCF-модель запускається на обох бекендах. Міграція йде через чотири етапи: встановити CPU-базову лінію, перевірити парність одного світу на GPU, розрахувати буфери контактів і обмежень, потім масштабувати до 2048 світів.

Встановлення CPU-базової лінії

Сцена завантажується з scene_pick_place.xml. Для коробки MuJoCo значення size — це напіврозміри. size="0.022 0.022 0.022" задає куб з ребрами 44 мм. База руки знаходиться в початку координат. Досіг простягається вздовж +X. Куби розташовані вздовж Y.

Керування працює на 50 Гц. Фізичні підкроки — 10 на кадр керування. Це дає крок фізики 0.002 секунди. Рядок mjm.opt.timestep = frame_dt / sim_substeps має виконатися до CPU-прогону і до завантаження моделі через mjw.put_model. Без цього кожне подальше вимірювання отримає невідповідність.

Успіх — не чистий вихід процесу. Це два вимірювані умови після встановлення кубів. Горизонтальна помилка між центрами кубів має бути xy_err ≤ 0.015 м. Вертикальна відстань між центрами — 0.035 м ≤ dz ≤ 0.055 м. Супутній репозиторій виводить ці числа в кінці запуску. Цей вивід — твердження, з яким порівнюється все інше.

Перевірка одного світу на GPU

Другий етап запускає один світ на GPU, хост залишається в ланці. Модель завантажується через m = mjw.put_model(mjm). Цей виклик одночасно перевіряє сумісність: він кидає помилку, якщо модель використовує непідтримувані функції, а не ігнорує їх тихо.

Пакетний стан виділяється через d = mjw.make_data(mjm, nworld=1, nconmax=spec.nconmax, njmax=spec.njmax). Профіль SO-101 використовує nconmax=128 і njmax=300 як стартові ємності. Хостовий стан наповнює пристрої. Кожен масив пристрою має провідну вимірність світу. Тому mjd.qpos[None, :] має форму (1, nq) замість (nq,).

Цикл кадрів повторює CPU-версію. Керування обчислюється на хості. Копіюється в d.ctrl кожен підкрок. mjw.step(m, d) просуває GPU-світ. Результати копіюються назад. mujoco.mj_forward(mjm, mjd) оновлює похідні хостові величини як mjd.xpos перед перевіркою складання. Читання цих полів після циклу не оновлює їх автоматично. Цей шлях перевіряє завдання. Це не бенчмарк пропускної здатності. Виклики .numpy() синхронізують і копіюють дані на хост кожен підкрок.

Розрахунок буферів перед масштабуванням

MJWarp виділяє буфери контактів і обмежень до кроків. Перевищення цих ємностей робить постраждалий прогін недійсним для верифікації або бенчмарку. Виконання продовжується з попередженням про переповнення, а не винятком. Попередження друкує необхідний бюджет: "narrowphase overflow - please increase nconmax to …". Також позначає постраждалі світи в Data.overflow.

Більші буфери споживають більше GPU-пам'яті. Правило: перевірити ємність на повному завданні перед зменшенням виділення. Розраховуємо їх проти моменту з найбільшою кількістю контактів. Для pick-and-place це миттєво, коли обидві пащі і стіл торкаються куба. Не коли рука висить у вільному просторі.

Інструмент mjwarp-testspeed --measure_alloc повідомляє реальне споживання контактів і обмежень. Він перериває прогін з ID проблемних світів, як тільки будь-який світ переповнюється. Ставимося до цих звітів як до невдач. Збільшуємо ліміт і перезапускаємо перед тим, як довіряти траєкторії або бенчмарку. Потім знову зменшуємо, коли змінюється модель, колізійна геометрія чи завдання.

Масштабування до 2048 світів

Масштабування змінює дві речі відносно другого етапу: 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,