💻 computer science

MuJoCoUni:Persistent Batched Runtime Primitives for MuJoCo

MuJoCoUni is a downstream MuJoCo distribution that introduces the C++/pybind11 `BatchEnvPool` primitive to enable high-throughput, stateful, and batched physics evaluation for online robot learning while strictly preserving upstream CPU MuJoCo semantics.

Original authors: Yufei Jia, Junzhe Wu

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yufei Jia, Junzhe Wu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a coach training a team of 10,000 robot athletes. Your goal is to teach them how to walk, grab objects, or run up stairs. To do this, you need a simulator—a digital gym where you can test their moves millions of times without breaking a sweat.

The paper introduces MuJoCoUni, a new tool designed to make this digital gym run much faster and more efficiently, specifically for robots that need to learn in real-time.

Here is the breakdown using simple analogies:

1. The Problem: The "One-at-a-Time" Bottleneck

Previously, if you wanted to train robots using the standard MuJoCo simulator (the industry standard for robot physics), you often had to treat them like a line of people waiting for a single cashier.

The Old Way: You send a command to Robot A, wait for the result, then send a command to Robot B, wait, and so on. Even if you had a super-fast computer, the software was designed to handle them one by one or in very small, rigid batches.
The Limitation: This is like trying to fill a swimming pool with a single water dropper. It works, but it's incredibly slow for large-scale training.

2. The Solution: The "Conveyor Belt" System

MuJoCoUni acts like a high-speed conveyor belt for your robot simulations.

The Batch: Instead of handling one robot at a time, MuJoCoUni grabs a whole "batch" (say, 1,000 robots) and processes them all at once.
The "Persistent" Feature: This is the key innovation. In the old system, every time you reset a robot to try again, the computer had to rebuild the robot's "brain" and "body" from scratch. MuJoCoUni keeps the robots "alive" in memory. It remembers their specific body shapes and settings. When a robot fails a task, the system simply "resets" that specific robot to the starting line while the others keep moving. This saves a massive amount of time.

3. How It Works: The "Specialized Workers"

The paper describes a core component called BatchEnvPool. Think of this as a factory floor manager:

The Models: It keeps a copy of every robot's blueprint (its physical structure) ready to go.
The Workers: It assigns a dedicated worker (a computer thread) to handle the physics calculations for a group of robots.
The Speed: Because these workers are already set up and waiting, they don't waste time building the robots from scratch. They just run the simulation, check the results, and reset the ones that fell down.

4. What It Can Do (The "Superpowers")

The paper highlights that MuJoCoUni isn't just about running faster; it adds specific tools for robot learning:

Smart Resets: If only 5 out of 1,000 robots fall down, the system only resets those 5. It doesn't stop the other 995. This is like a teacher only calling up the students who got a question wrong, while the rest of the class keeps working.
Randomization: To make robots robust, you often want to change their weight or the friction of the floor slightly every time they restart. MuJoCoUni can do this automatically and instantly during the reset, like a chef who subtly changes the spice level in a soup every time they serve a new bowl.
Instant Queries: Sometimes you need to know specific details, like "How high is the ground under the robot's foot?" or "What is the angle of the robot's arm?" MuJoCoUni can ask these questions for all 1,000 robots simultaneously without actually moving them forward in time.

5. The Result: Speed and Scale

The authors tested this on four different types of robots (a dog-like robot, a dexterous hand, an arm, and a human-like robot).

The Numbers: When they ran 4,000 robots at once, the new system was 15 to 500 times faster than the old Python-based methods.
The Sweet Spot: The system gets so efficient that it can handle hundreds of robots at once before the computer's processor becomes the bottleneck.

6. What It Is NOT

It is important to note what this paper doesn't claim:

It is not a GPU-based system (like those used in video games that run on graphics cards). It runs on standard computer processors (CPUs).
It does not change the actual physics rules. It uses the exact same physics engine as the original MuJoCo, ensuring that if a robot learns to walk here, it will walk the same way in the real world.
It is not a replacement for planning long, complex paths from start to finish. It is specifically designed for the "online learning" loop where a robot tries, fails, learns, and tries again quickly.

Summary

MuJoCoUni is a software upgrade that turns a slow, single-lane road into a multi-lane highway for robot training. It keeps the robots "parked" and ready to go, allowing researchers to run thousands of simulations in parallel, reset only the ones that fail, and tweak their settings instantly. This makes training complex robots much faster and more efficient, all while keeping the physics accurate and reliable.

Technical Summary: MuJoCoUni

Problem Statement

Robot learning systems increasingly integrate physics simulators directly into the training loop, requiring high-frequency, batched execution of physics steps, sensor reads, and environment resets. While upstream MuJoCo provides mature assets and solvers, its standard Python interface incurs overhead regarding object lifetime management, interface calls, and output shaping at high frequencies.

Existing solutions present trade-offs:

Upstream mujoco.rollout: Efficient for open-loop trajectory generation (planning, system identification) but lacks stateful environment persistence, sparse reset capabilities, and per-environment randomization lifecycles required for online reinforcement learning (RL).
GPU-Resident Simulators (e.g., Brax, MJX, Isaac Gym): Offer massive parallelism but often require models and data layouts to fit specific accelerator constraints, potentially sacrificing full feature coverage, complex contact handling, or the ability to reuse existing CPU-optimized MuJoCo assets.
Standard CPU Vectorized Runtimes: Often lack deep integration with the specific MuJoCo solver semantics or require significant engineering to maintain model persistence and domain randomization efficiently.

