🛠 NVIDIA Warp и MJWarp: как запустить тысячи симуляций для обучения роботов

MJWarp (MuJoCo Warp) переносит физику MuJoCo на GPU через NVIDIA Warp: MJCF-модель компилируется в CUDA-ядра, состояние сцены раскладывается по массивам на устройстве, а один вызов mjw.step продвигает весь пакет параллельных миров. В материале Hugging Face разбирают перенос манипулятора SO-101 из обычного сценария MuJoCo на CPU в 2048 параллельных сред и отделяют пропускную способность от задержки. Для одного мира MJWarp не обязательно быстрее — преимущество появляется в суммарной пропускной способности, когда сотни миров достаточно загружают GPU, чтобы выросло общее число шагов моделирования в секунду.

Механика переноса компактна: MuJoCo загружает и компилирует MJCF, MJWarp реализует физику в NVIDIA Warp, а тот компилирует ядра на лету для CUDA. На уровне API меняются три вызова: mujoco.MjModel на mjw.put_model, mujoco.MjData на mjw.make_data или put_data, mujoco.mj_step на mjw.step. Дальше подбирают буферы: nworld — число сред, nconmax — ожидаемое число контактов на один мир, njmax — верхний предел ограничений. Запас «на всякий случай» увеличивает потребление памяти и работу решателя, поэтому размеры выбирают впритык.

Авторы задают управление с частотой 50 Гц и 10 физических подшагов на кадр — шаг физического моделирования 0,002 секунды, одинаковый для CPU и GPU, чтобы замеры не искажались. Успех операции захвата и перемещения проверяют не кодом завершения, а числами: горизонтальная ошибка центров кубов xy_err ≤ 0,015 м и вертикальный зазор 0,035–0,055 м. При проверке одного мира каждый вызов `.numpy()` синхронизирует состояние с хостом, поэтому это не замер производительности; для оценки пропускной способности шаги оставляют на GPU и используют графы CUDA для повторного использования буферов. Warp 1.15 добавляет необязательный детерминизм — по умолчанию атомарные операции GPU зависят от планировщика, и повторные запуски могут слегка различаться.

Прирост тут не в физике, а в экономике GPU: MJWarp выигрывает, только когда пакет миров достаточно плотный, а nconmax и njmax подобраны под конкретную сцену. Чужой MJCF с запасом по контактам легко съест преимущество из-за раздутых буферов решателя.

#NVIDIAWarp #MJWarp #MuJoCo #робототехника #симуляции