mjbatch

Ejecuta miles de simulaciones de MuJoCo en paralelo en la CPU

Licencia

Apache-2.0

Resumen

mjbatch es una biblioteca de Python que ejecuta miles de simulaciones de MuJoCo en paralelo sobre la CPU. Un grupo de hilos en C++ hace el avance con el GIL liberado, y el lote se expone a Python como vistas NumPy vivas: bind te da un array sobre un campo de MjData a lo largo de todas las simulaciones, de modo que escribir en ctrl y leer qpos no requiere copias ni un bucle por simulación. expand hace lo mismo con los campos de MjModel, lo que permite que cada simulación lleve su propia física —fricciones distintas, longitudes de eslabón distintas— y set_const recalcula después las constantes derivadas. El repositorio incluye seis ejemplos autocontenidos que cubren aprendizaje por refuerzo, control predictivo, identificación de sistemas y codiseño de hardware.

El núcleo es deliberadamente pequeño. Construyes un Batch a partir de un único MjModel y un número de simulaciones, el grupo de hilos toma por defecto todas las CPU lógicas, y batch.step() las avanza todas. Lo que lo hace más que un bucle es que los arrays están vivos en lugar de recopilarse: qpos, ctrl = batch.bind("qpos"), batch.bind("ctrl") entrega vistas que un controlador puede escribir sobre todo el lote en una sola asignación, y se actualizan en el sitio conforme el lote avanza. Los parámetros de modelo por simulación son la otra mitad: batch.expand("geom_friction") devuelve un array que puedes aleatorizar por simulación, y eso es aleatorización de dominio e identificación de sistemas sin maquinaria adicional.

Los ejemplos son el argumento, porque cada uno es un método completo y no una demostración: un péndulo sobre carro levantado con iLQR, un MPC de péndulo sobre carro con muestreo predictivo, un humanoide G1 siguiendo un salto mortal de referencia con iLQR de horizonte deslizante, un controlador de joystick para el cuadrúpedo Go1 entrenado con PPO, CEM optimizando conjuntamente proporciones, engranajes y control de un brazo lanzador, y un ajuste de Gauss-Newton amortiguado de los parámetros inerciales de un brazo Rizon. El autor informa de que el controlador del Go1 aprende a caminar en menos de un minuto en un portátil M1 de hace cinco años. Se instala desde PyPI, fija mujoco 3.13.0, se construye con scikit-build-core y nanobind, y trae CI, una batería de pytest y un banco de escalado.

Por qué importa

La simulación masivamente paralela se ha vuelto una historia de GPU —MJX, Isaac, Warp— y eso ha subido en silencio el coste de entrada para quien no tiene una tarjeta de estación de trabajo. mjbatch pone ese mismo tipo de trabajo al alcance de un portátil usando bien la CPU en vez de un núcleo cada vez, y lo hace con MuJoCo a secas en lugar de una reimplementación, así que los modelos y la física siguen siendo los que ya te fías. Para un catálogo cuyo público suele tener un portátil y un brazo barato en vez de un clúster, ese es el lado útil del compromiso.

Qué problema resuelve

Conseguir miles de despliegues de MuJoCo suele obligar a pasarse a una pila de GPU y aceptar su factura de hardware y su reimplementación de la física. mjbatch mantiene MuJoCo de serie y extrae el paralelismo de la CPU que ya tienes, que basta para trabajo de control, identificación y codiseño a escala de portátil.

Puntos clave

• Un grupo de hilos en C++ avanza todo el lote con el GIL liberado, así que el paralelismo es real y no Python intercalado
• bind expone campos de MjData como arrays vivos sobre todas las simulaciones: el controlador escribe una vez en ctrl y lee qpos en el sitio, sin paso de recopilación
• expand da a cada simulación sus propios parámetros de MjModel, con set_const recalculando las constantes derivadas: aleatorización de dominio e identificación sin maquinaria extra
• Seis ejemplos trabajados, cada uno un método completo: iLQR, MPC por muestreo predictivo, PPO, codiseño de hardware con CEM y ajuste inercial de Gauss-Newton amortiguado
• El autor informa de un controlador de joystick del Go1 que aprende a caminar en menos de un minuto en un portátil M1 de hace cinco años
• Se instala desde PyPI, fija mujoco 3.13.0 y trae CI, batería de pytest y un banco de escalado

Para quién es

• Investigadores de MPC, optimización de trayectorias o CEM que necesitan muchos despliegues por paso
• Cualquiera que entrene controladores pequeños por refuerzo sin acceso a un clúster de GPU
• Profesionales que hacen identificación de sistemas o aleatorización de dominio sobre parámetros del modelo
• Usuarios de MuJoCo que quieren simulación por lotes sin pasarse a MJX o Warp

Limitaciones

• Solo CPU por diseño: para lotes a escala de GPU siguen siendo el camino MJX o MuJoCo Warp, y esto no los sustituye
• Seis días de vida al escribir esto, en v0.1.1 y clasificado como Beta; las 521 estrellas reflejan la trayectoria del autor tanto como la biblioteca
• Fija mujoco==3.13.0 exactamente, así que solo avanza cuando el autor sube la versión
• Necesita una cadena de herramientas C++ para compilar desde fuente, y los ejemplos que abren ventana necesitan pantalla salvo que se usen con --headless
• La superficie de API es deliberadamente estrecha —Batch, bind, expand, set_const— así que salirse de esa forma implica bajar directamente a MuJoCo
• En los clasificadores publicados solo aparecen macOS y Linux; no hay ruedas para Windows