mjlab/notebooks/create_new_task.ipynb

🤖 CartPole Tutorial with mjlab

This notebook demonstrates how to create a custom reinforcement learning task using mjlab. We'll build a CartPole environment from scratch, including:

1. Robot Definition - Define the CartPole model in MuJoCo XML
2. Task Configuration - Set up observations, actions, rewards, and terminations
3. Training - Train a policy using PPO
4. Evaluation - Visualize the simulation with the trained policy

Note

: This tutorial demonstrates how to create a new task in mjlab. For more context, see the official documentation.

📦 Setup and Installation

In [ ]:

# Clone the mjlab repository
!if [ ! -d "mjlab" ]; then git clone -q https://github.com/mujocolab/mjlab.git; fi
%cd /content/mjlab

# Install mjlab in editable mode
!pip install -e . -q

print("✓ Installation complete!")

🔑 WandB Setup

Configure Weights & Biases for experiment tracking. Add your WandB API key to Colab Secrets:

• WANDB_API_KEY: from wandb.ai/authorize
• WANDB_ENTITY: your wandb entity name

In [ ]:

import os

from google.colab import userdata

try:
  # Set this to use wandb logger
  os.environ["WANDB_API_KEY"] = userdata.get("WANDB_API_KEY")
  os.environ["WANDB_ENTITY"] = userdata.get("WANDB_ENTITY")

  print("✓ WandB configured successfully!")
except (AttributeError, KeyError):
  # Set this to disable wandb logger
  os.environ["WANDB_MODE"] = "disabled"

  print("⚠ WandB secrets not found. Training will proceed without WandB logging.")

---

🤖 Step 1: Define the Robot

We'll create a simple CartPole robot with:

• A sliding cart (1 DOF)
• A hinged pole (1 DOF)
• A velocity actuator to control the cart

📁 Structure Directories

In [ ]:

# Create the cartpole robot directory structure
!mkdir -p /content/mjlab/src/mjlab/asset_zoo/robots/cartpole/
!mkdir -p /content/mjlab/src/mjlab/asset_zoo/robots/cartpole/xmls

print("✓ Directory structure created")

📝 Create MuJoCo XML Model

This XML defines the CartPole physics:

• Ground plane for visualization
• Cart body with a sliding joint (±2m range)
• Pole body with a hinge joint (±90° range)
• Velocity actuator for cart control

In [ ]:

%%writefile /content/mjlab/src/mjlab/asset_zoo/robots/cartpole/xmls/cartpole.xml
<mujoco model="cartpole">
  <compiler angle="degree" coordinate="local" inertiafromgeom="true"/>
  <worldbody>
    <geom name="ground" type="plane" pos="0 0 0" size="5 5 0.1" rgba="0.8 0.9 0.8 1"/>
    <body name="cart" pos="0 0 0.1">
      <geom type="box" size="0.2 0.1 0.1" rgba="0.2 0.2 0.8 1" mass="1.0"/>
      <joint name="slide" type="slide" axis="1 0 0" limited="true" range="-2 2"/>
      <body name="pole" pos="0 0 0.1">
        <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0 0 1" rgba="0.8 0.2 0.2 1" mass="2.0"/>
        <joint name="hinge" type="hinge" axis="0 1 0" range="-90 90"/>
      </body>
    </body>
  </worldbody>
  <actuator>
    <velocity name="slide_velocity" joint="slide" ctrlrange="-20 20" kv="20"/>
  </actuator>
</mujoco>

⚙️ Create Robot Configuration

In [ ]:

%%writefile /content/mjlab/src/mjlab/asset_zoo/robots/cartpole/cartpole_constants.py
from pathlib import Path
import mujoco

from mjlab import MJLAB_SRC_PATH
from mjlab.entity import Entity, EntityCfg, EntityArticulationInfoCfg
from mjlab.actuator import XmlActuatorCfg

CARTPOLE_XML: Path = (
  MJLAB_SRC_PATH / "asset_zoo" / "robots" / "cartpole" / "xmls" / "cartpole.xml"
)
assert CARTPOLE_XML.exists(), f"XML not found: {CARTPOLE_XML}"

def get_spec() -> mujoco.MjSpec:
  return mujoco.MjSpec.from_file(str(CARTPOLE_XML))

def get_cartpole_robot_cfg() -> EntityCfg:
  """Get a fresh CartPole robot configuration instance."""
  actuators = (
    XmlActuatorCfg(
      target_names_expr=("slide",),
    ),
  )
  articulation = EntityArticulationInfoCfg(actuators=actuators)
  return EntityCfg(
    spec_fn=get_spec,
    articulation=articulation
  )

