mjbatch

Runs thousands of MuJoCo simulations in parallel on the CPU

License

Apache-2.0

Overview

mjbatch is a Python library that runs thousands of MuJoCo simulations in parallel on the CPU. A C++ thread pool does the stepping with the GIL released, and the batch is exposed to Python as live NumPy views: bind gives you an array over an MjData field across every simulation, so writing into ctrl and reading qpos needs no copying and no per-sim loop. expand does the same for MjModel fields, which makes each simulation carry its own physics — different frictions, different link lengths — with set_const recomputing the derived constants afterwards. The repository ships six self-contained examples covering reinforcement learning, model-predictive control, system identification and hardware co-design.

The core is deliberately small. You build a Batch from a single MjModel and a simulation count, the thread pool defaults to every logical CPU, and batch.step() advances all of them. What makes it more than a for-loop is that the arrays are live rather than gathered: qpos, ctrl = batch.bind("qpos"), batch.bind("ctrl") hands you views that a controller can write across the whole batch in one assignment, and they update in place as the batch steps. Per-simulation model parameters are the other half — batch.expand("geom_friction") returns an array you can randomise per simulation, which is domain randomisation and system identification with no extra machinery.

The examples are the argument for it, because each is a complete method rather than a demo: a cart-pole swung upright with iLQR, a cart-pole MPC using predictive sampling, a G1 humanoid tracking a reference backflip with receding-horizon iLQR, a Go1 quadruped joystick controller trained with PPO, CEM jointly optimising a throwing arm's proportions, gears and controls, and a damped Gauss-Newton fit of a Rizon arm's inertial parameters. The author reports the Go1 controller learning to walk in under a minute on a five-year-old M1 laptop. It installs from PyPI, pins mujoco 3.13.0, builds through scikit-build-core with nanobind, and carries CI, a pytest suite and a scaling benchmark.

Why it matters

Massively parallel simulation has become a GPU story — MJX, Isaac, Warp — and that quietly raised the cost of entry for anyone without a workstation card. mjbatch makes the same shape of work available on a laptop by using the CPU properly instead of one core at a time, and it does so through plain MuJoCo rather than a re-implementation, so the models and the physics are the ones you already trust. For a catalogue whose audience mostly owns a laptop and a low-cost arm rather than a cluster, that is the more useful end of the trade-off.

Problem it solves

Getting thousands of MuJoCo rollouts usually means moving to a GPU stack and accepting its hardware bill and its re-implementation of the physics. mjbatch keeps stock MuJoCo and extracts the parallelism from the CPU you already have, which is enough for control, identification and co-design work at laptop scale.

Key takeaways

• A C++ thread pool steps the whole batch with the GIL released, so parallelism is real rather than interleaved Python
• bind exposes MjData fields as live arrays across every simulation — a controller writes into ctrl once and reads qpos back in place, with no gather step
• expand gives each simulation its own MjModel parameters, with set_const recomputing derived constants — domain randomisation and system identification without extra machinery
• Six worked examples, each a complete method: iLQR, predictive-sampling MPC, PPO, CEM hardware co-design and damped Gauss-Newton inertial fitting
• The author reports a Go1 joystick controller learning to walk in under a minute on a five-year-old M1 laptop
• Installs from PyPI, pins mujoco 3.13.0, and carries CI, a pytest suite and a scaling benchmark

Who it is for

• Researchers doing MPC, trajectory optimisation or CEM who need many rollouts per step
• Anyone training small RL controllers without access to a GPU cluster
• Practitioners doing system identification or domain randomisation over model parameters
• MuJoCo users who want batched simulation without moving to MJX or Warp

Limitations

• CPU only by design — for GPU-scale batches MJX or MuJoCo Warp remain the route, and this does not replace them
• Six days old at the time of writing, at v0.1.1 and classified Beta; the 521 stars reflect the author's track record as much as the library's
• Pins mujoco==3.13.0 exactly, so it moves only when the author bumps it
• Needs a C++ toolchain to build from source, and the examples that open a window need a display unless run with --headless
• The API surface is deliberately narrow — Batch, bind, expand, set_const — so anything outside that shape means dropping to MuJoCo directly
• macOS and Linux only in the published classifiers; no Windows wheels