NVIDIA показала, как запустить сразу 2 048 симуляций робота на видеокарте

NVIDIA выпустила пошаговое руководство: как перенести симуляцию роборуки из физического движка MuJoCo¹, который считает на процессоре, на видеокарту с помощью MJWarp и запустить до 2 048 копий сцены одновременно.

Почему это важно: для обучения с подкреплением² важнее не скорость одной симуляции, а объём собранного опыта, то есть число копий сцены, которые работают параллельно.

Как это работает: MJWarp — версия MuJoCo, которую Google DeepMind и NVIDIA переписали на Warp, фреймворке NVIDIA для программ под видеокарту на Python. Все копии сцены хранятся в памяти видеокарты, и один вызов продвигает на шаг их все сразу.

Повод: переход показан на одной задаче: роборука SO-101³ ставит красный кубик с ребром 44 мм на синий. Руководство разбивает его на этапы:

• Сначала задачу решают в обычном MuJoCo и проверяют успех: центры кубиков по горизонтали должны разойтись не больше чем на 15 мм, а по высоте — на 35–55 мм.
• Затем запускают одну копию сцены на видеокарте и после каждого шага переносят её состояние на процессор, чтобы сверить те же два числа.
• Потом подбирают объём памяти под контакты по самому тяжёлому моменту, когда обе губки захвата и стол касаются кубика.
• В конце число копий поднимают до 2 048, а шаг записывают в CUDA-граф — готовую цепочку запусков, которую видеокарта повторяет без лишних затрат на вызовы.

Детали: отдельно разобраны ошибки, которые портят результат незаметно:

• Когда контактов больше, чем помещается в выделенную память, MJWarp не останавливает симуляцию, а только печатает предупреждение. Такой прогон нельзя брать ни для проверки, ни для замеров.
• Шаг физики, 0,002 секунды, нужно задать до переноса модели на видеокарту. Иначе расхождение попадёт во все замеры, а обученная модель управления будет отдавать команды не с той частотой, что на настоящем роботе.
• Видеокарта работает асинхронно. Поэтому таймер без синхронизации измеряет, как быстро Python поставил задачи в очередь, а не как быстро видеокарта их выполнила.

Но: у перехода есть оговорки:

• Одна копия сцены на MJWarp не обязательно считается быстрее, чем на процессоре: выигрыш — в суммарном числе шагов в секунду по всем копиям.
• Замеров скорости NVIDIA не приводит: результат зависит от сцены, настроек и железа, поэтому кривую предлагают построить скриптом на своей машине.
• Код примера лежит в сопутствующем репозитории, но вместо его адреса в статье пока стоит заглушка.

Что дальше: в следующей части серии ту же сцену перенесут в Newton — открытый физический движок, где MJWarp рассчитывает движение твёрдых тел, — а затем в Isaac Lab⁴.