# if __name__ == "__main__":
#   import mujoco.viewer as viewer
#   robot = Entity(get_cartpole_robot_cfg())
#   viewer.launch(robot.spec.compile())

In [ ]:

# Create __init__.py for the cartpole robot package
%%writefile /content/mjlab/src/mjlab/asset_zoo/robots/cartpole/__init__.py
# Empty __init__.py to mark the directory as a Python package

In [ ]:

import sys

# Append src dir to python path
mjlab_src = "/content/mjlab/src"
if mjlab_src not in sys.path:
  sys.path.insert(0, mjlab_src)
  print(f"✓ Added {mjlab_src} to Python path")

✅ Verify Robot Setup

Let's test that the robot can be loaded correctly.

In [ ]:

from mjlab.asset_zoo.robots.cartpole.cartpole_constants import get_cartpole_robot_cfg

from mjlab.entity import Entity

# Load the robot
robot = Entity(get_cartpole_robot_cfg())
model = robot.spec.compile()

# Display robot information
print("✓ CartPole robot loaded successfully!")
print(f"  • Degrees of Freedom (DOF): {model.nv}")
print(f"  • Number of Actuators: {model.nu}")
print(f"  • Bodies: {model.nbody}")
print(f"  • Joints: {model.njnt}")

📋 Register the Robot

Add the CartPole robot to the asset zoo registry.

In [ ]:

# Add CartPole import to robots __init__.py
with open("/content/mjlab/src/mjlab/asset_zoo/robots/__init__.py", "a") as f:
  f.write("\n# CartPole robot\n")
  f.write("from mjlab.asset_zoo.robots.cartpole.cartpole_constants import ")
  f.write("get_cartpole_robot_cfg as get_cartpole_robot_cfg\n")

print("✓ CartPole robot registered in asset zoo")

---

🎯 Step 2: Define the Task (MDP)

Now we'll define the Markov Decision Process:

• Observations: pole angle, angular velocity, cart position, cart velocity
• Actions: cart velocity commands
• Rewards: upright reward + effort penalty
• Terminations: pole tips over or timeout
• Events: random pushes for robustness

📁 Create Task Directory

In [ ]:

!mkdir -p /content/mjlab/src/mjlab/tasks/cartpole

print("✓ Task directory created")

📝 Create Environment Configuration

This file contains the MDP (Markov Decision Process) components:

1. Scene Config: 64 parallel environments
2. Actions: Joint velocity control with 20.0 scale
3. Observations: Normalized state variables
4. Rewards: Upright reward (5.0) + effort penalty (-0.01)
5. Events: Joint resets + random pushes
6. Terminations: Pole tipped (>30°) or timeout (10s)

In [ ]:

%%writefile /content/mjlab/src/mjlab/tasks/cartpole/env_cfg.py
"""CartPole task environment configuration."""

import math
import torch

from mjlab.envs import ManagerBasedRlEnvCfg
from mjlab.envs.mdp.actions import JointVelocityActionCfg
from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
from mjlab.managers.reward_manager import RewardTermCfg
from mjlab.managers.termination_manager import TerminationTermCfg
from mjlab.managers.event_manager import EventTermCfg
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.scene import SceneCfg
from mjlab.sim import MujocoCfg, SimulationCfg
from mjlab.viewer import ViewerConfig
from mjlab.asset_zoo.robots.cartpole.cartpole_constants import get_cartpole_robot_cfg
from mjlab.envs import mdp


