mjbatch

Esegue migliaia di simulazioni MuJoCo in parallelo sulla CPU

Licenza

Apache-2.0

Panoramica

mjbatch è una libreria Python che esegue migliaia di simulazioni MuJoCo in parallelo sulla CPU. Un thread pool C++ esegue i passi di simulazione con il GIL rilasciato, e il batch è esposto a Python come viste NumPy dirette: bind fornisce un array su un campo di MjData attraverso tutte le simulazioni, quindi scrivere in ctrl e leggere qpos non richiede copie né cicli per singola simulazione. expand fa lo stesso per i campi di MjModel, il che permette a ogni simulazione di avere la propria fisica (attriti diversi, lunghezze dei link diverse), con set_const che ricalcola poi le costanti derivate. Il repository include sei esempi autonomi che coprono apprendimento per rinforzo, controllo predittivo basato su modello (MPC), identificazione del sistema e co-progettazione dell'hardware.

Il nucleo è volutamente piccolo. Si costruisce un Batch da un singolo MjModel e un numero di simulazioni, il thread pool usa per impostazione predefinita tutte le CPU logiche, e batch.step() le fa avanzare tutte. Ciò che lo rende più di un ciclo for è che gli array sono diretti anziché raccolti: qpos, ctrl = batch.bind("qpos"), batch.bind("ctrl") restituisce viste che un controller può scrivere sull'intero batch con un'unica assegnazione, e che si aggiornano sul posto man mano che il batch avanza. I parametri del modello per singola simulazione sono l'altra metà: batch.expand("geom_friction") restituisce un array che si può randomizzare per ogni simulazione, il che equivale a domain randomization e identificazione del sistema senza meccanismi aggiuntivi.

Gli esempi sono l'argomento a suo favore, perché ciascuno è un metodo completo e non una demo: un carrello con pendolo portato in verticale con iLQR, un MPC del carrello con pendolo che usa il predictive sampling, un umanoide G1 che segue un salto mortale all'indietro di riferimento con iLQR a orizzonte mobile, un controller con joystick per il quadrupede Go1 addestrato con PPO, CEM che ottimizza congiuntamente le proporzioni, gli ingranaggi e i comandi di un braccio lanciatore, e un fit Gauss-Newton smorzato dei parametri inerziali di un braccio Rizon. L'autore riporta che il controller del Go1 impara a camminare in meno di un minuto su un portatile M1 di cinque anni. Si installa da PyPI, fissa mujoco 3.13.0, si compila tramite scikit-build-core con nanobind e include CI, una suite pytest e un benchmark di scalabilità.

Perché è importante

La simulazione massicciamente parallela è diventata una faccenda di GPU (MJX, Isaac, Warp), e questo ha alzato silenziosamente il costo d'ingresso per chiunque non abbia una scheda da workstation. mjbatch rende disponibile lo stesso tipo di lavoro su un portatile usando la CPU come si deve invece che un core alla volta, e lo fa attraverso MuJoCo standard anziché una reimplementazione, quindi i modelli e la fisica sono quelli di cui già ci si fida. Per un catalogo il cui pubblico possiede per lo più un portatile e un braccio a basso costo anziché un cluster, questo è il lato più utile del compromesso.

Il problema che risolve

Ottenere migliaia di rollout MuJoCo di solito significa passare a uno stack GPU e accettarne il costo dell'hardware e la reimplementazione della fisica. mjbatch mantiene MuJoCo standard ed estrae il parallelismo dalla CPU che si ha già, il che basta per lavori di controllo, identificazione e co-progettazione alla scala di un portatile.

Punti chiave

• Un thread pool C++ fa avanzare l'intero batch con il GIL rilasciato, quindi il parallelismo è reale e non Python interlacciato
• bind espone i campi di MjData come array diretti su tutte le simulazioni: un controller scrive in ctrl una sola volta e rilegge qpos sul posto, senza alcuna fase di raccolta
• expand dà a ogni simulazione i propri parametri di MjModel, con set_const che ricalcola le costanti derivate: domain randomization e identificazione del sistema senza meccanismi aggiuntivi
• Sei esempi elaborati, ciascuno un metodo completo: iLQR, MPC con predictive sampling, PPO, co-progettazione dell'hardware con CEM e fit inerziale Gauss-Newton smorzato
• L'autore riporta un controller con joystick per il Go1 che impara a camminare in meno di un minuto su un portatile M1 di cinque anni
• Si installa da PyPI, fissa mujoco 3.13.0 e include CI, una suite pytest e un benchmark di scalabilità

A chi si rivolge

• Ricercatori che si occupano di MPC, ottimizzazione della traiettoria o CEM e hanno bisogno di molti rollout per passo
• Chiunque addestri piccoli controller RL senza accesso a un cluster GPU
• Professionisti che fanno identificazione del sistema o domain randomization sui parametri del modello
• Utenti MuJoCo che vogliono la simulazione in batch senza passare a MJX o Warp

Limiti

• Solo CPU per scelta progettuale: per batch su scala GPU la strada restano MJX o MuJoCo Warp, e questa libreria non li sostituisce
• Ha sei giorni al momento della stesura, è alla v0.1.1 ed è classificata come Beta; le 521 stelle riflettono il curriculum dell'autore tanto quanto la libreria
• Fissa esattamente mujoco==3.13.0, quindi si aggiorna solo quando l'autore alza la versione
• Richiede una toolchain C++ per la compilazione dai sorgenti, e gli esempi che aprono una finestra richiedono un display a meno di essere eseguiti con --headless
• La superficie dell'API è volutamente ristretta (Batch, bind, expand, set_const), quindi tutto ciò che esce da questo schema richiede di scendere direttamente a MuJoCo
• Solo macOS e Linux nei classificatori pubblicati; nessuna wheel per Windows