In Silico

シミュレーション・制御・実機

MuJoCo MJXのJAX版とWarp版の違い：接触の扱い・速度・対応GPU

2026/9/30 （更新: 2026/10/11）

🤖 AIを使用して執筆

  • MuJoCo
  • ロボット
  • 強化学習
  • 物理シミュレーション
  • MJX
  • MuJoCo Warp
  • JAX
  • Warp
  • MuJoCo Playground
  • NVIDIA GPU
  • XLA
  • 接触
  • 拘束
  • 自動微分
  • スループット
  • 条件分岐
  • 自由度
  • スリープ
  • メッシュ
  • nconmax
  • njmax
  • naconmax
  • Data.overflow
  • TPU
  • Apple Silicon
  • Google DeepMind
  • Newton
  • 遅延
目次
背景・問い・要点
背景

ロボットの動かし方を強化学習で学習させるには、大量の試行が要る。実機だけでは試行の数が足りないので、GPU の上で物理シミュレーションを数千個同時に進めて試行を集める。同時に進めるシミュレーションの一つひとつを、世界（world）と呼ぶ。本稿は、物理シミュレータ MuJoCo の GPU 向け API である MJX を例に取る。MuJoCo は Google DeepMind が保守する物理シミュレータで、MJX には JAX 実装と Warp 実装の二つがある。作り手は、中規模から大規模の場面には Warp 実装を勧めている12。

問い

Warp 実装は、なぜ JAX 実装より大きな場面を速く計算できるのか。

要点

作り手の説明によると、JAX 実装はどの世界でも同じ個数の接触を生成しなければならず、Warp 実装は世界ごとに、実際に起きている個数だけ接触を生成できるからである2。 JAX 実装が生成する個数は、場面で起こりうる接触の数で決まるので、場面に物体を足すと、接触の少ない世界でも計算が増える。Warp 実装も接触の上限を前もって決めるが、上限は全世界の合計に対して置ける3。DeepMind の学習環境集 MuJoCo Playground の操作環境の多くで、Warp 実装は単位時間あたりに進むステップ数が1.5〜2倍になった2。代わりに Warp 実装は、高速に動く GPU が NVIDIA 製に限られ、自動微分に対応していない41。

モデル・例示

三つの世界で接触を数えると、二つの方式の差が分かる

※ この節の数値は説明のための仮定で、測定値ではありません。

接触は、物体どうしが触れている箇所である。シミュレータは1ステップごとに、触れている箇所を調べて接触の一覧を作る。これを接触の生成と呼ぶ。続いて、接触の一つひとつについて、物体がめり込まないための力を計算する。生成する接触が多いほど、1ステップの計算は増える。

三つの世界を同時に進める。この場面で起こりうる接触は、最大で40個とする。ある時点で実際に起きている接触の数は、世界ごとに違う。

  • 世界1：2個（物体がほとんど動いていない）
  • 世界2：5個
  • 世界3：40個（物体が積み重なっている）

この三つの世界を計算する方式は、二つある。

方式三つの世界で生成する接触の数
個数を揃える40 ＋ 40 ＋ 40 ＝ 120
世界ごとに変える2 ＋ 5 ＋ 40 ＝ 47

個数を揃える方式は、どの世界でも同じ個数の接触を生成する。その個数は、起こりうる接触を収められる数にする。ここでは40個である。世界1で実際に起きている接触は2個だが、40個分を生成する。残りの38個分は、実際には触れていない。

世界ごとに変える方式は、実際に起きている個数だけ接触を生成する。ただし、接触を入れる枠の大きさ（上限）は、この方式でも前もって決める。計算量は上限にも応じて増える。上限を世界ごとではなく全世界の合計に対して置くとし、世界あたりの見込みを16個にすると、合計の上限は 3 × 16 ＝ 48 になる。実際の接触の合計は47なので、上限に収まる。世界3の接触は40個で、世界あたりの見込みの16を超えている。それでも、世界1と世界2が使わなかった枠を世界3が使える。個数を揃える方式では、接触の枠が 40 × 3 ＝ 120 要る。

