CLAUDE.md

CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Repository Overview

This is a Hugging Face Space that provides a GPU-accelerated JupyterLab environment for training and simulating robots using the MuJoCo physics engine. The space covers a wide range of robotics applications including locomotion, manipulation, motion tracking, and general physics simulation. It is designed to run in a Docker container with NVIDIA GPU support for hardware-accelerated physics rendering.

What This Environment Supports

This is a general-purpose MuJoCo training environment with sample notebooks covering:

1. General MuJoCo Physics (tutorial.ipynb) - Comprehensive introduction to MuJoCo fundamentals including basic rendering, simulation loops, contacts, friction, tendons, actuators, sensors, and advanced rendering techniques

2. Locomotion (locomotion.ipynb) - Training quadrupedal and bipedal robots for walking, running, and acrobatic behaviors. Includes environments for Unitree Go1/G1, Boston Dynamics Spot, Google Barkour, Berkeley Humanoid, Unitree H1, and more

3. Manipulation (manipulation.ipynb) - Robot arm and dexterous hand control. Includes Franka Emika Panda pick-and-place tasks and Leap Hand dexterous manipulation with asymmetric actor-critic training

4. Motion Tracking (opentrack.ipynb) - Humanoid motion tracking and retargeting using the OpenTrack system with motion capture data

Architecture

Container Environment

• Base Image: nvidia/cuda:12.8.1-devel-ubuntu22.04
• Python: 3.13 (Miniconda)
• GPU Rendering: Uses EGL (OpenGL for headless rendering) with NVIDIA drivers
• Web Server: JupyterLab on port 7860

Key Components

1. GPU Initialization (init_gpu.py): Validates GPU setup before starting JupyterLab

   • Checks NVIDIA driver accessibility via nvidia-smi
   • Verifies EGL library availability (libEGL.so.1, libGL.so.1, libEGL_nvidia.so.0)
   • Tests EGL device initialization with multiple fallback methods (platform device, default display, surfaceless)
   • Validates MuJoCo rendering at multiple resolutions (64x64, 240x320, 480x640)
   • Critical environment variables: MUJOCO_GL=egl, PYOPENGL_PLATFORM=egl, EGL_PLATFORM=surfaceless

2. MuJoCo Playground Setup (init_mujoco.py): Downloads MuJoCo model assets

   • Imports mujoco_playground which automatically clones the mujoco_menagerie repository
   • This repository contains robot models (quadrupeds, bipeds, arms, hands, etc.)

3. Server Startup (start_server.sh): Container entrypoint

   • Sets up NVIDIA EGL library symlinks at runtime (searches /usr/local/nvidia/lib64, /usr/local/cuda/lib64, /usr/lib/nvidia)
   • Runs GPU validation (python init_gpu.py)
   • Downloads MuJoCo assets (python init_mujoco.py)
   • Disables JupyterLab announcements
   • Launches JupyterLab with iframe embedding support for Hugging Face Spaces

Sample Notebooks

Sample notebooks are organized in individual folders within samples/ and are automatically copied to /data/workspaces/ at container startup:

• samples/tutorial/ - Complete MuJoCo introduction (2258 lines) covering physics fundamentals, rendering, contacts, actuators, sensors, tendons, and camera control
• samples/locomotion/ - Quadrupedal and bipedal locomotion training (1762 lines) with PPO, domain randomization, curriculum learning, and policy fine-tuning
• samples/manipulation/ - Robot manipulation (649 lines) including pick-and-place (Panda arm) and dexterous manipulation (Leap Hand) with asymmetric actor-critic
• samples/opentrack/ - Humanoid motion tracking/retargeting (603 lines) including dataset download, training, checkpoint conversion, and video generation

Each sample is copied to its own workspace directory (/data/workspaces/<sample_name>/) at runtime. Notebooks are only copied if they don't already exist, preserving any user modifications.

Development Commands

Running Locally with Docker

# Build the container
docker build -t mujoco-training .

# Run with GPU support
docker run --gpus all -p 7860:7860 mujoco-training

Testing GPU Setup

# Validate GPU rendering capabilities (run inside container)
python init_gpu.py

# Check NVIDIA driver
nvidia-smi

# Test EGL libraries
ldconfig -p | grep EGL

JupyterLab Access

• Default port: 7860
• Default token: "huggingface" (set via JUPYTER_TOKEN environment variable)
• Default landing page: /lab/tree/workspaces/locomotion/locomotion.ipynb
• Notebook working directory: /data (when deployed as Hugging Face Space)

Persistent Storage and Workspaces

