NVIDIAのエンジニア、Warpを用いたMuJoCoロボティクスシミュレーションのGPUスケーリング手法を詳説

ロボットシミュレーションを単一CPU環境から大規模GPUバッチへ拡張

2026年9月23日、NVIDIAの研究者であるJohnny Nuñez Cano氏、Asier Arranz氏、Rishabh Chadha氏、Ben Oliveri氏は、ロボティクス開発者が従来のCPUベースのMuJoCoシミュレーションを高スループットなGPUワークロードへと移行するための実装ガイドを公開した。フィジカルAIや強化学習パイプラインが拡大する中、トレーニング速度の向上には、単一環境のレイテンシを単純に最適化すること以上に、数百から数千の環境を同時に実行することが重要性を増している。

このギャップを埋めるため、NVIDIAはNVIDIA Warp上で構築されたMuJoCo物理エンジンのGPU高速化実装であるMuJoCo Warp（MJWarp）を用いたワークフローを導入した。NVIDIAの「State of Simulation for Physical AI」シリーズの第2弾となる今回のチュートリアルでは、ブロック積みタスクを実行するSO-101フォロワーロボットアームを、標準的なPython駆動のCPU環境から、NVIDIA GPU上の2,048並列シミュレーション状態へと直接移行する手順が実演されている。

SO-101アームの移植：バッファサイズ設定、CUDA Graphs、そして等価性の検証

この構成の基盤となっているのがNVIDIA Warpだ。これは、静的型付けされたPythonカーネルをネイティブのCUDAコードへとコンパイルするPythonベースのフレームワークであり、ジャストインタイム（JIT）コンパイル、カーネルフュージョン、wp.Tapeによる自動微分、バージョン1.15で導入された決定論的実行モードをサポートしている。MJWarpはこのフレームワークをMuJoCoの物理エンジンに適用し、標準的なMJCF XMLシーン記述を処理しながら、先頭にバッチ次元を追加した配列をGPU上へと移行する。

移行作業は、制御周波数50Hz、1フレームあたり10回の物理サブステップ（物理タイムステップ0.002秒を想定）で構成されたSO-101のピック＆プレースタスクから始まる。アームのタスクは、44mmの赤い立方体を把持し、同じサイズの青い立方体の上に配置することだ。このワークフローをMJWarpへ移行するにはAPIの変更が伴う。開発者はmjw.put_model()を用いてモデルをアップロードし、mjw.make_data()またはmjw.put_data()でGPU常駐ステートを確保した上で、mjw.step()を介してすべての並列環境を同時に前進させる。

運用上極めて重要な要件となるのが、デバイス側の接触および拘束バッファの管理である。開発者はnconmax（環境あたりの最大接触数）、naconmax（全体の接触数上限）、njmax（環境あたりの最大拘束数）といった容量を定義しなければならない。MJWarpは狭相衝突（narrowphase collision）が容量を超えた際、処理を即座に中断するのではなく警告を発する仕様となっているため、適切に検証を行わないとシミュレーション結果が知らぬ間に無効化してしまうリスクがある。NVIDIAは、グリッパーの両爪とテーブルが同時に立方体に触れる瞬間など、タスク中で最も接触が増加するタイミングを基準にこれらの上限サイズを決定し、mjwarp-testspeed --measure_allocなどのツールを用いてメモリ割り当てを診断することを推奨している。

CPUベースラインとの単一環境における物理的等価性（パリティ）が確認できれば、環境を2,048以上へとスケールさせることが可能になる。数千規模のバッチ処理ステップにおけるディスパッチレイテンシを排除するため、開発者はmjw.step()をwp.ScopedCapture()でラップし、再実行可能なCUDA Graphsを作成する。プロファイリングにおいては、計測対象ループの前後にwp.synchronize()を挿入することや、ロールアウトステップ中に.numpy()を介して配列をホストメモリへ戻すのではなくシミュレーションデータをデバイス上に保持し続けることで、GPUの非同期実行を適切に考慮する必要がある。

フィジカルAIシミュレーションスタックをNewtonおよびIsaac Labへ展開

MJWarpへのシフトは、単一ロボットのモデル予測制御（MPC）とバッチ強化学習の間で広がりつつある業界の役割分担を浮き彫りにしている。従来のCPU版MuJoCoは遠隔操作や単一環境でのデバッグ、低遅延なMPCにおける標準であり続けている一方で、バッチGPU環境は大規模なデータ収集を主目的としている。

MJWarpは、最新の複数のフィジカルAIフレームワークと連携するように設計されている。MJWarpをPyTorchと直接結合するmjlabや、MJXのWarp実装を介したMuJoCo Playground、そしてより広範な学習パイプラインの実行バックエンドとして機能する。NVIDIAは、本シミュレーションシリーズの次回記事において、このSO-101環境を同社のマルチソルバーフレームワーク「Newton」へインポートする方法を実演すると明かした。そこではMJWarpが基盤となる剛体ソルバー（newton.solvers.SolverMuJoCo）として動作し、シーンをIsaac Labのトレーニングワークフローに直接接続することになる。