120 と 47 の開きを生む仕組みを、接触の生成について三点に分けて書く。

  1. 個数を揃える方式が生成する接触の数は、世界の数と、起こりうる接触の数の積で決まる。実際に起きている接触の数は関係しない。
  2. 場面に物体を足すと、起こりうる接触の数が増える。個数を揃える方式では、接触の少ない世界も含めて、すべての世界で生成する個数が増える。
  3. どの世界でも、起こりうる接触の数が少なく、実際に起きている接触の数とほぼ同じ場面では、二つの方式が生成する個数はほとんど変わらない。

JAX 実装は個数を揃える方式で、Warp 実装は世界ごとに変える方式である。

作り手の説明と実測は、接触の個数の扱いの違いと矛盾しない

Warp 実装は、MJX の API から MuJoCo Warp（以下 MJWarp）を呼ぶ2。MJWarp は Google DeepMind と NVIDIA が保守する MuJoCo の GPU 版で、NVIDIA の Warp で書かれている4。Warp は、GPU で実行する関数を書くための枠組みである5。

JAX 実装がどの世界でも同じ個数の接触を生成する理由を、作り手は GPU でのプログラムの動かし方で説明している。MuJoCo Playground は DeepMind の学習環境集である。その告知によると、JAX 実装では、1ステップごとに固定の個数の接触を生成するしかなかった。Warp は NVIDIA GPU を対象にしているので、GPU の中で並列に動く処理の一つひとつが、条件分岐で別々の処理へ進むプログラムを書ける。その結果、1ステップごとに個数の変わる接触を生成できるようになった2。

すべての世界が同じ処理を実行するなら、計算する接触の個数も全世界で同じになる。世界ごとに別の処理へ進めるなら、個数を世界ごとに変えられる。

場面に物体を足すと JAX 実装の計算が増える、ということと矛盾しない実測は、MJX の文書にある。humanoid（人型のモデル）の数を1体から10体まで増やすと、場面で起こりうる接触の数が増える。このとき JAX 実装のスループット（単位時間あたりに進むシミュレーションのステップ数）は、CPU で動く MuJoCo より急に落ちる1。文書はこの低下を、GPU や TPU が条件分岐を苦手とすることに結び付けている。衝突しうる物体の組を絞り込む処理は条件分岐を使うので、JAX 実装はこの処理を CPU の MuJoCo ほどうまく行えない1。この実測は、どの世界も同じ体数の場面で測っている。接触の少ない世界が混ざる場面を測ったものではない。

JAX 実装の利用者は、衝突を調べる物体の組を指定して、生成する個数を減らせる1。

起こりうる接触が少ない場面では二つの実装の差が小さい、ということと矛盾しない記述は、Playground の告知にある。JAX 実装は、衝突が少数で固定の場面や、球や箱のような単純な形どうしの衝突なら性能が良い。Warp 実装は、中規模から大規模の場面でよりよくスケールする2。Playground の操作環境の多くでは、Warp 実装のスループットが1.5〜2倍になった2。MJWarp の文書も、ほぼすべてのワークロードで MJX よりよくスケールし、複雑な場面で特にそうだと書く3。MJX の文書は、小さな場面で性能が出て勾配におおむね対応するシミュレータを求める利用者には JAX 実装がよい選択だとし、それ以外の利用者を Warp 実装へ向ける1。

作り手が挙げる Warp 実装の利点は、接触の個数だけではない。MJX の文書は、メッシュ（三角形の網で表した物体の表面）どうしの衝突に Warp 実装が全面的に対応している点も挙げる1。Playground の告知が Warp 実装を勧めるのも、メッシュどうしの衝突が中程度から多数ある環境である2。本稿は、このうち接触の個数だけを説明している。

JAX 実装から Warp 実装へ切り替えたときの違いも、告知は書いている。報酬の挙動は、JAX 実装と比べて「良好な一致」を得た。ただし PandPick 系と AlohaSinglePeg の環境では、報酬にわずかな回帰が残ると断っている2。Playground の Warp 対応はベータである2。

