kevinzakka/mjbatch

mjbatch – CPU 上の並列 MuJoCo シミュレーション

何であるか – mjbatch は、Python の GIL を解放する C++ スレッドプールを使用して、数千 の MuJoCo 物理シミュレーションを同時に実行できる小さな Python パッケージです。強化学習、モデル予測制御、システム同定、ハードウェア・コデザインに必要な大規模バッチ実行を必要とする研究者やエンジニアを対象としています。

主なアイデア

• スレッドプール実行 – 重い処理はネイティブ C++ スレッドで行われるため、各シミュレーションは Python のグローバルインタプリタロック（GIL）にブロッキングされず、別々のコアで実行されます。
• ライブ配列バインディング – qpos や ctrl などの MuJoCo データフィールドに直接バインドでき、バッチ全体にわたってインプレースで更新される NumPy 配列のように扱えます。
• シミュレーションごとのパラメータ – expand を使って geom_friction のようなモデル定数を個々のインスタンスごとに変化させられ、set_const で導出定数を再計算できます。

一般的なワークフロー

import mujoco, numpy as np
from mjbatch import Batch

# MuJoCo モデルを一度だけ読み込む
model = mujoco.MjModel.from_xml_path('scene.xml')

# 4096 個の並列シミュレーションのバッチを作成（デフォルトで論理CPUごとに1スレッド）
batch = Batch(model, num_sims=4096)

# 状態と制御配列にバインド – これらはバッチ全体への NumPy ビュー
qpos, ctrl = batch.bind('qpos'), batch.bind('ctrl')

# 各シミュレーションにモデルパラメータをランダム化
batch.expand('geom_friction')[:, :, 0] = np.random.uniform(0.4, 1.2, (4096, 1))

for _ in range(1000):
    ctrl[:] = policy(qpos)   # あなたのコントローラーがすべてのシミュレーションで同時に実行
    batch.step()             # すべてのシミュレーションを並列でステップ；qpos はインプレースで更新

上記のループは、単一シミュレーションのスクリプトと見分けがつかないようになっていますが、裏では 4,000 個の環境が同時にステップしています。

これで何ができるか

• 強化学習 – 大規模なリプレイバッファを生成したり、ポリシーを高速に評価したりできます（README には、古いラップトップで Go1 四足歩行ロボットが1分未満で歩行を学習する例が示されています）。
• モデル予測制御 / iLQR – 複数の初期条件にわたってリカレントホライズン最適化を実行できます。
• システム同定 – 複数の摂動モデルをシミュレーションし、合成データと比較することで慣性パラメータをフィットできます。
• コデザイン – 進化的または勾配ベースの手法を使って、ロボットの幾何形状、ギア比、制御ポリシーを同時に最適化できます。

リポジトリに同梱されている例

例	目的	技術
cartpole_swingup.py	ダブルポールカートをスイングアップ	iLQR
cartpole_mpc.py	予測サンプリング付きスイングアップ	MPC
g1_flip.py	人型ロボットのバックフリップ	リカレントホライズン iLQR
go1_joystick.py	四足歩行ロボットの遠隔操作	PPO 訓練済みポリシー
arm_throw.py	アームの形状と制御のコデザイン	CEM 最適化
rizon_inertia.py	慣性パラメータの同定	ダンプド・ガウス・ニュートン

uv run examples/<file>.py で任意の例を実行できます；一部はディスプレイが必要です（--headless で非表示でも動作します）。リポジトリには CI バッジと PyPI リリースも含まれており、パッケージがメンテナンスされており pip install mjbatch でインストール可能であることを示しています。

インストール

pip install mjbatch   # コンパイル済み拡張と純粋な Python ラッパーを取得

このライブラリは公式 MuJoCo Python バインディングに依存しているため、MuJoCo のライセンス（学術利用は無料）と mujoco パッケージのインストールが必要です。

ライセンス – Apache-2.0、商用およびオープンソース利用を許可しています。

---

要するに、mjbatch は GPU クラスタを必要とせずに MuJoCo を基盤とする実験をスケーリングしたいすべての人にとって実用的なツールです。スレッドの詳細を抽象化しつつ、バッチ全体のシミュレーション状態に直接 NumPy スタイルのアクセスを提供します。