NVIDIA Warp と MjWarp でロボティクスシミュレーションを GPU スケールへ移行する方法

記事のサマリー（TL;DR）

NVIDIA Warp は Python で書いた GPU カーネルを JIT コンパイルし、MJWarp は MuJoCo の物理演算をその上に実装して最大数千並列環境を 1 回の mjw.step で前進させる
SO-101 フォロワーアームを CPU の MuJoCo から 2,048 並列 MJWarp 環境へ移行する手順は「モデルアップロード → 1 世界での等価性検証 → バッファ容量確認 → バッチスケール」の 4 段階
CUDA Graph キャプチャとバッチサイズの適切な設定（nconmax・njmax）が GPU スループット計測の正確性を左右する

強化学習・ロボット制御向け並列シミュレーション環境を構築する際の実装ポイント

MJWarp が解決するのは「1 世界を速く回す」課題ではなく「何世界を同時に回せるか」という集計スループットの課題です。強化学習における経験収集フェーズでは、低遅延より大量サンプルのほうが重要であり、GPU デバイス上にシミュレーションと学習データを同居させることでPCIeバスのボトルネックを回避できます。日本でも自律ロボット・産業用マニピュレータ向けにIsaac Lab や MuJoCo Playground を利用したsim-to-real研究が活発化しており、MJWarp はその物理バックエンドとして直接利用可能です。kintone や Salesforce のような業務SaaSとの接続とは異なり、このスタックはNVIDIA GPU + Python 環境さえあればオンプレ・クラウドどちらでも構築できます。次回の Newton / Isaac Lab 記事で示されるマルチソルバー API や USD 対応に備えて、まず MJWarp ベースラインを固めておくことが移行コストの最小化につながります。

詳細

NVIDIA Warp とは何か

NVIDIA Warp は、Python で高性能 GPU カーネルを記述するためのフレームワークです。開発者は静的型付きカーネルを Python で書くだけで、CPU または CUDA 向けにコンパイルされます。最初の起動時にネイティブモジュルをビルドしてキャッシュし、以降の起動ではそれを再利用します。カーネル言語はパフォーマンス指向の Python サブセットであり、通常の Python が設定・メモリ確保・ディスパッチの統制を担います。

以下の小さなカーネルは、重力下で点の位置を前進させるロボティクス向けのサンプルです。1 論理スレッドが 1 点を担当するため、2 点から数百万点まで GPU 用語をコード中に持ち込まずにスケールします。

import numpy as np
import 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())

Warp の 3 つの価値軸は以下のとおりです。

柱	得られるもの
パフォーマンス	JIT コンパイル・カーネルフュージョン・CUDA Graph によるネイティブ CUDA 速度
使いやすさ	純 Python 記述、ベクトル・行列・クォータニオン・BVH・ハッシュグリッド・スパース行列・タイルプリミティブを内蔵
機能性	微分可能カーネルと DLPack 方式インターオプにより、シミュレーションを ML 学習ループ内に組み込める

ロボティクスで特に重要な 3 つの特性があります。

明示的な並列作業：wp.tid() で現在の論理スレッドが担当する点・接触・ボディ・世界を識別します。

明示的なデバイス配列：配列は選択したデバイス上に存在します。CUDA 配列に .numpy() を呼ぶと CPU メモリへの同期コピーが発生します。デバイス常駐の PyTorch・JAX パイプラインでは、Warp のフレームワークアダプタまたは DLPack 互換共有を使用してください。ゼロコピーパスではありません。

カーネル起動の合成：一連のカーネルを連続起動し、CUDA Graph にキャプチャして繰り返しのディスパッチオーバーヘッドを削減できます。Graph キャプチャは既存バッファに対してリプレイしますが、任意カーネルのフュージョンは行いません。

微分可能性と決定論的実行

wp.Tape でフォワードカーネル起動を記録し、backward() を呼ぶとアジョイントをリバース再生します。Warp 1.15 で導入された決定論的実行モードは、GPU アトミック命令のスケジューラ依存性をオプトインで抑制し、シミュレーション・検証・回帰テストでの再現性を保証します。ただしこれらは Warp 自体の機能であり、MJWarp ロールアウト全体の微分可能性や決定論を保証するものではありません。

インストール：pip install warp-lang（GPU 決定論は ≥ 1.15）、python -m warp.examples.browse でサンプルを確認できます。

MuJoCo Warp（MJWarp）とは何か

ロボットシミュレータは現在の関節位置・速度・制御・接触を受け取り、1 タイムステップだけ前進させる演算を繰り返します。ここでいう「世界（world）」とはそのシーンと状態の独立した 1 コピーです。SO-101 アームがキューブに手を伸ばす世界、わずかに異なるポーズから始まる同じアームの世界、というようにそれぞれ独立します。

