How mujocolab/mjlab works
1. Project shape
- mjlab is a Python robotics RL framework built on MuJoCo Warp for GPU simulation, with an Isaac Lab-style manager-based API for composing environments.
- The repository is organized as a package under src/mjlab, plus docs, tests, benchmark/train scripts, and a large set of robot/task assets.
2. Top-level runtime flow
- A task configuration selects a scene, assets, sensors, actuators, and MDP components.
- The environment steps through: actions → actuators → simulation → sensors/observations → rewards/terminations/metrics.
- RL entrypoints wrap this loop for training, playback, exporting, and evaluation.
3. Simulation core
- src/mjlab/sim/ contains the MuJoCo-facing runtime: simulation control, randomization, and shared simulation data.
- This is the low-level boundary to the external simulator and the place where domain randomization and reset-time state changes are applied.
4. Entity and scene modeling
- src/mjlab/entity/ defines entity abstractions and per-entity data/variants.
- src/mjlab/scene/ provides the scene container and scene XML, acting as the composition layer that assembles robots, terrain, cameras, and other objects into a runnable world.
5. Manager-based environment architecture
- src/mjlab/managers/ is the orchestration layer: action, command, curriculum, event, observation, reward, termination, metrics, and recorder managers.
- These managers are the main architectural boundary between task logic and runtime execution; they centralize how per-step signals are computed and how resets/events are triggered.
6. Sensors and actuators
- src/mjlab/sensor/ contains built-in sensors such as contact, camera, raycast, and terrain height sensing.
- src/mjlab/actuator/ implements actuator models including PD, DC motor, learned, XML-driven, and grouped/built-in variants.
- Together these are the I/O endpoints between policy outputs and simulated robot state.
7. Task layer and MDP definitions
- src/mjlab/tasks/ contains concrete task families: cartpole, manipulation, tracking, and velocity.
- Each task has its own environment configs, MDP logic (observations, rewards, terminations, commands, curriculums, metrics), and often task-specific RL runner code.
- This is where the repository’s behavioral differences live; the shared framework is reused across tasks.
8. Asset zoo and robot descriptions
- src/mjlab/asset_zoo/ stores robot-specific XMLs and mesh assets for robots like Unitree GO1, Unitree G1, and i2rt_yam.
- These assets are referenced by task configs and scenes, making the repository strongly asset-driven rather than code-only.
9. Terrain system
- src/mjlab/terrains/ provides terrain generation, primitive/heightfield terrain types, terrain entities, and configuration utilities.
- Terrain is a first-class scene component and is used heavily by velocity and locomotion-style tasks, affecting spawning, curriculum, and sensor inputs.
10. RL integration and export
- src/mjlab/rl/ contains runner/config/wrapper utilities for training and inference, including vectorized environment integration.
- There is also exporter support for deploying trained policies and utility code used by task-specific RL runners.
11. CLI and operational scripts
- src/mjlab/scripts/ provides command-line entrypoints for training, playing, demos, visualization, terrain inspection, scene export, and GCS-related utilities.
- The repo also includes scripts/cloud/ for cluster/cloud training workflows and scripts/benchmarks/ for throughput and nightly benchmark automation.
12. Visualization and inspection
- src/mjlab/viewer/ offers both native MuJoCo viewer support and a richer viser-based interface with overlays, reward panels, scene visualization, and policy/state updates.
- src/mjlab/utils/wrappers/video_recorder.py and related tooling support recording and offline inspection.
13. Utility layer
- src/mjlab/utils/ hosts reusable cross-cutting helpers: buffers, logging, randomness, torch/GPU helpers, XML utilities, noise models, NaN guards, and wrappers.
- src/mjlab/utils/lab_api/ is a notable compatibility layer forked from Isaac Lab, supporting the manager/config patterns used throughout the repository.
14. Validation and developer surface
- tests/ is extensive and maps closely to the architecture: unit tests for actuators, sensors, managers, terrains, tasks, viewer behavior, and RL export/runner flows.
- Docs under docs/ mirror the package structure, documenting architecture, API surface, tasks, training, viewers, sensors, and terrain generation; CI and release workflows automate testing, docs, Docker, and nightly benchmarks.
Main entry points: managers Managers · sim Simulation · tasks Tasks · scene Scene
Components 19
Core runtime
Simulation, scenes, managers, and I/O paths
- RL Env · Top-level step loop coordinating task execution
src/mjlab/envs/manager_based_rl_env.py - Managers · Action, command, curriculum, event, observation, reward, termination, metrics, recorder
src/mjlab/managers - Simulation · Low-level simulator control, randomization, and shared sim data
src/mjlab/sim - Scene · World assembly for robots, terrain, cameras, and other entities
src/mjlab/scene - Entities · Per-entity abstractions, data, and variants
src/mjlab/entity - Actuators · PD, DC motor, learned, XML-driven, and grouped actuators
src/mjlab/actuator - Sensors · Contact, camera, raycast, and terrain height sensors
src/mjlab/sensor - Terrains · Terrain generation, entities, and configuration utilities
src/mjlab/terrains
Tasks
Task families and their MDP definitions
- Tasks · Task family entrypoint and registration surface
src/mjlab/tasks - Cartpole · Simple baseline task and environment config
src/mjlab/tasks/cartpole - Manipulation · Robot manipulation tasks and task-specific MDP logic
src/mjlab/tasks/manipulation - Tracking · Tracking tasks with command, observation, reward, and metrics logic
src/mjlab/tasks/tracking - Velocity · Locomotion/velocity tasks with terrain-aware curriculum and rewards
src/mjlab/tasks/velocity
Support
RL, viewing, utilities, and assets
- Asset Zoo · Robot XMLs and meshes referenced by scenes and tasks
src/mjlab/asset_zoo - RL Runtime · Runner, config, wrappers, vectorized env integration, and export helpers
src/mjlab/rl - Viewer · Native MuJoCo and viser-based inspection, overlays, and policy/state views
src/mjlab/viewer - Scripts · Training, play, demo, export, visualization, and cloud tooling
src/mjlab/scripts - Utils · Buffers, logging, randomness, XML, torch/GPU, noise, wrappers, and compatibility helpers
src/mjlab/utils - Lab API · Isaac Lab-style helper layer used across configs and task composition
src/mjlab/utils/lab_api
Connections 25
- Tasks → Cartpole: includes · Task families provide concrete environments and MDP definitions
- Tasks → Manipulation: includes · Task families provide concrete environments and MDP definitions
- Tasks → Tracking: includes · Task families provide concrete environments and MDP definitions
- Tasks → Velocity: includes · Task families provide concrete environments and MDP definitions
- RL Env → Managers: uses · Environment delegates per-step logic to managers
- Managers → Actuators: actions · Action manager maps policy outputs to robot controls
- Managers → Sensors: observations · Observation manager gathers sensor/state features
- Managers → Simulation: reset/events · Events, domain randomization, and resets mutate sim state
- Managers → RL Runtime: signals · Rewards, terminations, metrics, and recorders feed RL loops
- Actuators → Simulation: commands · Actuators write control targets into simulation
- Simulation → Sensors: state · Sensors read simulated bodies, contacts, and camera data
- Scene → Simulation: builds · Scene composition becomes runnable simulation state
- Entities → Scene: composes · Entities are assembled into scenes and per-world variants
- Terrains → Scene: adds · Terrain entities and generators contribute to the scene
- Asset Zoo → Scene: references · Robot XMLs and meshes are consumed by task scenes
- Cartpole → Scene: configures · Task config selects scene and model assets
- Manipulation → Asset Zoo: uses · Manipulation tasks are robot/asset driven
- Tracking → Asset Zoo: uses · Tracking tasks are robot/asset driven
- Velocity → Terrains: depends · Velocity tasks rely on terrain generation and curriculum
- RL Runtime → RL Env: wraps · Training and inference run the manager-based environment
- Scripts → RL Runtime: launches · CLI entrypoints start training, playback, export, and evaluation
- Viewer → Simulation: inspects · Viewer reads and overlays simulation state
- Utils → Simulation: supports · Shared helpers back simulation, XML, logging, and GPU code
- Lab API → Managers: compatibility · Compatibility layer supports the config and manager patterns
- Lab API → Tasks: compatibility · Task configs and builders reuse Isaac Lab-style helpers