NVIDIA Warp y MJWarp simulan hasta 2.048 robots en paralelo sobre GPU
MuJoCo Warp, construido sobre el framework de kernels de NVIDIA, reutiliza el mismo MJCF de MuJoCo y reparte la física entre GPU para escalar a 2.048 entornos.
NVIDIA ha publicado una guía para llevar la física de MuJoCo de la CPU a la GPU con MuJoCo Warp (MJWarp), una implementación de ese motor construida sobre NVIDIA Warp. El caso que recorren es un brazo SO-101 que pasa del flujo clásico de MuJoCo a 2.048 entornos simulados en paralelo sobre GPU. No hay entrenamiento de políticas de por medio: la pieza cubre solo cómo preparar y escalar la escena.
MuJoCo en CPU ya reparte el muestreo entre núcleos, pero cuando la carga de aprendizaje crece la pregunta deja de ser cuánto tarda un mundo en avanzar y pasa a ser cuántos mundos caben a la vez. MuJoCo Warp mantiene el mismo MJCF, así que no hay que reescribir el modelo, y compila kernels CUDA con Warp para avanzar lotes grandes de estados con los datos de simulación y aprendizaje cerca de la GPU.
Warp, kernels tipados en Python
NVIDIA Warp es un framework de Python para escribir kernels acelerados por GPU. Los kernels se escriben en un subconjunto de Python con tipos estáticos y se compilan para CPU o CUDA; el primer lanzamiento construye y cachea un módulo nativo, y los siguientes lo reutilizan. El Python corriente se queda con la configuración, la reserva de memoria y la orquestación de lanzamientos.
El ejemplo del artículo es un kernel de pocas líneas que integra posiciones bajo gravedad: cada hilo lógico se ocupa de un punto, de modo que el mismo código escala de dos puntos a millones sin cambiar el flujo. Tres propiedades sostienen eso. wp.tid() identifica el punto, contacto, cuerpo o mundo del hilo actual. Los arrays viven en un dispositivo concreto, y llamar a .numpy() sobre uno en CUDA sincroniza y copia a memoria de CPU, no es una ruta sin copia. Y se pueden encadenar lanzamientos y capturarlos en un grafo CUDA para reducir el coste de despacho repetido.
Warp también trae kernels diferenciables, con wp.Tape grabando el forward y reproduciendo los adjuntos en orden inverso al llamar a backward(), y ejecución determinista desde la versión 1.15. Las atomicidades de GPU dependen del planificador por defecto, de forma que dos lanzamientos del mismo kernel pueden diferir ligeramente; los modos deterministas son opt-in y sacrifican algo de rendimiento a cambio de un orden reproducible, algo que pesa en validación y en tests de regresión. Se instala con pip install warp-lang, con la 1.15 o superior si se busca ese determinismo.
Qué motor elegir según el caso
La guía deja una tabla de atajos. Para teleoperación o control predictivo de un solo robot, MuJoCo en CPU. Para exprimir la física de MuJoCo en bruto, MJWarp o mjlab. Para recetas de entrenamiento en JAX, MuJoCo Playground con MJX. Y para integrar varios solvers con Isaac Lab, Newton, que llegará en otra entrega de la serie.
Lo relevante es que MJWarp no obliga a cambiar de motor ni de formato de modelo: se reutiliza el MJCF y se gana escalado por GPU, que es justo lo que necesita quien entrena políticas con muchas réplicas de la misma escena. Lo que todavía no cubre son las capas de sensores, gestores de escena y bucles de entrenamiento, que el propio texto sitúa en las entregas dedicadas a Newton e Isaac Lab.
