# mujocolab/mjlab architecture

> Architecture diagram of the GitHub repository mujocolab/mjlab, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-29). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/mujocolab/mjlab
- Repository: https://github.com/mujocolab/mjlab

## Overview

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: [`src/mjlab/managers`](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/managers) (Managers), [`src/mjlab/sim`](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/sim) (Simulation), [`src/mjlab/tasks`](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/tasks) (Tasks), [`src/mjlab/scene`](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/scene) (Scene)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_core["Core runtime"]
  node_env["RL Env<br/>manager-based env"]
  node_managers["Managers<br/>orchestration layer"]
  node_sim["Simulation<br/>MuJoCo runtime"]
  node_scene["Scene<br/>composition layer"]
  node_entity["Entities<br/>entity model"]
  node_actuator["Actuators<br/>control I/O"]
  node_sensor["Sensors<br/>state I/O"]
  node_terrains["Terrains<br/>terrain system"]
end

subgraph group_tasks["Tasks"]
  node_tasks["Tasks<br/>task registry"]
  node_cartpole["Cartpole<br/>task family"]
  node_manipulation["Manipulation<br/>task family"]
  node_tracking["Tracking<br/>task family"]
  node_velocity["Velocity<br/>task family"]
end

subgraph group_support["Support"]
  node_asset_zoo["Asset Zoo<br/>robot assets"]
  node_rl["RL Runtime<br/>training/inference"]
  node_viewer["Viewer<br/>visualization"]
  node_scripts["Scripts<br/>CLI entrypoints"]
  node_utils["Utils<br/>shared helpers"]
  node_lab_api["Lab API<br/>compat layer"]
end

node_tasks -->|"includes"| node_cartpole
node_tasks -->|"includes"| node_manipulation
node_tasks -->|"includes"| node_tracking
node_tasks -->|"includes"| node_velocity
node_env -->|"uses"| node_managers
node_managers -->|"actions"| node_actuator
node_managers -->|"observations"| node_sensor
node_managers -->|"reset/events"| node_sim
node_managers -->|"signals"| node_rl
node_actuator -->|"commands"| node_sim
node_sim -->|"state"| node_sensor
node_scene -->|"builds"| node_sim
node_entity -->|"composes"| node_scene
node_terrains -->|"adds"| node_scene
node_asset_zoo -->|"references"| node_scene
node_cartpole -->|"configures"| node_scene
node_manipulation -->|"uses"| node_asset_zoo
node_tracking -->|"uses"| node_asset_zoo
node_velocity -->|"depends"| node_terrains
node_rl -->|"wraps"| node_env
node_scripts -->|"launches"| node_rl
node_viewer -->|"inspects"| node_sim
node_utils -->|"supports"| node_sim
node_lab_api -->|"compatibility"| node_managers
node_lab_api -->|"compatibility"| node_tasks

click node_env "https://github.com/mujocolab/mjlab/blob/main/src/mjlab/envs/manager_based_rl_env.py"
click node_managers "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/managers"
click node_sim "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/sim"
click node_scene "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/scene"
click node_entity "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/entity"
click node_actuator "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/actuator"
click node_sensor "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/sensor"
click node_terrains "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/terrains"
click node_tasks "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/tasks"
click node_cartpole "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/tasks/cartpole"
click node_manipulation "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/tasks/manipulation"
click node_tracking "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/tasks/tracking"
click node_velocity "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/tasks/velocity"
click node_asset_zoo "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/asset_zoo"
click node_rl "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/rl"
click node_viewer "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/viewer"
click node_scripts "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/scripts"
click node_utils "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/utils"
click node_lab_api "https://github.com/mujocolab/mjlab/tree/main/src/mjlab/utils/lab_api"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_env,node_managers,node_sim,node_scene,node_entity,node_actuator,node_sensor,node_terrains toneBlue
class node_tasks,node_cartpole,node_manipulation,node_tracking,node_velocity toneAmber
class node_asset_zoo,node_rl,node_viewer,node_scripts,node_utils,node_lab_api toneMint
```

## Components

### Core runtime

Simulation, scenes, managers, and I/O paths

- [RL Env](https://github.com/mujocolab/mjlab/blob/main/src/mjlab/envs/manager_based_rl_env.py): manager-based env · Top-level step loop coordinating task execution · `src/mjlab/envs/manager_based_rl_env.py`
- [Managers](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/managers): orchestration layer · Action, command, curriculum, event, observation, reward, termination, metrics, recorder · `src/mjlab/managers`
- [Simulation](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/sim): MuJoCo runtime · Low-level simulator control, randomization, and shared sim data · `src/mjlab/sim`
- [Scene](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/scene): composition layer · World assembly for robots, terrain, cameras, and other entities · `src/mjlab/scene`
- [Entities](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/entity): entity model · Per-entity abstractions, data, and variants · `src/mjlab/entity`
- [Actuators](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/actuator): control I/O · PD, DC motor, learned, XML-driven, and grouped actuators · `src/mjlab/actuator`
- [Sensors](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/sensor): state I/O · Contact, camera, raycast, and terrain height sensors · `src/mjlab/sensor`
- [Terrains](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/terrains): terrain system · Terrain generation, entities, and configuration utilities · `src/mjlab/terrains`

### Tasks

Task families and their MDP definitions

- [Tasks](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/tasks): task registry · Task family entrypoint and registration surface · `src/mjlab/tasks`
- [Cartpole](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/tasks/cartpole): task family · Simple baseline task and environment config · `src/mjlab/tasks/cartpole`
- [Manipulation](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/tasks/manipulation): task family · Robot manipulation tasks and task-specific MDP logic · `src/mjlab/tasks/manipulation`
- [Tracking](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/tasks/tracking): task family · Tracking tasks with command, observation, reward, and metrics logic · `src/mjlab/tasks/tracking`
- [Velocity](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/tasks/velocity): task family · Locomotion/velocity tasks with terrain-aware curriculum and rewards · `src/mjlab/tasks/velocity`

### Support

RL, viewing, utilities, and assets

- [Asset Zoo](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/asset_zoo): robot assets · Robot XMLs and meshes referenced by scenes and tasks · `src/mjlab/asset_zoo`
- [RL Runtime](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/rl): training/inference · Runner, config, wrappers, vectorized env integration, and export helpers · `src/mjlab/rl`
- [Viewer](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/viewer): visualization · Native MuJoCo and viser-based inspection, overlays, and policy/state views · `src/mjlab/viewer`
- [Scripts](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/scripts): CLI entrypoints · Training, play, demo, export, visualization, and cloud tooling · `src/mjlab/scripts`
- [Utils](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/utils): shared helpers · Buffers, logging, randomness, XML, torch/GPU, noise, wrappers, and compatibility helpers · `src/mjlab/utils`
- [Lab API](https://github.com/mujocolab/mjlab/tree/main/src/mjlab/utils/lab_api): compat layer · Isaac Lab-style helper layer used across configs and task composition · `src/mjlab/utils/lab_api`

## Connections

- 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