MuJoCo と MJWarp の役割分担は明確です。

MuJoCo：1 ～少数の CPU 世界の開発・検査に適しています。
MJWarp：NVIDIA Warp 上に MuJoCo の物理パイプラインを実装し、モデルとバッチ状態を GPU に配置します。mjw.step 1 回でバッチ全体を前進させます。

MJWarp の価値は「1 世界の 1 ステップを速くすること」ではありません。数百〜数千の世界を同時に前進させることで GPU に十分な並列作業を与え、集計スループット（total world-steps/秒） を向上させることです。これは経験収集が重要な強化学習や大規模サンプリングに向いています。

判断の目安：

目的	選択肢
単体ロボット MPC / テレオペ	MuJoCo CPU
生の MuJoCo 物理での最大スループット	MJWarp（または mjlab）
JAX 学習レシピ	MuJoCo Playground / MJX（`impl='warp'`）
マルチソルバー + Isaac Lab 統合	Newton（本シリーズ次稿）

コア API の移行対照表

MuJoCo ホストワークフロー	MJWarp ワークフロー
`mujoco.MjModel`	`mjw.put_model(mjm)` でデバイスモデルを作成
`mujoco.MjData`	`mjw.put_data(mjm, mjd, ...)` で既存状態を保持・バッチ化
`mujoco.mj_step(mjm, mjd)`	`mjw.step(m, d)` で `d` 内の全世界を前進
`mjd.ctrl` などのホスト配列	`d.ctrl`（形状 `(nworld, nu)`）などのバッチデバイス配列

デフォルト状態が目的なら mjw.make_data()、初期化済み MuJoCo 状態をそのまま移行するなら mjw.put_data() を使います。

バッチリソースのアロケーションに必要なパラメータは以下のとおりです。

パラメータ	意味
`nworld`	並列環境の総数
`nconmax`	個別世界ごとの想定接触数（全体容量 ≈ nconmax × nworld）
`naconmax`	全環境合計の最大接触数（両方定義した場合はこちら優先）
`njmax`	世界ごとの制約数の上限

MuJoCo シーンを MJWarp へ移行するワークフロー

ステップ 1：MuJoCo CPU ベースラインの確立

SO-101 アーム・テーブル・2 つのキューブを MJCF で記述します。タスクは赤い 44mm キューブ（size="0.022 …" は半径値）をつかんで青いキューブの上に積み上げることです。

<mujoco model="so101_pick_place">
  <include file="so101.xml"/>
  <worldbody>
    <light pos="0.3 0 1.5" dir="0 0 -1" directional="true"/>
    <geom name="floor" type="plane" size="0 0 0.05"/>
    <geom name="table" type="box" pos="0.35 -0.04 0.012"
          size="0.16 0.26 0.012" rgba="0.32 0.32 0.32 1"
          friction="1 0.005 0.0005" condim="3"/>
    <body name="red_cube" pos="0.33 -0.13 0.046">
      <freejoint name="red_cube_joint"/>
      <geom type="box" size="0.022 0.022 0.022" mass="0.08"
            rgba="0.85 0.05 0.04 1" friction="1.2 0.005 0.0005" condim="3"/>
    </body>
    <body name="blue_cube" pos="0.33 0.06 0.046">
      <freejoint name="blue_cube_joint"/>
      <geom type="box" size="0.022 0.022 0.022" mass="0.08"
            rgba="0.05 0.20 0.90 1" friction="1.2 0.005 0.0005" condim="3"/>
    </body>
  </worldbody>
</mujoco>

ローディングと基本ループは通常の MuJoCo です。50 Hz 制御・10 サブステップで物理タイムステップを 0.002 秒に設定します。この値は mjw.put_model 前に設定しておく必要があります。

import mujoco
mjm = mujoco.MjModel.from_xml_path("scene_pick_place.xml")
mjd = mujoco.MjData(mjm)
fps = 50
sim_substeps = 10
frame_dt = 1.0 / fps
mjm.opt.timestep = frame_dt / sim_substeps  # 0.002 s

for _ in range(600):
    ctrl = controller.step(mjm, mjd, frame_dt)
    for _ in range(sim_substeps):
        mjd.ctrl[:mjm.nu] = ctrl
        mujoco.mj_step(mjm, mjd)

成功判定は 2 条件です。水平中心誤差 xy_err ≤ 0.015 m、垂直間隔 0.035 m ≤ dz ≤ 0.055 m。これが後続の等価性検証で使う基準値になります。

ステップ 2：1 世界での MJWarp 等価性検証

まず GPU 上で 1 世界を動かし、ホストをループ内に残したまま同じタスクを同じビューアで確認します。

wp.init()
import mujoco_warp as mjw
device = wp.get_device()
m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=1, nconmax=spec.nconmax, njmax=spec.njmax)
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)

