mjbatch

Führt Tausende MuJoCo-Simulationen parallel auf der CPU aus

Lizenz

Apache-2.0

Überblick

mjbatch ist eine Python-Bibliothek, die Tausende MuJoCo-Simulationen parallel auf der CPU ausführt. Ein C++-Threadpool übernimmt das Schrittmachen und gibt dabei den GIL frei, und der Stapel liegt in Python als lebende NumPy-Sicht vor: bind liefert ein Array über ein MjData-Feld hinweg für jede Simulation, sodass Schreiben in ctrl und Lesen von qpos ohne Kopieren und ohne Schleife je Simulation auskommt. expand tut dasselbe für MjModel-Felder, womit jede Simulation ihre eigene Physik trägt — andere Reibung, andere Gliedlängen — und set_const die abgeleiteten Konstanten danach neu berechnet. Das Repository liefert sechs eigenständige Beispiele zu Reinforcement Learning, modellprädiktiver Regelung, Systemidentifikation und Hardware-Co-Design.

Der Kern ist bewusst klein. Man baut einen Batch aus einem einzigen MjModel und einer Simulationszahl, der Threadpool nimmt standardmäßig jede logische CPU, und batch.step() bringt alle voran. Mehr als eine Schleife wird daraus, weil die Arrays leben statt eingesammelt zu werden: qpos, ctrl = batch.bind("qpos"), batch.bind("ctrl") gibt Sichten zurück, die ein Regler mit einer einzigen Zuweisung über den ganzen Stapel beschreiben kann, und sie aktualisieren sich beim Schrittmachen an Ort und Stelle. Modellparameter je Simulation sind die andere Hälfte — batch.expand("geom_friction") liefert ein Array, das sich je Simulation randomisieren lässt, und das ist Domain Randomization und Systemidentifikation ohne weitere Maschinerie.

Die Beispiele sind das Argument, denn jedes ist eine vollständige Methode und keine Demo: ein mit iLQR aufgeschwungenes Wagen-Pendel, eine Wagen-Pendel-MPC mit Predictive Sampling, ein G1-Humanoid, der einen Referenz-Rückwärtssalto mit gleitender iLQR nachfährt, ein mit PPO trainierter Go1-Joystickregler, CEM, das Proportionen, Getriebe und Ansteuerung eines Wurfarms gemeinsam optimiert, und eine gedämpfte Gauß-Newton-Anpassung der Trägheitsparameter eines Rizon-Arms. Der Autor berichtet, der Go1-Regler lerne auf einem fünf Jahre alten M1-Laptop in unter einer Minute laufen. Installation über PyPI, mujoco 3.13.0 fest gepinnt, Bau über scikit-build-core mit nanobind, dazu CI, eine pytest-Suite und ein Skalierungs-Benchmark.

Warum es wichtig ist

Massiv parallele Simulation ist zu einer GPU-Geschichte geworden — MJX, Isaac, Warp — und das hat die Einstiegskosten für alle ohne Workstation-Karte still erhöht. mjbatch macht dieselbe Art Arbeit auf einem Laptop verfügbar, indem es die CPU richtig nutzt statt einen Kern nach dem anderen, und zwar über schlichtes MuJoCo statt über eine Neuimplementierung — die Modelle und die Physik bleiben also die, denen man schon traut. Für einen Katalog, dessen Publikum meist einen Laptop und einen günstigen Arm besitzt und keinen Cluster, ist das die nützlichere Seite des Handels.

Welches Problem es löst

Tausende MuJoCo-Rollouts zu bekommen hieß bisher meist, auf einen GPU-Stack zu wechseln und dessen Hardwarerechnung samt eigener Physik-Neuimplementierung zu akzeptieren. mjbatch behält das originale MuJoCo und holt die Parallelität aus der CPU, die ohnehin dasteht — genug für Regelung, Identifikation und Co-Design in Laptop-Größenordnung.

Wichtigste Punkte

• Ein C++-Threadpool schrittet den ganzen Stapel bei freigegebenem GIL, die Parallelität ist also echt und nicht bloß verschränktes Python
• bind legt MjData-Felder als lebende Arrays über alle Simulationen offen — ein Regler schreibt einmal in ctrl und liest qpos an Ort und Stelle zurück, ohne Sammelschritt
• expand gibt jeder Simulation eigene MjModel-Parameter, set_const rechnet abgeleitete Konstanten nach — Domain Randomization und Systemidentifikation ohne Zusatzmaschinerie
• Sechs ausgearbeitete Beispiele, jedes eine vollständige Methode: iLQR, Predictive-Sampling-MPC, PPO, CEM-Hardware-Co-Design und gedämpfte Gauß-Newton-Trägheitsanpassung
• Der Autor berichtet, ein Go1-Joystickregler lerne auf einem fünf Jahre alten M1-Laptop in unter einer Minute laufen
• Installierbar über PyPI, pinnt mujoco 3.13.0, mit CI, pytest-Suite und Skalierungs-Benchmark

Für wen

• Forschende zu MPC, Trajektorienoptimierung oder CEM, die pro Schritt viele Rollouts brauchen
• Alle, die kleine RL-Regler ohne Zugang zu einem GPU-Cluster trainieren
• Praktiker, die Systemidentifikation oder Domain Randomization über Modellparameter betreiben
• MuJoCo-Nutzer, die gestapelte Simulation wollen, ohne zu MJX oder Warp zu wechseln

Einschränkungen

• Konstruktionsbedingt nur CPU — für Stapel in GPU-Größenordnung bleiben MJX oder MuJoCo Warp der Weg, dies ersetzt sie nicht
• Zum Zeitpunkt des Schreibens sechs Tage alt, bei v0.1.1 und als Beta eingestuft; die 521 Sterne spiegeln den Ruf des Autors ebenso wie die Bibliothek
• Pinnt mujoco==3.13.0 exakt, bewegt sich also nur, wenn der Autor die Version anhebt
• Braucht eine C++-Toolchain zum Bauen aus dem Quellcode, und die Beispiele mit Fenster brauchen ein Display, sofern nicht mit --headless gestartet
• Die API-Fläche ist bewusst schmal — Batch, bind, expand, set_const — alles außerhalb dieser Form heißt, direkt auf MuJoCo zurückzufallen
• In den veröffentlichten Classifiern nur macOS und Linux; keine Windows-Wheels