When deployed on Hugging Face Spaces, the /data directory is backed by persistent storage. At container startup, start_server.sh automatically:

1. Creates /data/workspaces/ if it doesn't exist
2. For each sample in samples/, creates /data/workspaces/<sample_name>/ if it doesn't exist
3. Copies the .ipynb file only if it doesn't already exist in the workspace (preserving user modifications)
4. Copies any additional files from the sample directory (datasets, scripts, etc.)

This ensures:

• User modifications to notebooks are preserved across container restarts
• Each sample has its own isolated workspace for generated data, models, and outputs
• Sample notebooks can include supporting files that are copied to the workspace
• Users can create additional workspaces in /data/workspaces/ for their own projects

Critical EGL Configuration

The container requires specific EGL configuration for headless GPU rendering:

1. NVIDIA EGL Vendor Config: Created at /usr/share/glvnd/egl_vendor.d/10_nvidia.json pointing to libEGL_nvidia.so.0
2. Library Path: LD_LIBRARY_PATH includes /usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/lib/x86_64-linux-gnu:/usr/local/cuda/lib64
3. Runtime Symlinks: start_server.sh creates symlinks to libEGL_nvidia.so.0 from mounted NVIDIA directories
4. Environment Variables: __EGL_VENDOR_LIBRARY_DIRS=/usr/share/glvnd/egl_vendor.d

Troubleshooting EGL Issues

If MuJoCo rendering fails:

1. Verify NVIDIA drivers: nvidia-smi should show GPU info
2. Check EGL vendor config: cat /usr/share/glvnd/egl_vendor.d/10_nvidia.json
3. Verify library loading: ldconfig -p | grep EGL
4. Run comprehensive diagnostic: python init_gpu.py
5. Check that MUJOCO_GL=egl is set: echo $MUJOCO_GL

Training Workflows

General MuJoCo Simulation (tutorial.ipynb)

Basic simulation loop:

import mujoco
model = mujoco.MjModel.from_xml_string(xml)
data = mujoco.MjData(model)

# Simulation loop
mujoco.mj_resetData(model, data)
while data.time < duration:
    mujoco.mj_step(model, data)
    # Read sensors, apply controls, etc.

Rendering:

with mujoco.Renderer(model, height, width) as renderer:
    mujoco.mj_forward(model, data)
    renderer.update_scene(data, camera="camera_name")
    pixels = renderer.render()

Locomotion Training (locomotion.ipynb)

Typical workflow using Brax + MuJoCo Playground:

1. Load environment: env = registry.load(env_name)
2. Get config: env_cfg = registry.get_default_config(env_name)
3. Configure PPO: ppo_params = locomotion_params.brax_ppo_config(env_name)
4. Apply domain randomization: randomizer = registry.get_domain_randomizer(env_name)
5. Train: Use brax.training.agents.ppo.train with the environment and randomization function
6. Save checkpoints: Policies saved to checkpoints/{env_name}/{step}/
7. Fine-tune: Restore from checkpoint and continue training with modified config

Available environments:

• Quadrupedal: Go1JoystickFlatTerrain, Go1JoystickRoughTerrain, Go1Getup, Go1Handstand, Go1Footstand, SpotFlatTerrainJoystick, SpotGetup, SpotJoystickGaitTracking, BarkourJoystick
• Bipedal: BerkeleyHumanoidJoystickFlatTerrain, BerkeleyHumanoidJoystickRoughTerrain, G1JoystickFlatTerrain, G1JoystickRoughTerrain, H1InplaceGaitTracking, H1JoystickGaitTracking, Op3Joystick, T1JoystickFlatTerrain, T1JoystickRoughTerrain

Full list: registry.locomotion.ALL_ENVS

Key training techniques:

• Domain Randomization: Randomizes friction, armature, center of mass, link masses for sim-to-real transfer
• Energy Penalties: energy_termination_threshold, reward_config.energy, reward_config.dof_acc to control power consumption and smoothness
• Curriculum Learning: Fine-tune from checkpoints with progressively modified reward configs
• Asymmetric Actor-Critic: Actor receives proprioception, critic receives privileged simulation state

Manipulation Training (manipulation.ipynb)

Similar to locomotion but focuses on:

• Pick-and-place tasks: PandaPickCubeOrientation (trains in ~3 minutes on RTX 4090)
• Dexterous manipulation: LeapCubeReorient (trains in ~33 minutes on RTX 4090)
• Asymmetric observations: Use policy_obs_key and value_obs_key in PPO params to train actor on sensor-like data while critic gets privileged state

Available environments: registry.manipulation.ALL_ENVS

Motion Tracking (opentrack.ipynb)