デバイス配列は先頭に世界次元を持ちます。mjd.qpos[None, :] でシェイプ (1, nq) になるのはそのためです。後でスケールアップしても先頭次元が変わるだけで呼び出しは変わりません。mjw.put_model() は非対応機能があればサイレントに無視するのではなく例外を投げるため、互換性チェックとして機能します。

フレームループではステップを GPU にリダイレクトし、結果をホストにミラーします。

def simulate_frame() -> None:
    ctrl = controller.step(mjm, mjd, frame_dt)
    for _ in range(sim_substeps):
        mjd.ctrl[:mjm.nu] = ctrl
        wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device))
        mjw.step(m, d)
        mjd.qpos[:] = d.qpos.numpy()[0]
        mjd.qvel[:] = d.qvel.numpy()[0]
        mujoco.mj_forward(mjm, mjd)

.numpy() の読み取りはサブステップごとに同期・コピーを発生させるため、これはタスク検証パスであり、スループットベンチマークではありません。d.qpos・d.qvel をコピーした後に mujoco.mj_forward(mjm, mjd) を呼んで mjd.xpos などの派生量を更新する必要があります。

ステップ 3：接触・制約バッファ容量の確認

MJWarp は mjw.step 前に接触・制約バッファを確保します。容量超過が発生すると、例外ではなく警告が出て Data.overflow フラグが立ちます。SO-101 プロファイルの初期値は nconmax=128・njmax=300 です。容量は「アーム・テーブル・キューブが同時に触れる最接触瞬間」で測定します。

mjwarp-testspeed --measure_alloc でシーンが実際に消費した接触数・制約数を計測し、超過があればロールアウト前に上限を引き上げてから再測定します。モデル・衝突ジオメトリ・タスクを変更するたびに再確認が必要です。

ステップ 4：2,048 世界へスケール

1 世界の等価性が確認できたら、バッチサイズを目標値に拡大し、初期状態を全世界にレプリケートします。

nworld = 2_048
d = mjw.make_data(mjm, 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)

with wp.ScopedCapture() as capture:
    mjw.step(m, d)
step_graph = capture.graph

np.tile で全世界に同一の初期状態を与えるのはスループット計測のベースラインとして正しい選択です。世界ごとのランダム化が必要な場合はデバイス上の d.qpos の各行に直接書き込みます。CUDA Graph はキャプチャ時のモデル・データバッファを再利用します。バッファ置き換え・nworld 変更・モデル再構築後は新たにキャプチャし直してください。

正確な計測方法：ウォームアップと同期の重要性

GPU 起動は非同期です。素朴なタイマーは Python がキューに積む速度を計測するに過ぎず、GPU が処理を完了した速度ではありません。ウォームアップ（カーネルコンパイル・アロケーション・キャッシュ）を済ませてから計測区間の前後で wp.synchronize() を呼びます。

import time
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
total = 200 * nworld
print(f"{total / elapsed:,.0f} world-steps/second")

バッチサイズとともに「集計 world-steps/秒」と「バッチステップあたりのミリ秒」の両方を報告します。1 世界のレイテンシ比較はバッチスループットを示しません。スケーリング曲線の測定には scaling_study.py（--worlds 1 64 1024 2048 8192 --steps 100）が利用できます。

スタック全体のレイヤー構成

レイヤー	スタック内の役割
NVIDIA Warp	Python カーネル言語：SIMT・自動微分・PyTorch/JAX インターオプ
MJWarp	Warp 上の MuJoCo 物理：同一 MJCF・バッチ GPU スループット
あなたのシーン（SO-101）	Menagerie / Robot Studio アセット + タスクジオメトリ
次（Newton / Isaac Lab）	マルチソルバー API・USD・センサ・マネージャ・学習ループ

MJWarp 上でポリシーを学習するには以下を使います。

Isaac Lab（Newton 経由）
mjlab（MJWarp + PyTorch の manager API を直接利用）
MuJoCo Playground（MJX の impl='warp'）

次のステップ

次稿では同一 MJCF 環境を Newton へ移行し、newton.solvers.SolverMuJoCo として MuJoCo Warp をリジッドボディソルバーに使います。Newton が追加するもの：マルチフォーマットアセット・交換可能ソルバー・センサ/IK ヘルパー・Isaac Lab パス。同じ SO-101 タスクとオプションの reBot プロファイルで説明が続きます。

インストールコマンド早見表：

Warp：pip install warp-lang → python -m warp.examples.browse
MJWarp：pip install mujoco-warp → mjwarp-viewer benchmarks/humanoid/humanoid.xml
SO-101 サンプル：git clone https://github.com/NVIDIA/accelerated-computing-hub.git

フィードバックは GitHub のリンク先リポジトリの Issue、または Discord（NVIDIA Omniverse）で受け付けています。