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 风格访问。