Warp実装の接触の上限は、世界どうしで融通できる

Warp 実装でも、接触の個数に上限を決める。GPU のメモリを前もって確保する必要があるからである。文書は「メモリと計算量はこれらのパラメータの値に応じて増える」と書き、上限をできるだけ小さく設定するよう勧める3。上限が費用になる点は、JAX 実装と Warp 実装に共通である。

JAX 実装との違いは、接触の上限を全世界の合計に対して置く点にある。文書は、全世界の接触の合計が nworld × nconmax を超えなければ、一つの世界の接触の数は nconmax を超えてよいと書く3。nworld は世界の数、nconmax は世界あたりの接触の見込みの数である。接触の上限は、全世界の合計の nworld × nconmax になる。MJX から Warp 実装を呼ぶときは、この合計を naconmax として直接渡す1。

そのため、接触の多い世界は、世界あたりの見込みの nconmax を超えて、接触の少ない世界が使わなかった枠を使える。JAX 実装では、どの世界にも、起こりうる接触を収められる枠が要る。

融通できるのは接触の枠だけである。シミュレータは接触ごとに、物体がめり込まないという条件を立てる。関節の可動域の限界なども、同じように条件として立てる。これらの条件をまとめて拘束（constraint）と呼ぶ。拘束の上限 njmax は世界ごとに決まっていて、ほかの世界の枠は使えない3。接触の多い世界の拘束を収められる njmax を決めると、同じ大きさの枠が接触の少ない世界にも確保される。メモリと計算量が枠の大きさに応じて増える点は、JAX 実装と同じである3。ただし、生成する拘束の個数そのものは、Warp 実装では世界ごとに変わる2。接触と拘束で扱いが違う理由を、文書は書いていない。

上限を超えると未定義動作になり、超過は世界ごとに Data.overflow に記録される3。新しい場面では、ビューアで読み込み、超過が出るたびに上限を上げて調整する1。

Warp 実装には、動いていない物体を計算から外す機構もある。スリープは「静止した物体を眠らせる」機構である3。場面の大きさは自由度で測る。自由度は、関節や物体が独立に動ける方向の数である。文書は、スリープを有効にし、場面を休止中の物体を含む独立した部分に分割できる限り、MJWarp は数百自由度の場面までスケールできると書く3。

Warp実装が高速に動くのはNVIDIA GPUだけで、自動微分に対応していない

Warp 実装には、制約が五つある。最初の二つは JAX 実装に無い制約で、Warp 実装を使えるかどうかを決める。次の二つは CPU の MuJoCo と比べたときの制約で、最後の一つは Warp を直接呼ぶときと比べたときの制約である。この三つは、用途によって問題になる。

高速に動くのは NVIDIA GPU だけである。 JAX 実装は、NVIDIA と AMD の GPU、Apple Silicon、Google Cloud TPU で動く1。JAX と、JAX が使うコンパイラの XLA の上に書かれているからである1。MJWarp の README は「高速なシミュレーションには NVIDIA GPU を要するが、開発とデバッグのために CPU にも対応する」と書く4。Warp が GPU での高速実行に使うのは NVIDIA の CUDA で、macOS 向けの配布は CPU での実行だけに対応する5。

自動微分に対応していない。 自動微分は、シミュレーションの出力を入力で微分した値（勾配）を自動で求める機能である。JAX 実装は、自動微分に「おおむね対応している」1。Warp 実装について、MJX の文書は「MJX-JAX と異なり、MJX-Warp は自動微分に対応せず、対応する当面の計画もない」と書く1。ただし、リポジトリでは実装が進んでいる。自動微分を扱う issue #500 で、NVIDIA で Warp と物理エンジンの Newton を担当するプロダクトマネージャが2026年8月4日に「この機能要望には積極的に取り組んでいる」と書き、実装の PR #1535 を示した。この PR は2026年9月19日時点で未マージである6。Warp 自体のカーネルは微分可能である5。この制約が続くかどうかは、まだ決まっていない。

