[software] 添加16DOF早期训练仿真与Sim2Real闭环

2026-07-21 16:15:14 +08:00
parent 9bd22225f9
commit e9e2c946b3
681 changed files with 137221 additions and 8 deletions
@@ -0,0 +1,76 @@
+# mjlab Nightly Benchmarks
+
+This directory contains scripts for automated nightly benchmarking of mjlab.
+
+## Overview
+
+The nightly benchmark system:
+1. Trains a tracking policy on the latest commit
+2. Evaluates the policy across 1024 trials
+3. Measures simulation throughput
+4. Generates an HTML report with historical trends
+5. Publishes results to GitHub Pages
+
+## Usage
+
+### Run the full nightly benchmark
+
+```bash
+./scripts/benchmarks/nightly_train.sh
+```
+
+### Skip training (regenerate report only)
+
+```bash
+SKIP_TRAINING=1 ./scripts/benchmarks/nightly_train.sh
+```
+
+### Skip training and throughput
+
+```bash
+SKIP_TRAINING=1 SKIP_THROUGHPUT=1 ./scripts/benchmarks/nightly_train.sh
+```
+
+### Regenerate report directly (no git operations)
+
+```bash
+uv run python scripts/benchmarks/generate_report.py \
+  --entity gcbc_researchers \
+  --tag nightly \
+  --output-dir benchmark_results
+```
+
+### Measure throughput only
+
+```bash
+uv run python scripts/benchmarks/measure_throughput.py \
+  --num-envs 4096 \
+  --output-dir benchmark_results
+```
+
+## Configuration
+
+Environment variables for `nightly_train.sh`:
+
+- `CUDA_DEVICE` - GPU device to use (default: 0)
+- `WANDB_TAGS` - Comma-separated tags for the run (default: nightly)
+- `SKIP_TRAINING` - Set to "1" to skip training
+- `SKIP_THROUGHPUT` - Set to "1" to skip throughput benchmarking
+
+## Automated Setup
+
+See [systemd/README.md](systemd/README.md) for instructions on setting up automated nightly runs using systemd timers.
+
+## Report Options
+
+The `generate_report.py` script supports:
+
+- `--eval-limit N` - Maximum number of NEW runs to evaluate per invocation (default: 10)
+  - Set to 0 for no limit
+  - Historical cached results are always preserved
+- `--tag TAG` - Filter runs by WandB tag (default: "nightly")
+- `--num-envs N` - Number of parallel environments for evaluation (default: 1024)
+
+## Viewing Results
+
+Reports are published to: https://mujocolab.github.io/mjlab/nightly/
@@ -0,0 +1,796 @@
+"""Generate benchmark report from evaluation metrics.
+
+This script runs policy evaluation on nightly runs and generates a static HTML
+dashboard for tracking policy performance over time.
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+from datetime import datetime
+from pathlib import Path
+
+import tyro
+import wandb
+
+import mjlab
+from mjlab.tasks.tracking.scripts.evaluate import EvaluateConfig, run_evaluate
+
+# Metrics to display: (key, label, unit, scale, higher_is_better)
+METRICS = [
+  ("success_rate", "Success Rate", "%", 100, True),
+  ("mpkpe", "MPKPE", "m", 1, False),
+  ("r_mpkpe", "R-MPKPE", "m", 1, False),
+  ("ee_pos_error", "EE Position Error", "m", 1, False),
+  ("ee_ori_error", "EE Orientation Error", "rad", 1, False),
+  ("joint_vel_error", "Joint Velocity Error", "rad/s", 1, False),
+]
+
+
+def evaluate_run(run_path: str, num_envs: int = 1024) -> dict:
+  """Evaluate a single run and return metrics with metadata."""
+  api = wandb.Api()
+  run = api.run(run_path)
+
+  print(f"Evaluating run: {run.name} ({run.id})")
+
+  cfg = EvaluateConfig(wandb_run_path=run_path, num_envs=num_envs)
+  metrics = run_evaluate("Mjlab-Tracking-Flat-Unitree-G1", cfg)
+
+  # Get commit SHA from run metadata.
+  commit = run.commit or run.config.get("commit", "unknown")
+
+  return {
+    "id": run.id,
+    "name": run.name,
+    "url": run.url,
+    "created_at": run.created_at,
+    "commit": commit[:7] if len(commit) > 7 else commit,
+    "metrics": metrics,
+  }
+
+
+def load_throughput_data(output_dir: Path) -> list[dict]:
+  """Load throughput benchmark data if available."""
+  data_file = output_dir / "throughput_data.json"
+  if not data_file.exists():
+    return []
+  with open(data_file) as f:
+    return json.load(f)
+
+
+def generate_html_report(runs: list[dict], output_dir: Path) -> None:
+  """Generate static HTML dashboard from evaluation data."""
+  output_dir.mkdir(parents=True, exist_ok=True)
+
+  # Save raw data.
+  with open(output_dir / "data.json", "w") as f:
+    json.dump(runs, f, indent=2, default=str)
+
+  # Copy task images for the throughput dashboard.
+  images_src = Path(__file__).parent / "nightly_images"
+  if images_src.is_dir():
+    images_dst = output_dir / "images"
+    if images_dst.exists():
+      shutil.rmtree(images_dst)
+    shutil.copytree(images_src, images_dst)
+
+  # Load throughput data if available.
+  throughput_data = load_throughput_data(output_dir)
+
+  html = generate_dashboard_html(runs, throughput_data)
+  with open(output_dir / "index.html", "w") as f:
+    f.write(html)
+
+  print(f"Report generated at {output_dir / 'index.html'}")
+
+
+def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> str:
+  """Generate the HTML dashboard content."""
+  runs_json = json.dumps(runs, default=str)
+  metrics_json = json.dumps(METRICS)
+  throughput_json = json.dumps(throughput_data, default=str)
+  timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
+  github_repo = "https://github.com/mujocolab/mjlab"
+
+  return f"""<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>mjlab Nightly Benchmark</title>
+    <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
+    <script src="https://cdn.jsdelivr.net/npm/chartjs-adapter-date-fns"></script>
+    <style>
+        :root {{
+            --bg: #ffffff;
+            --bg-card: #f6f8fa;
+            --text: #1f2328;
+            --text-dim: #656d76;
+            --border: #d0d7de;
+            --accent: #0969da;
+            --green: #1a7f37;
+            --red: #cf222e;
+        }}
+        @media (prefers-color-scheme: dark) {{
+            :root:not([data-theme="light"]) {{
+                --bg: #0d1117;
+                --bg-card: #161b22;
+                --text: #c9d1d9;
+                --text-dim: #8b949e;
+                --border: #30363d;
+                --accent: #58a6ff;
+                --green: #3fb950;
+                --red: #f85149;
+            }}
+        }}
+        :root[data-theme="dark"] {{
+            --bg: #0d1117;
+            --bg-card: #161b22;
+            --text: #c9d1d9;
+            --text-dim: #8b949e;
+            --border: #30363d;
+            --accent: #58a6ff;
+            --green: #3fb950;
+            --red: #f85149;
+        }}
+        * {{ margin: 0; padding: 0; box-sizing: border-box; }}
+        body {{
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
+            background: var(--bg);
+            color: var(--text);
+            padding: 2rem;
+            max-width: 1400px;
+            margin: 0 auto;
+        }}
+        header {{
+            display: flex;
+            justify-content: space-between;
+            align-items: center;
+            margin-bottom: 2rem;
+            padding-bottom: 1rem;
+            border-bottom: 1px solid var(--border);
+        }}
+        h1 {{ font-size: 1.5rem; }}
+        .subtitle {{
+            font-size: 0.875rem;
+            color: var(--text-dim);
+            margin-top: 0.25rem;
+        }}
+        .timestamp {{ color: var(--text-dim); font-size: 0.875rem; }}
+        .charts {{
+            display: grid;
+            grid-template-columns: 1fr;
+            gap: 1.5rem;
+            margin-bottom: 2rem;
+        }}
+        .chart-card {{
+            background: var(--bg-card);
+            border: 1px solid var(--border);
+            border-radius: 8px;
+            padding: 1rem;
+        }}
+        .chart-title {{
+            font-size: 0.875rem;
+            font-weight: 500;
+            margin-bottom: 0.5rem;
+            display: flex;
+            justify-content: space-between;
+        }}
+        .chart-value {{ color: var(--text-dim); }}
+        .chart-container {{ height: 350px; }}
+        a {{ color: var(--accent); text-decoration: none; }}
+        a:hover {{ text-decoration: underline; }}
+        .theme-toggle {{
+            background: var(--bg-card);
+            border: 1px solid var(--border);
+            border-radius: 6px;
+            padding: 0.5rem;
+            cursor: pointer;
+            color: var(--text);
+            font-size: 1rem;
+            line-height: 1;
+        }}
+        .theme-toggle:hover {{ border-color: var(--accent); }}
+        .header-right {{ display: flex; align-items: center; gap: 1rem; }}
+        .tabs {{
+            display: flex;
+            gap: 0.5rem;
+            margin-bottom: 1.5rem;
+        }}
+        .tab {{
+            background: var(--bg-card);
+            border: 1px solid var(--border);
+            border-radius: 6px;
+            padding: 0.5rem 1rem;
+            cursor: pointer;
+            color: var(--text);
+            font-size: 0.875rem;
+            font-weight: 500;
+        }}
+        .tab:hover {{ border-color: var(--accent); }}
+        .tab.active {{
+            background: var(--accent);
+            border-color: var(--accent);
+            color: white;
+        }}
+        .tab-content {{ display: none; }}
+        .tab-content.active {{ display: block; }}
+        .tab-description {{
+            font-size: 0.875rem;
+            color: var(--text-dim);
+            margin-bottom: 1.5rem;
+            line-height: 1.5;
+        }}
+        .task-grid {{
+            display: grid;
+            grid-template-columns: repeat(auto-fit, minmax(220px, 1fr));
+            gap: 1rem;
+            margin-bottom: 1.5rem;
+        }}
+        .task-card {{
+            background: var(--bg-card);
+            border: 2px solid var(--border);
+            border-radius: 10px;
+            padding: 0.75rem;
+            cursor: pointer;
+            transition: border-color 0.15s, box-shadow 0.15s;
+            text-align: center;
+        }}
+        .task-card:hover {{ border-color: var(--accent); }}
+        .task-card.active {{
+            border-color: var(--accent);
+            box-shadow: 0 0 0 1px var(--accent);
+        }}
+        .task-card img {{
+            width: 100%;
+            aspect-ratio: 16/10;
+            object-fit: cover;
+            border-radius: 6px;
+            margin-bottom: 0.5rem;
+            background: var(--border);
+        }}
+        .task-card .task-name {{
+            font-size: 0.85rem;
+            font-weight: 600;
+        }}
+        .task-card .task-stat {{
+            font-size: 0.75rem;
+            color: var(--text-dim);
+            margin-top: 0.2rem;
+        }}
+        .task-chart-area {{
+            display: none;
+        }}
+        .task-chart-area.active {{
+            display: block;
+        }}
+        footer {{
+            margin-top: 3rem;
+            padding-top: 1.5rem;
+            border-top: 1px solid var(--border);
+            font-size: 0.8rem;
+            color: var(--text-dim);
+            line-height: 1.6;
+        }}
+        @media (max-width: 600px) {{
+            body {{ padding: 1rem; }}
+            h1 {{ font-size: 1.25rem; }}
+            .tabs {{ flex-wrap: wrap; }}
+            .task-grid {{ grid-template-columns: 1fr; }}
+        }}
+    </style>
+</head>
+<body>
+    <header>
+        <div>
+            <h1><a href="{github_repo}" style="color:inherit;text-decoration:none">mjlab</a> Nightly Benchmark</h1>
+            <div class="subtitle">Performance tracking over time</div>
+        </div>
+        <div class="header-right">
+            <span class="timestamp">Updated: {timestamp}</span>
+            <button class="theme-toggle" id="theme-toggle" title="Toggle theme">
+                <span id="theme-icon"></span>
+            </button>
+        </div>
+    </header>
+
+    <div class="tabs">
+        <button class="tab active" data-tab="tracking">Tracking Eval</button>
+        <button class="tab" data-tab="throughput">Throughput</button>
+    </div>
+
+    <div id="tracking" class="tab-content active">
+        <p class="tab-description">Nightly motion imitation training and evaluation on Unitree G1 (1024 trials per run).</p>
+        <div class="charts" id="charts"></div>
+    </div>
+
+    <div id="throughput" class="tab-content">
+        <p class="tab-description">Physics simulation throughput across tasks (4096 parallel envs, NVIDIA RTX 5090).</p>
+        <div class="task-grid" id="task-grid"></div>
+        <div id="task-chart-panels"></div>
+    </div>
+
+    <footer>
+        These benchmarks run nightly using the latest commit.<br>
+        GPU: NVIDIA RTX 5090
+    </footer>
+
+    <script>
+        // Theme toggle logic
+        const themeToggle = document.getElementById('theme-toggle');
+        const themeIcon = document.getElementById('theme-icon');
+        const root = document.documentElement;
+
+        function getSystemTheme() {{
+            return window.matchMedia('(prefers-color-scheme: dark)').matches ? 'dark' : 'light';
+        }}
+
+        function getEffectiveTheme() {{
+            const stored = localStorage.getItem('theme');
+            if (stored === 'dark' || stored === 'light') return stored;
+            return getSystemTheme();
+        }}
+
+        function updateThemeIcon() {{
+            const stored = localStorage.getItem('theme');
+            if (!stored) {{
+                themeIcon.textContent = '\u2699\ufe0f'; // gear for auto
+                themeToggle.title = 'Theme: System (click to toggle)';
+            }} else if (stored === 'dark') {{
+                themeIcon.textContent = '\U0001f319'; // moon
+                themeToggle.title = 'Theme: Dark (click to toggle)';
+            }} else {{
+                themeIcon.textContent = '\u2600\ufe0f'; // sun
+                themeToggle.title = 'Theme: Light (click to toggle)';
+            }}
+        }}
+
+        function applyTheme() {{
+            const stored = localStorage.getItem('theme');
+            if (stored) {{
+                root.setAttribute('data-theme', stored);
+            }} else {{
+                root.removeAttribute('data-theme');
+            }}
+            updateThemeIcon();
+            updateChartColors();
+        }}
+
+        function cycleTheme() {{
+            const stored = localStorage.getItem('theme');
+            if (!stored) {{
+                // auto -> dark
+                localStorage.setItem('theme', 'dark');
+            }} else if (stored === 'dark') {{
+                // dark -> light
+                localStorage.setItem('theme', 'light');
+            }} else {{
+                // light -> auto
+                localStorage.removeItem('theme');
+            }}
+            applyTheme();
+        }}
+
+        themeToggle.addEventListener('click', cycleTheme);
+        window.matchMedia('(prefers-color-scheme: dark)').addEventListener('change', applyTheme);
+
+        const runs = {runs_json};
+        const METRICS = {metrics_json};
+        const GITHUB_REPO = '{github_repo}';
+
+        // Sort by date ascending for charts.
+        runs.sort((a, b) => new Date(a.created_at) - new Date(b.created_at));
+
+        const colors = {{
+            success_rate: '#3fb950',
+            mpkpe: '#58a6ff',
+            r_mpkpe: '#a371f7',
+            ee_pos_error: '#f0883e',
+            ee_ori_error: '#f85149',
+            joint_vel_error: '#79c0ff'
+        }};
+
+        let charts = [];
+
+        function updateChartColors() {{
+            const style = getComputedStyle(root);
+            const textDim = style.getPropertyValue('--text-dim').trim();
+            const border = style.getPropertyValue('--border').trim();
+            const isDark = getEffectiveTheme() === 'dark';
+            const gridColor = isDark ? '#3a424b' : '#d0d7de';
+            Chart.defaults.color = textDim;
+            Chart.defaults.borderColor = gridColor;
+            charts.forEach(c => c.update());
+        }}
+
+        // Initialize theme before creating charts so grid colors are correct.
+        applyTheme();
+
+        // Charts
+        const chartsContainer = document.getElementById('charts');
+
+        // Compute a rolling average over the last `window` points.
+        function rollingAvg(data, window) {{
+            return data.map((d, i) => {{
+                const start = Math.max(0, i - window + 1);
+                const slice = data.slice(start, i + 1);
+                const avg = slice.reduce((s, p) => s + p.y, 0) / slice.length;
+                return {{ x: d.x, y: avg }};
+            }});
+        }}
+
+        const AVG_WINDOW = 7;
+
+        METRICS.forEach(([key, label, unit, scale, higherIsBetter]) => {{
+            const data = runs.map(r => ({{
+                x: new Date(r.created_at),
+                y: r.metrics[key] * scale,
+                commit: r.commit,
+                name: r.name
+            }}));
+
+            const avgData = rollingAvg(data, AVG_WINDOW);
+            const color = colors[key] || '#58a6ff';
+
+            const latestVal = data[data.length - 1]?.y;
+            const arrow = higherIsBetter ? '\u2191' : '\u2193';
+            const tooltip = higherIsBetter ? 'Higher is better' : 'Lower is better';
+
+            const card = document.createElement('div');
+            card.className = 'chart-card';
+            card.innerHTML = `
+                <div class="chart-title">
+                    <span>${{label}} <span title="${{tooltip}}" style="cursor:help;opacity:0.6">${{arrow}}</span></span>
+                    <span class="chart-value">${{latestVal?.toFixed(3)}} ${{unit}}</span>
+                </div>
+                <div class="chart-container"><canvas></canvas></div>
+            `;
+            chartsContainer.appendChild(card);
+
+            charts.push(new Chart(card.querySelector('canvas'), {{
+                type: 'line',
+                data: {{
+                    datasets: [
+                        {{
+                            label: label,
+                            data: data,
+                            borderColor: color,
+                            backgroundColor: color + '20',
+                            borderWidth: 2,
+                            pointRadius: 4,
+                            tension: 0.1,
+                            fill: true
+                        }},
+                        {{
+                            label: `${{AVG_WINDOW}}-run avg`,
+                            data: avgData,
+                            borderColor: color,
+                            borderWidth: 2,
+                            borderDash: [6, 4],
+                            pointRadius: 0,
+                            tension: 0.3,
+                            fill: false
+                        }}
+                    ]
+                }},
+                options: {{
+                    responsive: true,
+                    maintainAspectRatio: false,
+                    onClick: (event, elements) => {{
+                        if (elements.length > 0 && elements[0].datasetIndex === 0) {{
+                            const d = data[elements[0].index];
+                            if (d?.commit && d.commit !== 'unknown') {{
+                                window.open(`${{GITHUB_REPO}}/commit/${{d.commit}}`, '_blank');
+                            }}
+                        }}
+                    }},
+                    plugins: {{
+                        legend: {{
+                            display: true,
+                            position: 'bottom',
+                            labels: {{ usePointStyle: true, pointStyle: 'line', boxHeight: 1 }}
+                        }},
+                        tooltip: {{
+                            filter: (item) => item.datasetIndex === 0,
+                            callbacks: {{
+                                title: (items) => {{
+                                    const d = items[0]?.raw;
+                                    return d ? `${{d.name}} (${{d.commit}})` : '';
+                                }},
+                                label: (item) => {{
+                                    const d = item.raw;
+                                    return `${{label}}: ${{d.y?.toFixed(4)}} ${{unit}}`;
+                                }},
+                                footer: (items) => {{
+                                    const d = items[0]?.raw;
+                                    return d?.commit && d.commit !== 'unknown' ? 'Click to view commit' : '';
+                                }}
+                            }}
+                        }}
+                    }},
+                    scales: {{
+                        x: {{
+                            type: 'time',
+                            time: {{ unit: 'day' }},
+                            ticks: {{ maxTicksLimit: 5 }},
+                            title: {{
+                                display: true,
+                                text: 'Date',
+                                font: {{ size: 11 }}
+                            }}
+                        }},
+                        y: {{
+                            ticks: {{ maxTicksLimit: 5 }},
+                            title: {{
+                                display: true,
+                                text: unit,
+                                font: {{ size: 11 }}
+                            }}
+                        }}
+                    }}
+                }}
+            }}));
+        }});
+
+        // Tab switching
+        document.querySelectorAll('.tab').forEach(tab => {{
+            tab.addEventListener('click', () => {{
+                document.querySelectorAll('.tab').forEach(t => t.classList.remove('active'));
+                document.querySelectorAll('.tab-content').forEach(c => c.classList.remove('active'));
+                tab.classList.add('active');
+                document.getElementById(tab.dataset.tab).classList.add('active');
+                // Update URL hash
+                history.replaceState(null, '', '#' + tab.dataset.tab);
+            }});
+        }});
+
+        // Handle URL hash on load
+        if (window.location.hash) {{
+            const tab = document.querySelector(`.tab[data-tab="${{window.location.hash.slice(1)}}"]`);
+            if (tab) tab.click();
+        }}
+
+        // Throughput data and task cards
+        const throughputData = {throughput_json};
+        const taskGrid = document.getElementById('task-grid');
+        const taskChartPanels = document.getElementById('task-chart-panels');
+
+        // Task metadata: display name and image path (relative to nightly/)
+        const taskMeta = {{
+            'Mjlab-Velocity-Flat-Unitree-Go1': {{ name: 'Velocity \u2014 Go1', img: 'images/velocity_go1.png' }},
+            'Mjlab-Tracking-Flat-Unitree-G1': {{ name: 'Tracking \u2014 G1', img: 'images/tracking_g1.png' }},
+            'Mjlab-Lift-Cube-Yam': {{ name: 'Lift Cube \u2014 Yam', img: 'images/lift_cube_yam.png' }}
+        }};
+
+        const throughputChartInstances = {{}};
+
+        if (throughputData.length > 0) {{
+            const tasks = [...new Set(throughputData.flatMap(d => d.results.map(r => r.task)))];
+            const latestRun = throughputData[throughputData.length - 1];
+
+            tasks.forEach((task, i) => {{
+                const meta = taskMeta[task] || {{ name: task, img: '' }};
+                const latestResult = latestRun?.results.find(r => r.task === task);
+                const latestSps = latestResult ? `${{(latestResult.env_sps / 1000).toFixed(0)}}K env steps/s` : '';
+
+                // Card
+                const card = document.createElement('div');
+                card.className = 'task-card' + (i === 0 ? ' active' : '');
+                card.dataset.task = task;
+                card.innerHTML = `
+                    <img src="${{meta.img}}" alt="${{meta.name}}" onerror="this.style.display='none'">
+                    <div class="task-name">${{meta.name}}</div>
+                    <div class="task-stat">${{latestSps}}</div>
+                `;
+                taskGrid.appendChild(card);
+
+                // Chart panel
+                const panel = document.createElement('div');
+                panel.className = 'task-chart-area' + (i === 0 ? ' active' : '');
+                panel.id = `task-panel-${{i}}`;
+                panel.innerHTML = `
+                    <div class="chart-card">
+                        <div class="chart-title">
+                            <span>${{meta.name}} \u2014 Throughput</span>
+                        </div>
+                        <div class="chart-container"><canvas></canvas></div>
+                    </div>
+                `;
+                taskChartPanels.appendChild(panel);
+
+                // Build datasets
+                const envData = [];
+                const physicsData = [];
+                throughputData.forEach(run => {{
+                    const result = run.results.find(r => r.task === task);
+                    if (!result) return;
+                    const point = {{ x: new Date(run.created_at), commit: run.commit }};
+                    envData.push({{ ...point, y: result.env_sps / 1000 }});
+                    physicsData.push({{ ...point, y: result.physics_sps / 1000 }});
+                }});
+
+                const chart = new Chart(panel.querySelector('canvas'), {{
+                    type: 'line',
+                    data: {{
+                        datasets: [
+                            {{
+                                label: 'Env SPS',
+                                data: envData,
+                                borderColor: '#58a6ff',
+                                backgroundColor: '#58a6ff20',
+                                borderWidth: 2,
+                                pointRadius: 4,
+                                tension: 0.1,
+                                fill: true
+                            }},
+                            {{
+                                label: 'Physics SPS',
+                                data: physicsData,
+                                borderColor: '#3fb950',
+                                backgroundColor: '#3fb95020',
+                                borderWidth: 2,
+                                pointRadius: 4,
+                                tension: 0.1,
+                                fill: true
+                            }}
+                        ]
+                    }},
+                    options: {{
+                        responsive: true,
+                        maintainAspectRatio: false,
+                        onClick: (event, elements) => {{
+                            if (elements.length > 0) {{
+                                const di = elements[0].datasetIndex;
+                                const idx = elements[0].index;
+                                const d = di === 0 ? envData[idx] : physicsData[idx];
+                                if (d?.commit && d.commit !== 'unknown') {{
+                                    window.open(`${{GITHUB_REPO}}/commit/${{d.commit}}`, '_blank');
+                                }}
+                            }}
+                        }},
+                        plugins: {{
+                            legend: {{ display: true, position: 'bottom' }},
+                            tooltip: {{
+                                mode: 'index',
+                                callbacks: {{
+                                    title: (items) => {{
+                                        const d = items[0]?.raw;
+                                        return d ? `Commit: ${{d.commit}}` : '';
+                                    }},
+                                    label: (item) => {{
+                                        return `${{item.dataset.label}}: ${{item.raw.y?.toFixed(0)}}K steps/s`;
+                                    }},
+                                    afterBody: (items) => {{
+                                        if (items.length >= 2) {{
+                                            const envY = items[0]?.raw?.y || 0;
+                                            const physY = items[1]?.raw?.y || 0;
+                                            const overhead = physY > 0 ? ((1 - envY / physY) * 100).toFixed(1) : '?';
+                                            return `Overhead: ${{overhead}}%`;
+                                        }}
+                                        return '';
+                                    }},
+                                    footer: (items) => {{
+                                        const d = items[0]?.raw;
+                                        return d?.commit && d.commit !== 'unknown' ? 'Click to view commit' : '';
+                                    }}
+                                }}
+                            }}
+                        }},
+                        scales: {{
+                            x: {{
+                                type: 'time',
+                                time: {{ unit: 'day' }},
+                                ticks: {{ maxTicksLimit: 5 }},
+                                title: {{ display: true, text: 'Date', font: {{ size: 11 }} }}
+                            }},
+                            y: {{
+                                ticks: {{ maxTicksLimit: 5 }},
+                                title: {{
+                                    display: true,
+                                    text: 'K steps/s',
+                                    font: {{ size: 11 }}
+                                }}
+                            }}
+                        }}
+                    }}
+                }});
+                charts.push(chart);
+                throughputChartInstances[task] = {{ chart, panelId: `task-panel-${{i}}` }};
+
+                // Card click handler
+                card.addEventListener('click', () => {{
+                    document.querySelectorAll('.task-card').forEach(c => c.classList.remove('active'));
+                    document.querySelectorAll('.task-chart-area').forEach(p => p.classList.remove('active'));
+                    card.classList.add('active');
+                    panel.classList.add('active');
+                    // Trigger resize so chart renders at correct size
+                    throughputChartInstances[task].chart.resize();
+                }});
+            }});
+        }} else {{
+            taskGrid.innerHTML = '<p style="color: var(--text-dim)">No throughput data available. Run measure_throughput.py to generate data.</p>';
+        }}
+    </script>
+</body>
+</html>
+"""
+
+
+def load_cached_results(output_dir: Path) -> dict[str, dict]:
+  """Load previously evaluated results from cache."""
+  data_file = output_dir / "data.json"
+  if not data_file.exists():
+    return {}
+
+  with open(data_file) as f:
+    runs = json.load(f)
+
+  return {run["id"]: run for run in runs}
+
+
+def main(
+  run_paths: list[str] | None = None,
+  entity: str = "gcbc_researchers",
+  project: str = "mjlab",
+  tag: str = "nightly",
+  eval_limit: int = 0,
+  num_envs: int = 1024,
+  output_dir: Path = Path("benchmark_results"),
+) -> None:
+  """Generate benchmark report by evaluating nightly runs.
+
+  Args:
+    run_paths: Specific run paths to evaluate (entity/project/run_id).
+    entity: WandB entity.
+    project: WandB project name.
+    tag: Filter runs by tag.
+    eval_limit: Maximum number of NEW runs to evaluate per invocation (0 = no limit).
+    num_envs: Number of envs for evaluation.
+    output_dir: Output directory for generated report.
+  """
+  # Load cached results to avoid re-evaluating old runs.
+  cached = load_cached_results(output_dir)
+  print(f"Loaded {len(cached)} cached evaluation results")
+
+  # Start with all cached results (preserves historical data).
+  eval_results_by_id: dict[str, dict] = dict(cached)
+  new_evals = 0
+
+  if run_paths:
+    for run_path in run_paths:
+      run_id = run_path.split("/")[-1]
+      if run_id in eval_results_by_id:
+        print(f"Using cached result for {run_id}")
+      else:
+        result = evaluate_run(run_path, num_envs)
+        eval_results_by_id[run_id] = result
+        new_evals += 1
+  else:
+    api = wandb.Api()
+    print(f"Fetching runs from {entity}/{project} with tag '{tag}'...")
+    runs = api.runs(f"{entity}/{project}", filters={"tags": tag}, order="-created_at")
+
+    for run in runs:
+      if run.state != "finished":
+        continue
+
+      if run.id in eval_results_by_id:
+        print(f"Using cached result for {run.name} ({run.id})")
+      else:
+        if eval_limit > 0 and new_evals >= eval_limit:
+          print(f"Reached eval limit ({eval_limit}), skipping remaining new runs")
+          break
+        run_path = f"{entity}/{project}/{run.id}"
+        result = evaluate_run(run_path, num_envs)
+        eval_results_by_id[run.id] = result
+        new_evals += 1
+
+  eval_results = list(eval_results_by_id.values())
+  print(f"Total runs: {len(eval_results)} ({new_evals} newly evaluated)")
+  generate_html_report(eval_results, output_dir)
+
+
+if __name__ == "__main__":
+  tyro.cli(main, config=mjlab.TYRO_FLAGS)
@@ -0,0 +1,246 @@
+"""Measure environment throughput for regression tracking.
+
+This script measures physics and environment step throughput across canonical tasks
+to catch performance regressions in the manager-based API.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+
+import torch
+import tyro
+import wandb
+
+import mjlab
+import mjlab.tasks  # noqa: F401 - registers tasks
+from mjlab.envs import ManagerBasedRlEnv
+from mjlab.tasks.registry import load_env_cfg
+from mjlab.tasks.tracking.mdp.commands import MotionCommandCfg
+
+
+@dataclass
+class BenchmarkResult:
+  """Results from a single benchmark run."""
+
+  task: str
+  num_envs: int
+  num_steps: int
+  decimation: int
+  physics_sps: float
+  env_sps: float
+  overhead_pct: float
+
+  def __str__(self) -> str:
+    return (
+      f"{self.task} (dec={self.decimation}):\n"
+      f"  Physics SPS: {self.physics_sps:,.0f}\n"
+      f"  Env SPS:     {self.env_sps:,.0f}\n"
+      f"  Overhead:    {self.overhead_pct:.1f}%"
+    )
+
+  def to_dict(self) -> dict:
+    return asdict(self)
+
+
+@dataclass
+class ThroughputConfig:
+  """Configuration for throughput benchmarking."""
+
+  num_envs: int = 4096
+  """Number of parallel environments."""
+
+  num_steps: int = 200
+  """Number of steps to measure (after warmup)."""
+
+  warmup_steps: int = 50
+  """Number of warmup steps before measuring."""
+
+  device: str = "cuda:0"
+  """Device to run on."""
+
+  tasks: list[str] = field(
+    default_factory=lambda: [
+      "Mjlab-Velocity-Flat-Unitree-Go1",
+      "Mjlab-Tracking-Flat-Unitree-G1",
+      "Mjlab-Lift-Cube-Yam",
+    ]
+  )
+  """Tasks to benchmark."""
+
+  tracking_motion: str = "rll_humanoid/wandb-registry-Motions/lafan_cartwheel:latest"
+  """W&B artifact path for tracking task motion (entity/project/name:alias)."""
+
+  output_dir: Path | None = None
+  """Output directory for JSON results. If None, results are only printed."""
+
+
+def measure_physics_sps(env: ManagerBasedRlEnv, num_steps: int) -> float:
+  """Measure raw physics stepping in env steps per second.
+
+  Runs num_steps worth of physics (i.e., num_steps * decimation sim.step calls)
+  and reports throughput in env steps/sec for direct comparison with env.step().
+  """
+  decimation = env.cfg.decimation
+  total_physics_steps = num_steps * decimation
+
+  torch.cuda.synchronize()
+  start = time.perf_counter()
+
+  for _ in range(total_physics_steps):
+    env.sim.step()
+
+  torch.cuda.synchronize()
+  elapsed = time.perf_counter() - start
+
+  # Report in env steps/sec (not physics steps/sec) for fair comparison.
+  return (num_steps * env.num_envs) / elapsed
+
+
+def measure_env_sps(env: ManagerBasedRlEnv, num_steps: int) -> float:
+  """Measure full environment step throughput in env steps per second."""
+  action_dim = sum(env.action_manager.action_term_dim)
+  action = torch.zeros(env.num_envs, action_dim, device=env.device)
+
+  torch.cuda.synchronize()
+  start = time.perf_counter()
+
+  for _ in range(num_steps):
+    env.step(action)
+
+  torch.cuda.synchronize()
+  elapsed = time.perf_counter() - start
+
+  return (num_steps * env.num_envs) / elapsed
+
+
+def benchmark_task(task: str, cfg: ThroughputConfig) -> BenchmarkResult:
+  """Benchmark a single task."""
+  print(f"\nBenchmarking {task}...")
+
+  env_cfg = load_env_cfg(task)
+  env_cfg.scene.num_envs = cfg.num_envs
+
+  # Handle tracking task motion file.
+  if len(env_cfg.commands) > 0:
+    motion_cmd = env_cfg.commands.get("motion")
+    if isinstance(motion_cmd, MotionCommandCfg):
+      api = wandb.Api()
+      artifact = api.artifact(cfg.tracking_motion)
+      motion_dir = artifact.download()
+      motion_cmd.motion_file = str(Path(motion_dir) / "motion.npz")
+
+  env = ManagerBasedRlEnv(cfg=env_cfg, device=cfg.device)
+  env.reset()
+
+  # Warmup.
+  action_dim = sum(env.action_manager.action_term_dim)
+  action = torch.zeros(env.num_envs, action_dim, device=env.device)
+  for _ in range(cfg.warmup_steps):
+    env.step(action)
+  torch.cuda.synchronize()
+
+  decimation = env.cfg.decimation
+  physics_sps = measure_physics_sps(env, cfg.num_steps)
+
+  env.reset()
+  torch.cuda.synchronize()
+
+  env_sps = measure_env_sps(env, cfg.num_steps)
+
+  overhead_pct = 100 * (1 - env_sps / physics_sps)
+
+  env.close()
+
+  return BenchmarkResult(
+    task=task,
+    num_envs=cfg.num_envs,
+    num_steps=cfg.num_steps,
+    decimation=decimation,
+    physics_sps=physics_sps,
+    env_sps=env_sps,
+    overhead_pct=overhead_pct,
+  )
+
+
+def get_git_commit() -> str:
+  """Get current git commit SHA."""
+  try:
+    result = subprocess.run(
+      ["git", "rev-parse", "HEAD"],
+      capture_output=True,
+      text=True,
+      check=True,
+    )
+    return result.stdout.strip()[:7]
+  except subprocess.CalledProcessError:
+    return "unknown"
+
+
+def save_results(results: list[BenchmarkResult], output_dir: Path) -> None:
+  """Save benchmark results to JSON, appending to existing data."""
+  output_dir.mkdir(parents=True, exist_ok=True)
+  data_file = output_dir / "throughput_data.json"
+
+  # Load existing data.
+  existing: list[dict] = []
+  if data_file.exists():
+    with open(data_file) as f:
+      existing = json.load(f)
+
+  # Create new run entry.
+  run_entry = {
+    "created_at": datetime.now(timezone.utc).isoformat(),
+    "commit": get_git_commit(),
+    "results": [r.to_dict() for r in results],
+  }
+
+  existing.append(run_entry)
+
+  with open(data_file, "w") as f:
+    json.dump(existing, f, indent=2)
+
+  print(f"\nResults saved to {data_file}")
+
+
+def main(cfg: ThroughputConfig) -> list[BenchmarkResult]:
+  """Run throughput benchmarks on all configured tasks."""
+  print("Throughput Benchmark")
+  print(f"  Envs: {cfg.num_envs}")
+  print(f"  Steps: {cfg.num_steps} (+ {cfg.warmup_steps} warmup)")
+  print(f"  Device: {cfg.device}")
+
+  results = []
+  for task in cfg.tasks:
+    result = benchmark_task(task, cfg)
+    results.append(result)
+    print(result)
+
+  print("\n" + "=" * 74)
+  print("Summary (all values in env steps per second):")
+  print("  Physics SPS: sim.step() only (×decimation per env step)")
+  print("  Env SPS: full env.step() including managers")
+  print("  Overhead: time spent on non-physics work (observations, rewards, etc.)")
+  print("=" * 74)
+  print(f"{'Task':<35} {'Dec':>4} {'Physics SPS':>12} {'Env SPS':>12} {'Overhead':>8}")
+  print("-" * 74)
+  for r in results:
+    task_short = r.task.replace("Mjlab-", "").replace("-Unitree-", "-")
+    print(
+      f"{task_short:<35} {r.decimation:>4} {r.physics_sps:>12,.0f} {r.env_sps:>12,.0f} {r.overhead_pct:>7.1f}%"
+    )
+
+  if cfg.output_dir:
+    save_results(results, cfg.output_dir)
+
+  return results
+
+
+if __name__ == "__main__":
+  cfg = tyro.cli(ThroughputConfig, config=mjlab.TYRO_FLAGS)
+  main(cfg)
@@ -0,0 +1,152 @@
+#!/bin/bash
+# Nightly training script for mjlab benchmarks
+#
+# This script clones mjlab fresh, runs the tracking benchmark, and generates a report.
+# It is designed to be called by a systemd timer or cron job.
+#
+# Usage:
+#   ./scripts/benchmarks/nightly_train.sh
+#
+# Environment variables:
+#   CUDA_DEVICE: GPU device to use (default: 0)
+#   WANDB_TAGS: Comma-separated tags for the run (default: nightly)
+#   SKIP_TRAINING: Set to "1" to skip training and only generate report
+#   SKIP_THROUGHPUT: Set to "1" to skip throughput benchmarking
+
+set -euo pipefail
+
+# Configuration
+CUDA_DEVICE="${CUDA_DEVICE:-0}"
+WANDB_TAGS="${WANDB_TAGS:-(\"nightly\",)}"
+SKIP_TRAINING="${SKIP_TRAINING:-0}"
+SKIP_THROUGHPUT="${SKIP_THROUGHPUT:-0}"
+
+# Training configuration
+TASK="Mjlab-Tracking-Flat-Unitree-G1"
+NUM_ENVS=4096
+MAX_ITERATIONS=6000
+REGISTRY_NAME="rll_humanoid/wandb-registry-Motions/side_kick_test"
+
+REPO_URL="git@github.com:mujocolab/mjlab.git"
+GH_PAGES_BRANCH="gh-pages"
+WORK_DIR="/tmp/mjlab-nightly-$$"
+GH_PAGES_DIR="/tmp/mjlab-gh-pages-$$"
+
+log() {
+    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"
+}
+
+error() {
+    log "ERROR: $*" >&2
+    exit 1
+}
+
+clear_gpu() {
+    local gpu_device="$1"
+    log "Clearing GPU $gpu_device..."
+    gpu_pids=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits -i "$gpu_device" 2>/dev/null || true)
+    if [[ -n "$gpu_pids" ]]; then
+        for pid in $gpu_pids; do
+            log "Killing process $pid on GPU $gpu_device"
+            kill -9 "$pid" 2>/dev/null || true
+        done
+        sleep 2  # Wait for processes to fully terminate
+    fi
+}
+
+cleanup() {
+    if [[ -d "$WORK_DIR" ]]; then
+        log "Cleaning up work directory..."
+        rm -rf "$WORK_DIR"
+    fi
+    if [[ -d "$GH_PAGES_DIR" ]]; then
+        log "Cleaning up gh-pages clone..."
+        rm -rf "$GH_PAGES_DIR"
+    fi
+}
+trap cleanup EXIT
+
+export GIT_SSH_COMMAND="ssh -i \"$HOME/.ssh/mjlab_nightly_ed25519\" \
+    -o IdentitiesOnly=yes \
+    -o StrictHostKeyChecking=accept-new"
+
+# Clone fresh copy of mjlab
+log "Cloning mjlab..."
+git clone "$REPO_URL" "$WORK_DIR"
+cd "$WORK_DIR"
+
+log "Starting nightly benchmark run"
+log "Task: $TASK"
+log "GPU: $CUDA_DEVICE"
+log "Commit: $(git rev-parse HEAD)"
+
+# Run training
+if [[ "$SKIP_TRAINING" != "1" ]]; then
+    log "Starting training..."
+
+    clear_gpu "$CUDA_DEVICE"
+
+    CUDA_VISIBLE_DEVICES="$CUDA_DEVICE" uv run train "$TASK" \
+        --env.scene.num-envs "$NUM_ENVS" \
+        --agent.max-iterations "$MAX_ITERATIONS" \
+        --registry-name "$REGISTRY_NAME" \
+        --agent.wandb-tags "$WANDB_TAGS"
+
+    log "Training completed"
+else
+    log "Skipping training (SKIP_TRAINING=1)"
+fi
+
+# Clone gh-pages branch (shallow clone for speed)
+log "Cloning gh-pages branch..."
+if git ls-remote --exit-code --heads origin "$GH_PAGES_BRANCH" > /dev/null 2>&1; then
+    git clone --branch "$GH_PAGES_BRANCH" --depth 1 "$REPO_URL" "$GH_PAGES_DIR"
+else
+    # Create new gh-pages branch
+    mkdir -p "$GH_PAGES_DIR"
+    cd "$GH_PAGES_DIR"
+    git init
+    git remote add origin "$REPO_URL"
+    git checkout -b "$GH_PAGES_BRANCH"
+    cd "$WORK_DIR"
+fi
+
+# Copy cached data if exists
+REPORT_DIR="$GH_PAGES_DIR/nightly"
+mkdir -p "$REPORT_DIR"
+
+# Run throughput benchmark
+if [[ "$SKIP_THROUGHPUT" != "1" ]]; then
+    log "Running throughput benchmark..."
+
+    clear_gpu "$CUDA_DEVICE"
+
+    CUDA_VISIBLE_DEVICES="$CUDA_DEVICE" uv run python scripts/benchmarks/measure_throughput.py \
+        --num-envs "$NUM_ENVS" \
+        --output-dir "$REPORT_DIR"
+    log "Throughput benchmark completed"
+else
+    log "Skipping throughput benchmark (SKIP_THROUGHPUT=1)"
+fi
+
+# Generate report (uses cached data.json if present, only evaluates new runs)
+log "Generating benchmark report..."
+uv run python scripts/benchmarks/generate_report.py \
+    --entity gcbc_researchers \
+    --tag nightly \
+    --output-dir "$REPORT_DIR"
+
+log "Report generated"
+
+# Commit and push
+cd "$GH_PAGES_DIR"
+git add -A
+if git diff --staged --quiet; then
+    log "No changes to commit"
+else
+    git commit -m "Update nightly tracking benchmark $(date '+%Y-%m-%d')"
+    git push origin "$GH_PAGES_BRANCH" || log "Failed to push"
+    log "Deployed to GitHub Pages"
+fi
+
+log "Nightly benchmark complete"
@@ -0,0 +1,65 @@
+# Systemd Setup for Nightly Benchmarks
+
+This directory contains systemd user service and timer files for running nightly
+mjlab benchmarks.
+
+## Setup
+
+```bash
+# 1. Create user systemd directory if it doesn't exist
+mkdir -p ~/.config/systemd/user
+
+# 2. Copy the service and timer files
+cp mjlab-nightly.service ~/.config/systemd/user/
+cp mjlab-nightly.timer ~/.config/systemd/user/
+
+# 3. Edit the service file to set your WANDB_API_KEY
+#    Get your key from: https://wandb.ai/authorize
+nano ~/.config/systemd/user/mjlab-nightly.service
+
+# 4. Reload systemd
+systemctl --user daemon-reload
+
+# 5. Enable and start the timer
+systemctl --user enable mjlab-nightly.timer
+systemctl --user start mjlab-nightly.timer
+
+# 6. Enable lingering (so timer runs even when you're not logged in)
+sudo loginctl enable-linger $USER
+```
+
+## Useful Commands
+
+```bash
+# Check timer status
+systemctl --user status mjlab-nightly.timer
+
+# List all timers and when they'll run next
+systemctl --user list-timers
+
+# Run the benchmark manually (without waiting for timer)
+systemctl --user start mjlab-nightly.service
+
+# View logs
+journalctl --user -u mjlab-nightly.service -f
+
+# View recent logs
+journalctl --user -u mjlab-nightly.service --since "1 hour ago"
+
+# Disable the timer
+systemctl --user disable mjlab-nightly.timer
+```
+
+## Configuration
+
+Edit `~/.config/systemd/user/mjlab-nightly.service` to customize:
+
+- `Environment="CUDA_DEVICE=0"` - Which GPU to use
+- `Environment="WANDB_API_KEY=..."` - Your WandB API key
+- `MemoryMax=32G` - Memory limit for the training job
+
+Edit `~/.config/systemd/user/mjlab-nightly.timer` to change the schedule:
+
+- `OnCalendar=*-*-* 02:00:00` - Default: 2 AM daily
+- `OnCalendar=Mon *-*-* 02:00:00` - Example: Mondays only at 2 AM
+- `OnCalendar=*-*-* 02,14:00:00` - Example: Twice daily at 2 AM and 2 PM
@@ -0,0 +1,38 @@
+# mjlab Nightly Benchmark Service
+#
+# Installation:
+#   1. Copy this file to ~/.config/systemd/user/mjlab-nightly.service
+#   2. Edit the paths below to match your setup
+#   3. Enable with: systemctl --user enable mjlab-nightly.timer
+#   4. Start with: systemctl --user start mjlab-nightly.timer
+#
+# To run manually: systemctl --user start mjlab-nightly.service
+# To check logs: journalctl --user -u mjlab-nightly.service
+
+[Unit]
+Description=mjlab Nightly Training Benchmark
+After=network-online.target
+Wants=network-online.target
+
+[Service]
+Type=oneshot
+
+# Update these paths to match your setup
+Environment="PATH=/home/kevin/.local/bin:/home/kevin/.cargo/bin:/usr/local/bin:/usr/bin:/bin"
+Environment="HOME=/home/kevin"
+Environment="CUDA_DEVICE=0"
+Environment="WANDB_API_KEY="
+
+WorkingDirectory=/home/kevin/dev/mjlab
+ExecStart=/home/kevin/dev/mjlab/scripts/benchmarks/nightly_train.sh
+
+# Logging
+StandardOutput=journal
+StandardError=journal
+
+# Resource limits (adjust as needed)
+Nice=10
+MemoryMax=32G
+
+[Install]
+WantedBy=default.target
@@ -0,0 +1,28 @@
+# mjlab Nightly Benchmark Timer
+#
+# This timer triggers the nightly training run.
+# Default: 2:00 AM daily
+#
+# Installation:
+#   1. Copy to ~/.config/systemd/user/mjlab-nightly.timer
+#   2. Enable: systemctl --user enable mjlab-nightly.timer
+#   3. Start: systemctl --user start mjlab-nightly.timer
+#
+# Check status: systemctl --user list-timers
+# Check next run: systemctl --user status mjlab-nightly.timer
+
+[Unit]
+Description=Run mjlab nightly benchmarks
+
+[Timer]
+# Run at 2:00 AM every day
+OnCalendar=*-*-* 02:00:00
+
+# Add randomized delay up to 30 minutes to avoid thundering herd
+RandomizedDelaySec=1800
+
+# If missed (e.g., machine was off), run on next boot
+Persistent=true
+
+[Install]
+WantedBy=timers.target
@@ -0,0 +1,3 @@
+# Cloud Training
+
+See the [Cloud Training](https://mjlab.readthedocs.io/en/latest/source/training/cloud.html) documentation for setup and usage.
@@ -0,0 +1,17 @@
+# SkyPilot job that runs a single W&B sweep agent on one GPU.
+#
+# Submit to an existing cluster provisioned by sweep-cluster.yaml:
+#   sky exec mjlab-sweep scripts/cloud/sweep-agent.yaml \
+#     --gpus A100:1 --env SWEEP_ID=<entity/project/sweep_id> -d
+
+resources:
+  accelerators: A100:1
+
+envs:
+  SWEEP_ID: ""
+  MUJOCO_GL: egl
+
+run: |
+  source "$HOME/.local/bin/env" 2>/dev/null || true
+  cd ~/sky_workdir
+  uv run wandb agent "$SWEEP_ID"
@@ -0,0 +1,34 @@
+# SkyPilot cluster definition for W&B sweep agents.
+#
+# Provisions a multi-GPU instance and installs dependencies. Does not
+# start any jobs. Use sweep-agent.yaml with sky exec for that.
+#
+# Usage:
+#   sky launch scripts/cloud/sweep-cluster.yaml -c mjlab-sweep --gpus A100:8
+
+name: mjlab-sweep-cluster
+
+resources:
+  cloud: lambda
+  accelerators: A100:8
+  autostop:
+    idle_minutes: 5
+    down: true
+
+workdir: .
+
+file_mounts:
+  ~/.netrc: ~/.netrc
+
+envs:
+  MUJOCO_GL: egl
+
+setup: |
+  # Install EGL for MuJoCo headless rendering.
+  sudo apt-get update && sudo apt-get install -y libegl-dev
+
+  # Install uv if not present.
+  command -v uv || curl -LsSf https://astral.sh/uv/install.sh | sh
+  source "$HOME/.local/bin/env" 2>/dev/null || true
+
+  uv sync --locked --no-dev
@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+# Launch a W&B sweep via SkyPilot.
+#
+# Provisions a single multi-GPU cluster and runs one sweep agent per GPU
+# using SkyPilot's job queue. Each agent pulls hyperparameters from the
+# W&B sweep controller independently.
+#
+# Usage:
+#   ./scripts/cloud/sweep-launch.sh [GPUS [CLOUD]]
+#
+# Examples:
+#   ./scripts/cloud/sweep-launch.sh A100:4          # 4 agents, default cloud
+#   ./scripts/cloud/sweep-launch.sh A100:8 gcp      # 8 agents on GCP
+#   ./scripts/cloud/sweep-launch.sh A100:8 lambda   # 8 agents on Lambda
+#   ./scripts/cloud/sweep-launch.sh                 # defaults to A100:4
+
+set -euo pipefail
+
+GPUS="${1:-A100:4}"
+CLOUD="${2:-}"
+GPU_TYPE="${GPUS%%:*}"
+NUM_AGENTS="${GPUS##*:}"
+CLUSTER_NAME="mjlab-sweep"
+
+echo "Creating W&B sweep..."
+SWEEP_ID=$(uv run wandb sweep scripts/cloud/sweep.yaml 2>&1 | grep "wandb agent" | awk '{print $NF}')
+
+if [ -z "$SWEEP_ID" ]; then
+  echo "Failed to create sweep."
+  exit 1
+fi
+
+echo "Sweep created: $SWEEP_ID"
+echo "Provisioning $GPUS cluster..."
+
+# Provision the cluster and run setup (no run section in this YAML).
+CLOUD_FLAG=${CLOUD:+--cloud "$CLOUD"}
+sky launch scripts/cloud/sweep-cluster.yaml \
+  -c "$CLUSTER_NAME" \
+  --gpus "$GPUS" \
+  ${CLOUD_FLAG} \
+  -y --retry-until-up
+
+echo "Submitting $NUM_AGENTS agents to job queue..."
+
+for i in $(seq 1 "$NUM_AGENTS"); do
+  echo "  Agent $i/$NUM_AGENTS"
+  sky exec "$CLUSTER_NAME" \
+    --gpus "${GPU_TYPE}:1" \
+    --env "SWEEP_ID=$SWEEP_ID" \
+    -d \
+    scripts/cloud/sweep-agent.yaml
+done
+
+echo ""
+echo "All agents launched. Monitor at:"
+echo "  sky queue $CLUSTER_NAME"
+echo "  sky logs $CLUSTER_NAME <JOB_ID>"
+echo "  W&B dashboard: https://wandb.ai/$SWEEP_ID"
+echo ""
+echo "When done: sky down $CLUSTER_NAME"
@@ -0,0 +1,37 @@
+# Example W&B sweep configuration. Customize the task, parameters, and
+# search space for your own experiment.
+#
+# Usage:
+#   ./scripts/cloud/sweep-launch.sh A100:4
+
+name: velocity-lr-entropy-sweep
+project: mjlab
+program: train
+method: random
+metric:
+  name: Train/mean_reward
+  goal: maximize
+
+parameters:
+  agent.algorithm.learning-rate:
+    distribution: log_uniform_values
+    min: 1e-4
+    max: 1e-2
+  agent.algorithm.entropy-coef:
+    distribution: log_uniform_values
+    min: 0.001
+    max: 0.1
+
+command:
+  - ${env}
+  - uv
+  - run
+  - ${program}
+  - Mjlab-Velocity-Flat-Unitree-G1
+  - --env.scene.num-envs
+  - "4096"
+  - --agent.max-iterations
+  - "6000"
+  - ${args}
+
+run_cap: 8
@@ -0,0 +1,47 @@
+# SkyPilot task for launching mjlab training on Lambda Cloud.
+#
+# Uses the pre-built Docker image from GHCR.
+#
+# Usage:
+#   sky launch scripts/cloud/train-docker.yaml \
+#     --env TASK=Mjlab-Velocity-Flat-Unitree-G1
+
+name: mjlab-train
+
+resources:
+  cloud: lambda
+  accelerators: A100:1
+  autostop:
+    idle_minutes: 5
+    down: true  # Terminates the instance when idle (stops billing).
+
+workdir: .
+
+file_mounts:
+  ~/.netrc: ~/.netrc
+
+envs:
+  TASK: Mjlab-Velocity-Flat-Unitree-G1
+  NUM_ENVS: "4096"
+  MAX_ITERATIONS: "6000"
+  MUJOCO_GL: egl
+
+setup: |
+  # Configure NVIDIA runtime for Docker if not already set up.
+  if ! sudo docker info 2>/dev/null | grep -q "nvidia"; then
+    sudo nvidia-ctk runtime configure --runtime=docker
+    sudo systemctl restart docker
+    sleep 3  # Wait for the daemon to be ready before pulling.
+  fi
+
+  sudo docker pull ghcr.io/mujocolab/mjlab:latest
+
+run: |
+  sudo docker run --rm --runtime=nvidia --gpus all \
+    -v "$HOME/.netrc:/root/.netrc:ro" \
+    -e MUJOCO_GL=egl \
+    ghcr.io/mujocolab/mjlab:latest \
+    uv run --no-dev train "$TASK" \
+      --env.scene.num-envs "$NUM_ENVS" \
+      --agent.max-iterations "$MAX_ITERATIONS" \
+      --gpu-ids all
@@ -0,0 +1,44 @@
+# SkyPilot task for launching mjlab training on Lambda Cloud.
+#
+# Installs mjlab directly with uv (no Docker).
+#
+# Usage:
+#   sky launch scripts/cloud/train.yaml \
+#     --env TASK=Mjlab-Velocity-Flat-Unitree-G1
+
+name: mjlab-train
+
+resources:
+  cloud: lambda
+  accelerators: A100:1
+  autostop:
+    idle_minutes: 5
+    down: true
+
+workdir: .
+
+file_mounts:
+  ~/.netrc: ~/.netrc
+
+envs:
+  TASK: Mjlab-Velocity-Flat-Unitree-G1
+  NUM_ENVS: "4096"
+  MAX_ITERATIONS: "6000"
+  MUJOCO_GL: egl
+
+setup: |
+  # Install EGL for MuJoCo headless rendering.
+  sudo apt-get update && sudo apt-get install -y libegl-dev
+
+  # Install uv if not present.
+  command -v uv || curl -LsSf https://astral.sh/uv/install.sh | sh
+  source "$HOME/.local/bin/env" 2>/dev/null || true
+
+  uv sync --locked --no-dev
+
+run: |
+  source "$HOME/.local/bin/env" 2>/dev/null || true
+  uv run train "$TASK" \
+    --env.scene.num-envs "$NUM_ENVS" \
+    --agent.max-iterations "$MAX_ITERATIONS" \
+    --gpu-ids all
@@ -0,0 +1,194 @@
+"""Body impulse demo with force visualization.
+
+A cylinder hangs from a ball joint like a punching bag. Random impulses
+swing it around while magenta arrows show the applied forces. Both native
+and Viser viewers render the arrows via ``apply_body_impulse``'s built-in
+debug visualization.
+
+Run with:
+  uv run mjpython scripts/demos/body_impulse.py                      # macOS
+  uv run python scripts/demos/body_impulse.py                        # Linux
+  uv run python scripts/demos/body_impulse.py --viewer viser         # Viser
+"""
+
+from __future__ import annotations
+
+import math
+import os
+
+import mujoco
+import torch
+import tyro
+
+import mjlab
+from mjlab.entity import EntityCfg
+from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
+from mjlab.envs.mdp.events import apply_body_impulse, reset_scene_to_default
+from mjlab.managers.event_manager import EventTermCfg
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.rl import RslRlVecEnvWrapper
+from mjlab.scene import SceneCfg
+from mjlab.utils.torch import configure_torch_backends
+from mjlab.viewer import NativeMujocoViewer, ViserPlayViewer
+
+BAG_RADIUS = 0.12  # punching bag radius
+BAG_HALF_HEIGHT = 0.25  # punching bag half-height
+BAG_DENSITY = 500.0  # ~5.7 kg
+ROPE_LEN = 0.3  # rope length in meters
+CEILING_Z = 1.0  # pivot height
+ROPE_RADIUS = 0.008  # visual rope thickness
+
+
+def create_pendulum_spec() -> mujoco.MjSpec:
+  spec = mujoco.MjSpec()
+  spec.modelname = "impulse_punching_bag"
+
+  # Ground plane.
+  ground = spec.worldbody.add_geom()
+  ground.type = mujoco.mjtGeom.mjGEOM_PLANE
+  ground.size[:] = (5.0, 5.0, 0.1)
+  ground.rgba[:] = (0.4, 0.5, 0.6, 1.0)
+
+  # Light.
+  light = spec.worldbody.add_light()
+  light.pos[:] = (0, 0, 4)
+  light.dir[:] = (0, 0, -1)
+  light.diffuse[:] = (0.8, 0.8, 0.8)
+
+  # Ceiling anchor (visual only).
+  anchor = spec.worldbody.add_geom()
+  anchor.type = mujoco.mjtGeom.mjGEOM_BOX
+  anchor.size[:] = (0.04, 0.04, 0.02)
+  anchor.pos[:] = (0, 0, CEILING_Z)
+  anchor.rgba[:] = (0.3, 0.3, 0.3, 1.0)
+  anchor.contype = 0
+  anchor.conaffinity = 0
+
+  # Pendulum body. Ball joint pivot is at (0, 0, CEILING_Z).
+  bag_body = spec.worldbody.add_body()
+  bag_body.name = "bag"
+  bag_body.pos[:] = (0, 0, CEILING_Z)
+
+  joint = bag_body.add_joint()
+  joint.name = "bag_joint"
+  joint.type = mujoco.mjtJoint.mjJNT_BALL
+  joint.damping[:] = 3.0
+  joint.frictionloss = 0.5
+
+  # Rope (visual only capsule from pivot to bag center).
+  rope = bag_body.add_geom()
+  rope.type = mujoco.mjtGeom.mjGEOM_CAPSULE
+  rope.size[:2] = (ROPE_RADIUS, ROPE_LEN / 2)
+  rope.pos[:] = (0, 0, -ROPE_LEN / 2)
+  rope.rgba[:] = (0.5, 0.4, 0.3, 1.0)
+  rope.contype = 0
+  rope.conaffinity = 0
+  rope.mass = 0.001  # negligible mass
+
+  # Punching bag cylinder hanging at the end of the rope.
+  geom = bag_body.add_geom()
+  geom.name = "bag_geom"
+  geom.type = mujoco.mjtGeom.mjGEOM_CYLINDER
+  geom.size[:2] = (BAG_RADIUS, BAG_HALF_HEIGHT)
+  geom.pos[:] = (0, 0, -ROPE_LEN - BAG_HALF_HEIGHT)
+  geom.density = BAG_DENSITY
+  geom.rgba[:] = (0.55, 0.15, 0.1, 0.35)
+
+  return spec
+
+
+def create_env_cfg() -> ManagerBasedRlEnvCfg:
+  bag_cfg = EntityCfg(
+    spec_fn=create_pendulum_spec,
+    init_state=EntityCfg.InitialStateCfg(
+      pos=(0.0, 0.0, 0.0),
+    ),
+  )
+
+  bag_mass = BAG_DENSITY * math.pi * BAG_RADIUS**2 * (2 * BAG_HALF_HEIGHT)
+  weight = bag_mass * 9.81
+  force_mag = weight * 0.8  # 0.8x body weight
+
+  cfg = ManagerBasedRlEnvCfg(
+    decimation=10,
+    scene=SceneCfg(
+      num_envs=1,
+      env_spacing=0.0,
+      extent=2.0,
+      entities={"bag": bag_cfg},
+    ),
+    events={
+      "reset_scene_to_default": EventTermCfg(
+        func=reset_scene_to_default,
+        mode="reset",
+      ),
+      "impulse": EventTermCfg(
+        func=apply_body_impulse,
+        mode="step",
+        params={
+          "force_range": (-force_mag, force_mag),
+          "torque_range": (-force_mag * 0.3, force_mag * 0.3),
+          "duration_s": (0.05, 0.1),
+          "cooldown_s": (1.0, 2.5),
+          "asset_cfg": SceneEntityCfg("bag", body_names=("bag",)),
+        },
+      ),
+    },
+  )
+
+  cfg.viewer.distance = 1.8
+  cfg.viewer.elevation = -10.0
+
+  return cfg
+
+
+class ZeroPolicy:
+  def __call__(self, obs: object) -> torch.Tensor:
+    del obs
+    return torch.zeros(1, 0)
+
+
+def main(device: str = "cpu", viewer: str = "auto") -> None:
+  configure_torch_backends()
+
+  env_cfg = create_env_cfg()
+  env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
+  env = RslRlVecEnvWrapper(env)
+
+  # Print force scaling info.
+  mjm = env.unwrapped.sim.mj_model
+  bag_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_BODY, "bag")
+  subtree_mass = mjm.body_subtreemass[bag_id]
+  weight = subtree_mass * 9.81
+  bag_mass = BAG_DENSITY * math.pi * BAG_RADIUS**2 * (2 * BAG_HALF_HEIGHT)
+  force_mag = bag_mass * 9.81 * 0.8
+  print("=" * 50)
+  print("Body Impulse Demo (punching bag)")
+  print(f"  Bag mass         : {bag_mass:.2f} kg")
+  print(f"  Bag weight       : {weight:.2f} N")
+  print(f"  Rope length      : {ROPE_LEN} m")
+  print(f"  Force range      : +/-{force_mag:.0f} N")
+  print(f"  Force/weight     : {force_mag / weight:.1f}x")
+  print("=" * 50)
+
+  if viewer == "auto":
+    has_display = bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"))
+    resolved = "native" if has_display else "viser"
+  else:
+    resolved = viewer
+
+  policy = ZeroPolicy()
+  if resolved == "native":
+    print("Launching native viewer...")
+    NativeMujocoViewer(env, policy).run()
+  elif resolved == "viser":
+    print("Launching Viser viewer...")
+    ViserPlayViewer(env, policy).run()
+  else:
+    raise ValueError(f"Unknown viewer: {viewer}")
+
+  env.close()
+
+
+if __name__ == "__main__":
+  tyro.cli(main, config=mjlab.TYRO_FLAGS)
@@ -0,0 +1,271 @@
+"""Demo: contact sensor history catches collisions missed by decimation.
+
+A bouncing ball with high restitution contacts the ground briefly during each
+bounce. With large decimation, the contact may start and end within
+intermediate substeps, so by the final substep there is no active contact and
+instantaneous sensor reads miss it. Setting ``history_length = decimation``
+captures every substep.
+
+Run with:
+  uv run python scripts/demos/contact_sensor_decimation.py
+  uv run python scripts/demos/contact_sensor_decimation.py --viewer
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import matplotlib.pyplot as plt
+import mujoco
+import numpy as np
+import torch
+
+from mjlab.entity import EntityCfg
+from mjlab.scene import Scene, SceneCfg
+from mjlab.sensor.contact_sensor import ContactMatch, ContactSensorCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+BOUNCING_BALL_XML = """
+<mujoco>
+  <option timestep="0.001"/>
+  <worldbody>
+    <body name="ground" pos="0 0 0">
+      <geom name="ground_geom" type="plane" size="5 5 0.1"/>
+    </body>
+    <body name="ball" pos="0 0 1">
+      <freejoint/>
+      <geom name="ball_geom" type="sphere" size="0.05" mass="0.1"
+            solref="-1000 0"/>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+DECIMATION = 20
+NUM_ENVS = 1
+NUM_POLICY_STEPS = 200
+PHYSICS_DT = 0.001
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def build(history_length: int) -> tuple[Scene, Simulation]:
+  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(BOUNCING_BALL_XML))
+  sensor_cfg = ContactSensorCfg(
+    name="ball_contact",
+    primary=ContactMatch(mode="geom", pattern="ball_geom", entity="ball"),
+    secondary=None,
+    fields=("found", "force"),
+    history_length=history_length,
+  )
+  scene_cfg = SceneCfg(
+    num_envs=NUM_ENVS,
+    env_spacing=3.0,
+    entities={"ball": entity_cfg},
+    sensors=(sensor_cfg,),
+  )
+  scene = Scene(scene_cfg, DEVICE)
+  model = scene.compile()
+  sim = Simulation(
+    num_envs=NUM_ENVS,
+    cfg=SimulationCfg(njmax=50),
+    model=model,
+    device=DEVICE,
+  )
+  scene.initialize(sim.mj_model, sim.model, sim.data)
+  return scene, sim
+
+
+def run_no_history():
+  """Read instantaneous contact at the end of each policy step."""
+  scene, sim = build(history_length=0)
+  sensor = scene["ball_contact"]
+
+  contact_detected = []
+  for _ in range(NUM_POLICY_STEPS):
+    for _ in range(DECIMATION):
+      sim.step()
+      scene.update(dt=PHYSICS_DT)
+    found = sensor.data.found[0, 0].item() > 0
+    contact_detected.append(found)
+
+  return contact_detected
+
+
+def run_with_history():
+  """Read full substep history to catch mid-decimation contacts."""
+  scene, sim = build(history_length=DECIMATION)
+  sensor = scene["ball_contact"]
+
+  contact_detected_instant = []
+  contact_detected_history = []
+  ball_height_substep = []
+  for _ in range(NUM_POLICY_STEPS):
+    for _ in range(DECIMATION):
+      sim.step()
+      scene.update(dt=PHYSICS_DT)
+      # qpos is always current after step; qpos[2] is z for a freejoint.
+      ball_height_substep.append(sim.data.qpos[0, 2].item())
+    data = sensor.data
+    found_instant = data.found[0, 0].item() > 0
+    # Check whether any substep in the decimation window had contact.
+    force_hist = data.force_history  # [B, N, H, 3]
+    found_history = (force_hist[0, 0].norm(dim=-1) > 1e-6).any().item()
+    contact_detected_instant.append(found_instant)
+    contact_detected_history.append(found_history)
+
+  return contact_detected_instant, contact_detected_history, ball_height_substep
+
+
+def run_viewer():
+  """Launch a Viser viewer showing the bouncing ball with contact forces."""
+  import viser
+
+  from mjlab.viewer.viser import ViserMujocoScene
+
+  scene, sim = build(history_length=0)
+
+  server = viser.ViserServer(label="Bouncing Ball")
+  viz = ViserMujocoScene(server, sim.mj_model, num_envs=NUM_ENVS)
+  viz.show_contact_forces = True
+  viz.show_contact_points = True
+  viz.create_scene_gui(
+    camera_distance=2.0,
+    camera_azimuth=90.0,
+    camera_elevation=20.0,
+  )
+
+  print("Open the Viser URL above to watch the bouncing ball.")
+  print("Contact forces and points are enabled by default.")
+  print("Press Ctrl+C to stop.\n")
+
+  try:
+    while True:
+      for _ in range(DECIMATION):
+        sim.step()
+        scene.update(dt=PHYSICS_DT)
+      viz.update(sim.data)
+      if viz.needs_update:
+        viz.refresh_visualization()
+      time.sleep(DECIMATION * PHYSICS_DT)
+  except KeyboardInterrupt:
+    print("\nShutting down...")
+    server.stop()
+
+
+def run_analysis():
+  """Run the analysis comparing instantaneous vs history contact detection."""
+  print("=" * 70)
+  print("Contact Sensor Decimation Demo")
+  print(f"  Ball dropped from 1m, restitution ~ 1, decimation = {DECIMATION}")
+  print(f"  Physics dt = {PHYSICS_DT}s, policy dt = {DECIMATION * PHYSICS_DT}s")
+  print("=" * 70)
+
+  no_hist = run_no_history()
+  instant, history, ball_height = run_with_history()
+
+  # Policy steps where history caught a contact that instant missed.
+  missed = []
+  for i in range(NUM_POLICY_STEPS):
+    if history[i] and not instant[i]:
+      missed.append(i)
+
+  total_contacts_instant = sum(instant)
+  total_contacts_history = sum(history)
+  total_contacts_no_hist = sum(no_hist)
+
+  print()
+  print(f"Total policy steps with contact (no history):   {total_contacts_no_hist}")
+  print(f"Total policy steps with contact (instant only): {total_contacts_instant}")
+  print(f"Total policy steps with contact (with history): {total_contacts_history}")
+  print()
+
+  if missed:
+    print(f"Contacts MISSED by instantaneous read but CAUGHT by history: {len(missed)}")
+    print(f"  Policy steps: {missed}")
+  else:
+    print("No missed contacts (try increasing decimation or adjusting drop height)")
+
+  print()
+  print("Step-by-step (showing first 60 policy steps):")
+  print(f"{'step':>6}  {'no_hist':>8}  {'instant':>8}  {'history':>8}  {'missed':>8}")
+  print("-" * 50)
+  for i in range(min(60, NUM_POLICY_STEPS)):
+    flag = " <<<" if (history[i] and not instant[i]) else ""
+    print(f"{i:>6}  {no_hist[i]!s:>8}  {instant[i]!s:>8}  {history[i]!s:>8}  {flag}")
+
+  # --- Plot ---
+  total_substeps = NUM_POLICY_STEPS * DECIMATION
+  t_substep = np.arange(total_substeps) * PHYSICS_DT
+
+  # Place markers at the minimum height within each policy step window.
+  min_height = []
+  min_time = []
+  for i in range(NUM_POLICY_STEPS):
+    start = i * DECIMATION
+    end = (i + 1) * DECIMATION
+    window = ball_height[start:end]
+    j = int(np.argmin(window))
+    min_height.append(window[j])
+    min_time.append(t_substep[start + j])
+
+  # Separate history detections into: caught by both, caught only by history.
+  idx_both = [i for i in range(NUM_POLICY_STEPS) if instant[i] and history[i]]
+  idx_history_only = missed  # history=True, instant=False
+
+  fig, ax = plt.subplots(figsize=(12, 4))
+  ax.plot(t_substep, ball_height, color="0.4", linewidth=0.8, label="Ball height")
+
+  if idx_both:
+    ax.scatter(
+      [min_time[i] for i in idx_both],
+      [min_height[i] for i in idx_both],
+      color="tab:green",
+      s=40,
+      zorder=3,
+      label="Detected by both",
+    )
+
+  if idx_history_only:
+    ax.scatter(
+      [min_time[i] for i in idx_history_only],
+      [min_height[i] for i in idx_history_only],
+      color="tab:red",
+      s=60,
+      marker="x",
+      linewidths=2,
+      zorder=4,
+      label="Caught only by history",
+    )
+
+  ax.set_xlabel("Time (s)")
+  ax.set_ylabel("Ball height (m)")
+  ax.set_title(
+    f"Contact sensor with decimation = {DECIMATION}: "
+    f"{len(missed)} collisions missed without history"
+  )
+  ax.legend(loc="upper right")
+  ax.set_ylim(bottom=-0.05)
+  fig.tight_layout()
+  fig.savefig("scripts/demos/contact_sensor_decimation.png", dpi=150)
+  print("\nPlot saved to scripts/demos/contact_sensor_decimation.png")
+  plt.close(fig)
+
+
+def main():
+  parser = argparse.ArgumentParser(description=__doc__)
+  parser.add_argument(
+    "--viewer",
+    action="store_true",
+    help="Launch a Viser viewer instead of running the analysis.",
+  )
+  args = parser.parse_args()
+
+  if args.viewer:
+    run_viewer()
+  else:
+    run_analysis()
+
+
+if __name__ == "__main__":
+  main()
@@ -0,0 +1,194 @@
+"""Interactive IK control demo.
+
+Drag the 3D transform control in the viser viewer to move the YAM end-effector.
+
+Run with:
+  MJLAB_WARP_QUIET=1 uv run scripts/demos/differential_ik.py
+"""
+
+from __future__ import annotations
+
+import time
+from types import SimpleNamespace
+
+import torch
+import viser
+
+from mjlab.asset_zoo.robots.i2rt_yam.yam_constants import get_yam_robot_cfg
+from mjlab.entity import Entity, EntityCfg
+from mjlab.envs.mdp.actions import DifferentialIKAction, DifferentialIKActionCfg
+from mjlab.sim.sim import MujocoCfg, Simulation, SimulationCfg
+from mjlab.utils.lab_api.math import quat_from_matrix
+from mjlab.viewer.viser import ViserMujocoScene
+
+DEMO_INIT_STATE = EntityCfg.InitialStateCfg(
+  pos=(0.0, 0.0, 0.01),
+  joint_pos={
+    "joint2": 0.6,
+    "joint3": 0.6,
+    "joint4": 0.0,
+    "left_finger": 0.037,
+    "right_finger": -0.037,
+  },
+  joint_vel={".*": 0.0},
+)
+
+IK_ITERATIONS = 10
+
+
+def main() -> None:
+  device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+  robot_cfg = get_yam_robot_cfg()
+  robot_cfg.init_state = DEMO_INIT_STATE
+  entity = Entity(robot_cfg)
+  model = entity.compile()
+  sim_cfg = SimulationCfg(mujoco=MujocoCfg(gravity=(0, 0, -9.81)))
+  sim = Simulation(num_envs=1, cfg=sim_cfg, model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+  entity.write_joint_position_to_sim(entity.data.default_joint_pos, joint_ids=None)
+  sim.forward()
+
+  env = SimpleNamespace(num_envs=1, device=device, scene={"robot": entity}, sim=sim)
+  ik_cfg = DifferentialIKActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    frame_name="grasp_site",
+    frame_type="site",
+    posture_weight=0.02,
+    joint_limit_weight=1e-1,
+    damping=1e-1,
+    use_relative_mode=False,
+  )
+  ik_action: DifferentialIKAction = ik_cfg.build(env)  # type: ignore[arg-type]
+  joint_ids = ik_action._joint_ids
+
+  grip_ids, _ = entity.find_joints("left_finger")
+  grip_joint_ids = torch.tensor(grip_ids, device=device, dtype=torch.long)
+  grip_open = torch.tensor([[0.037]], device=device)
+
+  server = viser.ViserServer(label="IK Control Demo")
+  scene = ViserMujocoScene(server, sim.mj_model, num_envs=1)
+  scene.create_scene_gui(
+    camera_distance=0.1,
+    camera_azimuth=135.0,
+    camera_elevation=30.0,
+  )
+
+  site_id = ik_action._frame_id
+  pos = sim.data.site_xpos[0, site_id].cpu().numpy()
+  xmat = sim.data.site_xmat[0, site_id]
+  quat = quat_from_matrix(xmat).cpu().numpy()
+
+  transform_ctrl = server.scene.add_transform_controls(
+    "/ik_target",
+    position=(float(pos[0]), float(pos[1]), float(pos[2])),
+    wxyz=(float(quat[0]), float(quat[1]), float(quat[2]), float(quat[3])),
+    scale=0.12,
+  )
+
+  needs_reset = [False]
+
+  with server.gui.add_folder("IK Control"):
+    reset_button = server.gui.add_button("Reset")
+    reset_button.on_click(lambda _: needs_reset.__setitem__(0, True))
+    iterations_slider = server.gui.add_slider(
+      "IK Iterations",
+      min=1,
+      max=50,
+      step=1,
+      initial_value=IK_ITERATIONS,
+    )
+
+  with server.gui.add_folder("IK Weights"):
+    damping_slider = server.gui.add_slider(
+      "Damping (λ)",
+      min=1e-2,
+      max=1.0,
+      step=1e-3,
+      initial_value=ik_cfg.damping,
+    )
+    pos_w_slider = server.gui.add_slider(
+      "Position Weight",
+      min=0.0,
+      max=10.0,
+      step=0.1,
+      initial_value=ik_cfg.position_weight,
+    )
+    ori_w_slider = server.gui.add_slider(
+      "Orientation Weight",
+      min=0.0,
+      max=10.0,
+      step=0.1,
+      initial_value=ik_cfg.orientation_weight,
+    )
+    jlim_w_slider = server.gui.add_slider(
+      "Joint Limit Weight",
+      min=0.0,
+      max=1.0,
+      step=0.01,
+      initial_value=ik_cfg.joint_limit_weight,
+    )
+    posture_w_slider = server.gui.add_slider(
+      "Posture Weight",
+      min=0.0,
+      max=1.0,
+      step=0.01,
+      initial_value=ik_cfg.posture_weight,
+    )
+
+  print("=" * 60)
+  print("IK Control Demo")
+  print("  Open the viser URL printed above")
+  print("  Drag the 3D transform control to move the end-effector")
+  print("=" * 60)
+
+  target_action = torch.zeros(1, 7, device=device)
+
+  def _reset() -> None:
+    entity.write_joint_position_to_sim(entity.data.default_joint_pos, joint_ids=None)
+    sim.forward()
+    ik_action.reset()
+    p = sim.data.site_xpos[0, site_id].cpu().numpy()
+    q = quat_from_matrix(sim.data.site_xmat[0, site_id]).cpu().numpy()
+    transform_ctrl.position = (float(p[0]), float(p[1]), float(p[2]))
+    transform_ctrl.wxyz = (float(q[0]), float(q[1]), float(q[2]), float(q[3]))
+
+  try:
+    while True:
+      if needs_reset[0]:
+        needs_reset[0] = False
+        _reset()
+
+      ik_cfg.damping = max(damping_slider.value, 1e-2)
+      ik_cfg.position_weight = max(pos_w_slider.value, 0.0)
+      ik_cfg.orientation_weight = max(ori_w_slider.value, 0.0)
+      ik_cfg.joint_limit_weight = max(jlim_w_slider.value, 0.0)
+      ik_cfg.posture_weight = max(posture_w_slider.value, 0.0)
+
+      p = transform_ctrl.position
+      w = transform_ctrl.wxyz
+      target_action[0, :3] = torch.tensor([p[0], p[1], p[2]], device=device)
+      target_action[0, 3:] = torch.tensor([w[0], w[1], w[2], w[3]], device=device)
+      ik_action.process_actions(target_action)
+
+      n_iter = int(iterations_slider.value)
+      for _ in range(n_iter):
+        dq = ik_action.compute_dq()
+        q = entity.data.joint_pos[:, joint_ids] + dq
+        entity.write_joint_position_to_sim(q, joint_ids=joint_ids)
+        entity.write_joint_position_to_sim(grip_open, joint_ids=grip_joint_ids)
+        sim.forward()
+
+      scene.update(sim.data)
+      if scene.needs_update:
+        scene.refresh_visualization()
+
+      time.sleep(1.0 / 30.0)
+  except KeyboardInterrupt:
+    print("\nShutting down...")
+    server.stop()
+
+
+if __name__ == "__main__":
+  main()
@@ -0,0 +1,125 @@
+"""Flat patch terrain demo.
+
+Spawns a Go1 on rough terrain with flat-patch sampling.
+On each reset, the robot lands on a flat patch.
+
+Run with:
+  uv run python scripts/demos/flat_patch_terrain.py [--viewer native|viser]
+
+Toggle visualization group 3 to see flat patch locations visualized as box sites.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import tyro
+
+import mjlab
+import mjlab.terrains as terrain_gen
+from mjlab.envs import ManagerBasedRlEnv
+from mjlab.envs.mdp import events as mdp
+from mjlab.managers.event_manager import EventTermCfg
+from mjlab.rl import RslRlVecEnvWrapper
+from mjlab.tasks.velocity.config.go1.env_cfgs import unitree_go1_rough_env_cfg
+from mjlab.terrains import FlatPatchSamplingCfg
+from mjlab.terrains.terrain_generator import TerrainGeneratorCfg
+from mjlab.utils.torch import configure_torch_backends
+from mjlab.viewer import NativeMujocoViewer, ViserPlayViewer
+
+
+def main(viewer: str = "auto") -> None:
+  configure_torch_backends()
+  device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+  cfg = unitree_go1_rough_env_cfg(play=True)
+
+  spawn_patch_cfg = FlatPatchSamplingCfg(
+    num_patches=100,
+    patch_radius=0.3,
+    max_height_diff=0.05,
+  )
+
+  # Override terrain: 1 row x 2 cols, curriculum mode so each column is deterministic.
+  # Column 0 = discrete obstacles, Column 1 = pyramid slope.
+  assert cfg.scene.terrain is not None
+  cfg.scene.terrain.terrain_generator = TerrainGeneratorCfg(
+    size=(4.0, 4.0),
+    num_rows=1,
+    num_cols=2,
+    border_width=1.0,
+    curriculum=True,
+    add_lights=True,
+    sub_terrains={
+      "discrete_obstacles": terrain_gen.HfDiscreteObstaclesTerrainCfg(
+        proportion=0.5,
+        obstacle_height_range=(0.05, 0.5),
+        obstacle_width_range=(0.4, 1.2),
+        num_obstacles=30,
+        platform_width=1.5,
+        border_width=0.25,
+        flat_patch_sampling={"spawn": spawn_patch_cfg},
+      ),
+      "pyramid_slope": terrain_gen.HfPyramidSlopedTerrainCfg(
+        proportion=0.5,
+        slope_range=(0.3, 0.8),
+        platform_width=1.5,
+        border_width=0.25,
+        flat_patch_sampling={"spawn": spawn_patch_cfg},
+      ),
+    },
+  )
+
+  # Remove all termination conditions except time limit.
+  for key in list(cfg.terminations):
+    if key != "time_out":
+      del cfg.terminations[key]
+
+  # Reset every 2 seconds to better showcase flat patch spawning.
+  cfg.episode_length_s = 2.0
+
+  # Replace reset_base event with flat-patch spawning.
+  cfg.events["reset_base"] = EventTermCfg(
+    func=mdp.reset_root_state_from_flat_patches,
+    mode="reset",
+    params={
+      "patch_name": "spawn",
+      "pose_range": {"z": (0.01, 0.05), "yaw": (-3.14, 3.14)},
+    },
+  )
+
+  print("=" * 60)
+  print("Flat Patch Terrain Demo")
+  print("  Toggle group 3 to see flat patch markers (orange spheres)")
+  print("  Press Enter in terminal to reset robot onto a flat patch")
+  print("=" * 60)
+
+  env = ManagerBasedRlEnv(cfg=cfg, device=device)
+  env = RslRlVecEnvWrapper(env)
+
+  class ZeroPolicy:
+    def __call__(self, obs) -> torch.Tensor:
+      del obs
+      return torch.zeros(env.unwrapped.action_space.shape, device=device)
+
+  policy = ZeroPolicy()
+
+  if viewer == "auto":
+    has_display = bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"))
+    resolved_viewer = "native" if has_display else "viser"
+  else:
+    resolved_viewer = viewer
+
+  if resolved_viewer == "native":
+    NativeMujocoViewer(env, policy).run()
+  elif resolved_viewer == "viser":
+    ViserPlayViewer(env, policy).run()
+  else:
+    raise ValueError(f"Unknown viewer: {viewer}")
+
+  env.close()
+
+
+if __name__ == "__main__":
+  tyro.cli(main, config=mjlab.TYRO_FLAGS)
@@ -0,0 +1,252 @@
+"""Raycast sensor demo.
+
+Run with:
+  uv run mjpython scripts/demos/raycast_sensor.py [--viewer native|viser]  # macOS
+  uv run python scripts/demos/raycast_sensor.py [--viewer native|viser]    # Linux
+
+Examples:
+  # Grid pattern (default)
+  uv run python scripts/demos/raycast_sensor.py --pattern grid
+
+  # Pinhole camera pattern
+  uv run python scripts/demos/raycast_sensor.py --pattern pinhole
+
+  # With yaw alignment (ignores pitch/roll)
+  uv run python scripts/demos/raycast_sensor.py --alignment yaw
+
+If using the native viewer, you can launch in interactive mode with:
+  uv run mjpython scripts/demos/raycast_sensor.py --viewer native --interactive
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Literal
+
+import mujoco
+import numpy as np
+import torch
+import tyro
+
+import mjlab
+import mjlab.terrains as terrain_gen
+from mjlab.entity import EntityCfg
+from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
+from mjlab.rl import RslRlVecEnvWrapper
+from mjlab.scene import SceneCfg
+from mjlab.sensor import (
+  GridPatternCfg,
+  ObjRef,
+  PinholeCameraPatternCfg,
+  RayCastSensorCfg,
+)
+from mjlab.terrains.terrain_entity import TerrainEntityCfg
+from mjlab.terrains.terrain_generator import TerrainGeneratorCfg
+from mjlab.utils.torch import configure_torch_backends
+from mjlab.viewer import NativeMujocoViewer, ViserPlayViewer
+
+
+def create_scanner_spec() -> mujoco.MjSpec:
+  spec = mujoco.MjSpec()
+  spec.modelname = "scanner"
+
+  mat = spec.add_material()
+  mat.name = "scanner_mat"
+  mat.rgba[:] = (1.0, 0.5, 0.0, 0.9)
+
+  scanner = spec.worldbody.add_body(mocap=True)
+  scanner.name = "scanner"
+  scanner.pos[:] = (0, 0, 2.0)
+
+  geom = scanner.add_geom()
+  geom.name = "scanner_geom"
+  geom.type = mujoco.mjtGeom.mjGEOM_BOX
+  geom.size[:] = (0.15, 0.15, 0.05)
+  geom.mass = 1.0
+  geom.material = "scanner_mat"
+
+  scanner.add_camera(name="scanner", fovy=58.0, resolution=(16, 12))
+
+  record_cam = scanner.add_camera(name="record_cam")
+  record_cam.pos[:] = (2, 0, 2)
+  record_cam.fovy = 40.0
+  record_cam.mode = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODY
+  record_cam.targetbody = "scanner"
+
+  return spec
+
+
+def create_env_cfg(
+  pattern: Literal["grid", "pinhole"],
+  alignment: Literal["base", "yaw", "world"],
+) -> ManagerBasedRlEnvCfg:
+  custom_terrain_cfg = TerrainGeneratorCfg(
+    size=(4.0, 4.0),
+    border_width=0.5,
+    num_rows=1,
+    num_cols=4,
+    curriculum=True,
+    sub_terrains={
+      "pyramid_stairs_inv": terrain_gen.BoxInvertedPyramidStairsTerrainCfg(
+        proportion=0.25,
+        step_height_range=(0.1, 0.25),
+        step_width=0.3,
+        platform_width=1.5,
+        border_width=0.25,
+      ),
+      "hf_pyramid_slope_inv": terrain_gen.HfPyramidSlopedTerrainCfg(
+        proportion=0.25,
+        slope_range=(0.6, 1.5),
+        platform_width=1.5,
+        border_width=0.25,
+        inverted=True,
+      ),
+      "random_rough": terrain_gen.HfRandomUniformTerrainCfg(
+        proportion=0.25,
+        noise_range=(0.05, 0.15),
+        noise_step=0.02,
+        border_width=0.25,
+      ),
+      "wave_terrain": terrain_gen.HfWaveTerrainCfg(
+        proportion=0.25,
+        amplitude_range=(0.15, 0.25),
+        num_waves=3,
+        border_width=0.25,
+      ),
+    },
+    add_lights=True,
+  )
+
+  terrain_cfg = TerrainEntityCfg(
+    terrain_type="generator",
+    terrain_generator=custom_terrain_cfg,
+    num_envs=1,
+  )
+
+  scanner_entity_cfg = EntityCfg(
+    spec_fn=create_scanner_spec,
+    init_state=EntityCfg.InitialStateCfg(pos=(0.65, -0.4, 0.5)),
+  )
+
+  if pattern == "grid":
+    pattern_cfg = GridPatternCfg(
+      size=(0.6, 0.6),
+      resolution=0.1,
+      direction=(0.0, 0.0, -1.0),
+    )
+  else:
+    assert pattern == "pinhole"
+    pattern_cfg = PinholeCameraPatternCfg.from_mujoco_camera("scanner/scanner")
+
+  raycast_cfg = RayCastSensorCfg(
+    name="terrain_scan",
+    frame=ObjRef(type="body", name="scanner", entity="scanner"),
+    pattern=pattern_cfg,
+    ray_alignment=alignment,
+    max_distance=5.0,
+    exclude_parent_body=True,
+    debug_vis=True,
+    viz=RayCastSensorCfg.VizCfg(
+      hit_color=(0.0, 1.0, 0.0, 0.9),
+      miss_color=(1.0, 0.0, 0.0, 0.5),
+      show_rays=False,
+      show_normals=True,
+    ),
+  )
+
+  cfg = ManagerBasedRlEnvCfg(
+    decimation=10,
+    scene=SceneCfg(
+      num_envs=1,
+      env_spacing=0.0,
+      extent=2.0,
+      terrain=terrain_cfg,
+      entities={"scanner": scanner_entity_cfg},
+      sensors=(raycast_cfg,),
+    ),
+  )
+
+  cfg.viewer.body_name = "scanner"
+  cfg.viewer.distance = 12.0
+  cfg.viewer.elevation = -25.0
+  cfg.viewer.azimuth = 135.0
+
+  return cfg
+
+
+def main(
+  viewer: str = "auto",
+  interactive: bool = False,
+  pattern: Literal["grid", "pinhole"] = "grid",
+  alignment: Literal["base", "yaw", "world"] = "base",
+) -> None:
+  configure_torch_backends()
+
+  device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+  print("=" * 60)
+  print("Raycast Sensor Demo - 4 Terrain Types")
+  print(f"  Pattern: {pattern}")
+  print(f"  Alignment: {alignment}")
+  print("=" * 60)
+  print()
+
+  env_cfg = create_env_cfg(pattern, alignment)
+  env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
+  env = RslRlVecEnvWrapper(env)
+
+  if viewer == "auto":
+    has_display = bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"))
+    resolved_viewer = "native" if has_display else "viser"
+  else:
+    resolved_viewer = viewer
+
+  use_auto_scan = (resolved_viewer == "viser") or (not interactive)
+
+  if use_auto_scan:
+
+    class AutoScanPolicy:
+      def __init__(self):
+        self.step_count = 0
+
+      def __call__(self, obs) -> torch.Tensor:
+        del obs
+        t = self.step_count * 0.005
+        y_period = 1000
+        y_normalized = (self.step_count % y_period) / y_period
+        y = -8.0 + 16.0 * y_normalized
+        x = 1.5 * np.sin(2 * np.pi * t * 0.3)
+        z = 1.0
+        env.unwrapped.sim.data.mocap_pos[0, 0, :] = torch.tensor(
+          [x, y, z], device=device, dtype=torch.float32
+        )
+        env.unwrapped.sim.data.mocap_quat[0, 0, :] = torch.tensor(
+          [1, 0, 0, 0], device=device, dtype=torch.float32
+        )
+        self.step_count += 1
+        return torch.zeros(env.unwrapped.action_space.shape, device=device)
+
+    policy = AutoScanPolicy()
+  else:
+
+    class PolicyZero:
+      def __call__(self, obs) -> torch.Tensor:
+        del obs
+        return torch.zeros(env.unwrapped.action_space.shape, device=device)
+
+    policy = PolicyZero()
+
+  if resolved_viewer == "native":
+    print("Launching native viewer...")
+    NativeMujocoViewer(env, policy).run()
+  elif resolved_viewer == "viser":
+    print("Launching viser viewer...")
+    ViserPlayViewer(env, policy).run()
+  else:
+    raise ValueError(f"Unknown viewer: {viewer}")
+
+  env.close()
+
+
+if __name__ == "__main__":
+  tyro.cli(main, config=mjlab.TYRO_FLAGS)
@@ -0,0 +1,36 @@
+#!/bin/bash
+# Fix mjpython on macOS by creating a symlink to libpython.
+# This is needed because mjpython expects libpython in .venv/lib/
+
+set -e
+
+VENV_DIR=".venv"
+
+if [[ "$(uname)" != "Darwin" ]]; then
+    echo "This script is only needed on macOS."
+    exit 0
+fi
+
+if [[ ! -d "$VENV_DIR" ]]; then
+    echo "Error: .venv directory not found. Run 'uv sync' first."
+    exit 1
+fi
+
+# Get Python version and prefix from the venv.
+PYTHON_VERSION=$("$VENV_DIR/bin/python" -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
+PYTHON_PREFIX=$("$VENV_DIR/bin/python" -c "import sys; print(sys.base_prefix)")
+DYLIB_NAME="libpython${PYTHON_VERSION}.dylib"
+
+# Find the dylib in the Python installation.
+DYLIB_PATH="$PYTHON_PREFIX/lib/$DYLIB_NAME"
+
+if [[ ! -f "$DYLIB_PATH" ]]; then
+    echo "Error: Could not find $DYLIB_PATH"
+    exit 1
+fi
+
+# Create the symlink.
+mkdir -p "$VENV_DIR/lib"
+ln -sf "$DYLIB_PATH" "$VENV_DIR/lib/$DYLIB_NAME"
+
+echo "Created symlink: $VENV_DIR/lib/$DYLIB_NAME -> $DYLIB_PATH"
@@ -0,0 +1,37 @@
+#!/usr/bin/env sh
+#
+# Injects useful arguments for running mjlab in docker.
+# See docs/source/installation.rst for usage.
+#
+# Patterned after the uv-in-docker example:
+# https://github.com/astral-sh/uv-docker-example/blob/5748835918ec293d547bbe0e42df34e140aca1eb/run.sh
+#
+# Key arguments:
+#   --rm                Remove the container after exiting
+#   --runtime=nvidia    Use NVIDIA Container runtime to give GPU access
+#   --gpus all          Expose all GPUs by default
+#   -v .:/app           Mount current directory to /app (code changes don't require rebuild)
+#   -v /app/.venv       Mount venv separately (keeps developer's environment out of container)
+#   -p 8080:8080        Publish port 8080 for viewing mjlab web interface on the host
+#   -it (conditional)   Launch in interactive mode if running in a terminal
+#                       (Note: if running training, there's a blocking wandb prompt before training begins)
+#   docker build        Build and launch the image (tag matches the Makefile)
+#   "$@"                Forward all arguments to the docker image
+
+
+if [ -t 1 ]; then
+    INTERACTIVE="-it"
+else
+    INTERACTIVE=""
+fi
+
+docker run \
+    --rm \
+    --runtime=nvidia \
+    --gpus all \
+    --volume .:/app \
+    --volume /app/.venv \
+    --publish 8080:8080 \
+    $INTERACTIVE \
+    $(docker build -qt mjlab .) \
+    "$@"
@@ -0,0 +1,305 @@
+"""Render images of each sub-terrain type for documentation.
+
+Generates one PNG per terrain type at a fixed difficulty level.
+Images are saved to docs/source/_static/terrains/.
+
+Run with:
+  uv run python scripts/tools/render_terrain_gallery.py
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import mujoco
+import numpy as np
+from PIL import Image
+
+import mjlab.terrains as terrain_gen
+from mjlab.terrains import TerrainEntity, TerrainEntityCfg
+from mjlab.terrains.terrain_generator import TerrainGeneratorCfg
+from mjlab.utils import spec_config as spec_cfg
+
+OUTPUT_DIR = Path("docs/source/_static/terrains")
+WIDTH = 1080
+HEIGHT = 1080
+DIFFICULTY = 0.85
+PATCH_SIZE = (8.0, 8.0)
+
+
+# Each entry: (filename, SubTerrainCfg instance).
+TERRAIN_TYPES: list[tuple[str, terrain_gen.SubTerrainCfg]] = [
+  # Primitive terrains.
+  (
+    "box_flat",
+    terrain_gen.BoxFlatTerrainCfg(proportion=1.0),
+  ),
+  (
+    "box_pyramid_stairs",
+    terrain_gen.BoxPyramidStairsTerrainCfg(
+      proportion=1.0,
+      step_height_range=(0.0, 0.2),
+      step_width=0.3,
+      platform_width=3.0,
+      border_width=1.0,
+    ),
+  ),
+  (
+    "box_inverted_pyramid_stairs",
+    terrain_gen.BoxInvertedPyramidStairsTerrainCfg(
+      proportion=1.0,
+      step_height_range=(0.0, 0.2),
+      step_width=0.3,
+      platform_width=3.0,
+      border_width=1.0,
+    ),
+  ),
+  (
+    "box_random_stairs",
+    terrain_gen.BoxRandomStairsTerrainCfg(
+      proportion=1.0,
+      step_width=0.8,
+      step_height_range=(0.1, 0.3),
+      platform_width=1.0,
+      border_width=0.25,
+    ),
+  ),
+  (
+    "box_open_stairs",
+    terrain_gen.BoxOpenStairsTerrainCfg(
+      proportion=1.0,
+      step_height_range=(0.1, 0.2),
+      step_width_range=(0.4, 0.8),
+      platform_width=1.0,
+      border_width=0.25,
+    ),
+  ),
+  (
+    "box_random_grid",
+    terrain_gen.BoxRandomGridTerrainCfg(
+      proportion=1.0,
+      grid_width=0.4,
+      grid_height_range=(0.0, 0.3),
+      platform_width=1.0,
+    ),
+  ),
+  (
+    "box_random_spread",
+    terrain_gen.BoxRandomSpreadTerrainCfg(
+      proportion=1.0,
+      num_boxes=80,
+      box_width_range=(0.1, 1.0),
+      box_length_range=(0.1, 2.0),
+      box_height_range=(0.05, 0.3),
+      platform_width=1.0,
+      border_width=0.25,
+    ),
+  ),
+  (
+    "box_stepping_stones",
+    terrain_gen.BoxSteppingStonesTerrainCfg(
+      proportion=1.0,
+      stone_size_range=(0.4, 0.8),
+      stone_distance_range=(0.2, 0.5),
+      stone_height=0.2,
+      stone_height_variation=0.1,
+      stone_size_variation=0.2,
+      displacement_range=0.1,
+      floor_depth=2.0,
+      platform_width=1.0,
+      border_width=0.25,
+    ),
+  ),
+  (
+    "box_narrow_beams",
+    terrain_gen.BoxNarrowBeamsTerrainCfg(
+      proportion=1.0,
+      num_beams=12,
+      beam_width_range=(0.2, 0.8),
+      beam_height=0.2,
+      spacing=0.8,
+      platform_width=1.0,
+      border_width=0.25,
+      floor_depth=2.0,
+    ),
+  ),
+  (
+    "box_tilted_grid",
+    terrain_gen.BoxTiltedGridTerrainCfg(
+      proportion=1.0,
+      grid_width=1.0,
+      tilt_range_deg=20.0,
+      height_range=0.3,
+      platform_width=1.0,
+      border_width=0.25,
+      floor_depth=2.0,
+    ),
+  ),
+  (
+    "box_nested_rings",
+    terrain_gen.BoxNestedRingsTerrainCfg(
+      proportion=1.0,
+      num_rings=8,
+      ring_width_range=(0.3, 0.6),
+      gap_range=(0.1, 0.4),
+      height_range=(0.1, 0.4),
+      platform_width=1.0,
+      border_width=0.25,
+      floor_depth=2.0,
+    ),
+  ),
+  # Heightfield terrains.
+  (
+    "hf_pyramid_slope",
+    terrain_gen.HfPyramidSlopedTerrainCfg(
+      proportion=1.0,
+      slope_range=(0.0, 0.7),
+      platform_width=2.0,
+      border_width=0.25,
+    ),
+  ),
+  (
+    "hf_random_uniform",
+    terrain_gen.HfRandomUniformTerrainCfg(
+      proportion=1.0,
+      noise_range=(0.02, 0.10),
+      noise_step=0.02,
+      border_width=0.25,
+    ),
+  ),
+  (
+    "hf_wave",
+    terrain_gen.HfWaveTerrainCfg(
+      proportion=1.0,
+      amplitude_range=(0.1, 0.5),
+      num_waves=6,
+      border_width=0.25,
+    ),
+  ),
+  (
+    "hf_discrete_obstacles",
+    terrain_gen.HfDiscreteObstaclesTerrainCfg(
+      proportion=1.0,
+      obstacle_width_range=(0.3, 1.0),
+      obstacle_height_range=(0.05, 0.3),
+      num_obstacles=40,
+      border_width=0.25,
+    ),
+  ),
+  (
+    "hf_perlin_noise",
+    terrain_gen.HfPerlinNoiseTerrainCfg(
+      proportion=1.0,
+      height_range=(0.0, 1.0),
+      octaves=4,
+      persistence=0.3,
+      lacunarity=2.0,
+      scale=10.0,
+      horizontal_scale=0.1,
+      border_width=0.50,
+    ),
+  ),
+]
+
+
+CAMERA_DISTANCE = 8.0
+CAMERA_ELEVATION_DEG = 50.0  # Degrees from horizontal (90 = top-down).
+CAMERA_AZIMUTH_DEG = 135.0  # Rotation around the vertical axis.
+FOV_PADDING = 1.1
+
+
+def render_terrain(
+  name: str,
+  sub_terrain_cfg: terrain_gen.SubTerrainCfg,
+) -> np.ndarray:
+  """Generate and render a single terrain type, return RGB array."""
+  terrain_cfg = TerrainEntityCfg(
+    terrain_type="generator",
+    terrain_generator=TerrainGeneratorCfg(
+      seed=42,
+      size=PATCH_SIZE,
+      num_rows=1,
+      num_cols=1,
+      border_width=0.0,
+      curriculum=False,
+      difficulty_range=(DIFFICULTY, DIFFICULTY),
+      color_scheme="height",
+      sub_terrains={name: sub_terrain_cfg},
+    ),
+    lights=(
+      spec_cfg.LightCfg(
+        name="sun",
+        type="directional",
+        dir=(-0.15, -0.15, -0.97),
+        castshadow=True,
+      ),
+    ),
+  )
+  terrain = TerrainEntity(terrain_cfg, device="cpu")
+
+  # Place camera at an isometric-ish angle to reveal 3D structure.
+  elev = np.deg2rad(CAMERA_ELEVATION_DEG)
+  azim = np.deg2rad(CAMERA_AZIMUTH_DEG)
+  cam_pos = CAMERA_DISTANCE * np.array(
+    [
+      np.cos(elev) * np.cos(azim),
+      np.cos(elev) * np.sin(azim),
+      np.sin(elev),
+    ]
+  )
+
+  # Look at a point slightly below the origin so the terrain centers
+  # in the frame instead of sitting in the bottom half.
+  lookat = np.array([0.0, 0.0, -2.0])
+  forward = lookat - cam_pos
+  forward /= np.linalg.norm(forward)
+  # MuJoCo camera convention: -z is forward, y is up.
+  up = np.array([0.0, 0.0, 1.0])
+  right = np.cross(forward, up)
+  right /= np.linalg.norm(right)
+  cam_up = np.cross(right, forward)
+
+  # Build rotation matrix (columns: right, cam_up, -forward).
+  rot = np.column_stack([right, cam_up, -forward])
+  # Convert to quaternion (wxyz).
+  from scipy.spatial.transform import Rotation
+
+  quat_xyzw = Rotation.from_matrix(rot).as_quat()
+  quat_wxyz = [quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]]
+
+  # FOV sized to fit the patch diagonal from this viewing angle.
+  apparent_size = max(PATCH_SIZE) * FOV_PADDING
+  fovy_rad = 2 * np.arctan2(apparent_size / 2, CAMERA_DISTANCE)
+
+  terrain.spec.worldbody.add_camera(
+    name="gallery",
+    pos=cam_pos.tolist(),
+    quat=quat_wxyz,
+    fovy=np.rad2deg(fovy_rad),
+  )
+
+  model = terrain.spec.compile()
+  model.vis.global_.offheight = HEIGHT
+  model.vis.global_.offwidth = WIDTH
+  data = mujoco.MjData(model)
+  mujoco.mj_forward(model, data)
+
+  with mujoco.Renderer(model, height=HEIGHT, width=WIDTH) as renderer:
+    renderer.update_scene(data, camera="gallery")
+    return renderer.render()
+
+
+def main() -> None:
+  OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+  for filename, sub_cfg in TERRAIN_TYPES:
+    print(f"Rendering {filename}...")
+    img = render_terrain(filename, sub_cfg)
+    Image.fromarray(img).save(OUTPUT_DIR / f"{filename}.png")
+    print(f"  Saved {OUTPUT_DIR / filename}.png")
+
+  print(f"\nDone. {len(TERRAIN_TYPES)} images saved to {OUTPUT_DIR}/")
+
+
+if __name__ == "__main__":
+  main()