def cartpole_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
  """Create CartPole environment configuration.

  Args:
    play: If True, disables corruption and extends episode length for evaluation.
  """

  # ==============================================================================
  # Scene Configuration
  # ==============================================================================

  scene_cfg = SceneCfg(
    num_envs=64 if not play else 16,  # Fewer envs for play mode
    extent=1.0,   # Spacing between environments
    entities={"robot": get_cartpole_robot_cfg()},
  )

  viewer_cfg = ViewerConfig(
    origin_type=ViewerConfig.OriginType.ASSET_BODY,
    entity_name="robot",
    body_name="pole",
    distance=3.0,
    elevation=10.0,
    azimuth=90.0,
  )

  sim_cfg = SimulationCfg(
    mujoco=MujocoCfg(
      timestep=0.02,  # 50 Hz control
      iterations=1,
    ),
  )

  # ==============================================================================
  # Actions
  # ==============================================================================

  actions = {
    "joint_pos": JointVelocityActionCfg(
      entity_name="robot",
      actuator_names=(".*",),
      scale=20.0,
      use_default_offset=False,
    ),
  }

  # ==============================================================================
  # Observations
  # ==============================================================================

  actor_terms = {
    "angle": ObservationTermCfg(
      func=lambda env: env.sim.data.qpos[:, 1:2] / math.pi
    ),
    "ang_vel": ObservationTermCfg(
      func=lambda env: env.sim.data.qvel[:, 1:2] / 5.0
    ),
    "cart_pos": ObservationTermCfg(
      func=lambda env: env.sim.data.qpos[:, 0:1] / 2.0
    ),
    "cart_vel": ObservationTermCfg(
      func=lambda env: env.sim.data.qvel[:, 0:1] / 20.0
    ),
  }

  observations = {
    "actor": ObservationGroupCfg(
      terms=actor_terms,
      concatenate_terms=True,
      enable_corruption=not play,  # Disable corruption in play mode
    ),
    "critic": ObservationGroupCfg(
      terms=actor_terms,  # Critic uses same observations
      concatenate_terms=True,
      enable_corruption=False,
    ),
  }

  # ==============================================================================
  # Rewards
  # ==============================================================================

  def compute_upright_reward(env):
    """Reward for keeping pole upright (cosine of angle)."""
    return env.sim.data.qpos[:, 1].cos()

  def compute_effort_penalty(env):
    """Penalty for control effort."""
    return -0.01 * (env.sim.data.ctrl[:, 0] ** 2)

  rewards = {
    "upright": RewardTermCfg(func=compute_upright_reward, weight=5.0),
    "effort": RewardTermCfg(func=compute_effort_penalty, weight=1.0),
  }

  # ==============================================================================
  # Events
  # ==============================================================================

  def random_push_cart(env, env_ids, force_range=(-5, 5)):
    """Apply random force to cart for robustness training."""
    n = len(env_ids)
    random_forces = (
      torch.rand(n, device=env.device) *
      (force_range[1] - force_range[0]) +
      force_range[0]
    )
    env.sim.data.qfrc_applied[env_ids, 0] = random_forces

  events = {
    "reset_robot_joints": EventTermCfg(
      func=mdp.reset_joints_by_offset,
      mode="reset",
      params={
        "asset_cfg": SceneEntityCfg("robot"),
        "position_range": (-0.1, 0.1),
        "velocity_range": (-0.1, 0.1),
      },
    ),
  }

  # Add random pushes only in training mode
  if not play:
    events["random_push"] = EventTermCfg(
      func=random_push_cart,
      mode="interval",
      interval_range_s=(1.0, 2.0),
      params={"force_range": (-20.0, 20.0)},
    )

  # ==============================================================================
  # Terminations
  # ==============================================================================

  def check_pole_tipped(env):
    """Check if pole has tipped beyond 30 degrees."""
    return env.sim.data.qpos[:, 1].abs() > math.radians(30)

  terminations = {
    "timeout": TerminationTermCfg(func=mdp.time_out, time_out=True),
    "tipped": TerminationTermCfg(func=check_pole_tipped, time_out=False),
  }

  # ==============================================================================
  # Environment Configuration
  # ==============================================================================

  return ManagerBasedRlEnvCfg(
    scene=scene_cfg,
    observations=observations,
    actions=actions,
    rewards=rewards,
    events=events,
    terminations=terminations,
    sim=sim_cfg,
    viewer=viewer_cfg,
    decimation=1,           # No action repeat
    episode_length_s=int(1e9) if play else 10.0,  # Infinite for play, 10s for training
  )

⚙️ Create RL Configuration

This file defines the PPO (Proximal Policy Optimization) training parameters.

In [ ]:

%%writefile /content/mjlab/src/mjlab/tasks/cartpole/rl_cfg.py
"""RL configuration for CartPole task."""

from mjlab.rl.config import (
  RslRlOnPolicyRunnerCfg,
  RslRlModelCfg,
  RslRlPpoAlgorithmCfg,
)


