Source-linked AI summary
mjlab: A Lightweight Framework for GPU-Accelerated Robot Learning
Kevin Zakka, Qiayuan Liao, Brent Yi, Louis Le Lay, Koushil Sreenath, Pieter Abbeel
TL;DR
Robot-learning frameworks must balance simulation fidelity and maintainable environment composition against installation and infrastructure complexity. mjlab combines Isaac Lab’s manager-based API with MuJoCo Warp GPU simulation and native MuJoCo access, yielding a lightweight framework with reference tasks and documented adoption. Its scope deliberately excludes cross-simulator portability and photorealistic rendering.
Problem
Existing frameworks trade off rich composable infrastructure against minimal setup and maintainability, leaving a gap for a lightweight, transparent, MuJoCo-native alternative.
Method
mjlab combines manager-based modular terms for observations, rewards, events, and curricula with MuJoCo Warp GPU simulation and direct MuJoCo-native data access.
Results
mjlab provides minimal dependencies, fast startup, a PyTorch-native training interface, three robot morphologies, three reference tasks, and reported adoption in education and open-source projects.
Takeaways & Limitations
The framework supports fast iteration and transparent MuJoCo-native workflows across locomotion, whole-body control, and manipulation tasks.
Takeaways & Limitations
mjlab targets a single MuJoCo Warp physics stack and excludes cross-simulator portability and photorealistic rendering.
Abstract
from arXiv · showhide
We present mjlab, a lightweight, open-source framework for robot learning that combines GPU-accelerated simulation with composable environments and minimal setup friction. mjlab adopts the manager-based API introduced by Isaac Lab, where users compose modular building blocks for observations, rewards, and events, and pairs it with MuJoCo Warp for GPU-accelerated physics. The result is a framework installable with a single command, requiring minimal dependencies, and providing direct access to native MuJoCo data structures. mjlab ships with reference implementations of velocity tracking, motion imitation, and manipulation tasks.
1 Introduction
Robot learning depends on simulation infrastructure that faithfully models actuators, contacts, sensing, and randomization while allowing researchers to iterate on policies. mjlab addresses this gap by combining Isaac Lab’s composable manager-based design with MuJoCo Warp’s GPU-accelerated physics and a lightweight, MuJoCo-native interface.
- Sim-to-real robot learning requires careful implementation of actuator dynamics, contact modeling, sensor noise, and domain randomization.
- Isaac Lab offers a rich GPU-accelerated manager-based API but adds Omniverse installation complexity and startup latency.
- MuJoCo Playground minimizes abstractions for rapid prototyping, but code duplication can make multi-robot and multi-task codebases difficult to maintain.
- mjlab combines composable observations, rewards, events, and commands with MuJoCo Warp, minimal dependencies, fast startup, native MuJoCo access, and a PyTorch-native training interface.It includes three robot morphologies and three reference tasks.
2 Design Philosophy and Scope
mjlab prioritizes minimal installation friction, transparent physics, and MuJoCo-native workflows over broad simulator support or feature completeness. Its scope includes composable rigid-body robot-learning infrastructure and sensors, while photorealistic rendering and cross-simulator portability remain outside scope.
- mjlab’s engineering commitments are minimal installation friction, transparent and inspectable physics, and tight MuJoCo integration.These commitments motivate explicit trade-offs in implementation and scope.
- The framework targets MuJoCo Warp and exposes native MjModel and MjData structures, rather than pursuing cross-simulator portability.
- mjlab provides ray-cast, depth, and RGB camera sensors, but photorealistic rendering is out of scope.The text describes training privileged policies and distilling them into vision-based controllers using external rendering as one possible approach.
- Manager-based composition uses reusable terms for observations, rewards, events, and curricula while keeping the implementation MuJoCo-native.This reduces code duplication across tasks.
- The infrastructure is intended for rigid-body robot learning and extension to custom robots, tasks, sensors, and actuators.Three example robots and tasks demonstrate the API.
- The design favors fast iteration, transparent physics, and MuJoCo-native workflows over feature completeness or broad simulator support.
3 Architecture
mjlab combines MuJoCo-native scene construction and GPU simulation with manager-based environment orchestration. Batched world state and CUDA graph replay support parallel stepping while modular managers organize rewards, observations, events, and diagnostics.
- MuJoCo Warp provides a GPU-accelerated MuJoCo backend built on NVIDIA Warp while preserving MuJoCo’s familiar model-and-data paradigm.
- A leading world dimension lets one MjData hold N independent simulation instances, enabling thousands of environments to step in parallel.Model parameters are shared by default and can be expanded per world for domain randomization.
- mjlab records the simulation step as a CUDA graph and replays it to eliminate CPU-side dispatch overhead.
- Entities are composed into an MjSpec, compiled into an MjModel on the CPU, and transferred to the GPU through MuJoCo Warp.
- The manager-based environment design composes self-contained reward, observation, and domain-randomization terms, with managers handling lifecycle, aggregation, and diagnostics.Terms may be plain functions or manager-term classes when state is needed.
- The structured manager-based API favors maintainability and extensibility over the minimal abstractions of MuJoCo Playground.
4 Components
mjlab’s components provide a unified, extensible foundation for robot learning, sensing, actuation, terrain generation, and visualization. The framework supports layered MuJoCo-native and custom GPU components, including configurable terrain difficulty and curriculum progression.
- mjlab’s core components are operated by managers during the simulation lifecycle.
- Entities: A single Entity class represents fixed or floating physical objects with or without joints through independent base-type and articulation properties.
- Sensors: Sensors can wrap MJCF-defined devices or be configured in Python, with custom ray-cast and contact sensors extending MuJoCo’s native capabilities.
- Actuators: Actuators can preserve MJCF behavior, instantiate MuJoCo-native types, or compute GPU-side torques using PD, DC motor, and learned MLP models.
- Terrain: Terrain generation combines primitive and heightfield patches whose difficulty scales from trivial to challenging through parameters such as step height, slope angle, and noise amplitude.
- Terrain: Curriculum terrain grids encode increasing difficulty along one axis so training can advance robots to harder terrain as performance improves.
- Viewer: mjlab provides native MuJoCo and browser-based Viser viewers for interactive or headless visualization.
5 Manager-Based API
mjlab structures environments through a Gym-compatible, manager-based pipeline in which modular managers coordinate actions, simulation, termination, rewards, resets, commands, events, and observations. Its MuJoCo-native design supports parallel worlds, configurable commands and observations, diagnostics, and physically consistent domain randomization.
- Environment lifecycle: The Gym-compatible environment exposes reset() and step(), with each step passing through a fixed manager-coordinated pipeline.The pipeline includes action, simulation, termination, reward, reset, command, event, and observation stages.
- Manager-based paradigm: Managers encapsulate specific MDP aspects, while users compose environments by registering modular terms such as rewards, observations, and events.Terms may be plain functions or stateful classes, and managers handle their lifecycle, aggregation, and diagnostics.
- Domain randomization: Its domain-randomization module preserves physical consistency when varying inertial parameters and supports applications beyond sim-to-real transfer.The module is also applicable to system identification and hardware co-design.
- Commands and observations: The command and observation managers expose configurable goals and processing pipelines, including clipping, scaling, noise, delay, history, and multiple observation groups.Commands include velocity targets, tracking waypoints, and object poses; observation groups can support asymmetric actor-critic architectures.
- Diagnostics and adaptation: mjlab includes managers for custom metrics, action routing, numerical-instability detection, weighted rewards, curricula, and domain-randomization events.The event system supports configurable parameter operations and distributions, while physically consistent inertial randomization covers mass, center of mass, and rotational inertia.
6 Tasks
mjlab ships a compact set of reference tasks spanning locomotion, whole-body control, and manipulation across humanoid, quadruped, and robot-arm platforms. These tasks demonstrate velocity tracking, motion imitation, and staged object lifting with task-specific observations and rewards.
- Task suite: mjlab ships three reference tasks spanning locomotion, whole-body control, and manipulation.The tasks use three morphologies: the Unitree G1 humanoid, Unitree Go1 quadruped, and YAM robot arm.
- Velocity tracking: Velocity tracking trains a controller to follow linear and angular velocity commands on flat or rough terrain.Observations include IMU readings, projected gravity, joint states, previous actions, and the commanded twist.
- Velocity tracking: A natural running gait emerged on the G1 humanoid in simulation and on real hardware under the velocity-tracking reward structure.The reward penalizes body angular velocity, angular momentum, joint-limit violations, action rate, and foot slip.
- Motion imitation: Motion imitation trains a Unitree G1 humanoid to track reference motion clips using DeepMimic with BeyondMimic extensions.Rewards penalize deviations in root motion, relative body poses, body velocities, and self-collision.
- Manipulation: The manipulation task trains a YAM robot arm to lift a cube to a target pose using observations of joint state, object-relative geometry, and previous actions.Its staged reward first guides the end-effector to the cube and then rewards lifting it to the target height.
7 Software Design
mjlab’s software design emphasizes composability, type safety, low boilerplate, and minimal setup friction. It connects GPU simulation to PyTorch through zero-copy data sharing and replaces fragile configuration inheritance with editable, typed instances.
- PyTorch-native interface: mjlab connects Warp simulation arrays to PyTorch tensors through a zero-copy TorchArray abstraction.Users write rewards, observations, and custom logic in pure PyTorch without directly handling Warp kernels.
- Instance-based configuration: Instance-based configuration replaces inherited dataclass hierarchies with typed dictionaries and mutable configuration instances.This avoids silent __post_init__ and misspelled-field override failures while making task variants easier to read, diff, and compose.
- CLI-first configuration: CLI-first configuration exposes typed parameters automatically through tyro, allowing single-flag overrides without configuration files or subclassing.Examples include reward weights, observation noise, and simulation timestep.
- Co-located definitions: Co-located definitions keep each configuration dataclass beside its implementation, reducing indirection and simplifying code navigation.For example, ActionTermCfg is kept alongside ActionTerm rather than separated into distinct files.
- Typing and testing: Static checking with pyright and ty, a py.typed marker, and tests across more than 40 files support maintainability and downstream use.The tests cover managers, sensors, actuators, domain randomization, and task execution.
- AI-assisted development: The lean, typed, tested codebase is described as amenable to AI-assisted development and autonomous verification of edits.The paper reports several recent contributions implemented end-to-end by AI agents.
- Minimal dependencies: mjlab is designed around uv and avoids heavyweight dependencies, enabling a fresh clone to launch a training job with a short command sequence.The framework also supports RSL-RL on-policy training and multi-GPU scaling through torchrunx.
8 Adoption
mjlab has been used in a UC Berkeley graduate robotics course, where students new to reinforcement learning trained their first legged locomotion policies. The framework has also been adopted by open-source projects and featured in a widely viewed tutorial.
- Course deployment: UC Berkeley’s graduate robotics course deployed mjlab, enabling students new to reinforcement learning to train their first legged locomotion policies.The framework was used in ME 292b/193b.
- External adoption: mjlab has been adopted by various open-source projects and featured in a tutorial from a YouTuber with over 1.5M subscribers.