1ステップにかかる時間は短くならない。 MJWarp の文書は「MJWarp はスループット、すなわち単位時間あたりのシミュレーションステップの総数に最適化されており、MuJoCo は遅延、すなわち一ステップにかかる時間に最適化されている」と書く3。モデル予測制御やテレオペレーションのように、1ステップの遅延が問題になる用途には、文書は CPU の MuJoCo を勧める3。JAX 実装も、一つの場面だけを進める用途には向かない。MJX の文書は、その場合に JAX 実装が CPU の MuJoCo より10倍遅くなることがあると書く1。

一続きの大きな機構では、CPU の MuJoCo ほど伸びない。 一続きの機構は、人型ロボットの全身のように、関節でつながった物体の集まりである。文書は、物体や自由度の多い場面では MJWarp が JAX 実装よりよくスケールするが、一続きの大きな機構では CPU の MuJoCo ほどはスケールしないと書く3。約60自由度を超える一続きの機構の性能改善を、文書は優先課題に挙げている3。

JAX から呼ぶと、直接呼ぶより遅くなる。 MJX の文書の表では、Humanoid の場面で、Warp を直接呼ぶと毎秒3.35M（M は百万）ステップ、JAX から呼ぶと2.96M である。Aloha Pot の場面では、2.45M と2.33M である1。低下は Humanoid で一割強、Aloha Pot で約5%である。

出典6件

  1. MuJoCo Documentation「MuJoCo XLA (MJX)」2026年9月19日取得. https://mujoco.readthedocs.io/en/stable/mjx.html — JAX 実装の対応ハードウェアと自動微分、Warp 実装の NVIDIA 特化と自動微分未対応、humanoid を増やしたときの JAX 実装の低下、毎秒ステップ数の表を載せる公式文書。 ↩ ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8 ↩9 ↩10 ↩11 ↩12 ↩13 ↩14 ↩15

  2. btaba「MuJoCo Playground works with MuJoCo Warp! (in Beta)」Discussion #197, 2025年8月28日. https://github.com/google-deepmind/mujoco_playground/discussions/197 — JAX は固定数、Warp はステップごとに可変数の接触を生成できるとし、操作環境の多くで1.5〜2倍のスループットと書く。 ↩ ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8 ↩9 ↩10 ↩11

  3. MuJoCo Documentation「MuJoCo Warp (MJWarp)」2026年9月19日取得. https://mujoco.readthedocs.io/en/latest/mjwarp/index.html — 接触の上限は全世界の合計で置け、拘束の上限 njmax は世界ごと、超過は未定義動作。スループット向けで、一続きの大きな機構では CPU 版ほど伸びないと書く公式文書。 ↩ ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8 ↩9 ↩10 ↩11 ↩12 ↩13

  4. google-deepmind/mujoco_warp README（GitHub）2026年9月19日取得. https://github.com/google-deepmind/mujoco_warp — NVIDIA 向けの MuJoCo の GPU 版で、Google DeepMind と NVIDIA が保守し、高速実行には NVIDIA GPU を要し、Warp 経由の微分可能性はまだ無いと書く。 ↩ ↩2 ↩3

  5. NVIDIA「Warp」README（GitHub）2026年9月19日時点. https://github.com/NVIDIA/warp — GPU で実行するカーネルを書く枠組みの説明書で、GPU での高速実行は CUDA によること、Warp のカーネルは微分可能で PyTorch・JAX などと組み合わせて使えることを書く。 ↩ ↩2 ↩3

  6. google-deepmind/mujoco_warp Issue #500「Differentiability」2025年7月13日作成. https://github.com/google-deepmind/mujoco_warp/issues/500 — NVIDIA の Warp・Newton 担当プロダクトマネージャが2026年8月4日に取り組み中と書き PR #1535 を示した。PR は2026年9月19日時点で未マージ。 https://github.com/google-deepmind/mujoco_warp/pull/1535 ↩

この記事はAIを使用して執筆されました。内容には誤りが含まれる可能性があります。事実や数字の出典は脚注から確かめられます。