DOCSEARCHbyAlgolia
Submit Docs

mjlab Documentation

repository·main·Indexed 25 days ago

https://github.com/mujocolab/mjlab

A lightweight, GPU-accelerated framework for robot learning and robotics research. mjlab combines a manager-based API inspired by Isaac Lab with MuJoCo Warp, providing composable building blocks for environment design and direct access to native MuJoCo data structures. It features PyTorch-native tensors with zero-copy GPU memory sharing and supports robots such as the Unitree G1, Unitree Go1, and i2rt YAM.

Tokens
69.4K
Snippets
152
Records
357
Agent score
83%

What's inside mjlab

  1. Overview of mjlab features

    main

    mjlab is a lightweight, open-source framework for robot learning that uses a manager-based API (inspired by Isaac Lab) and MuJoCo Warp for GPU-accelerated physics.

    Key technical features include:

    • Composable environments: Define observations, rewards, terminations, and other MDP terms as modular building blocks.
    • Direct MuJoCo data structures: Provides native access to MjModel and MjData without translation layers.
    • PyTorch-native: Observations, rewards, and actions are PyTorch tensors utilizing zero-copy GPU memory sharing.
    • Minimal dependencies: Designed for low startup latency and single-command installation via uv.

  3. Understand EntityData properties and GPU backing

    main

    The EntityData object provides access to simulation state via PyTorch tensors. These tensors are backed by MuJoCo Warp's GPU buffers, ensuring zero copy overhead.

    Key characteristics:

    • Shape: The first dimension of all properties is always num_envs, representing the number of parallel simulation worlds.
    • Data Freshness: Read properties reflect the state after sim.forward() has been called.

    Important Usage Note: If you write to the simulation state and then immediately attempt to read a derived property within the same event term, you must call sim.forward() between the write and the read to ensure the data is updated.

  4. Understand the Entity abstraction

    main

    An Entity represents a physical object in the simulation (e.g., a robot, a manipulated object, or a fixed fixture like a table). Unlike some other frameworks, mjlab uses a single Entity class for all variants, classified by two orthogonal properties:

    1. Base type:
      • is_fixed_base: The entity is welded to the world (e.g., a table).
      • is_floating_base: The entity has a free joint for 6-DOF movement (e.g., a humanoid).
    2. Articulation:
      • is_articulated: The entity has internal joints (revolute, prismatic, etc.).
      • is_non_articulated: The entity has no internal joints beyond a possible free joint.

    Important Note on Fixed-Base Entities: mjlab automatically wraps fixed-base entities in a mocap body to allow different positions across parallel environments. However, positioning only occurs when a reset event runs. You must include a reset event (e.g., reset_root_state_uniform) in your event configuration, otherwise, all fixed-base entities will remain at the origin. You can also reposition them at runtime using entity.write_mocap_pose_to_sim().

  6. Understand the mjlab architecture layers

    main

    mjlab is divided into two distinct layers:

    1. Simulation Layer: Models the robot and the world. It uses MjSpec to compose entities from MJCF files, which are then compiled and transferred to MuJoCo Warp for GPU-accelerated parallel simulation. This layer handles physics, entities, actuators, sensors, and scene composition.

    2. Manager Layer: Defines the reinforcement learning (RL) problem. It uses a manager-based design where small, self-contained terms (reward functions, observations, domain randomization) are registered with specific managers. These managers orchestrate the MDP (Markov Decision Process) lifecycle.

  7. Core design principles of mjlab

    main

    mjlab is built around three engineering commitments designed for efficient robot learning:

    1. Minimal installation friction: Uses lightweight dependencies and tools like uvx for fast setup.
    2. Transparent and inspectable physics: Uses the MuJoCo Warp stack to provide direct access to native MjModel and MjData structures for inspection and state access.
    3. Tight MuJoCo ecosystem integration: Works directly with MJCF files, MuJoCo Menagerie assets, and standard MuJoCo tooling without translation layers.

  8. Understand the scope of mjlab

    main

    mjlab is specialized for rigid-body robot learning.

    Included features:

    • Infrastructure for custom robots, tasks, sensors, and actuators.
    • Depth and raycast sensors for geometric perception.
    • Reference implementations for velocity tracking, motion imitation, and manipulation.

    Out of scope:

    • High-fidelity RGB rendering (though vision-based policies can be trained via distillation from privileged state).

  9. Understand Terrain Curriculum behavior

    main

    In a terrain grid (a num_rows x num_cols matrix of patches), columns represent terrain type variants and rows represent difficulty levels (row 0 is easiest, row num_rows - 1 is hardest).

    When TerrainGeneratorCfg.curriculum=True:

    1. Each column is assigned exactly one terrain type so difficulty increases monotonically along rows.
    2. At construction, each environment is assigned a random starting row within [0, max_init_terrain_level].
    3. The terrain_levels_vel curriculum term promotes or demotes environments on each reset based on distance traveled during the episode.
    4. Environments that reach the maximum level are randomly reassigned to any row to maintain coverage across all difficulty levels.

  10. Use the mjlab.scene module for environment management

    main
    The mjlab.scene module provides the core classes for managing simulation scenes in mjlab. It primarily consists of the Scene class for handling the simulation environment and the SceneCfg class for configuring scene parameters.

  11. Launch a viewer using the play script

    main

    You can launch interactive viewers using the --viewer flag with the play script.

    • Use --viewer native for a desktop window (MuJoCo native viewer).
    • Use --viewer viser for a browser-based viewer (opens at localhost:8080).
    • The default is auto, which selects native if a display server is available and falls back to viser on headless machines.

    For quick exploration without a trained checkpoint, use --agent zero or --agent random to use a dummy policy.

    # Desktop window (MuJoCo native viewer).
uv run play Mjlab-Velocity-Flat-Unitree-G1 --viewer native \
    --wandb-run-path your-entity/your-project/run_id

# Browser-based viewer (opens localhost:8080).
uv run play Mjlab-Velocity-Flat-Unitree-G1 --viewer viser \
    --wandb-run-path your-entity/your-project/run_id

# Quick exploration with a dummy policy
uv run play Mjlab-Velocity-Flat-Unitree-G1 --agent zero --viewer viser

  12. Play back a trained policy

    main

    Use the play command to evaluate a policy. You can load a policy from a Weights & Biases (W&B) run path or a local checkpoint file.

    Key arguments:

    • --agent: Policy mode. Options: "trained" (default), "zero" (zero actions), or "random" (uniform random).
    • --viewer: Viewer backend. Options: "native" (MuJoCo viewer) or "viser" (browser-based).
    • --no-terminations: Disable termination conditions so the policy runs indefinitely.
    # From W&B
uv run play Mjlab-Velocity-Flat-Unitree-G1 \
    --wandb-run-path your-entity/mjlab/run-id

# From a local checkpoint
uv run play Mjlab-Velocity-Flat-Unitree-G1 \
    --checkpoint-file logs/rsl_rl/g1_velocity/2025-01-27_14-30-00/model_1000.pt