Kevin Zakka releases mjbatch, a Python library for running thousands of MuJoCo simulations in parallel on CPU, and the repo hit 331 stars in one day

Kevin Zakka, a UC Berkeley PhD, has published mjbatch, a Python library that runs thousands of MuJoCo simulations in parallel on CPU. The repository was created on 10 September and has 331 stars and 26 forks in about a day. It is Apache-2.0 licensed and built to make batched physics experiments cheap on hardware people already have.

What it is

MuJoCo is the physics engine most robotics research uses for training and control. It runs a single simulation quickly, but running hundreds of them at once is what a reinforcement learning loop, a system-identification sweep, or a model-predictive controller actually needs. mjbatch handles the batching in a C++ thread pool with the Python GIL released, so the parallelism scales with the CPU rather than being blocked on the interpreter.

What is in it

The library exposes two primitives that make batched simulation ergonomic. bind gives live array access to MjData fields, so the state and controls of every parallel simulation are visible as numpy arrays. expand and set_const do the same for MjModel fields, so each simulation can carry its own model parameters, which is what you need for randomised training environments or parameter sweeps.

The example workloads shipped with the repo are the kind that would otherwise motivate reaching for a GPU: cart-pole controllers, a humanoid learning backflips, a quadruped tracking a joystick input, and system-identification runs. There is also a co-design demonstration that varies hardware parameters across the batch to search for a better robot design at the same time as a controller.

The audience is specific and the number matters. A robotics or RL team that has been treating batched MuJoCo as a job for a machine with a GPU can now run the same experiments on a workstation or a cheap cloud CPU instance. The GIL-released thread pool is the piece that made this awkward to write from scratch; a library that packages it, with numpy-shaped access to state, removes the reason to keep the batching inside a specific RL framework.

The author is a well-known robotics researcher, which is part of why 331 stars in a day is a plausible number rather than an inflated one. mjbatch is early: v0.x, one week of pushes, a small set of examples. For a team already using MuJoCo through its Python bindings, this is worth a read and a benchmark against the current setup; for anyone whose workflow is bottlenecked on how many rollouts fit in a training run, it is the kind of library that changes the shape of the loop rather than the speed of one step.

Source

Primary source: kevinzakka/mjbatch on GitHub.