傳統 MuJoCo 提供快速的 CPU 機器人模擬，用於開發、測試和控制機器人，並能透過 CPU 核心平行取樣。然而，隨著學習工作負載的增長，問題從單一世界運行速度轉變為能同時運行多少個世界。GPU 加速使得這些世界能夠以大型批次推進，同時保持模擬和學習資料靠近裝置。

MuJoCo Warp (MJWarp) 建基於 NVIDIA Warp，能將相容的 MuJoCo 模型帶入 GPU 規模的運算。在本文中，我們將把 SO-101 追隨者機械手臂從熟悉的 MuJoCo 工作流程，轉移到多達 2,048 個平行的 MJWarp 環境中，並探討實現此轉變的技術和驗證步驟。

圖 1 顯示 MJWarp 如何將 Python 連接到 GPU 模擬。MuJoCo 載入並編譯 MJCF 模型；MJWarp 則在 NVIDIA Warp 中實現物理運算，後者編譯 CUDA 核心以在 NVIDIA GPU 上推進模擬狀態。這是我們「實體 AI 模擬現況」系列文章的第二篇。

本系列第一篇文章描繪了機器人模擬的現況。在本文中，我們準備並擴展模擬環境；我們不訓練策略。後續的 Newton 和 Isaac Lab 文章將涵蓋下一層的整合。

文章概述了整個技術堆疊，其中 NVIDIA Warp 是用於編寫高效能 GPU 加速核心的 Python 框架，支援自動微分與 PyTorch/JAX 互通。MJWarp 則是在 Warp 上實現 MuJoCo 物理，提供批次 GPU 吞吐量。對於單一機器人控制，MuJoCo CPU 仍是首選；而若追求原始 MuJoCo 物理的最大吞吐量，MJWarp 或 mjlab 則是理想選擇。

NVIDIA Warp 是一個 Python 框架，用於編寫高效能、GPU 加速的核心。Warp 讓開發者能以 Python 編寫靜態型別的核心，並將其編譯為 CPU 或 CUDA 執行。首次啟動會建構並快取原生模組；後續啟動則重複使用它。

核心語言是 Python 中一個注重效能的子集，而一般的 Python 則負責配置、分配和啟動協調。這個小型、以機器人為導向的核心在重力作用下推進點位。一個邏輯執行緒處理一個點，因此相同的程式碼可以從兩個點擴展到數百萬個點，而無需在控制流程中引入 GPU 術語。

Warp 的三大價值主張為：效能、易用性和功能。效能方面，它透過即時編譯 (JIT compilation)、核心融合 (kernel fusion) 和 CUDA Graphs 提供原生 CUDA 速度。易用性方面，它提供純 Python 編寫，內建向量、矩陣、四元數、BVH、雜湊網格、稀疏矩陣和圖塊原語。

功能方面，它支援可微分核心和 DLPack 風格的互通性，使模擬能夠融入機器學習訓練迴圈中。

文章中提供了一個簡短的 Python 程式碼範例，展示如何使用 wp.kernel 裝飾器定義一個簡單的重力積分核心，並透過 wp.launch 在 GPU 上執行，以更新點位和速度。

三個特性使其在機器人領域中非常有用：明確的平行工作、明確的裝置陣列和可組合的核心啟動。wp.tid() 識別由當前邏輯執行緒擁有的點、接觸、物體或世界，實現明確的平行工作。陣列存在於選定的裝置上，透過 .numpy() 呼叫 CUDA 陣列會同步並複製到 CPU 記憶體，而非零複製路徑。

對於駐留在裝置上的 PyTorch 或 JAX 管道，應使用 Warp 的框架轉接器或 DLPack 相容的共享方式。程式可以啟動一系列專注的核心，並將支援的 CUDA 工作捕獲到圖形中，以減少重複的調度開銷。圖形捕獲會針對現有緩衝區重播啟動，但不會融合任意核心。

可微分性和確定性是 Warp 另外兩個值得了解的功能，儘管在本文的 SO-101 工作流程中並未使用。Warp 核心是可微分的：wp.Tape 會記錄在其上下文內進行的前向核心啟動，並在呼叫 backward() 時反向重播其伴隨，這也是團隊在 Warp 中建構可微分幾何、CFD 和客製化物理的原因，包括用於模擬和設計最佳化的 CAE 工作流程。

Warp 也支援確定性執行，此功能在 Warp 1.15 中引入：GPU 原子操作預設是依賴排程器的，因此重複啟動相同核心可能會略有不同。選擇性啟用的確定性模式會犧牲一些效能，以換取模擬、驗證和迴歸測試中的可重現排序。這些是 Warp 的功能，而非整個 MJWarp 運行結果的可微分性或確定性保證。

您可以透過 pip install warp-lang (版本 ≥ 1.15 支援 GPU 確定性) 試用 Warp，然後執行 python -m warp.examples.browse 或參考教學筆記本。