def cartpole_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
  """Create RL runner configuration for CartPole task."""
  return RslRlOnPolicyRunnerCfg(
    actor=RslRlModelCfg(
      hidden_dims=(256, 128, 64),  # Smaller network for simpler task
      activation="elu",
      obs_normalization=True,
      distribution_cfg={
        "class_name": "GaussianDistribution",
        "init_std": 1.0,
        "std_type": "scalar",
      },
    ),
    critic=RslRlModelCfg(
      hidden_dims=(256, 128, 64),
      activation="elu",
      obs_normalization=True,
    ),
    algorithm=RslRlPpoAlgorithmCfg(
      value_loss_coef=1.0,
      use_clipped_value_loss=True,
      clip_param=0.2,
      entropy_coef=0.01,
      num_learning_epochs=5,
      num_mini_batches=4,
      learning_rate=1.0e-3,
      schedule="adaptive",
      gamma=0.99,
      lam=0.95,
      desired_kl=0.01,
      max_grad_norm=1.0,
    ),
    experiment_name="cartpole",
    save_interval=50,
    num_steps_per_env=24,
    max_iterations=5_000,  # Fewer iterations for simpler task
  )

📋 Register the Task Environment

Register the CartPole task with mjlab registry.

In [ ]:

%%writefile /content/mjlab/src/mjlab/tasks/cartpole/__init__.py
"""CartPole task registration."""

from mjlab.tasks.registry import register_mjlab_task
from mjlab.rl.runner import MjlabOnPolicyRunner

from .env_cfg import cartpole_env_cfg
from .rl_cfg import cartpole_ppo_runner_cfg

register_mjlab_task(
  task_id="Mjlab-Cartpole",
  env_cfg=cartpole_env_cfg(),
  play_env_cfg=cartpole_env_cfg(play=True),
  rl_cfg=cartpole_ppo_runner_cfg(),
  runner_cls=MjlabOnPolicyRunner,
)

---

🚀 Step 3: Train the Agent

Now let's train a PPO policy to balance the CartPole!

In [ ]:

!python -m mjlab.scripts.train Mjlab-Cartpole --agent.max-iterations 201 --agent.save-interval 20

📁 Locate Training Checkpoints

After training, checkpoints are saved locally.

In [ ]:

import os
import re
from pathlib import Path

# Find the most recent training run
log_dir = Path("/content/mjlab/logs/rsl_rl/cartpole")
if log_dir.exists():
  runs = sorted(log_dir.glob("*"), key=os.path.getmtime, reverse=True)
  if runs:
    latest_run = runs[0]
    print(f"✓ Latest training run: {latest_run.name}\n")

    # List checkpoints - sorted by iteration number
    checkpoints = list(latest_run.glob("model_*.pt"))
    if checkpoints:
      # Extract iteration number and sort numerically
      def get_iteration(ckpt):
        match = re.search(r"model_(\d+)\.pt", ckpt.name)
        return int(match.group(1)) if match else 0

      checkpoints = sorted(checkpoints, key=get_iteration)

      print(f"Found {len(checkpoints)} checkpoints:")
      for ckpt in checkpoints[-5:]:  # Show last 5
        size_mb = ckpt.stat().st_size / (1024 * 1024)
        iteration = get_iteration(ckpt)
        print(f"  • {ckpt.name} (iteration {iteration}, {size_mb:.2f} MB)")

      # Store the last checkpoint path
      last_checkpoint = str(checkpoints[-1])
      print(f"\n💾 Last checkpoint: {last_checkpoint}")
    else:
      print("⚠ No checkpoints found yet")
  else:
    print("⚠ No training runs found")
else:
  print("⚠ Log directory not found. Have you run training yet?")

---

🎮 Step 4: Visualize the Trained Policy

Let's see the trained policy in action!

🌐 Launch Viser API

In [ ]:

import subprocess
import sys

process = subprocess.Popen(
  [
    "python",
    "-m",
    "mjlab.scripts.play",
    "Mjlab-Cartpole",
    "--checkpoint_file",
    last_checkpoint,
    "--num_envs",
    "4",
  ],
  stdout=subprocess.PIPE,
  stderr=subprocess.STDOUT,
  universal_newlines=True,
  bufsize=1,
)

detected_port = None

for line in process.stdout:
  print(line, end="")
  sys.stdout.flush()

  # Extract port number from viser output
  port_match = re.search(r":(\d{4})", line)
  if port_match and "viser" in line.lower():
    detected_port = int(port_match.group(1))
    print("\n" + "=" * 34)
    print(f"✅ Server is running on port {detected_port}!")
    print("=" * 34)
    break

🖥️ Embed Client as iframe

In [ ]:

from google.colab import output

port = detected_port if detected_port else 8081
output.serve_kernel_port_as_iframe(port=port, height=700)