mjlab example: go2_velocity

This example demonstrates how to integrate a custom robot into an existing mjlab velocity-tracking task. The setup trains a Unitree Go2 quadruped to walk and track commanded body velocities on flat or rough terrain.

📁 Project Structure

src/go2_velocity/
├── __init__.py                 # Task registration (entry point)
├── env_cfgs.py                 # Environment configs (sensors, rewards, terminations)
├── rl_cfg.py                   # RL hyperparameters (PPO)
└── go2/
    ├── go2_constants.py        # Robot definition (actuators, collisions, init state)
    └── xmls/
        ├── go2.xml             # MuJoCo MJCF model
        └── assets/             # Meshes and textures

---

🧠 How It Works

1. Depend on mjlab

Add mjlab as a dependency and register your task using Python entry points so that mjlab can auto-discover it.

[build-system]
requires = ["uv_build>=0.8.18,<0.9.0"]
build-backend = "uv_build"

[project]
dependencies = ["mjlab"]

[project.entry-points."mjlab.tasks"]
go2_velocity = "go2_velocity"

⚠️ The [build-system] section is required. Without it, the package will not be installed and the entry point will not be registered.

If developing locally against a checkout of mjlab:

[tool.uv.sources]
mjlab = { path = "../mjlab", editable = true }

---

2. Define Your Robot

The file go2_constants.py defines everything specific to the robot hardware:

• MuJoCo spec loaded from MJCF XML
• Actuators (PD gains, torque limits, armature)
• Collision geometry (feet, illegal contacts)
• Initial joint configuration

Key components:

• get_spec() – loads the MJCF model and assets
• BuiltinPositionActuatorCfg – PD controller configuration
• EntityCfg – bundles the spec, actuators, collisions, and initial state

This is the only place you need to touch when swapping in a different robot.

---

3. Configure the Environment

env_cfgs.py builds on make_velocity_env_cfg() (mjlab’s default velocity task):

• Registers the Go2 robot entity
• Defines foot contact sensors and illegal contacts
• Adjusts reward weights and termination conditions
• Configures viewer and terrain curriculum

A play=True variant disables domain randomization for evaluation and debugging.

---

4. Configure Reinforcement Learning

rl_cfg.py defines PPO hyperparameters via RslRlOnPolicyRunnerCfg:

• Policy / value network architecture
• Learning rate and optimizer settings
• PPO clip range and entropy coefficient

You can start with defaults and tune progressively.

---

5. Register Tasks

In __init__.py, tasks are registered using register_mjlab_task():

register_mjlab_task(
    task_id="Mjlab-Velocity-Rough-Unitree-Go2",
    env_cfg=unitree_go2_rough_env_cfg(),
    play_env_cfg=unitree_go2_rough_env_cfg(play=True),
    rl_cfg=unitree_go2_ppo_runner_cfg(),
    runner_cls=VelocityOnPolicyRunner,
)

Multiple variants (flat / rough terrain) can be registered in the same file.

---

▶️ Usage

Sanity Check

Watch the robot stand and fall under zero actions:

uv run play Mjlab-Velocity-Rough-Unitree-Go2 --agent zero

Training

CUDA_VISIBLE_DEVICES=0 uv run train Mjlab-Velocity-Rough-Unitree-Go2 \
  --env.scene.num-envs 4096 \
  --agent.max-iterations 3_000

Play a Trained Policy

uv run play Mjlab-Velocity-Rough-Unitree-Go2 \
  --wandb-run-path <wandb-run-path>

📚 References

• anymal_c_velocity
  • https://github.com/mujocolab/anymal_c_velocity
• Unitree RL Mjlab
  • https://github.com/unitreerobotics/unitree_rl_mjlab

---