MuJoCoUni

English | 简体中文

MuJoCoUni is the standalone UniLab batch-executor layer for official MuJoCo. It provides the BatchEnvPool API used by UniLab without modifying MuJoCo solver, contact, integrator, or source-tree internals.

System Boundary

MuJoCoUni sits between UniLab and official MuJoCo.

MJCF / XML asset
        |
        v
official MuJoCo compiler
        |
        v
mjModel
        |
        v
MuJoCoUni BatchEnvPool
  - captures/copies mjModel
  - owns model pool
  - owns worker mjData
  - calls official MuJoCo C APIs
        ^
        |
UniLab backend
  - receives task/training command
  - owns rollout pace
  - packs state/control/reset arrays
  - unpacks state/sensor arrays

Responsibility split:

UniLab:
  task config, commands, rewards, rollout pace, CPU/GPU bridge,
  training orchestration, logging

MuJoCoUni:
  BatchEnvPool, model cloning, worker mjData, local thread pool,
  batched step/forward/reset/query execution

MuJoCo:
  MJCF compiler, mjModel/mjData definitions, solver, contact,
  integrator, sensor layout, official C API

For larger CPU/GPU training, UniLab remains the bridge above local MuJoCoUni executors:

CPU rollout side
  MuJoCoUni BatchEnvPool instances
  produce observation/reward/done/sensor data
  request action tensors
        |
        v
UniLab bridge
  batches rollout traffic
  routes data toward GPU learner/action service
  routes actions/control commands back to CPU workers
        |
        v
GPU side
  consumes training data
  performs learner/inference work
  returns action tensors or policy-side feedback through UniLab

Version Policy

MuJoCoUni has its own package version, independent of the MuJoCo solver version.

Current release:

mujoco-uni-runtime==0.5.0
mujoco>=3.5,<3.12

The public metadata is available from Python:

import mujoco_uni

print(mujoco_uni.__version__)
print(mujoco_uni.MUJOCO_VERSION)
print(mujoco_uni.MUJOCO_VERSION_SPEC)

MuJoCoUni supports one official MuJoCo solver version per Python environment. The native extension records the MuJoCo version used at build time, and the runtime fails fast if the loaded mujoco package does not match that native build target.

Version switching is active but process-level: MuJoCoUni selects an existing versioned uv environment, then runs the target command in that environment before Python imports mujoco. Normal training launch does not create, install, or rebuild environments.

env-mj35  -> mujoco==3.5.x  -> build/install mujoco-uni-runtime
env-mj36  -> mujoco==3.6.x  -> build/install mujoco-uni-runtime
env-mj37  -> mujoco==3.7.x  -> build/install mujoco-uni-runtime
env-mj38  -> mujoco==3.8.x  -> build/install mujoco-uni-runtime
env-mj39  -> mujoco==3.9.x  -> build/install mujoco-uni-runtime
env-mj310 -> mujoco==3.10.x -> build/install mujoco-uni-runtime
env-mj311 -> mujoco==3.11.x -> build/install mujoco-uni-runtime

Default and fallback selection prefer discovered environments in this order:

3.11 > 3.8 > 3.10 > 3.9 > 3.7 > 3.6 > 3.5

If the requested version is not found, MuJoCoUni prints a warning and falls back to the preferred existing environment. If no MuJoCo environment exists, launch fails with a clear setup error.

For UniLab training, the normal task command is the only user-facing launcher. Set the process selector when running the command:

MUJOCO_UNI_VERSION=3.8 uv run train --algo ppo --task go2_joystick_flat --sim mujoco
MUJOCO_UNI_VERSION=3.10 uv run train --algo ppo --task go2_joystick_flat --sim mujoco

Unset MUJOCO_UNI_VERSION keeps the active Python environment behavior. Exact requests such as 3.8.0 require that exact runtime; minor requests such as 3.8 accept a compatible 3.8.x runtime.

MuJoCoUni owns the internal discovery and spawning services used by UniLab. Explicit environment preparation is a setup operation, not part of the normal training launch path.

The required MuJoCo Python model pointer helpers, _address and _from_model_ptr, are checked at import time.

Package Layout

The structure mirrors the DrakeUni split between runtime code, native source, and compiled artifacts:

