NVIDIA Warp と MuJoCo Warp (MJWarp) でロボット用 2,048 ワールド GPU シミュレーションを実現

TL;DR

NVIDIA Warp とその MuJoCo 互換レイヤーである MJWarp を使用すると、1 つの GPU 上で数千ものロボットシミュレーションワールドをバッチ処理でき、従来の CPU 基準 MuJoCo と比べて桁違いに高い集約スループットを達成できます。

スタック概要

レイヤー 役割
NVIDIA Warp CUDA にコンパイルされる Python カーネル言語。自動微分を提供し、PyTorch/JAX と相互運用可能
MJWarp MuJoCo 物理の GPU 実装。MJCF 形式を保持し、バッチステップをサポート
あなたのシーン (SO‑101) 標準的な Menagerie / Robot Studio アセットとタスクジオメトリ
Next (Newton / Isaac Lab) マルチソルバー API、USD アセット、センサー、マネージャー、トレーニングループ

各コンポーネントの使用タイミング

必要な機能 使用するもの
単一ロボットの MPC や遠隔操作 従来の MuJoCo CPU
最大の MuJoCo 物理スループット MJWarp (または mjlab)
JAX をベースにしたトレーニングレシピ MuJoCo Playground / MJX (実装 = warp)
マルチソルバー + Isaac Lab 統合 Newton (今後公開予定)

NVIDIA Warp: 最小限の GPU カーネル

Warp を使用すると、静的に型指定された Python カーネルを JIT コンパイルしてネイティブ CUDA に変換できます。以下の例では、重力下での点の位置を統合しており、2 点から数百万点まで、カーネルロジックを変更せずにスケーリング可能です。

import numpy as np, warp as wp

@wp.kernel
def integrate(positions: wp.array[wp.vec3],
                  velocities: wp.array[wp.vec3],
                  dt: float):
    i = wp.tid()
    velocities[i] += wp.vec3(0.0, 0.0, -9.81) * dt
    positions[i] += velocities[i] * dt

wp.init()
device = "cuda:0" if wp.is_cuda_available() else "cpu"
start = np.array([[0.0, 0.0, 0.5], [0.2, 0.0, 0.5]], dtype=np.float32)
positions = wp.array(start, dtype=wp.vec3, device=device)
velocities = wp.zeros_like(positions)

wp.launch(integrate, dim=len(start), inputs=[positions, velocities, 0.01], device=device)
wp.synchronize_device(device)
print(positions.numpy())

ロボット工学における重要性

  • 明示的な並列処理 – wp.tid() は論理スレッド（点、接触、剛体、またはワールド）を識別します。
  • 明示的なデバイス配列 – 配列は選択されたデバイス上に存在します。.numpy() はホストへのコピーをトリガーし、DLPack アダプタにより PyTorch/JAX とのゼロコピー共有が可能になります。
  • 合成可能な起動 – 複数のカーネルを CUDA グラフにキャプチャでき、起動オーバーヘッドを削減できます。

微分可能性と決定性（Warp の機能）

  • 微分可能なカーネル – wp.Tape は前方起動を記録し、backward() で自動的に随伴を計算します。これにより、微分可能なジオメトリ、CFD、カスタム物理パイプラインが可能になります。
  • 決定的実行 – Warp 1.15 で導入。決定的モードはわずかなパフォーマンスの低下を犠牲に、再現可能なアトミック順序を提供し、レグレッションテストに有用です。これらの機能はオプションであり、SO‑101 のウォークスルーには必須ではありません。

MuJoCo Warp (MJWarp) とは？

MJWarp は Warp 上に構築された MuJoCo 物理パイプラインの GPU ベース実装です。同じ MJCF モデルを受け入れ、モデルと複数の独立した状態を GPU に配置し、単一の mjw.step 呼び出しで全バッチを進めることが可能です。

  • レイテンシ vs スループット – MJWarp は単一ワールドステップの高速化を保証しません。その利点は 集約スループット にあります。すなわち、多数のワールドを並列で実行した場合の 1 秒あたりのワールドステップ総数です。これは強化学習や大規模サンプリングの主要な指標です。

MuJoCo シーンを MJWarp に移行する方法

ガイドでは、3 つの段階を経て説明します。CPU ベースラインの確立、1 ワールド GPU ロールアウトの検証、そして 2,048 ワールドへのスケーリングです。

1. MuJoCo CPU ベースラインの確立

  • mujoco.MjModel.from_xml_path を使用して、SO‑101 ピックアンドプレース MJCF（標準的な Menagerie アセット）を読み込みます。
  • 50 Hz の制御ループと 10 回の物理サブステップ（timestep = 0.002 s）を使用します。
  • キューブが安定した後、水平誤差 ≤ 0.015 m かつ垂直方向の分離が 0.035 m から 0.055 m の範囲内であることを確認して、タスクの成功を検証します。

