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 机械臂惯性参数。作者称 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 硬件协同设计，以及阻尼高斯-牛顿惯性拟合
作者称一个 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