mjbatch
在 CPU 上平行執行數千個 MuJoCo 模擬
授權條款
Apache-2.0
概述
mjbatch 是一個在 CPU 上平行執行數千個 MuJoCo 模擬的 Python 函式庫。C++ 執行緒池在釋放 GIL 的情況下執行步進,批次則以即時 NumPy 視圖的形式提供給 Python:bind 會提供一個涵蓋所有模擬中某個 MjData 欄位的陣列,因此寫入 ctrl 與讀取 qpos 都不需要複製,也不需要逐一模擬的迴圈。expand 對 MjModel 欄位做同樣的事,讓每個模擬都帶有自己的物理參數——不同的摩擦係數、不同的連桿長度——之後再由 set_const 重新計算衍生常數。儲存庫附有六個獨立完整的範例,涵蓋強化學習、模型預測控制、系統識別與硬體協同設計。
核心刻意保持精簡。你以單一 MjModel 與模擬數量建立一個 Batch,執行緒池預設使用所有邏輯 CPU,而 batch.step() 會推進所有模擬。讓它不只是一個 for 迴圈的關鍵,在於陣列是即時的,而非收集而來:qpos, ctrl = batch.bind("qpos"), batch.bind("ctrl") 會交給你一組視圖,控制器只需一次賦值就能寫入整個批次,而且這些視圖會隨批次步進原地更新。逐模擬的模型參數則是另一半——batch.expand("geom_friction") 會回傳一個可針對每個模擬隨機化的陣列,無需額外機制就能實現領域隨機化與系統識別。
範例正是支持它的論據,因為每個範例都是完整的方法,而非展示:以 iLQR 將倒單擺台車甩至直立、使用預測取樣的倒單擺台車模型預測控制(MPC)、以滾動時域 iLQR 追蹤參考後空翻的 G1 人形機器人、以 PPO 訓練的 Go1 四足機器人搖桿控制器、以 CEM 聯合最佳化投擲手臂的比例、齒輪與控制,以及對 Rizon 手臂慣性參數進行的阻尼 Gauss-Newton 擬合。作者表示,Go1 控制器在一台五年前的 M1 筆電上不到一分鐘就學會走路。它可從 PyPI 安裝,鎖定 mujoco 3.13.0,透過 scikit-build-core 與 nanobind 建置,並附有 CI、pytest 測試套件與擴展性基準測試。
為何重要
大規模平行模擬已成為 GPU 的天下——MJX、Isaac、Warp——這也悄悄墊高了沒有工作站顯示卡者的入門成本。mjbatch 藉由妥善運用 CPU、而非一次只用一個核心,讓同樣類型的工作也能在筆電上進行,而且它是透過原生 MuJoCo 而非重新實作來做到的,因此模型與物理都是你早已信任的那一套。對於一個讀者大多擁有筆電與低成本手臂、而非運算叢集的目錄來說,這是取捨中更實用的那一端。
解決的問題
要取得數千次 MuJoCo 推演,通常意味著得轉向 GPU 軟體堆疊,並接受其硬體開銷與對物理引擎的重新實作。mjbatch 保留原生 MuJoCo,並從你現有的 CPU 中榨取平行能力,這已足以應付筆電規模的控制、識別與協同設計工作。
重點
- C++ 執行緒池在釋放 GIL 的情況下步進整個批次,因此是真正的平行處理,而非交錯執行的 Python
- bind 將 MjData 欄位以涵蓋所有模擬的即時陣列形式提供——控制器只需寫入 ctrl 一次,並原地讀回 qpos,無需收集步驟
- expand 讓每個模擬都擁有自己的 MjModel 參數,並由 set_const 重新計算衍生常數——無需額外機制即可實現領域隨機化與系統識別
- 六個實作範例,每個都是完整的方法:iLQR、預測取樣模型預測控制(MPC)、PPO、CEM 硬體協同設計,以及阻尼 Gauss-Newton 慣性擬合
- 作者表示,Go1 搖桿控制器在一台五年前的 M1 筆電上不到一分鐘就學會走路
- 可從 PyPI 安裝,鎖定 mujoco 3.13.0,並附有 CI、pytest 測試套件與擴展性基準測試
適用對象
- 進行模型預測控制(MPC)、軌跡最佳化或 CEM、每一步都需要大量推演的研究人員
- 無法使用 GPU 叢集、但要訓練小型強化學習控制器的任何人
- 針對模型參數進行系統識別或領域隨機化的實務工作者
- 想要批次模擬、又不想轉向 MJX 或 Warp 的 MuJoCo 使用者
限制
- 設計上僅支援 CPU——要進行 GPU 規模的批次運算,MJX 或 MuJoCo Warp 仍是正途,本專案無法取代它們
- 撰寫本文時僅推出六天,版本為 v0.1.1,分類為 Beta;521 顆星反映的是作者過往的成績,也同樣反映函式庫本身
- 精確鎖定 mujoco==3.13.0,因此只有在作者升級時才會跟著更新
- 從原始碼建置需要 C++ 工具鏈,而會開啟視窗的範例需要顯示器,除非以 --headless 執行
- API 介面刻意保持精簡——Batch、bind、expand、set_const——因此超出此範圍的需求都得直接改用 MuJoCo
- 已發布的分類標記中僅支援 macOS 與 Linux;沒有 Windows wheel