mjbatch

Exécute des milliers de simulations MuJoCo en parallèle sur le CPU

Licence

Apache-2.0

Aperçu

mjbatch est une bibliothèque Python qui exécute des milliers de simulations MuJoCo en parallèle sur le CPU. Un pool de fils en C++ assure l'avancement en relâchant le GIL, et le lot est exposé à Python sous forme de vues NumPy vivantes : bind renvoie un tableau couvrant un champ MjData sur toutes les simulations, si bien qu'écrire dans ctrl et lire qpos se passe de copie et de boucle par simulation. expand fait de même pour les champs MjModel, ce qui donne à chaque simulation sa propre physique — frottements différents, longueurs de segment différentes — set_const recalculant ensuite les constantes dérivées. Le dépôt livre six exemples autonomes couvrant l'apprentissage par renforcement, la commande prédictive, l'identification de systèmes et la co-conception matérielle.

Le cœur est délibérément réduit. On construit un Batch à partir d'un seul MjModel et d'un nombre de simulations, le pool de fils prend par défaut tous les CPU logiques, et batch.step() les fait toutes avancer. Ce qui dépasse la simple boucle, c'est que les tableaux sont vivants au lieu d'être rassemblés : qpos, ctrl = batch.bind("qpos"), batch.bind("ctrl") rend des vues qu'un régulateur peut écrire sur tout le lot en une affectation, et qui se mettent à jour sur place à mesure que le lot avance. Les paramètres de modèle par simulation sont l'autre moitié — batch.expand("geom_friction") renvoie un tableau que l'on randomise par simulation, et c'est là de la randomisation de domaine et de l'identification de systèmes sans machinerie supplémentaire.

Les exemples sont l'argument, car chacun est une méthode complète et non une démonstration : un pendule sur chariot redressé par iLQR, une commande prédictive de pendule sur chariot par échantillonnage prédictif, un humanoïde G1 suivant un salto arrière de référence par iLQR à horizon glissant, un régulateur de manette pour le quadrupède Go1 entraîné par PPO, un CEM optimisant conjointement les proportions, les réducteurs et la commande d'un bras lanceur, et un ajustement de Gauss-Newton amorti des paramètres inertiels d'un bras Rizon. L'auteur rapporte que le régulateur du Go1 apprend à marcher en moins d'une minute sur un portable M1 vieux de cinq ans. Installation depuis PyPI, mujoco 3.13.0 épinglé, construction via scikit-build-core et nanobind, avec CI, suite pytest et banc de montée en charge.

Pourquoi c'est important

La simulation massivement parallèle est devenue une affaire de GPU — MJX, Isaac, Warp — et cela a discrètement relevé le ticket d'entrée pour qui n'a pas de carte de station de travail. mjbatch remet ce type de travail à portée d'un portable en utilisant vraiment le CPU plutôt qu'un cœur à la fois, et il le fait via MuJoCo tel quel plutôt qu'une réimplémentation : les modèles et la physique restent ceux auxquels on se fie déjà. Pour un catalogue dont le lectorat possède surtout un portable et un bras bon marché plutôt qu'une grappe, c'est le côté utile du compromis.

Le problème résolu

Obtenir des milliers de déroulés MuJoCo obligeait le plus souvent à passer sur une pile GPU et à en accepter la facture matérielle et la réimplémentation de la physique. mjbatch garde MuJoCo d'origine et tire le parallélisme du CPU déjà présent, ce qui suffit pour de la commande, de l'identification et de la co-conception à l'échelle d'un portable.

Points clés

• Un pool de fils C++ fait avancer tout le lot avec le GIL relâché : le parallélisme est réel, pas du Python entrelacé
• bind expose les champs MjData comme des tableaux vivants sur toutes les simulations — le régulateur écrit une fois dans ctrl et relit qpos sur place, sans étape de rassemblement
• expand donne à chaque simulation ses propres paramètres MjModel, set_const recalculant les constantes dérivées — randomisation de domaine et identification sans machinerie supplémentaire
• Six exemples aboutis, chacun une méthode complète : iLQR, MPC par échantillonnage prédictif, PPO, co-conception matérielle par CEM et ajustement inertiel de Gauss-Newton amorti
• L'auteur rapporte qu'un régulateur de manette du Go1 apprend à marcher en moins d'une minute sur un portable M1 vieux de cinq ans
• S'installe depuis PyPI, épingle mujoco 3.13.0, avec CI, suite pytest et banc de montée en charge

Pour qui

• Chercheurs en MPC, optimisation de trajectoire ou CEM ayant besoin de nombreux déroulés par pas
• Quiconque entraîne de petits régulateurs par renforcement sans accès à une grappe de GPU
• Praticiens de l'identification de systèmes ou de la randomisation de domaine sur les paramètres du modèle
• Utilisateurs de MuJoCo qui veulent de la simulation par lots sans migrer vers MJX ou Warp

Limites

• CPU uniquement par conception — pour des lots à l'échelle du GPU, MJX ou MuJoCo Warp restent la voie, et ceci ne les remplace pas
• Six jours d'existence au moment d'écrire, en v0.1.1 et classé Beta ; les 521 étoiles tiennent autant au parcours de l'auteur qu'à la bibliothèque
• Épingle mujoco==3.13.0 exactement, elle ne bouge donc que quand l'auteur relève la version
• Exige une chaîne d'outils C++ pour compiler depuis les sources, et les exemples qui ouvrent une fenêtre demandent un affichage sauf avec --headless
• La surface d'API est délibérément étroite — Batch, bind, expand, set_const — sortir de cette forme suppose de redescendre directement dans MuJoCo
• Les classificateurs publiés ne mentionnent que macOS et Linux ; pas de roues Windows