NVIDIA 工程師詳解運用 Warp 實現 MuJoCo 機器人模擬的 GPU 規模化擴展

將機器人模擬從單一 CPU 世界擴展至龐大的 GPU 批次處理

2026 年 9 月 23 日，NVIDIA 研究員 Johnny Nuñez Cano、Asier Arranz、Rishabh Chadha 與 Ben Oliveri 發表了一份實作指南，展示機器人開發者如何將傳統基於 CPU 的 MuJoCo 模擬轉移至高吞吐量的 GPU 工作負載。隨著實體 AI（Physical AI）與強化學習管線的不斷擴展，訓練速度日益取決於同時執行數百乃至數千個環境，而非僅僅最佳化單一世界的延遲。

為了弭平此落差，NVIDIA 提出了一套運用 MuJoCo Warp（MJWarp）的工作流程，這是一個建構在 NVIDIA Warp 之上、支援 GPU 加速的 MuJoCo 物理引擎實作版本。作為 NVIDIA「實體 AI 模擬現況」（State of Simulation for Physical AI）系列的第二篇專題，該指南示範了如何將執行積木堆疊任務的 SO-101 從動機械手臂，從標準的 Python 驅動 CPU 環境直接遷移至 NVIDIA GPU 上的 2,048 個平行模擬狀態。

移植 SO-101 機械手臂：緩衝區大小設定、CUDA Graphs 與一致性驗證

這套架構的核心基礎是 NVIDIA Warp，這是一個以 Python 為基礎的框架，能將靜態型別的 Python 核心（kernels）編譯為原生 CUDA 程式碼，並支援即時編譯（JIT）、核心融合（kernel fusion）、透過 wp.Tape 實現的自動微分，以及在 1.15 版本中引入的確定性執行（deterministic execution）模式。MJWarp 將此框架應用於 MuJoCo 物理引擎，在處理標準 MJCF XML 場景描述的同時，將陣列轉移至 GPU 並加上前置的批次維度（leading batch dimension）。

遷移工作從 SO-101 的夾取放置（pick-and-place）任務開始，設定控制頻率為 50 Hz，每影格進行 10 個物理子步驟，目標物理時間步長為 0.002 秒。該機械手臂的任務是抓取一個 44 毫米的紅色立方體，並放置在同樣大小的藍色立方體上方。將此工作流程轉移至 MJWarp 涉及 API 的切換：開發者使用 mjw.put_model() 上傳模型，透過 mjw.make_data() 或 mjw.put_data() 配置駐留於 GPU 上的狀態，並經由 mjw.step() 同步推進所有平行世界的進度。

一項關鍵的運作要求是管理裝置端的接觸（contact）與約束（constraint）緩衝區。開發者必須定義各項容量上限，例如 nconmax（每個環境的最大接觸數）、naconmax（全域接觸上限）以及 njmax（每個環境的最大約束數）。由於窄相位（narrowphase）碰撞超過容量上限時，MJWarp 僅會發出警告而不會立即終止執行，若未經檢查，執行結果可能面臨默默失效的風險。NVIDIA 建議在任務中接觸最繁雜的時刻（例如夾爪兩側與桌面同時接觸到立方體時）來估算這些上限值，並利用 mjwarp-testspeed --measure_alloc 等工具來診斷記憶體配置。

一旦單一世界的物理行為經過與 CPU 基準的一致性（parity）驗證後，環境便能擴展至 2,048 個或更多世界。為了消除跨數千個批次步驟的調度延遲（dispatch latency），開發者可將 mjw.step() 包裹在 wp.ScopedCapture() 中，以建立可重複播放的 CUDA Graphs。效能分析時必須考量 GPU 的非同步執行特性，在計時迴圈前後插入 wp.synchronize()，並在推展步驟（rollout steps）中將模擬數據保留在裝置端，而非透過 .numpy() 將陣列拉回主機端記憶體。

拓展實體 AI 模擬技術堆疊：邁向 Newton 與 Isaac Lab

轉向 MJWarp 突顯了業界在「單一機器人模型預測控制（MPC）」與「批次強化學習」之間日益明顯的分工。傳統 CPU 版 MuJoCo 仍是用於遠距遙控操作（teleoperation）、單一世界除錯與低延遲 MPC 的標準工具，而批次 GPU 環境則專門鎖定海量數據的收集需求。

MJWarp 旨在與多種現代實體 AI 框架無縫串接。它可作為 mjlab（直接結合 MJWarp 與 PyTorch）的執行後端，亦可透過 MJX 的 Warp 實作支援 MuJoCo Playground 及更廣泛的學習管線。NVIDIA 表示，該模擬系列的下一篇將展示如何將這個相同的 SO-101 環境匯入其多解算器框架 Newton 中；屆時 MJWarp 將作為底層剛體解算器（newton.solvers.SolverMuJoCo），將場景直接串接至 Isaac Lab 的訓練工作流程。