mjbatch

CPU 上で数千の MuJoCo シミュレーションを並列に走らせる

ライセンス

Apache-2.0

概要

mjbatch は、数千の MuJoCo シミュレーションを CPU 上で並列に走らせる Python ライブラリです。ステップ実行は C++ のスレッドプールが GIL を解放した状態で担い、バッチは実時間の NumPy ビューとして Python 側に公開されます。bind は全シミュレーションにまたがる MjData フィールドの配列を返すので、ctrl への書き込みと qpos の読み出しにコピーもシミュレーションごとのループも要りません。expand は MjModel フィールドについて同じことを行い、各シミュレーションが自前の物理——異なる摩擦、異なるリンク長——を持てるようにし、その後 set_const が派生定数を再計算します。リポジトリには強化学習、モデル予測制御、システム同定、ハードウェア協調設計をカバーする六つの自己完結した例が同梱されています。

中核は意図的に小さく保たれています。単一の MjModel とシミュレーション数から Batch を作ると、スレッドプースは既定で全論理 CPU を使い、batch.step() がすべてを進めます。これが単なる for ループを超えるのは、配列が集約されたものではなく生きているからです。qpos, ctrl = batch.bind("qpos"), batch.bind("ctrl") が返すのはビューであり、コントローラは一度の代入でバッチ全体に書き込め、バッチが進むにつれてその場で更新されます。シミュレーションごとのモデルパラメータがもう半分です——batch.expand("geom_friction") はシミュレーションごとにランダム化できる配列を返し、これがそのままドメインランダム化とシステム同定になります。追加の仕掛けは要りません。

例こそがこのライブラリの論拠です。どれもデモではなく完結した手法だからです。iLQR で振り上げるカートポール、予測サンプリングによるカートポール MPC、後退ホライズン iLQR で参照バク転を追従する G1 ヒューマノイド、PPO で学習した Go1 四足のジョイスティック制御、投擲アームの比率とギアと制御を CEM が同時最適化する例、そして減衰ガウス・ニュートン法で Rizon アームの慣性パラメータを合わせ込む例です。著者は、五年前の M1 ノート PC で Go1 の制御器が一分足らずで歩行を学習したと報告しています。PyPI から導入でき、mujoco 3.13.0 を固定し、scikit-build-core と nanobind でビルドされ、CI、pytest スイート、スケーリングベンチマークを備えます。

重要な理由

大規模並列シミュレーションは GPU の物語になりました——MJX、Isaac、Warp——そしてそれは、ワークステーション用カードを持たない人の参入コストを静かに引き上げました。mjbatch は、一度に一コアではなく CPU をきちんと使い切ることで、同じ形の作業をノート PC の上に取り戻します。しかも再実装ではなく素の MuJoCo を通すので、モデルも物理もすでに信頼しているものそのままです。読者の多くがクラスタではなくノート PC と安価なアームを持っているこのカタログにとって、こちらがこのトレードオフの有用な側です。

解決する課題

数千回の MuJoCo ロールアウトを得るには、たいてい GPU スタックへ移り、そのハードウェア代と物理の再実装を受け入れることを意味してきました。mjbatch は素の MuJoCo のまま、すでに手元にある CPU から並列性を引き出します。ノート PC 規模の制御・同定・協調設計の仕事にはそれで十分です。

主なポイント

• C++ スレッドプールが GIL を解放したままバッチ全体を進めるので、並列性は Python の交互実行ではなく本物である
• bind は MjData フィールドを全シミュレーションにまたがる実時間配列として公開する——コントローラは ctrl に一度書き、qpos をその場で読み返す。集約ステップは存在しない
• expand は各シミュレーションに独自の MjModel パラメータを与え、set_const が派生定数を再計算する——追加の仕掛けなしでドメインランダム化とシステム同定ができる
• 六つの作り込まれた例はどれも完結した手法である：iLQR、予測サンプリング MPC、PPO、CEM によるハードウェア協調設計、減衰ガウス・ニュートン慣性同定
• 著者は、五年前の M1 ノート PC で Go1 のジョイスティック制御器が一分足らずで歩行を学習したと報告している
• PyPI から導入でき、mujoco 3.13.0 を固定し、CI・pytest スイート・スケーリングベンチマークを備える

対象

• 各ステップで多数のロールアウトを要する MPC・軌道最適化・CEM の研究者
• GPU クラスタを使えない環境で小規模な強化学習制御器を訓練する人
• モデルパラメータ上でシステム同定やドメインランダム化を行う実務者
• MJX や Warp に移らずにバッチシミュレーションが欲しい MuJoCo ユーザ

制限事項

• 設計上 CPU 専用——GPU 規模のバッチには MJX や MuJoCo Warp が依然として道であり、本ライブラリはそれらを置き換えない
• 執筆時点で生後六日、v0.1.1 で Beta 分類。521 個のスターはライブラリと同じくらい著者の実績を映している
• mujoco==3.13.0 を厳密に固定しているため、著者が上げない限り動かない
• ソースからのビルドには C++ ツールチェーンが要り、ウィンドウを開く例は --headless を付けない限りディスプレイが必要
• API 面は意図的に狭い——Batch、bind、expand、set_const——ため、その形から外れることは MuJoCo に直接降りることを意味する
• 公開されている分類子は macOS と Linux のみで、Windows 用ホイールはない