There is a need for a CPU-based, batched runtime that retains the full semantic fidelity of upstream MuJoCo (models, sensors, contact, constraints) while providing high-throughput parallel execution and stateful environment management for online RL and batched physics evaluation.

Methodology

MuJoCoUni is a lightweight downstream distribution of MuJoCo that concentrates its engineering additions strictly within the Python binding layer, leaving the core physics kernel, solver, integrator, and contact model unchanged.

Core Architecture: `BatchEnvPool`

The central component is BatchEnvPool, a C++/pybind11 executor that manages:

Persistent Resources: It owns $N$ copies of mjModel (one per environment) and $W$ mjData workers (one per thread).
Thread Pool: An internal thread pool dispatches environment indices to worker threads for parallel execution.
Stateful Execution: Unlike rollout, which is stateless per call, BatchEnvPool maintains the state of environments across calls, allowing for sparse resets and persistent model variants.

Key Primitives

The API exposes several primitives designed for online robot training and evaluation:

step: Performs $n$ integration steps (mj_step) on all environments. It returns only the final state (and optionally final sensor data), avoiding the overhead of returning full trajectories.
forward: Executes mj_forward to update sensors based on the current state without advancing dynamics, useful for observation construction.
reset: Implements sparse reset, acting only on a subset of terminated environments (env_ids). It supports reset-lifecycle domain randomization, applying field patches (e.g., mass, inertia, friction) and calling mj_setConst only for affected environments.
compute_site_jacobians: Batched computation of translational and rotational Jacobians for specific sites, running the minimal kinematic prefix (mj_kinematics, mj_comPos) followed by mj_jacSite.
sample_hfield_height: Bilinearly samples MuJoCo height-field geometries in a batch, supporting various frame alignments (world, body, yaw).

Design Boundaries

The implementation avoids modifying the MuJoCo source tree. It relies on mj_copyModel for environment instantiation and mj_setConst for randomization. The design supports two modes:

Single Shared Model: All environments share a model copy (memory efficient, less randomization flexibility).
Model Variants: Each environment owns a distinct model copy, allowing for geometry-level randomization (e.g., link lengths, mesh scales) pre-compiled at construction time.

Key Contributions

Stateful Batched Runtime: Introduces BatchEnvPool to bridge the gap between open-loop trajectory generation and the requirements of online, stateful RL environments.
Efficient Sparse Reset & Randomization: Provides a mechanism to reset only terminated environments and apply domain randomization (both field patches and model variants) with minimal overhead, scaling linearly with the number of reset environments.
Batched Physics Queries: Exposes high-performance, batched implementations of sensor forward passes, site Jacobians, and height-field sampling without advancing dynamics.
Preservation of Semantics: Ensures that the physics behavior remains identical to upstream CPU MuJoCo, maintaining compatibility with existing XML/MJB assets, debugging tools, and sensor definitions.

Results

The authors validated MuJoCoUni on four robot models (Unitree Go1, Wonik Allegro, Franka Panda, CMU Humanoid) using an Intel i9-14900HX CPU with 16 threads.

Throughput: Throughput saturates around 256–512 environments. At saturation, the system achieves:
- Wonik Allegro: ~1.8M steps/second.
- Unitree Go1: ~1.2M steps/second.
- Franka Panda: ~410k steps/second.
- CMU Humanoid: ~290k steps/second.
Reset Performance: The C++ BatchEnvPool path is significantly faster than Python loops. For a full reset of 4096 environments, it takes 3.5 ms compared to 53 ms for a Python for-loop (~15× speedup).
Jacobian Performance: Computing site Jacobians for 4096 environments takes 0.53 ms in C++ versus 11.9 ms in Python (~22× speedup).
Height-Field Sampling: Sampling height fields for 4096 environments takes 0.52 ms in C++ versus 290 ms in Python (~555× speedup).
Model Variants: The overhead of using distinct model copies per environment (model-variant mode) versus a single shared model is negligible at saturation, with throughput becoming essentially identical.

Significance and Scope

The paper positions MuJoCoUni as a complementary engineering solution rather than a replacement for GPU-resident simulators or upstream rollout.

Complement to GPU: It targets workloads where feature coverage (full MuJoCo semantics, complex contacts, debugging) and cross-platform portability are prioritized over raw accelerator residency. It allows systems to leverage modern multicore CPUs for high-throughput data generation while retaining the stability of the CPU MuJoCo solver.
Complement to rollout: It addresses the specific needs of online robot RL and stateful evaluation loops, which require persistent environments, sparse resets, and current-state queries, tasks for which rollout (designed for full-trajectory generation) is not optimized.
System Integration: It serves as a simulation backend that decouples physics execution from higher-level training logic, logging, and scheduling. It is particularly suited for sim-to-real training requiring domain randomization, terrain-aware locomotion, and batch optimization (e.g., evolutionary strategies).

The authors emphasize that the contribution is engineering-oriented, providing a concrete, reproducible path for CPU-batched execution that preserves the upstream MuJoCo ecosystem.