От CPU к GPU: почему важна пропускная способность
Классический MuJoCo обеспечивает быструю симуляцию на CPU, но при масштабировании задач машинного обучения вопрос смещается от скорости одного мира к количеству одновременно обрабатываемых миров. NVIDIA Warp и новый модуль MJWarp переносят физику MuJoCo на GPU, позволяя продвигать состояния симуляции большими батчами. Это сохраняет данные на устройстве и значительно повышает агрегатную пропускную способность (total world-steps per second), что важнее для RL, чем минимизация задержки одного шага.
Архитектура и ключевые компоненты
Решение строится на стеке, где NVIDIA Warp выступает языком ядра Python (SIMT, автодифференцирование), а MJWarp реализует физику MuJoCo поверх Warp. Это позволяет использовать одни и те же модели MJCF, но получать пакетную обработку на GPU. Для интеграции с ML-пайплайнами поддерживается DLPack и совместимость с PyTorch/JAX.
| Компонент | Роль в стеке |
|---|---|
| NVIDIA Warp | Язык Python-ядер: JIT-компиляция в CUDA, автодифференцирование, совместимость с PyTorch/JAX |
| MJWarp | Физика MuJoCo на Warp: те же модели MJCF, пакетная пропускная способность на GPU |
| Сцена (SO-101) | Активы Menagerie/Robot Studio + геометрия задачи |
| Next (Newton/Isaac Lab) | Мульти-решатель, USD, сенсоры, менеджеры, циклы обучения |
Технические детали и производительность
В статье демонстрируется переход робота-манипулятора SO-101 от стандартного workflow MuJoCo к MJWarp. Ключевое достижение — возможность масштабирования до 2 048 параллельных сред (worlds). Это достигается за счет явного параллелизма: один логический поток обрабатывает одну точку/тело/мир, что позволяет коду масштабироваться от двух точек до миллионов без изменения логики управления.
Важно различать метрики:
- Latency (Задержка): время в секундах для одного шага симуляции одного мира.
- Aggregate Throughput (Агрегатная пропускная способность): общее количество выполненных шагов мира в секунду. Именно этот показатель оптимизирует MJWarp для задач сбора опыта (experience collection).
Установка и совместимость
Для использования требуется пакет warp-lang версии не ниже 1.15 (для поддержки детерминированного выполнения на GPU). Установка осуществляется через pip:
pip install warp-lang
Разработчики отмечают, что MJWarp не заменяет CPU-версию для отладки одного робота или телеуправления, но становится стандартом для large-scale sampling и обучения политик, где критичен объем собранных данных, а не отклик отдельной симуляции.
Источник: Hugging Face blog ↗