Skip to content
kevinzakkaPublic

About

A Python library for running thousands of MuJoCo simulations in parallel on CPU

Topics

Resources

Stars

607 stars

Watchers

4 watching

Forks

Repository files navigation

mjbatch

Build PyPI version

mjbatch is a Python library for running thousands of MuJoCo simulations in parallel on CPU.

Note

This repository is archived. Batched simulation is now part of MuJoCo itself, which this project inspired. Please use the mjbatch API in the main repository instead.

Features include:

  • C++ thread pool execution, with the GIL released;
  • Live array access to simulation state and controls across the batch, with bind for MjData fields;
  • Per-simulation model parameters, with expand for MjModel fields and set_const to recompute derived constants.

For example:

import mujoco, numpy as np
from mjbatch import Batch

model = mujoco.MjModel.from_xml_path("scene.xml")
batch = Batch(model, num_sims=4096)  # threads default to every logical 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)             # your controller, all 4096 at once
  batch.step()                       # step them in parallel; qpos updates in place

Examples

We showcase a range of applications built using mjbatch: RL, MPC, SysID, and hardware co-design. Each example is a self-contained, performant implementation. For instance, the Go1 RL controller learns to walk in under a minute on a five-year-old M1 laptop.

cart-pole swing-up cart-pole MPC
A two-pole cart swung upright with iLQR A cart-pole swing-up controller using predictive sampling
G1 backflip Go1 joystick
A G1 humanoid tracking a reference backflip with receding-horizon iLQR A Go1 quadruped joystick controller trained with PPO
throwing arm co-design Rizon inertia identification
CEM jointly optimizes a robot arm's proportions, gears, and controls Damped Gauss–Newton fits a Rizon arm's inertial parameters to synthetic motion data

Run with uv run examples/<file>.py; some need uv sync --group examples. The ones that open a window need a display; --headless runs the solver without one.

License

Apache-2.0.

About

A Python library for running thousands of MuJoCo simulations in parallel on CPU

Topics

Resources

Stars

607 stars

Watchers

4 watching

Forks

Releases

Packages

Contributors

Languages