kevinzakka/mjbatch

mjbatch – 在 CPU 上進行並行 MuJoCo 模擬

是什麼 – mjbatch 是一個小型 Python 套件，可使用釋放 Python GIL 的 C++ 線程池，同時執行 數千 個 MuJoCo 物理模擬。它針對需要大量批次推演的科研人員與工程師，適用於強化學習、模型預測控制、系統辨識或硬體共設計。

核心理念

• 線程池執行 – 重負載由原生 C++ 線程處理，因此每個模擬在獨立核心上執行，不會被 Python 的全域解釋器鎖（GIL）阻塞。
• 即時陣列綁定 – 可直接綁定到 MuJoCo 資料欄位（如 qpos、ctrl），並將其視為 NumPy 陣列，整個批次的資料會原地更新。
• 每模擬參數 – 透過 expand 可對模型常數（如 geom_friction）進行每實例變更，透過 set_const 重新計算衍生常數。

典型工作流程

import mujoco, numpy as np
from mjbatch import Batch

# 一次性載入 MuJoCo 模型
model = mujoco.MjModel.from_xml_path('scene.xml')

# 建立 4096 個並行模擬批次（預設每個邏輯 CPU 一個線程）
batch = Batch(model, num_sims=4096)

# 綁定狀態與控制陣列 – 這些是整個批次的 NumPy 視圖
qpos, ctrl = batch.bind('qpos'), batch.bind('ctrl')

# 為每個模擬隨機化模型參數
batch.expand('geom_friction')[:, :, 0] = np.random.uniform(0.4, 1.2, (4096, 1))

for _ in range(1000):
    ctrl[:] = policy(qpos)   # 你的控制器在所有模擬上同時執行
    batch.step()             # 並行推進每個模擬；qpos 原地更新

上述迴圈看起來就像單一模擬腳本，但底層有 4,000 個環境同時推進。

你可以用它做什麼

• 強化學習 – 快速產生大規模回放快取或評估策略（README 展示了在舊筆電上 Go1 四足機器人不到一分鐘內學會行走）。
• 模型預測控制 / iLQR – 在多個初始條件下執行滾動時域最佳化器。
• 系統辨識 – 透過模擬多個擾動模型並對比合成資料來擬合慣性參數。
• 共設計 – 使用進化或梯度基礎方法聯合優化機器人幾何、齒輪比與控制策略。

倉儲內建範例

範例	目標	技術
cartpole_swingup.py	將雙擺小車擺起	iLQR
cartpole_mpc.py	帶預測採樣的擺起	MPC
g1_flip.py	使人形機器人後空翻	滾動時域 iLQR
go1_joystick.py	遙控四足機器人	PPO 訓練策略
arm_throw.py	協同設計手臂比例與控制	CEM 最佳化
rizon_inertia.py	識別慣性參數	阻尼高斯-牛頓

使用 uv run examples/<file>.py 執行任意範例；部分需要顯示（--headless 可無顯示運行）。倉儲還提供 CI 指示燈與 PyPI 發布，顯示該套件已維護且可透過 pip install mjbatch 安裝。

安裝

pip install mjbatch   # 拉取編譯擴展與純 Python 包裝

此套件依賴官方 MuJoCo Python 綁定，因此需要 MuJoCo 許可證（學術用途免費）與已安裝的 mujoco 套件。

許可證 – Apache-2.0，允許商業與開源使用。

---

簡而言之，mjbatch 是任何希望在不依賴 GPU 集群的情況下擴展 MuJoCo 基礎實驗的人的實用工具。它抽象了執行緒細節，同時提供對整個模擬批次狀態的直接 NumPy 風格存取。