什麼是 MuJoCo Warp (MJWarp)？機器人模擬器會重複計算接下來會發生什麼：給定當前的關節位置、速度、控制和接觸，它會將場景推進一個小時間步。在本文中，「一個世界」指的是該場景及其狀態的一個獨立副本。一個世界可能包含 SO-101 機械手臂伸向一個方塊；另一個世界則可能包含從略微不同姿勢開始的相同手臂。

MuJoCo 和 MJWarp 可以運行相同的相容機器人和任務，但它們組織工作的方式不同。MuJoCo 自然適合開發和檢查一個或幾個 CPU 世界。MJWarp 是 MuJoCo 物理管線的 NVIDIA Warp 實現，它將模型和一批獨立狀態放置在 NVIDIA GPU 上；一次呼叫 mjw.step 即可推進整個批次。

MJWarp 的價值不一定在於單一世界更快的步進速度。它的能力在於能夠同時推進數百或數千個世界，為 GPU 提供足夠的平行工作，以提高總體吞吐量，即每秒完成的總世界步數。這有利於強化學習和大規模取樣，在這些應用中，收集經驗比最小化單一環境的延遲更重要。

本部落格涵蓋以下內容：驗證一個 MuJoCo 世界，將其轉移到 MJWarp，形成批次，驗證並正確測量它。求解器調整、雅可比表示以及專門的多 GPU 或確定性主題對於此次遷移並非必需，可以單獨討論。

因此，區別很明確：延遲是單次模擬步進的實際時間。總體吞吐量是在測量到的實際時間每秒內完成的世界步進總數。

基本用法：結構、批次大小和最小步進。核心 API 轉換很小：MuJoCo 主機工作流程中的 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 則是每個世界約束的硬性上限。

效能調優。1. CUDA 圖形捕獲：mjw.step 包含許多核心啟動；捕獲一次，重複播放：with wp.ScopedCapture() as capture: mjw.step(m, d) wp.capture_launch(capture.graph)。

2. 緊密調整 nconmax / naconmax / njmax：記憶體和工作量會隨之擴展。使用 mjwarp-testspeed --measure_alloc 進行調整，並在 mjwarp-viewer 中觀察溢位。

其他調優考量。在確定接觸和約束緩衝區大小後，測試求解器迭代限制，而無需改變任務行為。網格和 CCD 設定會增加記憶體使用；當測量的接觸計數允許時，nccdmax / naccdmax 可以減少 CCD 緩衝區分配。MJWarp 的緊湊求解器使用 MuJoCo 的 Newton 約束求解器和休眠功能，而非獨立的 Newton 物理引擎框架。

緊湊求解器和多 GPU 配置超出了本教學的範圍；請查閱 MJWarp 效能調優文件。要在 MJWarp 物理上訓練策略，可以使用 Isaac Lab 透過 Newton、mjlab (直接在 MJWarp + PyTorch 上使用管理器 API) 或 MuJoCo Playground 透過 MJX (impl='warp')。

安裝/試用：pip install mujoco-warp，然後執行 mjwarp-viewer path/to/scene.xml 或參考 Colab 教學。

將 MuJoCo 場景遷移到 MJWarp 的工作流程。

建立 MuJoCo CPU 基準。

場景。這裡還沒有任何 MJWarp 特定的內容：一個 SO-101 機械手臂、一張桌子和兩個要堆疊的方塊，以普通的 MJCF 格式編寫。圖 2 顯示了從 MuJoCo CPU 模擬渲染的 SO-101 拾取放置場景。任務是抓取紅色的 44 毫米方塊並將其堆疊在藍色方塊上；MJWarp 驗證也使用相同的機器人和場景。

文章中提供了一個 MJCF XML 程式碼片段，定義了一個名為 so101_pick_place 的 MuJoCo 模型，其中包含 SO-101 機械手臂、地板、桌子以及兩個可自由移動的紅色和藍色方塊，每個方塊邊長為 44 毫米，用於演示拾取和放置任務。

對於 MJCF 方塊，size 值是半延伸量：size="0.022 ..." 定義了一個邊長為 44 毫米的方塊。任務使用此尺寸作為其成功閾值。機械手臂的底座位於原點，其伸展方向沿 +X 軸，方塊則沿 Y 軸排列。

在配套的儲存庫中，此檔案是自動生成的，而非手寫：resolve_pick_place_scene() 將 Menagerie 機械手臂複製到 .generated/ 中，從機器人設定檔填寫桌子和方塊座標，並寫入 scene_pick_place.xml。本教學使用 SO-101 設定檔；可選的 reBot 變體將在下方描述。

載入。編譯和步進都是普通的 MuJoCo 流程。文章中提供了一個 Python 程式碼片段，展示如何載入 scene_pick_place.xml，初始化 MjModel 和 MjData，設定控制器速率和物理步進，然後在一個迴圈中執行 600 個控制幀，每個幀內多次呼叫 mujoco.mj_step 來推進模擬。

請記住這個模式：每個幀計算一次控制，然後物理模擬步進 sim_substeps 次。第二道門檻只改變內部迴圈。