src/mujoco_uni/
  __init__.py
  metadata.py              # MuJoCoUni package metadata and supported range
  batch_env.py              # Stable public BatchEnvPool API
  mujoco_runtime/
    api.py                  # MuJoCoUni-owned access to official mujoco
    version_control.py      # MuJoCo solver-version control
  runtime/
    batch.py                # Python API, validation, compatibility behavior
  compiled/
    __init__.py             # Native extension loader and diagnostics
    _batch_env*.so          # Generated extension after local build
  native/
    batch_env.cc            # Native pybind11 executor
    threadpool.h/.cc        # Local thread pool

The stable import path is:

from mujoco_uni.batch_env import BatchEnvPool, SUPPORTED_FIELDS

Execution Model

BatchEnvPool accepts one mujoco.MjModel or a compatible sequence of mujoco.MjModel objects. At construction time it reads the official MuJoCo model pointer from Python, copies the model with mj_copyModel, and stores pool-owned models internally.

The hot path is native C++ calling the official MuJoCo C API. The pool uses:

• one logical pool-owned mjModel assignment per environment slot,
• one reusable mjData per worker thread,
• disjoint environment chunks,
• disjoint output slots,
• one synchronization point per batch operation.

There is no MPI or OpenMP inside the base BatchEnvPool executor. Large-scale multi-process, multi-socket, or multi-node collection composes multiple local executors from a layer above BatchEnvPool.

Worker CPU affinity (Linux)

BatchEnvPool accepts an optional cpu_ids sequence that pins each native worker thread to an explicit CPU, so multiple pools can be kept on disjoint CPU sets:

pool = BatchEnvPool(model, nbatch=64, cpu_ids=[0, 1, 2, 3])

cpu_ids[i] pins worker thread i, so the mapping is 1:1 with nthread: its length must equal nthread, and when nthread is omitted it is inferred from len(cpu_ids). CPU ids are validated at construction (non-negative, unique, available to the process); invalid input raises ValueError. The configured mapping is exposed read-only via pool.cpu_ids, and pool.worker_cpu_ids() reports the CPU each worker was actually observed on after pinning. When cpu_ids is not given, scheduling behavior is unchanged. CPU pinning is supported on Linux only.

Public API

The stable public API is mujoco_uni.batch_env.

Exported symbols:

• BatchEnvPool
• SUPPORTED_FIELDS
• batch_available
• batch_import_error

Supported randomization/model fields:

body_mass
body_ipos
body_iquat
body_inertia
dof_armature
gravity
geom_friction
kp
kd

Core BatchEnvPool behavior:

• construct from one model or a sequence of length 1 / nbatch,
• step over the full pool and return final state, optionally with final sensor data,
• forward over the full pool and return sensor data,
• sparse reset with optional model-field randomization,
• site Jacobian queries,
• hfield height sampling,
• non-owning model views through get_model, get_models, and get_all_models.

Returned model views remain valid only while the pool is alive.

Installation

MuJoCoUni is published as mujoco-uni-runtime. The default install path is a prebuilt wheel bound to the default MuJoCo version (3.11.0), published for linux x86_64 / aarch64 and macOS arm64 × Python 3.10–3.13:

pip install "mujoco==3.11.0" mujoco-uni-runtime

One released runtime version carries exactly one MuJoCo binding (PyPI filename identity), and the native extension refuses to load against any MuJoCo version other than its build-time one (exact-version watchdog at import time).

For any non-default MuJoCo version, the sdist is the fallback and the only path: build from source against your environment's mujoco with --no-build-isolation (so the build does not see a throwaway isolated environment). Windows users take this path as well — no Windows wheels.

pip install "mujoco>=3.5,<3.12" pybind11 numpy setuptools wheel
pip install mujoco-uni-runtime --no-binary mujoco-uni-runtime --no-build-isolation

See docs/release-coordination.md for the cross-repo release coordination rules.

Prerequisites

Building from source requires:

• a C++17 toolchain — macOS: Xcode Command Line Tools (xcode-select --install); Debian/Ubuntu: sudo apt install build-essential; Windows: MSVC Build Tools
• Python development headers — Debian/Ubuntu system Python: sudo apt install python3-dev (or python3.X-dev matching your version). Not needed when using a uv-managed Python (uv python install), which bundles the headers and is the recommended option.

