Ingenieros de NVIDIA detallan el escalado en GPU para la simulación robótica en MuJoCo mediante Warp

Escalando la simulación robótica de mundos individuales en CPU a lotes masivos en GPU

El 23 de septiembre de 2026, los investigadores de NVIDIA Johnny Nuñez Cano, Asier Arranz, Rishabh Chadha y Ben Oliveri publicaron una guía de implementación que muestra cómo los desarrolladores de robótica pueden migrar las simulaciones tradicionales de MuJoCo basadas en CPU a cargas de trabajo de alto rendimiento en GPU. A medida que se expanden los flujos de trabajo de IA física y aprendizaje por refuerzo, la velocidad de entrenamiento depende cada vez más de ejecutar cientos o miles de entornos en simultáneo, en lugar de simplemente optimizar la latencia de un único mundo.

Para acortar esta brecha, NVIDIA introdujo un flujo de trabajo que utiliza MuJoCo Warp (MJWarp), una implementación acelerada por GPU del motor de física MuJoCo construida sobre NVIDIA Warp. Como segunda entrega de la serie "State of Simulation for Physical AI" de NVIDIA, la guía práctica demuestra cómo trasladar un brazo robótico seguidor SO-101 que ejecuta una tarea de apilamiento de bloques desde un entorno estándar en CPU basado en Python directamente a 2.048 estados de simulación paralelizados en una GPU de NVIDIA.

Migración del brazo SO-101: dimensionamiento de buffers, CUDA Graphs y validación de paridad

En la base de esta configuración se encuentra NVIDIA Warp, un framework basado en Python que compila kernels de Python con tipado estático a código nativo de CUDA, con soporte para compilación just-in-time (JIT), fusión de kernels, diferenciación automática mediante wp.Tape y modos de ejecución deterministas introducidos en la versión 1.15. MJWarp aplica este framework al motor de física de MuJoCo, procesando descripciones de escenas estándar en XML MJCF y transfiriendo los arreglos a la GPU con una dimensión de lote (batch) inicial añadida.

La migración comienza con una tarea de recogida y colocación (pick-and-place) del SO-101 configurada a una frecuencia de control de 50 Hz con 10 subpasos de física por cuadro, apuntando a un paso temporal de física de 0,002 segundos. La tarea del brazo consiste en sujetar un cubo rojo de 44 mm y colocarlo sobre un cubo azul de idéntico tamaño. Trasladar este flujo de trabajo a MJWarp implica una transición de API: los desarrolladores cargan el modelo mediante mjw.put_model(), asignan el estado residente en la GPU usando mjw.make_data() o mjw.put_data(), y hacen avanzar todos los mundos paralelos simultáneamente a través de mjw.step().

Un requisito operativo crucial es gestionar los buffers de contactos y restricciones en el dispositivo. Los desarrolladores deben definir capacidades como nconmax (máximo de contactos por entorno), naconmax (límite global de contactos) y njmax (máximo de restricciones por entorno). Dado que MJWarp emite advertencias en lugar de terminar la ejecución de inmediato cuando las colisiones en fase estrecha (narrowphase) superan la capacidad, las ejecuciones corren el riesgo de invalidarse silenciosamente a menos que se verifiquen. NVIDIA recomienda dimensionar estos límites durante el momento de mayor contacto de una tarea —como cuando tanto las mordazas de la pinza como la mesa tocan un cubo— y diagnosticar las asignaciones con herramientas como mjwarp-testspeed --measure_alloc.

Una vez verificada la paridad física en un solo mundo frente a la referencia en CPU, los entornos pueden escalarse a 2.048 mundos o más. Para eliminar la latencia de despacho en miles de pasos por lotes, los desarrolladores envuelven mjw.step() dentro de wp.ScopedCapture() para crear CUDA Graphs reproducibles. El análisis de rendimiento (profiling) debe considerar la ejecución asíncrona de la GPU insertando wp.synchronize() antes y después de los bucles cronometrados, y manteniendo los datos de simulación en el dispositivo en lugar de transferir los arreglos a la memoria del host mediante .numpy() durante los pasos de despliegue (rollout).

Expandiendo la pila de simulación de IA física hacia Newton e Isaac Lab

La transición a MJWarp subraya una creciente división en la industria entre el control predictivo basado en modelos (MPC) para un solo robot y el aprendizaje por refuerzo por lotes. Mientras que el clásico MuJoCo en CPU sigue siendo el estándar para teleoperación, depuración en mundos individuales y MPC de baja latencia, los entornos por lotes en GPU apuntan a la recolección masiva de datos.

MJWarp está diseñado para interactuar con varios frameworks modernos de IA física. Funciona como backend de ejecución para mjlab (que vincula MJWarp directamente con PyTorch), MuJoCo Playground a través de la implementación en Warp de MJX, y pipelines de aprendizaje más amplios. NVIDIA afirmó que la próxima entrega de la serie de simulación demostrará cómo importar este mismo entorno del SO-101 a su framework multi-solver, Newton, donde MJWarp operará como el solver de cuerpos rígidos subyacente (newton.solvers.SolverMuJoCo) para conectar las escenas directamente con los flujos de trabajo de entrenamiento de Isaac Lab.