2. 1 ワールド MJWarp の同等性検証

import warp as wp, mujoco_warp as mjw
wp.init()
device = wp.get_device()

# モデルのアップロードと単一ワールドバッチの割り当て
m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=1, nconmax=spec.nconmax, njmax=spec.njmax)

# CPU ロールアウトからの状態を初期化
wp.copy(d.qpos, wp.array(mjd.qpos[None, :], dtype=wp.float32, device=device))
wp.copy(d.qvel, wp.array(mjd.qvel[None, :], dtype=wp.float32, device=device))
wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device))

mjw.forward(m, d)
  • ホストループは CPU ワークフローを模倣しており、各サブステップでコントロールをデバイスにコピーし、qpos/qvel を後続計算のために取得します。
  • このパスはパフォーマンスベンチマークではありません。GPU シミュレーションが CPU 結果を再現していることを検証するためのものです。

3. 接触および制約バッファのサイズ調整

  • MJWarp は各ワールドに対して接触（nconmax）および制約（njmax）バッファを事前に割り当てます。容量が不足するとオーバーフロー警告が発生し、トラジェクトリが破損する可能性があります。
  • mjwarp-testspeed --measure_alloc を使用して実際の使用量を測定し、必要に応じて制限を引き上げます（例：SO‑101 タスクでは nconmax=128、njmax=300）。

4. 2,048 ワールドへのスケーリング

nworld = 2_048
d = mjw.make_data(m, nworld=nworld, nconmax=spec.nconmax, njmax=spec.njmax)

# 初期状態をバッチ全体に複製
wp.copy(d.qpos, wp.array(np.tile(mjd.qpos, (nworld, 1)), dtype=wp.float32, device=device))
wp.copy(d.qvel, wp.array(np.tile(mjd.qvel, (nworld, 1)), dtype=wp.float32, device=device))
wp.copy(d.ctrl, wp.array(np.tile(mjd.ctrl, (nworld, 1)), dtype=wp.float32, device=device))

mjw.forward(m, d)

# ステップを CUDA グラフにキャプチャして低オーバーヘッドの再実行を可能に
with wp.ScopedCapture() as cap:
    mjw.step(m, d)
step_graph = cap.graph
  • キャプチャされたグラフは同じバッファを再利用します。再実行間で更新が必要なのは d.ctrl のみです。

5. 集約スループットの測定

# ワームアップ（コンパイル、割り当て）
for _ in range(10):
    wp.capture_launch(step_graph)
wp.synchronize()


t0 = time.perf_counter()
for _ in range(200):
    wp.capture_launch(step_graph)
wp.synchronize()
elapsed = time.perf_counter() - t0
world_steps = 200 * nworld
print(f"{world_steps/elapsed:,.0f} world‑steps/second")
  • バッチサイズとともに、ワールドステップ/秒 と バッチステップあたりのミリ秒 の両方を報告してください。
  • 提供された scaling_study.py スクリプトを使用して、バッチサイズ（1、64、1,024、2,048、8,192）をスイープし、スループットが飽和するポイントを観察してください。

はじめ方

  • Warp – pip install warp-lang（決定的モードには ≥ 1.15 必須）。python -m warp.examples.browse で例を探索。
  • MJWarp – pip install mujoco-warp；mjwarp-viewer path/to/scene.xml でビューアを起動。
  • SO‑101 リソース – NVIDIA の Sim‑to‑Real コースと Menagerie アセット。
  • MJWarp 上でのトレーニング – mjlab、MuJoCo Playground（MJX impl='warp'）、または近日公開予定の Newton 統合による Isaac Lab。

次回予告

次回の投稿では、MJWarp を Newton 物理エンジン内に統合し、マルチフォーマットアセット、交換可能なソルバー、センサー、および Isaac Lab トレーニングパイプラインを追加します。同じ SO‑101 タスク（およびオプションの reBot バリアント）を再利用して、追加の統合ステップを説明します。

参考文献

  • Physical AI におけるシミュレーションの現状 – このシリーズの最初のブログ。
  • NVIDIA Warp GitHub とドキュメント（v1.15.0 で GPU の決定性が追加）。
  • MuJoCo Warp リポジトリと公式ドキュメント。
  • mjlab（arXiv:2601.22074）と MuJoCo Playground。
  • NVIDIA SO‑101 シミュレーションから実世界への学習パス。
  • 今後公開予定の Newton リポジトリ。

Sources