ESC
开源 1 分钟阅读

mjbatch:在CPU上并行运行数千个MuJoCo仿真的Python库

开源库mjbatch发布:可在CPU上并行运行数千个MuJoCo物理仿真,特性包括释放GIL的C++线程池执行、跨批次实时数组访问状态与控制量、以及逐仿真模型参数配置。示例涵盖强化学习、MPC、系统辨识与硬件协同设计,Go1四足机器人RL控制器在五年前的M1笔记本上不到一分钟即可学会走路。采用Apache-2.0许可。

来源:GitHub

mjbatch

mjbatch 是一个用于在CPU上并行运行数千个MuJoCo仿真的Python库。

功能特性包括:

  • 采用C++线程池执行,并释放GIL;
  • 可对整个批次内所有仿真的状态和控制量进行实时数组访问;
  • 支持逐仿真设置模型参数,并可通过 set_const 重新计算派生常量。

例如:

import mujoco, numpy as np
from mjbatch import Batch

model = mujoco.MjModel.from_xml_path("scene.xml")
batch = Batch(model, num_sims=4096)  # 线程数默认为所有逻辑CPU核心
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)             # 你的控制器,一次处理全部4096个仿真
  batch.step()                       # 并行推进所有仿真;qpos就地更新

示例

我们展示了一系列基于 mjbatch 构建的应用:强化学习(RL)、模型预测控制(MPC)、系统辨识(SysID)以及硬件协同设计。每个示例都是独立完整、性能高效的实现。例如,Go1强化学习控制器在一台五年前的M1笔记本电脑上不到一分钟即可学会走路。

cart-pole swing-upcart-pole MPC
使用 iLQR 将双杆小车摆正直立基于 predictive sampling 的小车摆杆起摆控制器
G1 backflipGo1 joystick
G1人形机器人通过滚动时域iLQR跟踪参考后空翻动作使用PPO训练的Go1四足机器人遥杆控制器
throwing arm co-designRizon inertia identification
CEM联合优化机械臂的尺寸比例、齿轮与控制阻尼Gauss–Newton法将Rizon机械臂的惯性参数拟合到合成运动数据

使用 uv run examples/<file>.py 运行;部分示例需要先执行 uv sync --group examples。需要打开窗口的示例必须有显示设备;使用 --headless 可在没有显示器的情况下运行求解器。

许可证

Apache-2.0。