uv projects declare the same setup:

[project.optional-dependencies]
mujoco = ["mujoco~=3.11.0", "mujoco-uni-runtime==0.5.0", "pybind11>=2.12", "wheel"]

[tool.uv]
no-build-isolation-package = ["mujoco-uni-runtime"]

For development beside UniLab:

cd /path/to/mujoco_uni
uv sync
uv pip install --force-reinstall --no-deps --no-build-isolation -e .

The common development commands are also wrapped in the Makefile:

make sync       # uv sync
make install    # editable install (builds the native extension)
make check      # ruff + pytest
make test       # pytest only
make matrix     # version-matrix checks across supported MuJoCo versions

UniLab imports through its compatibility/backend layer, which in turn imports:

from mujoco_uni.batch_env import BatchEnvPool, SUPPORTED_FIELDS

Rebuilding The Native Extension

Editable installs are useful while developing MuJoCoUni in one environment. They generate a local extension such as:

src/mujoco_uni/compiled/_batch_env.cpython-313-darwin.so

That artifact is tied to the active Python environment, platform, and MuJoCo patch version. Rebuild after switching virtual environments, Python versions, or MuJoCo versions:

uv pip install "mujoco==3.10.0" pybind11 wheel
uv pip install --force-reinstall --no-deps --no-build-isolation -e .

After selecting a solver version, run local checks without an automatic sync if you want to preserve the already-built native target:

uv run --no-sync pytest -q

Validation

Standalone checks:

uv run ruff check .
uv run pytest -q

or simply make check. After switching the MuJoCo version with make mujoco MJ=<version>, run tests with make test-no-sync so the freshly built native target is preserved.

Version-matrix checks:

uv run python tools/version_matrix.py --pytest

or make matrix. Full UniLab validation is available as make test-unilab and make test-unilab-train.

The default matrix covers:

3.5.0 3.6.0 3.7.0 3.8.0 3.8.1 3.9.0 3.10.0 3.11.0

Full UniLab task validation is separate from the quick package matrix:

uv run python tools/version_matrix.py --versions 3.5.0 3.10.0 --unilab
uv run python tools/version_matrix.py --versions 3.5.0 3.10.0 --unilab-train

The --unilab-train mode runs a short one-iteration training smoke in each selected MuJoCo environment.

UniLab integration checks:

cd ../UniLab
uv run pytest \
  tests/base/test_mujoco_batch_env_randomization.py \
  tests/base/test_mujoco_batch_env_jacobian.py \
  tests/base/backend/test_mujoco_site_jacobian.py \
  tests/envs/locomotion/test_go2_rough_height_scan.py \
  tests/envs/locomotion/test_go2_footstand.py \
  tests/envs/locomotion/test_go2_terrain_spawn.py \
  -q

Training smoke:

cd ../UniLab
uv run python scripts/train_rsl_rl.py \
  task=go2_joystick_flat/mujoco \
  algo.seed=1 \
  algo.num_envs=256 \
  algo.num_steps_per_env=24 \
  algo.max_iterations=2 \
  algo.save_interval=100 \
  training.no_play=true \
  training.logger=tensorboard \
  training.device=cpu \
  training.log_root=logs/mujoco_uni_smoke

Thread Safety

Built-in MuJoCo sensors are read from mjData.sensordata and are supported by the batch executor.

Custom sensors, plugins, and global MuJoCo callbacks are thread-safety-sensitive. If nthread > 1, any global mutable callback/plugin state is the caller's responsibility.

Scope Boundaries

MuJoCoUni stays outside these responsibilities:

• MuJoCo source-code changes,
• MuJoCo solver/contact/integrator forks,
• distributed factorization of one MuJoCo solve across MPI ranks,
• OpenMP inside the MuJoCoUni executor,
• UniLab task YAMLs, rollout code, reward functions, and DrakeUni behavior.

Citation

Paper page: MuJoCoUni

@article{jia2026mujocouni,
  title   = {MuJoCoUni: Persistent Batched Runtime Primitives for MuJoCo},
  author  = {Jia, Yufei and Wu, Junzhe},
  journal = {arXiv preprint arXiv:2605.24922},
  year    = {2026},
  url     = {https://arxiv.org/abs/2605.24922}
}

License

Apache-2.0.