NVIDIA Warp 和 MuJoCo Warp (MJWarp) 讓機器人實現 2,048 個世界之 GPU 模擬

TL;DR

NVIDIA Warp 及其與 MuJoCo 兼容的層 MJWarp 可讓您在單一 GPU 上批量處理數千個機器人模擬世界，實現比傳統 CPU-based MuJoCo 高出數個數量級的總吞吐量。

系統概覽

層級 功能
NVIDIA Warp 編譯為 CUDA 的 Python 核心語言，提供自動微分，並與 PyTorch/JAX 互操作
MJWarp 基於 Warp 的 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 核心，可即時編譯為原生 CUDA。以下範例在重力作用下整合點位置，從兩個點擴展至百萬點，無需改變核心邏輯。

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 不保證單一世界的步進更快；其優勢在於 總吞吐量：當多個世界並行運行時，每秒完成的世界步數總和。這是強化學習與大規模採樣的關鍵指標。

將 MuJoCo 場景遷移至 MJWarp

本指南介紹三個階段：建立 CPU 基準、驗證單一世界 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. 驗證單一世界 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 變體）來說明額外的整合步驟。

參考資料

  • 物理 AI 模擬現況 – 系列首篇部落格。
  • NVIDIA Warp GitHub 與文件（v1.15.0 新增 GPU 確定性）。
  • MuJoCo Warp 倉儲與官方文件。
  • mjlab（arXiv:2601.22074）與 MuJoCo Playground。
  • NVIDIA SO‑101 模擬到真實學習路徑。
  • 即將推出的 Newton 倉儲。

Sources