OpenTrack workflow for humanoid motion tracking:

1. Clone repository: git clone https://github.com/GalaxyGeneralRobotics/OpenTrack.git
2. Download mocap data: From huggingface.co/datasets/robfiras/loco-mujoco-datasets (Lafan1/UnitreeG1)
3. Train policy: python train_policy.py --exp_name debug --terrain_type flat_terrain
4. Convert checkpoint: python brax2torch.py --exp_name <exp_name> (Brax → PyTorch)
5. Generate videos: python play_policy.py --exp_name <exp_name> --use_renderer

Python Dependencies

Core stack (see requirements.txt):

• JupyterLab: 4.4.3 (with tornado 6.2 for compatibility)
• JAX: CUDA 12 support via jax[cuda12]
• MuJoCo: 3.3+ with MuJoCo MJX (JAX-based physics)
• Brax: JAX-based RL framework for massively parallel training
• MuJoCo Playground: Collection of robot environments and training utilities
• Supporting libraries: mediapy (video rendering), ipywidgets, nvidia-cusparse-cu12

File Structure

/
├── Dockerfile                      # Container with CUDA 12.8 + EGL setup
├── start_server.sh                 # Container entrypoint
├── init_gpu.py                     # GPU validation script (comprehensive EGL tests)
├── init_mujoco.py                  # MuJoCo Playground asset downloader
├── requirements.txt                # Python dependencies
├── packages.txt                    # System packages (currently empty)
├── on_startup.sh                   # Custom startup commands (placeholder)
├── login.html                      # Custom JupyterLab login page
└── samples/                        # Example notebooks (organized by topic)
    ├── tutorial/
    │   └── tutorial.ipynb          # MuJoCo fundamentals (2258 lines)
    ├── locomotion/
    │   └── locomotion.ipynb        # Robot locomotion (1762 lines)
    ├── manipulation/
    │   └── manipulation.ipynb      # Robot manipulation (649 lines)
    └── opentrack/
        └── opentrack.ipynb         # Motion tracking (603 lines)

When deployed as a Hugging Face Space with persistent storage:

/data/                              # Persistent storage volume (mounted at runtime)
└── workspaces/                     # Sample workspaces (created by start_server.sh)
    ├── tutorial/
    │   ├── tutorial.ipynb          # Copied from samples/, preserves user edits
    │   └── ...                     # User-generated data, models, outputs
    ├── locomotion/
    │   ├── locomotion.ipynb
    │   ├── checkpoints/            # Training checkpoints
    │   └── ...
    ├── manipulation/
    │   ├── manipulation.ipynb
    │   └── ...
    └── opentrack/
        ├── opentrack.ipynb
        ├── datasets/               # Downloaded mocap data
        ├── models/                 # Trained models
        └── videos/                 # Generated videos

Performance Notes

• Physics simulation: Can achieve 50,000+ Hz on single GPU with JAX/MJX (much faster than rendering)
• Rendering: Typically 30-60 Hz, much slower than physics
• Training times (on RTX 4090 / L40S):
  • Simple manipulation: 3 minutes
  • Quadrupedal joystick: 7 minutes
  • Bipedal locomotion: 17 minutes
  • Dexterous manipulation: 33 minutes
• Brax parallelization: Uses thousands of parallel environments for fast training
• Checkpointing: Critical for curriculum learning and fine-tuning

Common Patterns

Visualization Options

scene_option = mujoco.MjvOption()
scene_option.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True           # Show joints
scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True   # Show contacts
scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = True   # Show forces
scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True    # Transparency
scene_option.flags[mujoco.mjtVisFlag.mjVIS_PERTFORCE] = True      # Show perturbations

Named Access Pattern

# Instead of using indices
model.geom_rgba[geom_id, :]

# Use named access
model.geom('green_sphere').rgba
data.geom('box').xpos
data.joint('swing').qpos
data.sensor('accelerometer').data

Rendering Modes

• RGB rendering: renderer.render() - returns pixels
• Depth rendering: renderer.enable_depth_rendering() then renderer.render()
• Segmentation: renderer.enable_segmentation_rendering() - returns object IDs and types

Important Notes

• This is designed for Hugging Face Spaces with GPU instances (NVIDIA L40S or similar)
• All training uses JAX/Brax for massive parallelization across thousands of environments
• Policies are typically saved using Orbax checkpointing for fine-tuning
• Domain randomization is critical for sim-to-real transfer
• The environment supports multiple RL algorithms (PPO, SAC) through Brax
• Asymmetric actor-critic (different observations for policy and value function) is commonly used