diff --git a/CHANGELOG.md b/CHANGELOG.md index b67cdf3..bd25d59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,10 @@ Format: each entry includes the date, a short title, and a description of what c ## Unreleased +### Learn with the GR00T, MuJoCo and MCAP notebook + +The supplied GR00T N1.7 lab is available alongside the ABC/XDOF notebook in a shared notebook index, with dedicated kernel setup and local artifact guidance. LIBERO task selection follows the chosen suite, and before/after comparisons require distinct checkpoints saved during training. ABC video conversion also supports current FFmpeg; GPU execution remains unvalidated. *(PR #36)* + ### Compare ABC-VLA and adapted pi0.5 on bimanual episodes Run compatible policies through repeated ABC simulator episodes using existing Trials, Campaigns, evidence, traces and reports. A guided comparison and matching CLI preserve scene resets, checkpoint identities, task outcomes and policy-call measurements; a separate ABC/XDOF conversion and pi0.5 training recipe prepares adapted checkpoints. Real data conversion and CPU simulation are validated; GPU training and learned-policy comparison remain pending. *(PR #35)* diff --git a/examples/abc-pi05/dataset.py b/examples/abc-pi05/dataset.py index ad0d223..e466ec9 100644 --- a/examples/abc-pi05/dataset.py +++ b/examples/abc-pi05/dataset.py @@ -248,8 +248,8 @@ def frames(ep: Episode): "error", "-i", str(ep.path / "combined_camera-images-rgb.mp4"), - "-vsync", - "0", + "-fps_mode", + "passthrough", "-f", "rawvideo", "-pix_fmt", diff --git a/notebooks/README.md b/notebooks/README.md new file mode 100644 index 0000000..de15afb --- /dev/null +++ b/notebooks/README.md @@ -0,0 +1,13 @@ +# ROVE notebooks + +Clone ROVE to get these Jupyter learning labs. Each notebook has its own setup; +use its dedicated Python environment rather than the ROVE application environment. + +| Notebook | What you can explore | Setup | +| --- | --- | --- | +| [ABC and XDOF](abc_xdof/hello_abc.ipynb) | Demonstrations, action predictions, simulator episodes and optional ABC-VLA fine-tuning. | [ABC instructions](abc_xdof/README.md) | +| [GR00T N1.7, MuJoCo and MCAP](gr00t/gr00t_n17_finetuning_mujoco_mcap_lab_v2.ipynb) | SO100 data inspection, MuJoCo, optional GR00T fine-tuning and LIBERO evaluation, and synthetic MCAP cleaning. | [GR00T instructions](gr00t/README.md) | + +The notebooks are standalone exercises. GPU training and inference require the +runtime described in each guide. Notebook outputs do not automatically become +ROVE Trials or Campaigns. diff --git a/notebooks/gr00t/.gitignore b/notebooks/gr00t/.gitignore new file mode 100644 index 0000000..2229b72 --- /dev/null +++ b/notebooks/gr00t/.gitignore @@ -0,0 +1,4 @@ +.venv/ +.work/ +gr00t_lab/ +*.executed.ipynb diff --git a/notebooks/gr00t/README.md b/notebooks/gr00t/README.md new file mode 100644 index 0000000..984a3a0 --- /dev/null +++ b/notebooks/gr00t/README.md @@ -0,0 +1,72 @@ +# GR00T N1.7, MuJoCo and MCAP lab + +Open [the notebook](gr00t_n17_finetuning_mujoco_mcap_lab_v2.ipynb) for a learning +exercise covering SO100 demonstration inspection, MuJoCo fundamentals, optional +GR00T fine-tuning and LIBERO evaluation, and synthetic ROS 2 MCAP cleaning. +Browse the [notebook index](../README.md) for the ABC and XDOF lab. + +This is a standalone notebook, not a GR00T integration with ROVE's evaluation +engine. Its source is included without saved outputs. Notebook structure and +code-cell syntax are checked, and regression tests exercise suite selection, +checkpoint scheduling and comparison guards. Its installations, training and +GPU evaluation cells have not been executed as part of this addition. + +## Open with a dedicated kernel + +Install Python 3.12, Git, uv and FFmpeg. Git LFS may be needed for upstream demo +assets. From the ROVE repository root: + +```bash +uv venv --python 3.12 --seed notebooks/gr00t/.venv +uv pip install --python notebooks/gr00t/.venv/bin/python jupyterlab ipykernel +notebooks/gr00t/.venv/bin/python -m ipykernel install --sys-prefix \ + --name rove-gr00t --display-name "ROVE GR00T (Python 3.12)" +GR00T_LAB_ROOT="$PWD/notebooks/gr00t/.work" \ + notebooks/gr00t/.venv/bin/jupyter lab \ + notebooks/gr00t/gr00t_n17_finetuning_mujoco_mcap_lab_v2.ipynb --ip=127.0.0.1 +``` + +Choose **ROVE GR00T (Python 3.12)** in Jupyter. In VS Code, use **Select Kernel → +Python Environments** and select `notebooks/gr00t/.venv/bin/python`. Selecting +Homebrew's system Python instead can produce the missing `ipykernel` message. +The seeded environment includes pip because the notebook installs its data and +simulation dependencies into the selected kernel. + +## Choose the execution scope + +The notebook separates macOS/CPU data, MuJoCo and MCAP exercises from the full +Linux/NVIDIA model path. Work through the prerequisite checks first. The full +path needs a compatible CUDA environment, appropriate model access and graphics +libraries for rendered simulation. Docker also needs NVIDIA GPU access. + +Fine-tuning, larger dataset downloads and LIBERO evaluation have explicit flags +in the first configuration cell. Keep them disabled until their prerequisites are +ready. Some earlier setup cells still install packages, clone the upstream GR00T +repository and retrieve assets. Dependencies and the upstream checkout are not +pinned by this notebook; inspect the commands before execution and record the +versions used for any comparison. + +The MCAP exercise generates synthetic data and exports a staging episode; it does +not by itself create a validated production training dataset. Short training runs +demonstrate mechanics and do not establish model quality or real-robot success. + +## Compare checkpoints + +Choose `LIBERO_SUITE` once. The notebook selects a matching example task from the +[upstream task list](https://github.com/NVIDIA/Isaac-GR00T/blob/main/examples/LIBERO/README.md) +and uses it in both rollout command templates. The default `goal` suite evaluates +putting the bowl on the plate; suite `10` uses the stove/moka task. + +The 100-step learning exercise now saves checkpoints at steps 50 and 100. The +comparison requires two existing, distinct checkpoint directories, including +when paths are symlink aliases. Missing or identical checkpoints block comparison +command generation and the before/after plot. With valid checkpoints, enable +`RUN_MUJOCO_BEFORE_AFTER` to print each server/client command pair for manual +execution with the same task, seed and episode count, and separate video folders. + +## Keep generated files local + +The launch command puts downloads, model checkpoints and generated data under +the ignored `.work/` folder. The notebook's default `gr00t_lab/` folder is also +ignored when it is created beside the notebook. Keep credentials outside notebook +cells, and clear outputs before committing or sharing an executed copy. diff --git a/notebooks/gr00t/gr00t_n17_finetuning_mujoco_mcap_lab_v2.ipynb b/notebooks/gr00t/gr00t_n17_finetuning_mujoco_mcap_lab_v2.ipynb new file mode 100644 index 0000000..b16998d --- /dev/null +++ b/notebooks/gr00t/gr00t_n17_finetuning_mujoco_mcap_lab_v2.ipynb @@ -0,0 +1,2647 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "91aa799a", + "metadata": {}, + "source": [ + "# GR00T N1.7 VLA Lab: Mac data work → MuJoCo → Linux GPU fine-tuning → MCAP cleaning\n", + "\n", + "**Updated for macOS and Linux/Docker workflows; public APIs/docs checked 2026-09-15.**\n", + "\n", + "This notebook deliberately has two execution modes:\n", + "\n", + "### Mode A — macOS / CPU development\n", + "You can run:\n", + "- GR00T / LeRobot dataset inspection;\n", + "- plotting and data QA;\n", + "- the standalone **MuJoCo fundamentals lab**;\n", + "- MCAP generation, decoding, synchronization and cleaning;\n", + "- LIBERO dataset download/inspection.\n", + "\n", + "You **do not** run the current NVIDIA GR00T N1.7 CUDA training stack on macOS.\n", + "\n", + "### Mode B — Linux + NVIDIA CUDA, including Docker\n", + "You can additionally run:\n", + "- GR00T N1.7 model loading;\n", + "- fine-tuning;\n", + "- open-loop inference;\n", + "- LIBERO closed-loop evaluation in MuJoCo;\n", + "- before/after checkpoint rollout videos.\n", + "\n", + "The learning progression is:\n", + "\n", + "```text\n", + "SO100 data inspection\n", + " ↓\n", + "MuJoCo fundamentals\n", + " ↓\n", + "GR00T fine-tune + training metrics\n", + " ↓\n", + "open-loop action metrics\n", + " ↓\n", + "LIBERO / MuJoCo closed-loop rollouts\n", + " ↓\n", + "early checkpoint vs later checkpoint videos\n", + " ↓\n", + "MCAP raw episode cleaning\n", + " ↓\n", + "robot episode staged for validation\n", + "```\n", + "\n", + "> **Important:** `LIBERO_PANDA` is a post-training embodiment in GR00T N1.7. The raw base model cannot be used directly as a LIBERO policy. Therefore the closed-loop \"before vs after\" experiment compares an **early fine-tuning checkpoint** with a **later/final checkpoint** using identical simulator seeds.\n", + "\n", + "### Key public references\n", + "\n", + "- GR00T: https://github.com/NVIDIA/Isaac-GR00T\n", + "- GR00T policy/inference guide: https://github.com/NVIDIA/Isaac-GR00T/blob/main/getting_started/policy.md\n", + "- GR00T fine-tuning: https://github.com/NVIDIA/Isaac-GR00T/blob/main/getting_started/finetune_new_embodiment.md\n", + "- GR00T LIBERO benchmark: https://github.com/NVIDIA/Isaac-GR00T/blob/main/examples/LIBERO/README.md\n", + "- MuJoCo Python: https://mujoco.readthedocs.io/en/stable/python.html\n", + "- LIBERO: https://github.com/Lifelong-Robot-Learning/LIBERO\n", + "- MCAP ROS 2: https://mcap.dev/guides/python/ros2" + ] + }, + { + "cell_type": "markdown", + "id": "fb67f463", + "metadata": {}, + "source": [ + "## 0. Platform expectations\n", + "\n", + "### macOS\n", + "\n", + "Use macOS for the **data + simulation learning parts** of this notebook.\n", + "\n", + "The official `mujoco` Python package installs with:\n", + "\n", + "```bash\n", + "pip install mujoco\n", + "```\n", + "\n", + "and bundles the MuJoCo runtime. A special `mjpython` launcher is only needed for some interactive viewer patterns on macOS; this notebook uses normal stepping/offscreen rendering and falls back gracefully if rendering is unavailable.\n", + "\n", + "The current GR00T N1.7 root environment is NVIDIA CUDA-oriented, so the notebook **does not call `uv sync` for GR00T on macOS**.\n", + "\n", + "### Linux / Docker + NVIDIA GPU\n", + "\n", + "For full GR00T training/evaluation you need:\n", + "- Linux;\n", + "- NVIDIA GPU with appropriate VRAM;\n", + "- NVIDIA driver / CUDA runtime;\n", + "- for Docker: NVIDIA Container Toolkit and `--gpus all`;\n", + "- EGL/OpenGL system libraries for headless MuJoCo/LIBERO rendering.\n", + "\n", + "NVIDIA currently recommends about **40 GB+ GPU VRAM** for the standard GR00T N1.7 fine-tuning path.\n", + "\n", + "A typical GPU-container pattern is:\n", + "\n", + "```bash\n", + "docker run --rm -it \\\n", + " --gpus all \\\n", + " --ipc=host \\\n", + " -p 8888:8888 \\\n", + " -v \"$PWD\":/workspace \\\n", + " -w /workspace \\\n", + " \n", + "```\n", + "\n", + "Inside a Debian/Ubuntu-based container, make sure these are available before LIBERO evaluation:\n", + "\n", + "```bash\n", + "apt-get update\n", + "apt-get install -y git git-lfs ffmpeg libegl1-mesa-dev libglu1-mesa\n", + "```\n", + "\n", + "### Hugging Face access\n", + "\n", + "GR00T N1.7 loads the gated `nvidia/Cosmos-Reason2-2B` backbone. For training/inference:\n", + "1. request/accept access on Hugging Face;\n", + "2. authenticate via `HF_TOKEN` or the Hugging Face CLI." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "56ca08bc", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import math\n", + "import os\n", + "import platform\n", + "import re\n", + "import shlex\n", + "import shutil\n", + "import subprocess\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "WORK_ROOT = Path(os.environ.get(\"GR00T_LAB_ROOT\", Path.cwd() / \"gr00t_lab\")).resolve()\n", + "REPO_DIR = WORK_ROOT / \"Isaac-GR00T\"\n", + "\n", + "SYSTEM = platform.system()\n", + "IS_MAC = SYSTEM == \"Darwin\"\n", + "IS_LINUX = SYSTEM == \"Linux\"\n", + "# Container runtimes expose the lowercase 'container' variable.\n", + "IN_DOCKER = Path(\"/.dockerenv\").exists() or bool(os.environ.get(\"container\")) # noqa: SIM112\n", + "\n", + "\n", + "def command_ok(cmd):\n", + " try:\n", + " return (\n", + " subprocess.run(\n", + " cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=10\n", + " ).returncode\n", + " == 0\n", + " )\n", + " except Exception:\n", + " return False\n", + "\n", + "\n", + "HAS_NVIDIA_SMI = shutil.which(\"nvidia-smi\") is not None\n", + "HAS_VISIBLE_NVIDIA_GPU = HAS_NVIDIA_SMI and command_ok([\"nvidia-smi\", \"-L\"])\n", + "FULL_GROOT_SUPPORTED = IS_LINUX and HAS_VISIBLE_NVIDIA_GPU\n", + "\n", + "# Heavy operations are opt-in.\n", + "RUN_SIMPLE_FINETUNE = False\n", + "RUN_SIMPLE_EVAL = False\n", + "RUN_COMPLEX_DOWNLOAD = False\n", + "RUN_COMPLEX_FINETUNE = False\n", + "RUN_LIBERO_SIM_SETUP = False\n", + "RUN_MUJOCO_BEFORE_AFTER = False\n", + "\n", + "# A smoke test proves the pipeline; it does not produce a strong policy.\n", + "SIMPLE_MAX_STEPS = 50\n", + "COMPLEX_MAX_STEPS = 100\n", + "GLOBAL_BATCH_SIZE = 32\n", + "LIBERO_SUITE = \"goal\" # goal, object, spatial, 10\n", + "\n", + "WORK_ROOT.mkdir(parents=True, exist_ok=True)\n", + "\n", + "print(\"System :\", SYSTEM)\n", + "print(\"Docker :\", IN_DOCKER)\n", + "print(\"NVIDIA GPU visible :\", HAS_VISIBLE_NVIDIA_GPU)\n", + "print(\"Full GR00T mode :\", FULL_GROOT_SUPPORTED)\n", + "print(\"WORK_ROOT :\", WORK_ROOT)\n", + "\n", + "if IS_MAC:\n", + " print(\"\\nMode: macOS data / MuJoCo / MCAP lab. CUDA GR00T cells will skip cleanly.\")\n", + "elif FULL_GROOT_SUPPORTED:\n", + " print(\"\\nMode: Linux CUDA. Full GR00T training and LIBERO evaluation are available.\")\n", + "else:\n", + " print(\n", + " \"\\nMode: Linux without visible NVIDIA GPU. Data/MuJoCo/MCAP cells work; GR00T training skips.\"\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "45b02787", + "metadata": {}, + "source": [ + "## 1. Check the machine and rendering prerequisites" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3cb27ab5", + "metadata": {}, + "outputs": [], + "source": [ + "usage = shutil.disk_usage(WORK_ROOT)\n", + "print(\"OS:\", platform.platform())\n", + "print(\"Python kernel:\", sys.version.split()[0])\n", + "print(f\"Free disk: {usage.free / 2**30:.1f} GiB\")\n", + "\n", + "\n", + "def show_command(cmd):\n", + " if not shutil.which(cmd[0]):\n", + " print(f\"\\n{cmd[0]}: not found\")\n", + " return\n", + " p = subprocess.run(cmd, text=True, capture_output=True)\n", + " print(f\"\\n$ {' '.join(cmd)}\")\n", + " print((p.stdout or p.stderr).strip()[:8000])\n", + "\n", + "\n", + "show_command([\"nvidia-smi\"])\n", + "show_command([\"ffmpeg\", \"-version\"])\n", + "\n", + "if IN_DOCKER and FULL_GROOT_SUPPORTED:\n", + " print(\"\\nDocker GPU looks visible. LIBERO additionally needs EGL/OpenGL libraries.\")\n", + "elif IN_DOCKER and not FULL_GROOT_SUPPORTED:\n", + " print(\"\\nDocker detected but NVIDIA GPU is not visible.\")\n", + " print(\"Start the container with NVIDIA Container Toolkit and --gpus all.\")" + ] + }, + { + "cell_type": "markdown", + "id": "8f8615ce", + "metadata": {}, + "source": [ + "### System packages\n", + "\n", + "The notebook does not run OS package managers automatically.\n", + "\n", + "**macOS** (Homebrew, if needed):\n", + "\n", + "```bash\n", + "brew install git git-lfs ffmpeg\n", + "```\n", + "\n", + "**Debian/Ubuntu Linux or container**:\n", + "\n", + "```bash\n", + "apt-get update\n", + "apt-get install -y git git-lfs ffmpeg libegl1-mesa-dev libglu1-mesa\n", + "```\n", + "\n", + "MuJoCo itself comes from the Python package; there is no separate MuJoCo binary download required for the normal Python API." + ] + }, + { + "cell_type": "markdown", + "id": "e50a3760", + "metadata": {}, + "source": [ + "## 2. Install cross-platform notebook dependencies" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "03e3bd12", + "metadata": {}, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "pkgs = [\n", + " \"numpy\",\n", + " \"pandas\",\n", + " \"pyarrow\",\n", + " \"matplotlib\",\n", + " \"pillow\",\n", + " \"imageio\",\n", + " \"imageio-ffmpeg\",\n", + " \"huggingface_hub\",\n", + " \"mujoco\",\n", + " \"mcap\",\n", + " \"mcap-ros2-support\",\n", + "]\n", + "subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", *pkgs])\n", + "print(\"Cross-platform notebook dependencies installed.\")" + ] + }, + { + "cell_type": "markdown", + "id": "af9dcfa7", + "metadata": {}, + "source": [ + "## 3. Clone GR00T; create the CUDA environment only where supported\n", + "\n", + "Cloning is useful on **both** Mac and Linux because the repository contains the small SO100 demo dataset and the reference configs.\n", + "\n", + "`uv sync` for the GR00T root environment is only attempted on **Linux with a visible NVIDIA GPU**." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2a883b6c", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "if not REPO_DIR.exists():\n", + " subprocess.check_call(\n", + " [\n", + " \"git\",\n", + " \"clone\",\n", + " \"--recurse-submodules\",\n", + " \"https://github.com/NVIDIA/Isaac-GR00T.git\",\n", + " str(REPO_DIR),\n", + " ]\n", + " )\n", + "else:\n", + " print(\"Repository already exists:\", REPO_DIR)\n", + "\n", + "# Pull LFS objects when git-lfs is available.\n", + "if shutil.which(\"git-lfs\") or command_ok([\"git\", \"lfs\", \"version\"]):\n", + " try:\n", + " subprocess.check_call([\"git\", \"lfs\", \"pull\"], cwd=REPO_DIR)\n", + " except Exception as e:\n", + " print(\"git lfs pull warning:\", e)\n", + "else:\n", + " print(\"git-lfs not found. Install it if demo media/parquet files appear as tiny pointer files.\")\n", + "\n", + "# uv is cross-platform; install it if absent, but only sync the GR00T CUDA env on supported Linux.\n", + "uv = shutil.which(\"uv\")\n", + "if uv is None:\n", + " try:\n", + " subprocess.check_call([\"bash\", \"-lc\", \"curl -LsSf https://astral.sh/uv/install.sh | sh\"])\n", + " except Exception as e:\n", + " print(\"uv installer warning:\", e)\n", + "\n", + " candidates = [\n", + " shutil.which(\"uv\"),\n", + " str(Path.home() / \".local/bin/uv\"),\n", + " str(Path.home() / \".cargo/bin/uv\"),\n", + " ]\n", + " uv = next((x for x in candidates if x and Path(x).exists()), None)\n", + "\n", + "print(\"uv:\", uv)\n", + "\n", + "if FULL_GROOT_SUPPORTED:\n", + " if uv is None:\n", + " raise RuntimeError(\"uv is required for full GR00T mode.\")\n", + " subprocess.check_call([uv, \"sync\", \"--python\", \"3.12\"], cwd=REPO_DIR)\n", + " subprocess.check_call(\n", + " [uv, \"run\", \"python\", \"-c\", \"import gr00t; print('GR00T import OK')\"],\n", + " cwd=REPO_DIR,\n", + " )\n", + "else:\n", + " print(\"Skipping GR00T `uv sync`: full model stack requires Linux + visible NVIDIA GPU.\")" + ] + }, + { + "cell_type": "markdown", + "id": "0c7f8b2b", + "metadata": {}, + "source": [ + "## 4. Hugging Face authentication check" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f9252a05", + "metadata": {}, + "outputs": [], + "source": [ + "from huggingface_hub import HfApi\n", + "\n", + "HF_TOKEN = os.environ.get(\"HF_TOKEN\")\n", + "try:\n", + " who = HfApi(token=HF_TOKEN).whoami()\n", + " print(\"Hugging Face user:\", who.get(\"name\", who))\n", + "except Exception as e:\n", + " print(\"Hugging Face auth not active:\", type(e).__name__)\n", + " print(\"Dataset downloads may still work when public; GR00T gated model access requires login.\")\n", + "\n", + "if FULL_GROOT_SUPPORTED and not HF_TOKEN:\n", + " print(\"\\nFor model training/inference you can run:\")\n", + " print(\" huggingface-cli login\")\n", + " print(\"or export HF_TOKEN before starting Jupyter.\")" + ] + }, + { + "cell_type": "markdown", + "id": "5d513ea1", + "metadata": {}, + "source": [ + "# Part A — simple manipulation: SO100 cube → bowl\n", + "\n", + "The bundled `cube_to_bowl_5` dataset exposes the GR00T data contract without a large download.\n", + "\n", + "Its modality config uses front + wrist RGB, current robot state, **5 arm values + 1 gripper value**, a **16-step future action horizon**, relative arm actions, absolute gripper actions and language.\n", + "\n", + "Five demonstrations are useful for learning the mechanics, not for claiming a strong policy." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21029681", + "metadata": {}, + "outputs": [], + "source": [ + "SIMPLE_DATASET = REPO_DIR / \"demo_data\" / \"cube_to_bowl_5\"\n", + "assert SIMPLE_DATASET.exists(), f\"Missing {SIMPLE_DATASET}\"\n", + "print(SIMPLE_DATASET)\n", + "\n", + "for p in sorted(SIMPLE_DATASET.rglob(\"*\")):\n", + " if p.is_file():\n", + " rel = p.relative_to(SIMPLE_DATASET)\n", + " print(f\"{rel!s:85s} {p.stat().st_size / 2**20:8.2f} MiB\")" + ] + }, + { + "cell_type": "markdown", + "id": "8f0f48fb", + "metadata": {}, + "source": [ + "## 5. Inspect the metadata contract\n", + "\n", + "`meta/modality.json` is the GR00T-specific bridge between flat arrays in parquet and meaningful robot fields." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0f2d9b31", + "metadata": {}, + "outputs": [], + "source": [ + "from pprint import pprint\n", + "\n", + "\n", + "def load_json(path):\n", + " with open(path) as f:\n", + " return json.load(f)\n", + "\n", + "\n", + "def read_jsonl(path):\n", + " with open(path) as f:\n", + " return [json.loads(line) for line in f if line.strip()]\n", + "\n", + "\n", + "modality = load_json(SIMPLE_DATASET / \"meta\" / \"modality.json\")\n", + "info = load_json(SIMPLE_DATASET / \"meta\" / \"info.json\")\n", + "episodes = read_jsonl(SIMPLE_DATASET / \"meta\" / \"episodes.jsonl\")\n", + "tasks = read_jsonl(SIMPLE_DATASET / \"meta\" / \"tasks.jsonl\")\n", + "\n", + "print(\"modality.json\")\n", + "pprint(modality)\n", + "print(\"\\nFirst episode metadata:\")\n", + "pprint(episodes[0])\n", + "print(\"\\nTasks:\")\n", + "pprint(tasks[:10])" + ] + }, + { + "cell_type": "markdown", + "id": "ce3c8dfe", + "metadata": {}, + "source": [ + "## 6. Inspect one episode at timestep level\n", + "\n", + "A GR00T/LeRobot row is not an `(instruction, answer)` pair. It contains timestamped state, action, episode identity and task annotation. Video lives in parallel MP4 files." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9e7e1675", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "simple_parquets = sorted(SIMPLE_DATASET.glob(\"data/chunk-*/episode_*.parquet\"))\n", + "assert simple_parquets\n", + "ep0 = pd.read_parquet(simple_parquets[0])\n", + "\n", + "display(ep0.head())\n", + "print(\"Columns:\", list(ep0.columns))\n", + "print(\"Rows:\", len(ep0))\n", + "\n", + "state0 = np.asarray(ep0.iloc[0][\"observation.state\"], dtype=float)\n", + "action0 = np.asarray(ep0.iloc[0][\"action\"], dtype=float)\n", + "print(\"State dimension :\", state0.shape)\n", + "print(\"Action dimension:\", action0.shape)\n", + "print(\"First state :\", state0)\n", + "print(\"First action :\", action0)" + ] + }, + { + "cell_type": "markdown", + "id": "be11f418", + "metadata": {}, + "source": [ + "## 7. Reusable dataset QA\n", + "\n", + "Checks required columns, finite state/action data, timestamp order, gaps, jump norms and video counts." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4bedd146", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "REQUIRED_COLUMNS = {\n", + " \"observation.state\",\n", + " \"action\",\n", + " \"timestamp\",\n", + " \"episode_index\",\n", + " \"index\",\n", + " \"task_index\",\n", + " \"next.done\",\n", + "}\n", + "\n", + "\n", + "def _stack_array_col(series):\n", + " return np.stack([np.asarray(x, dtype=float).reshape(-1) for x in series])\n", + "\n", + "\n", + "def robust_threshold(x, z=8.0):\n", + " x = np.asarray(x, dtype=float)\n", + " if len(x) == 0:\n", + " return np.inf\n", + " med = np.nanmedian(x)\n", + " mad = np.nanmedian(np.abs(x - med))\n", + " if mad < 1e-12:\n", + " return med + max(1e-6, np.nanstd(x) * 6)\n", + " return med + z * 1.4826 * mad\n", + "\n", + "\n", + "def audit_gr00t_dataset(root: Path):\n", + " root = Path(root)\n", + " rows = []\n", + " parquets = sorted(root.glob(\"data/chunk-*/episode_*.parquet\"))\n", + "\n", + " for pq in parquets:\n", + " df = pd.read_parquet(pq)\n", + " missing = sorted(REQUIRED_COLUMNS - set(df.columns))\n", + " if missing:\n", + " rows.append({\"episode\": pq.name, \"missing_columns\": \",\".join(missing)})\n", + " continue\n", + "\n", + " ts = df[\"timestamp\"].to_numpy(dtype=float)\n", + " dt = np.diff(ts)\n", + " states = _stack_array_col(df[\"observation.state\"])\n", + " actions = _stack_array_col(df[\"action\"])\n", + " state_jump = (\n", + " np.linalg.norm(np.diff(states, axis=0), axis=1) if len(states) > 1 else np.array([])\n", + " )\n", + " action_jump = (\n", + " np.linalg.norm(np.diff(actions, axis=0), axis=1) if len(actions) > 1 else np.array([])\n", + " )\n", + " positive_dt = dt[dt > 0]\n", + " median_dt = np.median(positive_dt) if len(positive_dt) else np.nan\n", + " gap_threshold = 3 * median_dt if np.isfinite(median_dt) else np.nan\n", + "\n", + " rows.append(\n", + " {\n", + " \"episode\": pq.name,\n", + " \"rows\": len(df),\n", + " \"state_dim\": states.shape[1],\n", + " \"action_dim\": actions.shape[1],\n", + " \"duration_s\": float(ts[-1] - ts[0]) if len(ts) else 0.0,\n", + " \"median_dt_s\": float(median_dt) if np.isfinite(median_dt) else np.nan,\n", + " \"nonmonotonic_or_duplicate_ts\": int(np.sum(dt <= 0)),\n", + " \"large_gaps\": int(np.sum(dt > gap_threshold)) if np.isfinite(gap_threshold) else 0,\n", + " \"nonfinite_state_values\": int((~np.isfinite(states)).sum()),\n", + " \"nonfinite_action_values\": int((~np.isfinite(actions)).sum()),\n", + " \"state_jump_flags\": int(np.sum(state_jump > robust_threshold(state_jump))),\n", + " \"action_jump_flags\": int(np.sum(action_jump > robust_threshold(action_jump))),\n", + " \"done_on_last_row\": bool(df.iloc[-1][\"next.done\"]) if len(df) else False,\n", + " \"missing_columns\": \"\",\n", + " }\n", + " )\n", + "\n", + " result = pd.DataFrame(rows)\n", + " videos = sorted(root.glob(\"videos/chunk-*/*/episode_*.mp4\"))\n", + " camera_dirs = sorted({v.parent.name for v in videos})\n", + " print(f\"Episodes: {len(parquets)}\")\n", + " print(f\"Videos : {len(videos)} across cameras {camera_dirs}\")\n", + " return result\n", + "\n", + "\n", + "simple_audit = audit_gr00t_dataset(SIMPLE_DATASET)\n", + "display(simple_audit)" + ] + }, + { + "cell_type": "markdown", + "id": "3b934430", + "metadata": {}, + "source": [ + "## 8. Plot state and action traces" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e72ac561", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "ts = ep0[\"timestamp\"].to_numpy(float)\n", + "states = _stack_array_col(ep0[\"observation.state\"])\n", + "actions = _stack_array_col(ep0[\"action\"])\n", + "\n", + "plt.figure(figsize=(12, 5))\n", + "for i in range(states.shape[1]):\n", + " plt.plot(ts, states[:, i], label=f\"state[{i}]\")\n", + "plt.xlabel(\"time (s)\")\n", + "plt.ylabel(\"state\")\n", + "plt.title(\"SO100 episode 0 — observation.state\")\n", + "plt.legend(ncol=3)\n", + "plt.show()\n", + "\n", + "plt.figure(figsize=(12, 5))\n", + "for i in range(actions.shape[1]):\n", + " plt.plot(ts, actions[:, i], label=f\"action[{i}]\")\n", + "plt.xlabel(\"time (s)\")\n", + "plt.ylabel(\"action\")\n", + "plt.title(\"SO100 episode 0 — action\")\n", + "plt.legend(ncol=3)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "03b27e4a", + "metadata": {}, + "source": [ + "## 9. View the camera episode" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0d1edfec", + "metadata": {}, + "outputs": [], + "source": [ + "from IPython.display import Video, display\n", + "\n", + "simple_videos = sorted(SIMPLE_DATASET.glob(\"videos/chunk-*/*/episode_000000.mp4\"))\n", + "for video in simple_videos:\n", + " print(video.relative_to(SIMPLE_DATASET))\n", + " display(Video(str(video), embed=False, width=480))" + ] + }, + { + "cell_type": "markdown", + "id": "3b6f4255", + "metadata": {}, + "source": [ + "## 10. Read the actual SO100 modality config\n", + "\n", + "Separate **dataset metadata** (`meta/modality.json`) from the **model sampling config** (`so100_config.py`). Changing the action horizon requires recomputing statistics." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6143e753", + "metadata": {}, + "outputs": [], + "source": [ + "so100_cfg = REPO_DIR / \"examples\" / \"SO100\" / \"so100_config.py\"\n", + "print(so100_cfg.read_text())" + ] + }, + { + "cell_type": "markdown", + "id": "8430f255", + "metadata": {}, + "source": [ + "## 11. Fine-tune the simple task\n", + "\n", + "NVIDIA's documented single-GPU example uses 2000 steps and global batch size 32. This notebook defaults to **50 steps** only to prove the pipeline works.\n", + "\n", + "Set `RUN_SIMPLE_FINETUNE = True` on a 40 GB+ VRAM GPU." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "79a9096b", + "metadata": {}, + "outputs": [], + "source": [ + "SIMPLE_OUT = WORK_ROOT / \"checkpoints\" / \"so100_cube_to_bowl\"\n", + "SIMPLE_OUT.mkdir(parents=True, exist_ok=True)\n", + "\n", + "simple_cmd = [\n", + " uv,\n", + " \"run\",\n", + " \"python\",\n", + " \"gr00t/experiment/launch_finetune.py\",\n", + " \"--base-model-path\",\n", + " \"nvidia/GR00T-N1.7-3B\",\n", + " \"--dataset-path\",\n", + " str(SIMPLE_DATASET),\n", + " \"--embodiment-tag\",\n", + " \"NEW_EMBODIMENT\",\n", + " \"--modality-config-path\",\n", + " \"examples/SO100/so100_config.py\",\n", + " \"--num-gpus\",\n", + " \"1\",\n", + " \"--output-dir\",\n", + " str(SIMPLE_OUT),\n", + " \"--max-steps\",\n", + " str(SIMPLE_MAX_STEPS),\n", + " \"--save-steps\",\n", + " str(max(1, SIMPLE_MAX_STEPS)),\n", + " \"--save-total-limit\",\n", + " \"2\",\n", + " \"--global-batch-size\",\n", + " str(GLOBAL_BATCH_SIZE),\n", + " \"--dataloader-num-workers\",\n", + " \"4\",\n", + "]\n", + "\n", + "print(\"Command:\\n\", \" \".join(map(str, simple_cmd)))\n", + "\n", + "if RUN_SIMPLE_FINETUNE and FULL_GROOT_SUPPORTED:\n", + " subprocess.check_call(simple_cmd, cwd=REPO_DIR)\n", + "elif RUN_SIMPLE_FINETUNE and not FULL_GROOT_SUPPORTED:\n", + " print(\"Requested training, but this runtime is not Linux + NVIDIA CUDA.\")\n", + "else:\n", + " print(\"\\nNot executed. Set RUN_SIMPLE_FINETUNE=True on Linux + NVIDIA GPU.\")" + ] + }, + { + "cell_type": "markdown", + "id": "ddc7a27d", + "metadata": {}, + "source": [ + "## 12. Optional open-loop evaluation\n", + "\n", + "Open-loop evaluation compares predicted action chunks with recorded ground truth. It is useful for debugging but does **not** prove closed-loop task success." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a4a1a750", + "metadata": {}, + "outputs": [], + "source": [ + "def newest_checkpoint(folder: Path):\n", + " ckpts = [p for p in Path(folder).glob(\"checkpoint-*\") if p.is_dir()]\n", + "\n", + " def step(p):\n", + " m = re.search(r\"checkpoint-(\\d+)$\", p.name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + " return max(ckpts, key=step) if ckpts else None\n", + "\n", + "\n", + "simple_ckpt = newest_checkpoint(SIMPLE_OUT)\n", + "print(\"Newest checkpoint:\", simple_ckpt)\n", + "\n", + "if RUN_SIMPLE_EVAL and FULL_GROOT_SUPPORTED:\n", + " if simple_ckpt is None:\n", + " raise RuntimeError(\"No simple checkpoint found. Run fine-tuning first.\")\n", + " eval_cmd = [\n", + " uv,\n", + " \"run\",\n", + " \"python\",\n", + " \"gr00t/eval/open_loop_eval.py\",\n", + " \"--dataset-path\",\n", + " str(SIMPLE_DATASET),\n", + " \"--embodiment-tag\",\n", + " \"NEW_EMBODIMENT\",\n", + " \"--model-path\",\n", + " str(simple_ckpt),\n", + " \"--traj-ids\",\n", + " \"0\",\n", + " \"--execution-horizon\",\n", + " \"16\",\n", + " \"--steps\",\n", + " \"200\",\n", + " ]\n", + " print(\" \".join(map(str, eval_cmd)))\n", + " subprocess.check_call(eval_cmd, cwd=REPO_DIR)\n", + "elif RUN_SIMPLE_EVAL and not FULL_GROOT_SUPPORTED:\n", + " print(\"Requested GR00T inference, but this runtime is not Linux + NVIDIA CUDA.\")\n", + "else:\n", + " print(\"Set RUN_SIMPLE_EVAL=True after training to run open-loop evaluation.\")" + ] + }, + { + "cell_type": "markdown", + "id": "06d5755e", + "metadata": {}, + "source": [ + "# Training and inference metrics\n", + "\n", + "For this VLA lab, separate four questions:\n", + "\n", + "| Question | Metric |\n", + "|---|---|\n", + "| Is optimization progressing? | training loss, learning rate, gradient norm |\n", + "| Does the action predictor match demonstrations? | open-loop action MSE / MAE |\n", + "| Which control dimensions are wrong? | per-action-dimension traces/errors |\n", + "| Does the robot complete the task? | closed-loop MuJoCo success rate |\n", + "\n", + "A conventional `eval_loss` is possible in GR00T's trainer when a real evaluation dataset is configured and `eval_strategy` is enabled. Do **not** manufacture a frame-level random split of robot trajectories just to get an `eval_loss`: adjacent frames leak heavily.\n", + "\n", + "The bundled five-episode SO100 demo is too small for a statistically meaningful validation split. We therefore emphasize:\n", + "- local training-loss curves;\n", + "- episode-level open-loop action metrics;\n", + "- closed-loop LIBERO success rate.\n", + "\n", + "If `eval_loss` exists in `trainer_state.json`, the plotting code below will show it automatically." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1199e7db", + "metadata": {}, + "outputs": [], + "source": [ + "def load_trainer_history(output_dir: Path):\n", + " output_dir = Path(output_dir)\n", + " ckpts = sorted(\n", + " [p for p in output_dir.glob(\"checkpoint-*\") if p.is_dir()],\n", + " key=lambda p: int(p.name.split(\"-\")[-1]) if p.name.split(\"-\")[-1].isdigit() else -1,\n", + " )\n", + " if not ckpts:\n", + " return pd.DataFrame()\n", + "\n", + " state_path = ckpts[-1] / \"trainer_state.json\"\n", + " if not state_path.exists():\n", + " print(\"No trainer_state.json:\", state_path)\n", + " return pd.DataFrame()\n", + "\n", + " state = json.loads(state_path.read_text())\n", + " hist = pd.DataFrame(state.get(\"log_history\", []))\n", + " print(\"Loaded:\", state_path)\n", + " print(\"Global step:\", state.get(\"global_step\"))\n", + " return hist\n", + "\n", + "\n", + "def plot_training_history(output_dir: Path, title=\"GR00T training\"):\n", + " hist = load_trainer_history(output_dir)\n", + " if hist.empty:\n", + " print(\"No metrics yet.\")\n", + " return hist\n", + "\n", + " display(hist.tail(20))\n", + "\n", + " metric_groups = [\n", + " [\"loss\", \"eval_loss\"],\n", + " [\"learning_rate\"],\n", + " [\"grad_norm\"],\n", + " [\"train_accuracy\", \"eval_accuracy\"],\n", + " ]\n", + "\n", + " for metrics in metric_groups:\n", + " available = [m for m in metrics if m in hist.columns and hist[m].notna().any()]\n", + " if not available:\n", + " continue\n", + "\n", + " plt.figure(figsize=(10, 4))\n", + " for metric in available:\n", + " subset = hist[[\"step\", metric]].dropna()\n", + " plt.plot(subset[\"step\"], subset[metric], marker=\".\", label=metric)\n", + " plt.xlabel(\"training step\")\n", + " plt.ylabel(\"metric\")\n", + " plt.title(f\"{title}: {', '.join(available)}\")\n", + " plt.legend()\n", + " plt.show()\n", + "\n", + " return hist\n", + "\n", + "\n", + "# After running SO100 fine-tuning:\n", + "simple_history = plot_training_history(SIMPLE_OUT, \"SO100 cube → bowl\")" + ] + }, + { + "cell_type": "markdown", + "id": "7a12acbb", + "metadata": {}, + "source": [ + "### Open-loop checkpoint comparison\n", + "\n", + "GR00T's `open_loop_eval.py` reports unnormalized **MSE** and **MAE** and saves plots comparing predicted and demonstrated actions.\n", + "\n", + "For a useful learning experiment, save multiple checkpoints and evaluate the **same episode IDs** against each checkpoint.\n", + "\n", + "Example:\n", + "\n", + "```text\n", + "checkpoint-250 MSE 0.12\n", + "checkpoint-1000 MSE 0.05\n", + "checkpoint-2000 MSE 0.03\n", + "```\n", + "\n", + "Then check whether those improvements also produce higher **closed-loop success**. The two metrics can diverge." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9731e3f2", + "metadata": {}, + "outputs": [], + "source": [ + "def list_checkpoints(output_dir: Path):\n", + " def step(p):\n", + " m = re.fullmatch(r\"checkpoint-(\\d+)\", p.name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + " return sorted(\n", + " [\n", + " p\n", + " for p in Path(output_dir).glob(\"checkpoint-*\")\n", + " if p.is_dir() and re.fullmatch(r\"checkpoint-\\d+\", p.name)\n", + " ],\n", + " key=step,\n", + " )\n", + "\n", + "\n", + "def checkpoint_save_interval(max_steps):\n", + " if isinstance(max_steps, bool) or not isinstance(max_steps, int) or max_steps < 2:\n", + " raise ValueError(\"Use at least 2 training steps to save an early and a later checkpoint.\")\n", + " return max_steps // 2\n", + "\n", + "\n", + "def validate_checkpoint_pair(early, late):\n", + " if early is None or late is None:\n", + " raise ValueError(\n", + " \"Two checkpoints are required. Fine-tune with the configured save interval \"\n", + " \"or select two existing checkpoints.\"\n", + " )\n", + " early, late = Path(early).resolve(), Path(late).resolve()\n", + " if not early.is_dir() or not late.is_dir():\n", + " raise ValueError(\"Both checkpoint directories must exist before comparing them.\")\n", + " if early == late:\n", + " raise ValueError(\"Early and late must be different checkpoints, not the same directory.\")\n", + " return early, late\n", + "\n", + "\n", + "print(\"SO100 checkpoints:\")\n", + "for p in list_checkpoints(SIMPLE_OUT):\n", + " print(\" \", p.name)\n", + "\n", + "print(\"\"\"\n", + "For each checkpoint, run open_loop_eval.py on identical episode IDs and save a separate plot.\n", + "The official evaluator prints:\n", + " - per-trajectory MSE\n", + " - per-trajectory MAE\n", + " - average MSE / MAE\n", + "and draws ground-truth vs predicted action traces.\n", + "\"\"\")" + ] + }, + { + "cell_type": "markdown", + "id": "f66529a4", + "metadata": {}, + "source": [ + "# Part B — MuJoCo fundamentals (works on Mac and Linux)\n", + "\n", + "Before using LIBERO, spend a few minutes with MuJoCo directly.\n", + "\n", + "MuJoCo has two central objects:\n", + "\n", + "```text\n", + "MjModel = static model: bodies, joints, geometry, actuators, timestep\n", + "MjData = changing simulation state: qpos, qvel, controls, contacts, time\n", + "```\n", + "\n", + "The essential loop is:\n", + "\n", + "```python\n", + "model = mujoco.MjModel.from_xml_string(MJCF)\n", + "data = mujoco.MjData(model)\n", + "\n", + "data.ctrl[:] = ...\n", + "mujoco.mj_step(model, data)\n", + "```\n", + "\n", + "This section uses a small actuated 2-link arm so you can see:\n", + "- MJCF;\n", + "- joints and actuators;\n", + "- control input;\n", + "- physics stepping;\n", + "- `qpos` / `qvel`;\n", + "- optional camera rendering.\n", + "\n", + "It does **not** depend on GR00T or CUDA." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8eb74adb", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import mujoco\n", + "import numpy as np\n", + "\n", + "MJCF = r\"\"\"\n", + "\n", + " \n", + "\"\"\"\n", + "\n", + "mj_model = mujoco.MjModel.from_xml_string(MJCF)\n", + "mj_data = mujoco.MjData(mj_model)\n", + "\n", + "print(\"MuJoCo version :\", mujoco.__version__)\n", + "print(\"nq / nv / nu :\", mj_model.nq, mj_model.nv, mj_model.nu)\n", + "print(\"timestep :\", mj_model.opt.timestep)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "607c0414", + "metadata": {}, + "outputs": [], + "source": [ + "# Command two joint-position targets and record the response.\n", + "duration_s = 3.0\n", + "target = np.array([0.8, -1.0], dtype=float)\n", + "\n", + "times, qpos_log, qvel_log = [], [], []\n", + "\n", + "while mj_data.time < duration_s:\n", + " mj_data.ctrl[:] = target\n", + " mujoco.mj_step(mj_model, mj_data)\n", + "\n", + " if len(times) == 0 or mj_data.time - times[-1] >= 0.01:\n", + " times.append(mj_data.time)\n", + " qpos_log.append(mj_data.qpos.copy())\n", + " qvel_log.append(mj_data.qvel.copy())\n", + "\n", + "qpos_log = np.asarray(qpos_log)\n", + "qvel_log = np.asarray(qvel_log)\n", + "\n", + "plt.figure(figsize=(10, 4))\n", + "plt.plot(times, qpos_log[:, 0], label=\"shoulder qpos\")\n", + "plt.plot(times, qpos_log[:, 1], label=\"elbow qpos\")\n", + "plt.axhline(target[0], linestyle=\"--\", label=\"shoulder target\")\n", + "plt.axhline(target[1], linestyle=\"--\", label=\"elbow target\")\n", + "plt.xlabel(\"simulation time (s)\")\n", + "plt.ylabel(\"joint position (rad)\")\n", + "plt.title(\"MuJoCo: position actuators driving a 2-link arm\")\n", + "plt.legend()\n", + "plt.show()\n", + "\n", + "print(\"Final qpos:\", mj_data.qpos.copy())\n", + "print(\"Final qvel:\", mj_data.qvel.copy())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8b1c1f58", + "metadata": {}, + "outputs": [], + "source": [ + "# Optional offscreen render.\n", + "# Rendering depends on a working graphics backend. Physics stepping above is independent of rendering.\n", + "\n", + "try:\n", + " renderer = mujoco.Renderer(mj_model, height=360, width=480)\n", + " renderer.update_scene(mj_data)\n", + " rgb = renderer.render()\n", + "\n", + " plt.figure(figsize=(8, 6))\n", + " plt.imshow(rgb)\n", + " plt.axis(\"off\")\n", + " plt.title(\"MuJoCo offscreen render\")\n", + " plt.show()\n", + " renderer.close()\n", + "except Exception as e:\n", + " print(\"Offscreen rendering unavailable in this runtime:\", repr(e))\n", + " if IS_MAC:\n", + " print(\n", + " \"On macOS, interactive viewer workflows may require `mjpython`; physics/data cells still work.\"\n", + " )\n", + " if IN_DOCKER:\n", + " print(\"In Docker, check EGL/OpenGL libraries and NVIDIA GPU passthrough.\")" + ] + }, + { + "cell_type": "markdown", + "id": "bc24a58a", + "metadata": {}, + "source": [ + "### What changes when we move from this toy arm to LIBERO?\n", + "\n", + "Almost nothing about the **simulation loop**:\n", + "\n", + "```text\n", + "toy MuJoCo LIBERO / MuJoCo\n", + "----------- ----------------\n", + "mj_step() env.step(action)\n", + "qpos/qvel Panda proprioception\n", + "simple camera task cameras\n", + "2 joint controls 7-DoF Panda + gripper\n", + "manual target GR00T action chunk\n", + "no task condition language-conditioned task\n", + "```\n", + "\n", + "LIBERO adds robot/task assets, object scenes, observations, task-reset logic and success conditions around MuJoCo." + ] + }, + { + "cell_type": "markdown", + "id": "4441f31c", + "metadata": {}, + "source": [ + "# Part C — more complex manipulation: LIBERO Panda\n", + "\n", + "LIBERO adds a Franka/Panda-style setting, multiple tasks, language-conditioned behavior, richer scene goals and closed-loop simulation.\n", + "\n", + "This notebook defaults to **LIBERO Goal**. Change `LIBERO_SUITE = \"10\"` for longer multi-step tasks.\n", + "The evaluation task is selected from the same suite automatically. Dataset downloads, fine-tuning output, published checkpoint selection and rollout commands all use `LIBERO_SUITE`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8fc13c9e", + "metadata": {}, + "outputs": [], + "source": [ + "LIBERO_REPOS = {\n", + " \"goal\": \"IPEC-COMMUNITY/libero_goal_no_noops_1.0.0_lerobot\",\n", + " \"object\": \"IPEC-COMMUNITY/libero_object_no_noops_1.0.0_lerobot\",\n", + " \"spatial\": \"IPEC-COMMUNITY/libero_spatial_no_noops_1.0.0_lerobot\",\n", + " \"10\": \"IPEC-COMMUNITY/libero_10_no_noops_1.0.0_lerobot\",\n", + "}\n", + "# One example from each suite in NVIDIA's LIBERO task list.\n", + "LIBERO_DEFAULT_ENVS = {\n", + " \"goal\": \"libero_sim/put_the_bowl_on_the_plate\",\n", + " \"object\": \"libero_sim/pick_up_the_alphabet_soup_and_place_it_in_the_basket\",\n", + " \"spatial\": \"libero_sim/pick_up_the_black_bowl_from_table_center_and_place_it_on_the_plate\",\n", + " \"10\": \"libero_sim/KITCHEN_SCENE3_turn_on_the_stove_and_put_the_moka_pot_on_it\",\n", + "}\n", + "\n", + "\n", + "def libero_env_for_suite(suite):\n", + " if suite not in LIBERO_DEFAULT_ENVS:\n", + " raise ValueError(f\"Unknown LIBERO suite {suite!r}; choose goal, object, spatial or 10.\")\n", + " return LIBERO_DEFAULT_ENVS[suite]\n", + "\n", + "\n", + "LIBERO_ENV = libero_env_for_suite(LIBERO_SUITE)\n", + "\n", + "repo_id = LIBERO_REPOS[LIBERO_SUITE]\n", + "suite_dir_name = repo_id.split(\"/\")[-1]\n", + "LIBERO_DATASET = REPO_DIR / \"examples\" / \"LIBERO\" / suite_dir_name\n", + "\n", + "print(\"Suite:\", LIBERO_SUITE)\n", + "print(\"Evaluation task:\", LIBERO_ENV)\n", + "print(\"HF dataset:\", repo_id)\n", + "print(\"Local path:\", LIBERO_DATASET)" + ] + }, + { + "cell_type": "markdown", + "id": "edb049f2", + "metadata": {}, + "source": [ + "## 13. Download LIBERO\n", + "\n", + "The official GR00T recipe downloads a LeRobot dataset from Hugging Face and copies NVIDIA's `modality.json` into `meta/`.\n", + "\n", + "For LIBERO Goal, NVIDIA also ships a replacement for one known corrupted wrist-camera episode." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a0fc1d70", + "metadata": {}, + "outputs": [], + "source": [ + "from huggingface_hub import snapshot_download\n", + "\n", + "print(\"Dataset:\", repo_id)\n", + "print(\"Destination:\", LIBERO_DATASET)\n", + "\n", + "if RUN_COMPLEX_DOWNLOAD:\n", + " snapshot_download(\n", + " repo_id=repo_id,\n", + " repo_type=\"dataset\",\n", + " local_dir=str(LIBERO_DATASET),\n", + " token=HF_TOKEN,\n", + " )\n", + "\n", + " meta_dir = LIBERO_DATASET / \"meta\"\n", + " meta_dir.mkdir(parents=True, exist_ok=True)\n", + " shutil.copy2(REPO_DIR / \"examples/LIBERO/modality.json\", meta_dir / \"modality.json\")\n", + "\n", + " if LIBERO_SUITE == \"goal\":\n", + " patch = REPO_DIR / \"examples/LIBERO/patches/episode_000082.mp4\"\n", + " target = (\n", + " LIBERO_DATASET / \"videos/chunk-000/observation.images.wrist_image/episode_000082.mp4\"\n", + " )\n", + " if patch.exists():\n", + " shutil.copy2(patch, target)\n", + " print(\"Applied NVIDIA wrist-video patch:\", target)\n", + "else:\n", + " print(\"Not executed. Set RUN_COMPLEX_DOWNLOAD=True.\")\n", + " print(\"This dataset download works on both macOS and Linux.\")" + ] + }, + { + "cell_type": "markdown", + "id": "796c1000", + "metadata": {}, + "source": [ + "## 14. Audit LIBERO" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0613eff6", + "metadata": {}, + "outputs": [], + "source": [ + "if LIBERO_DATASET.exists():\n", + " libero_audit = audit_gr00t_dataset(LIBERO_DATASET)\n", + " display(libero_audit.head(20))\n", + " libero_tasks = read_jsonl(LIBERO_DATASET / \"meta/tasks.jsonl\")\n", + " print(\"\\nSample task descriptions:\")\n", + " for row in libero_tasks[:20]:\n", + " print(row)\n", + "else:\n", + " print(\"LIBERO dataset not downloaded yet.\")" + ] + }, + { + "cell_type": "markdown", + "id": "43a7c1ec", + "metadata": {}, + "source": [ + "## 15. Fine-tune LIBERO — learning run versus reference run\n", + "\n", + "NVIDIA's published recipe uses about 20,000 steps, 8 GPUs in the reference command, global batch size 640 and `state_dropout_prob=0.2`.\n", + "\n", + "The one-GPU 100-step setting below is only for learning the mechanics.\n", + "The notebook saves halfway through the run and at subsequent intervals, retaining two checkpoints. With the default 100 steps this produces checkpoints at steps 50 and 100. At least two training steps are required.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9bd01194", + "metadata": {}, + "outputs": [], + "source": [ + "COMPLEX_OUT = WORK_ROOT / \"checkpoints\" / f\"libero_{LIBERO_SUITE}\"\n", + "COMPLEX_OUT.mkdir(parents=True, exist_ok=True)\n", + "\n", + "complex_cmd = [\n", + " uv,\n", + " \"run\",\n", + " \"python\",\n", + " \"gr00t/experiment/launch_finetune.py\",\n", + " \"--base-model-path\",\n", + " \"nvidia/GR00T-N1.7-3B\",\n", + " \"--dataset-path\",\n", + " str(LIBERO_DATASET),\n", + " \"--embodiment-tag\",\n", + " \"LIBERO_PANDA\",\n", + " \"--num-gpus\",\n", + " \"1\",\n", + " \"--output-dir\",\n", + " str(COMPLEX_OUT),\n", + " \"--max-steps\",\n", + " str(COMPLEX_MAX_STEPS),\n", + " \"--save-steps\",\n", + " str(checkpoint_save_interval(COMPLEX_MAX_STEPS)),\n", + " \"--save-total-limit\",\n", + " \"2\",\n", + " \"--global-batch-size\",\n", + " str(GLOBAL_BATCH_SIZE),\n", + " \"--dataloader-num-workers\",\n", + " \"4\",\n", + " \"--state-dropout-prob\",\n", + " \"0.2\",\n", + "]\n", + "\n", + "print(\"Learning/smoke-test command:\\n\", \" \".join(map(str, complex_cmd)))\n", + "\n", + "if RUN_COMPLEX_FINETUNE and FULL_GROOT_SUPPORTED:\n", + " if not LIBERO_DATASET.exists():\n", + " raise RuntimeError(\"Download LIBERO first.\")\n", + " subprocess.check_call(complex_cmd, cwd=REPO_DIR)\n", + "elif RUN_COMPLEX_FINETUNE and not FULL_GROOT_SUPPORTED:\n", + " print(\"Requested fine-tuning, but this runtime is not Linux + NVIDIA CUDA.\")\n", + "else:\n", + " print(\"\\nNot executed. Set RUN_COMPLEX_FINETUNE=True on Linux + NVIDIA GPU.\")\n", + "\n", + "print(\"\\nReference-scale NVIDIA recipe uses NUM_GPUS=8, MAX_STEPS=20000, GLOBAL_BATCH_SIZE=640.\")" + ] + }, + { + "cell_type": "markdown", + "id": "b23f851e", + "metadata": {}, + "source": [ + "## 16. Download NVIDIA's already-fine-tuned LIBERO checkpoint\n", + "\n", + "If your goal is to learn evaluation rather than pay for a full benchmark fine-tune, use NVIDIA's published checkpoint." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "deb8ca22", + "metadata": {}, + "outputs": [], + "source": [ + "PRETRAINED_LIBERO_ROOT = WORK_ROOT / \"checkpoints\" / \"GR00T-N1.7-LIBERO\"\n", + "checkpoint_subdir = f\"libero_{LIBERO_SUITE}\"\n", + "\n", + "allow_patterns = [\n", + " f\"{checkpoint_subdir}/config.json\",\n", + " f\"{checkpoint_subdir}/embodiment_id.json\",\n", + " f\"{checkpoint_subdir}/model-*.safetensors\",\n", + " f\"{checkpoint_subdir}/model.safetensors.index.json\",\n", + " f\"{checkpoint_subdir}/processor_config.json\",\n", + " f\"{checkpoint_subdir}/statistics.json\",\n", + "]\n", + "\n", + "print(\"Hugging Face model: nvidia/GR00T-N1.7-LIBERO\")\n", + "print(\"Subset:\", checkpoint_subdir)\n", + "print(\"Destination:\", PRETRAINED_LIBERO_ROOT)\n", + "print(\"\\nTo download, run:\")\n", + "print(\"\"\"\n", + "snapshot_download(\n", + " repo_id=\"nvidia/GR00T-N1.7-LIBERO\",\n", + " local_dir=str(PRETRAINED_LIBERO_ROOT),\n", + " allow_patterns=allow_patterns,\n", + " token=HF_TOKEN,\n", + ")\n", + "\"\"\")" + ] + }, + { + "cell_type": "markdown", + "id": "8bb5d10e", + "metadata": {}, + "source": [ + "## 17. Closed-loop LIBERO evaluation = GR00T policy + MuJoCo environment\n", + "\n", + "This is where MuJoCo becomes the **robot evaluation world** rather than a toy simulator.\n", + "\n", + "The NVIDIA path is:\n", + "\n", + "```text\n", + "GR00T policy server (CUDA)\n", + " │\n", + " │ actions\n", + " ▼\n", + "LIBERO client\n", + " │\n", + " ▼\n", + "MuJoCo Panda + objects\n", + " │\n", + " ├── RGB observations\n", + " ├── robot state\n", + " └── success / failure\n", + "```\n", + "\n", + "The current NVIDIA LIBERO setup script is Linux-oriented and configures **EGL** headless rendering. Therefore:\n", + "\n", + "- **Mac:** run the standalone MuJoCo lab and data work locally; use Linux GPU/Docker for GR00T+LIBERO.\n", + "- **Linux GPU/Docker:** full closed-loop LIBERO evaluation is supported.\n", + "\n", + "LIBERO uses sparse task success: a rollout matters because the robot actually completes the manipulation, not because its action MSE is low." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ed26f72b", + "metadata": {}, + "outputs": [], + "source": [ + "setup_cmd = [\"bash\", \"gr00t/eval/sim/LIBERO/setup_libero.sh\"]\n", + "print(\"One-time LIBERO setup:\", \" \".join(setup_cmd))\n", + "\n", + "if RUN_LIBERO_SIM_SETUP and FULL_GROOT_SUPPORTED:\n", + " subprocess.check_call(setup_cmd, cwd=REPO_DIR)\n", + "elif RUN_LIBERO_SIM_SETUP and not FULL_GROOT_SUPPORTED:\n", + " print(\"LIBERO+GR00T closed-loop setup requires Linux + NVIDIA CUDA.\")\n", + "else:\n", + " print(\"Not executed. Set RUN_LIBERO_SIM_SETUP=True on the Linux GPU machine.\")\n", + "\n", + "print(f\"\"\"\n", + "Official server-client pattern:\n", + "\n", + "Terminal 1 — policy server:\n", + " uv run python gr00t/eval/run_gr00t_server.py \\\\\n", + " --model-path \\\\\n", + " --embodiment-tag LIBERO_PANDA \\\\\n", + " --use-sim-policy-wrapper\n", + "\n", + "Terminal 2 — MuJoCo/LIBERO rollout:\n", + " gr00t/eval/sim/LIBERO/libero_uv/.venv/bin/python gr00t/eval/rollout_policy.py \\\\\n", + " --n-episodes 10 \\\\\n", + " --policy-client-host 127.0.0.1 \\\\\n", + " --policy-client-port 5555 \\\\\n", + " --max-episode-steps 720 \\\\\n", + " --env-name {LIBERO_ENV} \\\\\n", + " --n-action-steps 8 \\\\\n", + " --n-envs 1 \\\\\n", + " --seed 42 \\\\\n", + " --video-dir /tmp/libero_rollout_video\n", + "\"\"\")" + ] + }, + { + "cell_type": "markdown", + "id": "f55a5714", + "metadata": {}, + "source": [ + "## 18. Before/after inference in MuJoCo\n", + "\n", + "For LIBERO, **do not compare the raw GR00T base model to a LIBERO checkpoint**. `LIBERO_PANDA` is not a base-model embodiment.\n", + "\n", + "Instead compare:\n", + "\n", + "```text\n", + "same MuJoCo task\n", + "same random seed\n", + "same episode count\n", + "same action horizon\n", + "\n", + "EARLY checkpoint ──► rollout videos + success rate\n", + "LATE checkpoint ──► rollout videos + success rate\n", + "```\n", + "\n", + "The default learning run saves at steps 50 and 100. Comparison requires two existing, distinct checkpoint directories; a missing checkpoint or two paths resolving to the same directory blocks the comparison.\n", + "\n", + "The rollout evaluator can save videos with `--video-dir`, and it prints the final success rate." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5919c7c9", + "metadata": {}, + "outputs": [], + "source": [ + "LIBERO_ENV = libero_env_for_suite(LIBERO_SUITE)\n", + "LIBERO_EPISODES = 10\n", + "LIBERO_SEED = 42\n", + "\n", + "complex_ckpts = list_checkpoints(COMPLEX_OUT)\n", + "EARLY_CKPT = complex_ckpts[0] if complex_ckpts else None\n", + "LATE_CKPT = complex_ckpts[-1] if len(complex_ckpts) >= 2 else None\n", + "\n", + "# You can also set LATE_CKPT to NVIDIA's published fully trained suite checkpoint:\n", + "# LATE_CKPT = PRETRAINED_LIBERO_ROOT / f\"libero_{LIBERO_SUITE}\"\n", + "\n", + "\n", + "def rollout_client_command(video_dir, port=5555):\n", + " libero_python = REPO_DIR / \"gr00t/eval/sim/LIBERO/libero_uv/.venv/bin/python\"\n", + " return [\n", + " str(libero_python),\n", + " \"gr00t/eval/rollout_policy.py\",\n", + " \"--n-episodes\",\n", + " str(LIBERO_EPISODES),\n", + " \"--policy-client-host\",\n", + " \"127.0.0.1\",\n", + " \"--policy-client-port\",\n", + " str(port),\n", + " \"--max-episode-steps\",\n", + " \"720\",\n", + " \"--env-name\",\n", + " LIBERO_ENV,\n", + " \"--n-action-steps\",\n", + " \"8\",\n", + " \"--n-envs\",\n", + " \"1\",\n", + " \"--seed\",\n", + " str(LIBERO_SEED),\n", + " \"--video-dir\",\n", + " str(video_dir),\n", + " ]\n", + "\n", + "\n", + "COMPARISON_PAIR = None\n", + "try:\n", + " COMPARISON_PAIR = validate_checkpoint_pair(EARLY_CKPT, LATE_CKPT)\n", + "except ValueError as error:\n", + " if RUN_MUJOCO_BEFORE_AFTER:\n", + " raise\n", + " print(\"Comparison not ready:\", error)\n", + "\n", + "if COMPARISON_PAIR is not None:\n", + " EARLY_CKPT, LATE_CKPT = COMPARISON_PAIR\n", + " print(\"Early checkpoint:\", EARLY_CKPT)\n", + " print(\"Late checkpoint :\", LATE_CKPT)\n", + " if RUN_MUJOCO_BEFORE_AFTER:\n", + " print(\"Run these commands from\", REPO_DIR, \"for each checkpoint, one at a time.\")\n", + " for label, checkpoint in zip((\"early\", \"late\"), COMPARISON_PAIR, strict=True):\n", + " server = [\n", + " uv or \"uv\",\n", + " \"run\",\n", + " \"python\",\n", + " \"gr00t/eval/run_gr00t_server.py\",\n", + " \"--model-path\",\n", + " str(checkpoint),\n", + " \"--embodiment-tag\",\n", + " \"LIBERO_PANDA\",\n", + " \"--use-sim-policy-wrapper\",\n", + " \"--host\",\n", + " \"127.0.0.1\",\n", + " \"--port\",\n", + " \"5555\",\n", + " ]\n", + " video_dir = WORK_ROOT / \"videos\" / f\"libero_{LIBERO_SUITE}_{label}\"\n", + " print(f\"\\n{label.title()} — terminal 1 (stop the previous server first):\")\n", + " print(shlex.join(server))\n", + " print(f\"{label.title()} — terminal 2:\")\n", + " print(shlex.join(rollout_client_command(video_dir)))\n", + " else:\n", + " print(\"Set RUN_MUJOCO_BEFORE_AFTER=True to generate the two rollout command pairs.\")" + ] + }, + { + "cell_type": "markdown", + "id": "aeed8c77", + "metadata": {}, + "source": [ + "### Run both checkpoints\n", + "\n", + "Set `RUN_MUJOCO_BEFORE_AFTER=True` and rerun the preceding cell to generate commands for both checkpoints. The flag generates commands; it does not start GPU processes automatically. For each checkpoint:\n", + "\n", + "1. start the GR00T policy server with that checkpoint;\n", + "2. run the MuJoCo client with the **same seed**;\n", + "3. write videos to a separate directory;\n", + "4. save stdout to a log file.\n", + "\n", + "Example server:\n", + "\n", + "```bash\n", + "uv run python gr00t/eval/run_gr00t_server.py --model-path /path/to/checkpoint-500 --embodiment-tag LIBERO_PANDA --use-sim-policy-wrapper --host 127.0.0.1 --port 5555\n", + "```\n", + "\n", + "The client already prints:\n", + "\n", + "```text\n", + "results: ...\n", + "success rate: 0.7\n", + "Video saved to: ...\n", + "```\n", + "\n", + "Run again for the late checkpoint using the same environment and `--seed 42`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fa5090bc", + "metadata": {}, + "outputs": [], + "source": [ + "def parse_rollout_log(text):\n", + " rate = None\n", + " m = re.search(r\"success rate:\\s*([0-9.eE+-]+)\", text)\n", + " if m:\n", + " rate = float(m.group(1))\n", + "\n", + " video = None\n", + " m = re.search(r\"Video saved to:\\s*(.+)\", text)\n", + " if m:\n", + " video = m.group(1).strip()\n", + "\n", + " return {\"success_rate\": rate, \"video_dir\": video}\n", + "\n", + "\n", + "def compare_success_rates(early_log_text, late_log_text):\n", + " validate_checkpoint_pair(EARLY_CKPT, LATE_CKPT)\n", + " early = parse_rollout_log(early_log_text)\n", + " late = parse_rollout_log(late_log_text)\n", + "\n", + " frame = pd.DataFrame(\n", + " [\n", + " {\"checkpoint\": \"early\", \"success_rate\": early[\"success_rate\"]},\n", + " {\"checkpoint\": \"late\", \"success_rate\": late[\"success_rate\"]},\n", + " ]\n", + " )\n", + " display(frame)\n", + "\n", + " if frame[\"success_rate\"].notna().all():\n", + " plt.figure(figsize=(6, 4))\n", + " plt.bar(frame[\"checkpoint\"], frame[\"success_rate\"])\n", + " plt.ylim(0, 1)\n", + " plt.ylabel(\"closed-loop task success rate\")\n", + " plt.title(f\"MuJoCo / LIBERO before vs after — seed {LIBERO_SEED}\")\n", + " plt.show()\n", + "\n", + " return frame\n", + "\n", + "\n", + "# Example after saving logs:\n", + "# early_text = Path(\"early_rollout.log\").read_text()\n", + "# late_text = Path(\"late_rollout.log\").read_text()\n", + "# compare_success_rates(early_text, late_text)" + ] + }, + { + "cell_type": "markdown", + "id": "d992c69d", + "metadata": {}, + "source": [ + "### What to look for in the videos\n", + "\n", + "Do not only ask \"did it succeed?\"\n", + "\n", + "Watch for:\n", + "- whether it reaches the correct object;\n", + "- grasp approach/orientation;\n", + "- oscillation;\n", + "- overshoot;\n", + "- whether the gripper closes at the right phase;\n", + "- recovery after small errors;\n", + "- repeated action patterns;\n", + "- collisions;\n", + "- whether failure occurs because of perception, planning or low-level action execution.\n", + "\n", + "That qualitative failure analysis is often more informative than another decimal place of MSE." + ] + }, + { + "cell_type": "markdown", + "id": "bcb55cc3", + "metadata": {}, + "source": [ + "# Part D — MCAP data cleaning exercise\n", + "\n", + "This is the part that matters when demonstrations come from an industrial arm through ROS 2.\n", + "\n", + "Typical raw rates differ:\n", + "\n", + "```text\n", + "joint state 50–500 Hz\n", + "action/command 20–100 Hz\n", + "RGB camera 15–60 Hz\n", + "force/torque 100–1000 Hz\n", + "```\n", + "\n", + "Do **not** concatenate messages by row number.\n", + "\n", + "Use:\n", + "\n", + "```text\n", + "ROS 2 / MCAP episode\n", + " ↓\n", + "topic inventory\n", + " ↓\n", + "timestamp QA\n", + " ↓\n", + "remove duplicates / invalid messages\n", + " ↓\n", + "detect sensor gaps + state/action jumps\n", + " ↓\n", + "choose a common training timebase\n", + " ↓\n", + "interpolate/nearest-sample each modality\n", + " ↓\n", + "trim idle head/tail\n", + " ↓\n", + "manual replay / task-success filter\n", + " ↓\n", + "GR00T LeRobot episode\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "85d3f8ee", + "metadata": {}, + "source": [ + "## 18. Create a deliberately dirty MCAP episode\n", + "\n", + "The synthetic recording contains robot state, action, two camera-frame streams and task text. It intentionally has duplicate timestamps, an out-of-order timestamp, a camera gap and one implausible robot-state jump." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "48c246cd", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "from mcap_ros2.writer import Writer as McapWriter\n", + "\n", + "MCAP_DIR = WORK_ROOT / \"mcap\"\n", + "MCAP_DIR.mkdir(parents=True, exist_ok=True)\n", + "SYNTH_MCAP = MCAP_DIR / \"dirty_pick_place_episode.mcap\"\n", + "\n", + "STATE_SCHEMA = \"\"\"\\\n", + "float64[5] arm\n", + "float64 gripper\n", + "\"\"\"\n", + "ACTION_SCHEMA = \"\"\"\\\n", + "float64[5] arm\n", + "float64 gripper\n", + "\"\"\"\n", + "FRAME_SCHEMA = \"\"\"\\\n", + "uint32 frame_id\n", + "\"\"\"\n", + "TASK_SCHEMA = \"\"\"\\\n", + "string text\n", + "\"\"\"\n", + "\n", + "base_ns = 1_700_000_000_000_000_000\n", + "\n", + "with open(SYNTH_MCAP, \"wb\") as f:\n", + " w = McapWriter(f)\n", + " state_schema = w.register_msgdef(\"demo_msgs/RobotState\", STATE_SCHEMA)\n", + " action_schema = w.register_msgdef(\"demo_msgs/RobotAction\", ACTION_SCHEMA)\n", + " frame_schema = w.register_msgdef(\"demo_msgs/FrameMeta\", FRAME_SCHEMA)\n", + " task_schema = w.register_msgdef(\"demo_msgs/Task\", TASK_SCHEMA)\n", + "\n", + " w.write_message(\n", + " topic=\"/task\",\n", + " schema=task_schema,\n", + " message={\"text\": \"pick the red cube and place it in the bowl\"},\n", + " log_time=base_ns,\n", + " publish_time=base_ns,\n", + " sequence=0,\n", + " )\n", + "\n", + " n_state = 300\n", + " t_state = np.arange(n_state) / 50.0\n", + " ts_state = base_ns + (t_state * 1e9).astype(np.int64)\n", + " ts_state[65] = ts_state[64]\n", + " ts_state[155] = ts_state[154] - 5_000_000\n", + "\n", + " for i, (t, ts) in enumerate(zip(t_state, ts_state, strict=False)):\n", + " arm = np.array(\n", + " [\n", + " 0.30 * np.sin(0.7 * t),\n", + " 0.25 * np.sin(0.9 * t + 0.3),\n", + " 0.20 * np.cos(0.6 * t),\n", + " 0.15 * np.sin(1.2 * t),\n", + " 0.10 * np.cos(0.8 * t),\n", + " ]\n", + " )\n", + " if i == 190:\n", + " arm = arm + 3.0\n", + " gripper = 1.0 if t < 2.4 else 0.0\n", + " w.write_message(\n", + " topic=\"/robot/state\",\n", + " schema=state_schema,\n", + " message={\"arm\": arm.tolist(), \"gripper\": float(gripper)},\n", + " log_time=int(ts),\n", + " publish_time=int(ts),\n", + " sequence=i,\n", + " )\n", + "\n", + " n_action = 120\n", + " t_action = np.arange(n_action) / 20.0\n", + " ts_action = base_ns + (t_action * 1e9).astype(np.int64)\n", + " ts_action[48] = ts_action[47]\n", + " for i, (t, ts) in enumerate(zip(t_action, ts_action, strict=False)):\n", + " arm = np.array(\n", + " [\n", + " 0.30 * np.sin(0.7 * (t + 0.08)),\n", + " 0.25 * np.sin(0.9 * (t + 0.08) + 0.3),\n", + " 0.20 * np.cos(0.6 * (t + 0.08)),\n", + " 0.15 * np.sin(1.2 * (t + 0.08)),\n", + " 0.10 * np.cos(0.8 * (t + 0.08)),\n", + " ]\n", + " )\n", + " gripper = 1.0 if t < 2.3 else 0.0\n", + " w.write_message(\n", + " topic=\"/robot/action\",\n", + " schema=action_schema,\n", + " message={\"arm\": arm.tolist(), \"gripper\": float(gripper)},\n", + " log_time=int(ts),\n", + " publish_time=int(ts),\n", + " sequence=i,\n", + " )\n", + "\n", + " frame_times = np.arange(90) / 15.0\n", + " front_times = [t for t in frame_times if not (3.0 < t < 3.45)]\n", + " for i, t in enumerate(front_times):\n", + " ts = base_ns + int(t * 1e9)\n", + " w.write_message(\n", + " topic=\"/camera/front/frame\",\n", + " schema=frame_schema,\n", + " message={\"frame_id\": i},\n", + " log_time=ts,\n", + " publish_time=ts,\n", + " sequence=i,\n", + " )\n", + "\n", + " for i, t in enumerate(frame_times):\n", + " ts = base_ns + int(t * 1e9)\n", + " if i == 40:\n", + " ts = base_ns + int(frame_times[39] * 1e9)\n", + " w.write_message(\n", + " topic=\"/camera/wrist/frame\",\n", + " schema=frame_schema,\n", + " message={\"frame_id\": i},\n", + " log_time=ts,\n", + " publish_time=ts,\n", + " sequence=i,\n", + " )\n", + "\n", + " w.finish()\n", + "\n", + "print(\"Created:\", SYNTH_MCAP)\n", + "print(\"Size:\", f\"{SYNTH_MCAP.stat().st_size / 1024:.1f} KiB\")" + ] + }, + { + "cell_type": "markdown", + "id": "90f85877", + "metadata": {}, + "source": [ + "## 19. Generic MCAP topic inventory — use this on real bags too" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f036893d", + "metadata": {}, + "outputs": [], + "source": [ + "from collections import defaultdict\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "from mcap.reader import make_reader\n", + "\n", + "\n", + "def inventory_mcap(path):\n", + " rows = defaultdict(list)\n", + " schemas = {}\n", + "\n", + " with open(path, \"rb\") as f:\n", + " reader = make_reader(f)\n", + " for schema, channel, message in reader.iter_messages(log_time_order=False):\n", + " rows[channel.topic].append(int(message.log_time))\n", + " schemas[channel.topic] = schema.name if schema else None\n", + "\n", + " out = []\n", + " for topic, ts_ns in rows.items():\n", + " ts = np.asarray(ts_ns, dtype=np.int64) / 1e9\n", + " dt_record_order = np.diff(ts)\n", + " ts_sorted = np.sort(ts)\n", + " dt_sorted = np.diff(ts_sorted)\n", + " positive = dt_sorted[dt_sorted > 0]\n", + " med_dt = np.median(positive) if len(positive) else np.nan\n", + "\n", + " out.append(\n", + " {\n", + " \"topic\": topic,\n", + " \"schema\": schemas[topic],\n", + " \"messages\": len(ts),\n", + " \"span_s\": float(ts_sorted[-1] - ts_sorted[0]) if len(ts) > 1 else 0.0,\n", + " \"median_rate_hz\": float(1 / med_dt)\n", + " if np.isfinite(med_dt) and med_dt > 0\n", + " else np.nan,\n", + " \"duplicates\": int(len(ts_sorted) - len(np.unique(ts_sorted))),\n", + " \"record_order_nonmonotonic\": int(np.sum(dt_record_order < 0)),\n", + " \"gaps_gt_3x_median\": int(np.sum(dt_sorted > 3 * med_dt))\n", + " if np.isfinite(med_dt)\n", + " else 0,\n", + " \"max_gap_s\": float(np.max(dt_sorted)) if len(dt_sorted) else 0.0,\n", + " }\n", + " )\n", + "\n", + " return pd.DataFrame(out).sort_values(\"topic\").reset_index(drop=True)\n", + "\n", + "\n", + "MCAP_PATH = SYNTH_MCAP\n", + "display(inventory_mcap(MCAP_PATH))" + ] + }, + { + "cell_type": "markdown", + "id": "d7c63a46", + "metadata": {}, + "source": [ + "For a real recording:\n", + "\n", + "```python\n", + "MCAP_PATH = Path(\"/data/episodes/episode_000123.mcap\")\n", + "display(inventory_mcap(MCAP_PATH))\n", + "```\n", + "\n", + "Look for inconsistent camera FPS, controller pauses, duplicates, cross-camera drift, force/torque saturation and missing TF." + ] + }, + { + "cell_type": "markdown", + "id": "ed39e276", + "metadata": {}, + "source": [ + "## 20. Decode ROS 2 fields without installing ROS 2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb19d316", + "metadata": {}, + "outputs": [], + "source": [ + "from mcap_ros2.decoder import DecoderFactory\n", + "\n", + "with open(MCAP_PATH, \"rb\") as f:\n", + " reader = make_reader(f, decoder_factories=[DecoderFactory()])\n", + " seen = set()\n", + "\n", + " for schema, channel, _message, ros_msg in reader.iter_decoded_messages(log_time_order=False):\n", + " if channel.topic not in seen:\n", + " print(\"\\nTOPIC:\", channel.topic)\n", + " print(\"TYPE :\", schema.name)\n", + " print(\"MSG :\", ros_msg)\n", + " seen.add(channel.topic)\n", + " if len(seen) >= 5:\n", + " break" + ] + }, + { + "cell_type": "markdown", + "id": "8a1b4d59", + "metadata": {}, + "source": [ + "## 21. Extract robot streams\n", + "\n", + "Customize this mapping for your robot. Keep source timestamps." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "90a4f5b3", + "metadata": {}, + "outputs": [], + "source": [ + "def decode_demo_mcap(path):\n", + " state_rows, action_rows, front_rows, wrist_rows, task_rows = [], [], [], [], []\n", + "\n", + " with open(path, \"rb\") as f:\n", + " reader = make_reader(f, decoder_factories=[DecoderFactory()])\n", + " for _schema, channel, message, ros_msg in reader.iter_decoded_messages(\n", + " log_time_order=False\n", + " ):\n", + " t = message.log_time / 1e9\n", + " topic = channel.topic\n", + "\n", + " if topic == \"/robot/state\":\n", + " vals = list(ros_msg.arm)\n", + " state_rows.append(\n", + " {\n", + " \"t\": t,\n", + " **{f\"arm_{i}\": float(v) for i, v in enumerate(vals)},\n", + " \"gripper\": float(ros_msg.gripper),\n", + " }\n", + " )\n", + " elif topic == \"/robot/action\":\n", + " vals = list(ros_msg.arm)\n", + " action_rows.append(\n", + " {\n", + " \"t\": t,\n", + " **{f\"arm_{i}\": float(v) for i, v in enumerate(vals)},\n", + " \"gripper\": float(ros_msg.gripper),\n", + " }\n", + " )\n", + " elif topic == \"/camera/front/frame\":\n", + " front_rows.append({\"t\": t, \"frame_id\": int(ros_msg.frame_id)})\n", + " elif topic == \"/camera/wrist/frame\":\n", + " wrist_rows.append({\"t\": t, \"frame_id\": int(ros_msg.frame_id)})\n", + " elif topic == \"/task\":\n", + " task_rows.append({\"t\": t, \"text\": str(ros_msg.text)})\n", + "\n", + " return {\n", + " \"state\": pd.DataFrame(state_rows),\n", + " \"action\": pd.DataFrame(action_rows),\n", + " \"front\": pd.DataFrame(front_rows),\n", + " \"wrist\": pd.DataFrame(wrist_rows),\n", + " \"task\": pd.DataFrame(task_rows),\n", + " }\n", + "\n", + "\n", + "streams = decode_demo_mcap(MCAP_PATH)\n", + "for name, df in streams.items():\n", + " print(name, df.shape)\n", + " display(df.head(3))" + ] + }, + { + "cell_type": "markdown", + "id": "2b1aa1f7", + "metadata": {}, + "source": [ + "## 22. Timestamp cleaning\n", + "\n", + "Preserve raw records, sort, remove exact duplicates, report changes and never silently interpolate across a large outage." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a3d9301b", + "metadata": {}, + "outputs": [], + "source": [ + "def clean_timestamps(df, name, gap_factor=3.0):\n", + " x = df.copy()\n", + " before = len(x)\n", + " original_nonmono = int(np.sum(np.diff(x[\"t\"].to_numpy()) <= 0)) if len(x) > 1 else 0\n", + "\n", + " x = x.sort_values(\"t\").reset_index(drop=True)\n", + " dup = int(x[\"t\"].duplicated(keep=\"first\").sum())\n", + " x = x.drop_duplicates(\"t\", keep=\"first\").reset_index(drop=True)\n", + "\n", + " dt = np.diff(x[\"t\"].to_numpy())\n", + " pos = dt[dt > 0]\n", + " med = np.median(pos) if len(pos) else np.nan\n", + " gap_limit = gap_factor * med if np.isfinite(med) else np.nan\n", + " gap_idx = np.where(dt > gap_limit)[0] if np.isfinite(gap_limit) else np.array([], dtype=int)\n", + "\n", + " report = {\n", + " \"stream\": name,\n", + " \"rows_before\": before,\n", + " \"rows_after_dedupe\": len(x),\n", + " \"nonmonotonic_or_duplicate_in_record_order\": original_nonmono,\n", + " \"duplicates_removed\": dup,\n", + " \"median_dt_s\": med,\n", + " \"gap_limit_s\": gap_limit,\n", + " \"large_gap_count\": int(len(gap_idx)),\n", + " \"large_gaps_after_row\": gap_idx.tolist(),\n", + " }\n", + " return x, report\n", + "\n", + "\n", + "cleaned = {}\n", + "reports = []\n", + "for name in [\"state\", \"action\", \"front\", \"wrist\"]:\n", + " cleaned[name], rep = clean_timestamps(streams[name], name)\n", + " reports.append(rep)\n", + "\n", + "display(pd.DataFrame(reports))" + ] + }, + { + "cell_type": "markdown", + "id": "1bd9ec96", + "metadata": {}, + "source": [ + "## 23. Detect implausible state jumps\n", + "\n", + "For real hardware, replace the generic robust threshold with real joint-position, velocity, acceleration and EEF limits." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "346695fc", + "metadata": {}, + "outputs": [], + "source": [ + "state = cleaned[\"state\"].copy()\n", + "joint_cols = [c for c in state.columns if c.startswith(\"arm_\")]\n", + "\n", + "dt = np.diff(state[\"t\"].to_numpy())\n", + "dq = np.diff(state[joint_cols].to_numpy(), axis=0)\n", + "speed_norm = np.linalg.norm(dq / dt[:, None], axis=1)\n", + "\n", + "speed_limit = robust_threshold(speed_norm, z=8.0)\n", + "bad_after = np.where(speed_norm > speed_limit)[0] + 1\n", + "\n", + "print(\"Robust speed-norm threshold:\", speed_limit)\n", + "print(\"Flagged row indices:\", bad_after.tolist())\n", + "\n", + "plt.figure(figsize=(12, 4))\n", + "plt.plot(state[\"t\"].iloc[1:], speed_norm)\n", + "plt.axhline(speed_limit, linestyle=\"--\", label=\"robust threshold\")\n", + "plt.xlabel(\"time (s)\")\n", + "plt.ylabel(\"joint-speed vector norm\")\n", + "plt.title(\"State jump detection\")\n", + "plt.legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "56f2730a", + "metadata": {}, + "source": [ + "## 24. Remove isolated bad samples with an audit trail\n", + "\n", + "The demo removes a one-frame spike only if neighboring samples are mutually consistent. Long suspicious regions should trigger manual review." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4514c990", + "metadata": {}, + "outputs": [], + "source": [ + "def remove_isolated_state_outliers(df, joint_cols, threshold):\n", + " x = df.copy().reset_index(drop=True)\n", + " drop = []\n", + " arr = x[joint_cols].to_numpy(float)\n", + " t = x[\"t\"].to_numpy(float)\n", + "\n", + " for i in range(1, len(x) - 1):\n", + " dt_prev = t[i] - t[i - 1]\n", + " dt_next = t[i + 1] - t[i]\n", + " dt_bridge = t[i + 1] - t[i - 1]\n", + " if min(dt_prev, dt_next, dt_bridge) <= 0:\n", + " continue\n", + "\n", + " v_prev = np.linalg.norm((arr[i] - arr[i - 1]) / dt_prev)\n", + " v_next = np.linalg.norm((arr[i + 1] - arr[i]) / dt_next)\n", + " v_bridge = np.linalg.norm((arr[i + 1] - arr[i - 1]) / dt_bridge)\n", + "\n", + " if v_prev > threshold and v_next > threshold and v_bridge < threshold:\n", + " drop.append(i)\n", + "\n", + " return x.drop(index=drop).reset_index(drop=True), drop\n", + "\n", + "\n", + "state_clean, dropped_state_rows = remove_isolated_state_outliers(\n", + " cleaned[\"state\"], joint_cols, speed_limit\n", + ")\n", + "cleaned[\"state\"] = state_clean\n", + "print(\"Dropped isolated state rows:\", dropped_state_rows)" + ] + }, + { + "cell_type": "markdown", + "id": "105bb702", + "metadata": {}, + "source": [ + "## 25. Synchronize modalities on a visual timebase\n", + "\n", + "We use front-camera timestamps, nearest-match the wrist camera, reject poor camera sync, then interpolate continuous state/action values." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "066ce015", + "metadata": {}, + "outputs": [], + "source": [ + "def interpolate_numeric(df, target_t, columns):\n", + " src_t = df[\"t\"].to_numpy(float)\n", + " out = {}\n", + " for col in columns:\n", + " out[col] = np.interp(target_t, src_t, df[col].to_numpy(float))\n", + " return pd.DataFrame(out)\n", + "\n", + "\n", + "front = cleaned[\"front\"]\n", + "wrist = cleaned[\"wrist\"]\n", + "action = cleaned[\"action\"]\n", + "state = cleaned[\"state\"]\n", + "\n", + "start_t = max(front.t.min(), wrist.t.min(), action.t.min(), state.t.min())\n", + "end_t = min(front.t.max(), wrist.t.max(), action.t.max(), state.t.max())\n", + "\n", + "front_use = front[(front.t >= start_t) & (front.t <= end_t)].copy().reset_index(drop=True)\n", + "\n", + "wrist_t = wrist[\"t\"].to_numpy(float)\n", + "front_t = front_use[\"t\"].to_numpy(float)\n", + "nearest_idx = np.abs(wrist_t[:, None] - front_t[None, :]).argmin(axis=0)\n", + "nearest_wrist_t = wrist_t[nearest_idx]\n", + "cam_sync_error = np.abs(nearest_wrist_t - front_t)\n", + "\n", + "nominal_cam_dt = np.median(np.diff(np.sort(front_t)))\n", + "MAX_CAMERA_SYNC_ERROR = 0.5 * nominal_cam_dt\n", + "keep = cam_sync_error <= MAX_CAMERA_SYNC_ERROR\n", + "\n", + "timeline = front_use.loc[keep, \"t\"].to_numpy(float)\n", + "matched_wrist = wrist.iloc[nearest_idx[keep]].reset_index(drop=True)\n", + "\n", + "state_cols = [c for c in state.columns if c != \"t\"]\n", + "action_cols = [c for c in action.columns if c != \"t\"]\n", + "\n", + "state_aligned = interpolate_numeric(state, timeline, state_cols)\n", + "action_aligned = interpolate_numeric(action, timeline, action_cols)\n", + "\n", + "aligned = pd.DataFrame(\n", + " {\n", + " \"t\": timeline,\n", + " \"front_frame_id\": front_use.loc[keep, \"frame_id\"].to_numpy(),\n", + " \"wrist_frame_id\": matched_wrist[\"frame_id\"].to_numpy(),\n", + " \"camera_sync_error_s\": cam_sync_error[keep],\n", + " }\n", + ")\n", + "for c in state_cols:\n", + " aligned[f\"state.{c}\"] = state_aligned[c].to_numpy()\n", + "for c in action_cols:\n", + " aligned[f\"action.{c}\"] = action_aligned[c].to_numpy()\n", + "\n", + "print(\"Frames before camera-sync filter:\", len(front_use))\n", + "print(\"Frames after camera-sync filter:\", len(aligned))\n", + "print(\"Max accepted camera sync error:\", aligned.camera_sync_error_s.max())\n", + "display(aligned.head())" + ] + }, + { + "cell_type": "markdown", + "id": "d686653d", + "metadata": {}, + "source": [ + "## 26. Visualize raw → cleaned → aligned data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "88bf285c", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure(figsize=(12, 5))\n", + "plt.plot(streams[\"state\"][\"t\"], streams[\"state\"][\"arm_0\"], \".\", alpha=0.4, label=\"raw state arm_0\")\n", + "plt.plot(cleaned[\"state\"][\"t\"], cleaned[\"state\"][\"arm_0\"], \"-\", label=\"cleaned state arm_0\")\n", + "plt.plot(aligned[\"t\"], aligned[\"state.arm_0\"], \"o\", markersize=3, label=\"aligned @ camera time\")\n", + "plt.xlabel(\"absolute time (s)\")\n", + "plt.ylabel(\"joint/state value\")\n", + "plt.title(\"Raw → cleaned → camera-aligned\")\n", + "plt.legend()\n", + "plt.show()\n", + "\n", + "plt.figure(figsize=(12, 4))\n", + "plt.plot(aligned[\"t\"], aligned[\"camera_sync_error_s\"] * 1000)\n", + "plt.xlabel(\"absolute time (s)\")\n", + "plt.ylabel(\"front↔wrist timestamp error (ms)\")\n", + "plt.title(\"Camera synchronization error\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "c96c6636", + "metadata": {}, + "source": [ + "## 27. Head/tail idle trimming\n", + "\n", + "This uses action motion as a simple first pass. Production pipelines should combine it with task-state/success annotations." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e3c09804", + "metadata": {}, + "outputs": [], + "source": [ + "a_cols = [c for c in aligned.columns if c.startswith(\"action.arm_\")]\n", + "a = aligned[a_cols].to_numpy(float)\n", + "t = aligned[\"t\"].to_numpy(float)\n", + "\n", + "if len(aligned) >= 3:\n", + " da = np.linalg.norm(np.diff(a, axis=0), axis=1)\n", + " dt = np.diff(t)\n", + " action_speed = da / np.maximum(dt, 1e-9)\n", + " move_threshold = max(np.percentile(action_speed, 20), 1e-5)\n", + " moving = np.r_[False, action_speed > move_threshold]\n", + "\n", + " idx = np.where(moving)[0]\n", + " if len(idx):\n", + " margin_frames = 3\n", + " lo = max(0, idx[0] - margin_frames)\n", + " hi = min(len(aligned), idx[-1] + margin_frames + 1)\n", + " aligned_trimmed = aligned.iloc[lo:hi].reset_index(drop=True)\n", + " else:\n", + " aligned_trimmed = aligned.copy()\n", + "else:\n", + " aligned_trimmed = aligned.copy()\n", + "\n", + "print(\"Before trim:\", len(aligned))\n", + "print(\"After trim :\", len(aligned_trimmed))" + ] + }, + { + "cell_type": "markdown", + "id": "1185f929", + "metadata": {}, + "source": [ + "## 28. Export a GR00T-shaped staging episode\n", + "\n", + "The export mirrors the SO100 5-arm+gripper dimensional convention and creates tiny synthetic MP4s only so the final directory shape is inspectable.\n", + "\n", + "**The generated `info.json` is staging metadata, not a complete production LeRobot export.** For real training, use real images and regenerate/validate complete LeRobot metadata and GR00T statistics." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "69a0cc14", + "metadata": {}, + "outputs": [], + "source": [ + "import shutil\n", + "\n", + "import imageio.v2 as imageio\n", + "from PIL import Image, ImageDraw\n", + "\n", + "EXPORT_ROOT = WORK_ROOT / \"mcap\" / \"cleaned_episode_gr00t_staging\"\n", + "if EXPORT_ROOT.exists():\n", + " shutil.rmtree(EXPORT_ROOT)\n", + "\n", + "meta_dir = EXPORT_ROOT / \"meta\"\n", + "data_dir = EXPORT_ROOT / \"data\" / \"chunk-000\"\n", + "front_dir = EXPORT_ROOT / \"videos\" / \"chunk-000\" / \"observation.images.front\"\n", + "wrist_dir = EXPORT_ROOT / \"videos\" / \"chunk-000\" / \"observation.images.wrist\"\n", + "\n", + "for d in [meta_dir, data_dir, front_dir, wrist_dir]:\n", + " d.mkdir(parents=True, exist_ok=True)\n", + "\n", + "df = aligned_trimmed.copy()\n", + "t0 = df[\"t\"].iloc[0]\n", + "rel_t = df[\"t\"].to_numpy(float) - t0\n", + "\n", + "state_cols_export = [f\"state.arm_{i}\" for i in range(5)] + [\"state.gripper\"]\n", + "action_cols_export = [f\"action.arm_{i}\" for i in range(5)] + [\"action.gripper\"]\n", + "\n", + "episode_df = pd.DataFrame(\n", + " {\n", + " \"observation.state\": [\n", + " np.asarray(row, dtype=np.float32) for row in df[state_cols_export].to_numpy()\n", + " ],\n", + " \"action\": [np.asarray(row, dtype=np.float32) for row in df[action_cols_export].to_numpy()],\n", + " \"timestamp\": rel_t.astype(np.float64),\n", + " \"annotation.human.task_description\": np.zeros(len(df), dtype=np.int64),\n", + " \"task_index\": np.zeros(len(df), dtype=np.int64),\n", + " \"episode_index\": np.zeros(len(df), dtype=np.int64),\n", + " \"index\": np.arange(len(df), dtype=np.int64),\n", + " \"next.reward\": np.zeros(len(df), dtype=np.float32),\n", + " \"next.done\": np.r_[np.zeros(max(0, len(df) - 1), dtype=bool), True],\n", + " }\n", + ")\n", + "episode_df.to_parquet(data_dir / \"episode_000000.parquet\", index=False)\n", + "\n", + "modality_json = {\n", + " \"state\": {\n", + " \"single_arm\": {\"start\": 0, \"end\": 5},\n", + " \"gripper\": {\"start\": 5, \"end\": 6},\n", + " },\n", + " \"action\": {\n", + " \"single_arm\": {\"start\": 0, \"end\": 5},\n", + " \"gripper\": {\"start\": 5, \"end\": 6},\n", + " },\n", + " \"video\": {\n", + " \"front\": {\"original_key\": \"observation.images.front\"},\n", + " \"wrist\": {\"original_key\": \"observation.images.wrist\"},\n", + " },\n", + " \"annotation\": {\n", + " \"human.task_description\": {\"original_key\": \"task_index\"},\n", + " },\n", + "}\n", + "(meta_dir / \"modality.json\").write_text(json.dumps(modality_json, indent=2))\n", + "\n", + "task = \"pick the red cube and place it in the bowl\"\n", + "(meta_dir / \"tasks.jsonl\").write_text(json.dumps({\"task_index\": 0, \"task\": task}) + \"\\n\")\n", + "(meta_dir / \"episodes.jsonl\").write_text(\n", + " json.dumps(\n", + " {\n", + " \"episode_index\": 0,\n", + " \"tasks\": [task],\n", + " \"length\": len(episode_df),\n", + " }\n", + " )\n", + " + \"\\n\"\n", + ")\n", + "\n", + "fps = 1.0 / np.median(np.diff(rel_t)) if len(rel_t) > 1 else 15.0\n", + "info_staging = {\n", + " \"codebase_version\": \"staging-created-by-gr00t-lab\",\n", + " \"robot_type\": \"industrial_arm_demo\",\n", + " \"total_episodes\": 1,\n", + " \"total_frames\": len(episode_df),\n", + " \"fps\": float(fps),\n", + " \"note\": \"Staging metadata only. Regenerate/validate with official LeRobot/GR00T tooling before training.\",\n", + "}\n", + "(meta_dir / \"info.json\").write_text(json.dumps(info_staging, indent=2))\n", + "\n", + "video_fps = max(1, round(fps))\n", + "\n", + "\n", + "def make_frame(i, n, label, wrist=False):\n", + " img = Image.new(\"RGB\", (320, 180), (245, 245, 245))\n", + " d = ImageDraw.Draw(img)\n", + " frac = i / max(1, n - 1)\n", + " x = int(40 + frac * 220)\n", + " y = int(90 + (30 if wrist else 45) * math.sin(frac * 2 * math.pi))\n", + " d.ellipse((x - 10, y - 10, x + 10, y + 10), fill=(60, 80, 180))\n", + " d.text((10, 10), f\"{label} frame {i}\", fill=(0, 0, 0))\n", + " return np.asarray(img)\n", + "\n", + "\n", + "for path, label, wrist_flag in [\n", + " (front_dir / \"episode_000000.mp4\", \"front\", False),\n", + " (wrist_dir / \"episode_000000.mp4\", \"wrist\", True),\n", + "]:\n", + " writer = imageio.get_writer(str(path), fps=video_fps, codec=\"libx264\", quality=7)\n", + " for i in range(len(episode_df)):\n", + " writer.append_data(make_frame(i, len(episode_df), label, wrist=wrist_flag))\n", + " writer.close()\n", + "\n", + "print(\"Exported staging dataset:\", EXPORT_ROOT)\n", + "for p in sorted(EXPORT_ROOT.rglob(\"*\")):\n", + " if p.is_file():\n", + " print(p.relative_to(EXPORT_ROOT))" + ] + }, + { + "cell_type": "markdown", + "id": "72f4c0df", + "metadata": {}, + "source": [ + "## 29. Audit the cleaned staging episode" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "eb2a9ab0", + "metadata": {}, + "outputs": [], + "source": [ + "staging_audit = audit_gr00t_dataset(EXPORT_ROOT)\n", + "display(staging_audit)\n", + "\n", + "display(pd.read_parquet(EXPORT_ROOT / \"data/chunk-000/episode_000000.parquet\").head())" + ] + }, + { + "cell_type": "markdown", + "id": "b0917b6a", + "metadata": {}, + "source": [ + "## 30. Generate GR00T statistics before real training\n", + "\n", + "Once a real converted dataset has complete LeRobot v2 metadata, compute GR00T statistics:\n", + "\n", + "```bash\n", + "uv run python gr00t/data/stats.py --dataset-path /path/to/your_dataset --embodiment-tag NEW_EMBODIMENT --modality-config-path examples/SO100/so100_config.py\n", + "```\n", + "\n", + "If your industrial arm has a different state/action space, create a matching modality config instead of reusing SO100." + ] + }, + { + "cell_type": "markdown", + "id": "892055b9", + "metadata": {}, + "source": [ + "# Part E — translate this to an industrial arm\n", + "\n", + "A practical architecture:\n", + "\n", + "```text\n", + " task instruction\n", + " │\n", + " ▼\n", + "front camera ────────────── GR00T\n", + "wrist camera ────────────────┤\n", + "joint state / EEF pose ──────┤\n", + " │\n", + " future action chunk\n", + " │\n", + " ▼\n", + " safety / rate limiter\n", + " │\n", + " ▼\n", + " robot controller / ROS2\n", + " │\n", + " physical arm\n", + "```\n", + "\n", + "For a UR5e / ABB / FANUC-style experiment:\n", + "\n", + "1. choose joint-space or Cartesian EEF control;\n", + "2. record cameras, state, demonstrated target, gripper, task text and optionally force/torque;\n", + "3. prefer one MCAP per episode;\n", + "4. QA timestamps, camera sync, holes, physical limits, stale frames, task success and idle periods;\n", + "5. convert to GR00T LeRobot;\n", + "6. write a custom `modality.json` and modality config;\n", + "7. generate statistics;\n", + "8. start with a narrow workspace/distribution;\n", + "9. fine-tune;\n", + "10. do open-loop checks, then supervised closed-loop evaluation.\n", + "\n", + "### Safety boundary\n", + "\n", + "Do not connect a freshly fine-tuned policy directly to an industrial robot at unrestricted speed. Keep normal safety envelopes, conservative rate/position limits, an emergency stop and supervised low-speed evaluation." + ] + }, + { + "cell_type": "markdown", + "id": "dad31305", + "metadata": {}, + "source": [ + "# Part F — suggested learning sequence\n", + "\n", + "1. **Understand the data:** run sections 5–10.\n", + "2. **Prove fine-tuning works:** run the 50-step SO100 smoke test.\n", + "3. **Richer manipulation:** inspect LIBERO Goal, then fine-tune briefly or use NVIDIA's checkpoint.\n", + "4. **Data engineering:** run the MCAP exercise and deliberately inject larger gaps, stale frames, clock offsets and gripper discontinuities.\n", + "5. **Bring your own ROS 2 episode:** point `MCAP_PATH` at it, inventory topics, inspect schemas, then replace `decode_demo_mcap()` with your real mappings." + ] + }, + { + "cell_type": "markdown", + "id": "f4587d14", + "metadata": {}, + "source": [ + "# Runtime matrix\n", + "\n", + "| Section | macOS | Linux CPU | Linux/Docker + NVIDIA GPU |\n", + "|---|---:|---:|---:|\n", + "| Clone repo / inspect SO100 | ✅ | ✅ | ✅ |\n", + "| LeRobot parquet/video QA | ✅ | ✅ | ✅ |\n", + "| MuJoCo fundamentals | ✅ | ✅ | ✅ |\n", + "| MCAP cleaning | ✅ | ✅ | ✅ |\n", + "| LIBERO dataset download/inspection | ✅ | ✅ | ✅ |\n", + "| GR00T N1.7 fine-tuning | — | — | ✅ |\n", + "| GR00T open-loop model inference | — | — | ✅ |\n", + "| LIBERO + GR00T closed-loop MuJoCo | — | — | ✅ |\n", + "| Before/after rollout videos | — | — | ✅ |\n", + "\n", + "`—` means the notebook intentionally skips that path rather than treating the platform as broken." + ] + }, + { + "cell_type": "markdown", + "id": "34d03456", + "metadata": {}, + "source": [ + "# Appendix — public downloads\n", + "\n", + "| Purpose | Resource |\n", + "|---|---|\n", + "| GR00T source + bundled demos | `https://github.com/NVIDIA/Isaac-GR00T` |\n", + "| GR00T N1.7 base | `nvidia/GR00T-N1.7-3B` |\n", + "| Required VLM backbone | `nvidia/Cosmos-Reason2-2B` |\n", + "| LIBERO Goal data | `IPEC-COMMUNITY/libero_goal_no_noops_1.0.0_lerobot` |\n", + "| LIBERO Object data | `IPEC-COMMUNITY/libero_object_no_noops_1.0.0_lerobot` |\n", + "| LIBERO Spatial data | `IPEC-COMMUNITY/libero_spatial_no_noops_1.0.0_lerobot` |\n", + "| LIBERO-10 data | `IPEC-COMMUNITY/libero_10_no_noops_1.0.0_lerobot` |\n", + "| NVIDIA LIBERO checkpoints | `nvidia/GR00T-N1.7-LIBERO` |\n", + "| DROID sample generator | bundled `scripts/download_droid_sample.py` |\n", + "| Full DROID source | `lerobot/droid_1.0.1` |\n", + "| MCAP Python | `mcap` |\n", + "| ROS 2 MCAP decoding without ROS | `mcap-ros2-support` |\n", + "\n", + "The full DROID dataset is very large; NVIDIA's example script can fetch a small sample for pipeline testing instead.\n", + "\n", + "| MuJoCo Python runtime/docs | `https://mujoco.readthedocs.io/en/stable/python.html` |\n", + "| LIBERO benchmark | `https://github.com/Lifelong-Robot-Learning/LIBERO` |\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/test_gr00t_notebook.py b/tests/test_gr00t_notebook.py new file mode 100644 index 0000000..b57af64 --- /dev/null +++ b/tests/test_gr00t_notebook.py @@ -0,0 +1,182 @@ +"""Exercise notebook comparison controls without installing or running GR00T.""" + +import ast +import json +import re +import shlex +from pathlib import Path + +import pytest + +NOTEBOOK = ( + Path(__file__).resolve().parents[1] + / "notebooks/gr00t/gr00t_n17_finetuning_mujoco_mcap_lab_v2.ipynb" +) + + +def code_cells(): + return [ + "".join(cell["source"]) + for cell in json.loads(NOTEBOOK.read_text())["cells"] + if cell["cell_type"] == "code" + ] + + +def cell_containing(marker): + return next(cell for cell in code_cells() if marker in cell) + + +def notebook_functions(): + names = { + "list_checkpoints", + "checkpoint_save_interval", + "validate_checkpoint_pair", + "compare_success_rates", + } + definitions = [ + statement + for cell in code_cells() + for statement in ast.parse(cell).body + if isinstance(statement, ast.FunctionDef) and statement.name in names + ] + namespace = {"Path": Path, "re": re, "shlex": shlex} + exec(compile(ast.Module(body=definitions, type_ignores=[]), str(NOTEBOOK), "exec"), namespace) + return namespace + + +def comparison_context(tmp_path, suite="goal", requested=False): + namespace = notebook_functions() + output = tmp_path / "checkpoints" + output.mkdir(exist_ok=True) + namespace.update( + LIBERO_SUITE=suite, + REPO_DIR=tmp_path / "upstream with spaces", + WORK_ROOT=tmp_path / "workspace with spaces", + COMPLEX_OUT=output, + RUN_MUJOCO_BEFORE_AFTER=requested, + uv="uv", + ) + exec(cell_containing("LIBERO_REPOS ="), namespace) + return namespace + + +def test_notebook_has_valid_code_and_no_saved_outputs(): + notebook = json.loads(NOTEBOOK.read_text()) + assert notebook["nbformat"] == 4 + assert len(notebook["cells"]) == 87 + for cell in notebook["cells"]: + if cell["cell_type"] == "code": + assert cell["execution_count"] is None + assert cell["outputs"] == [] + ast.parse("".join(cell["source"])) + + +@pytest.mark.parametrize( + ("suite", "task"), + [ + ("goal", "put_the_bowl_on_the_plate"), + ("object", "pick_up_the_alphabet_soup_and_place_it_in_the_basket"), + ("spatial", "pick_up_the_black_bowl_from_table_center_and_place_it_on_the_plate"), + ("10", "KITCHEN_SCENE3_turn_on_the_stove_and_put_the_moka_pot_on_it"), + ], +) +def test_selected_suite_reaches_both_rollout_templates(tmp_path, capsys, suite, task): + namespace = comparison_context(tmp_path, suite, requested=True) + for step in (50, 100): + (namespace["COMPLEX_OUT"] / f"checkpoint-{step}").mkdir() + namespace.update(RUN_LIBERO_SIM_SETUP=False, FULL_GROOT_SUPPORTED=False) + exec( + cell_containing('setup_cmd = ["bash", "gr00t/eval/sim/LIBERO/setup_libero.sh"]'), namespace + ) + setup_text = capsys.readouterr().out + assert f"--env-name libero_sim/{task}" in setup_text + exec(cell_containing("COMPARISON_PAIR = None"), namespace) + output = capsys.readouterr().out + commands = [ + shlex.split(line) + for line in output.splitlines() + if "run_gr00t_server.py" in line or "rollout_policy.py" in line + ] + assert len(commands) == 4 + servers, clients = commands[::2], commands[1::2] + assert [Path(cmd[cmd.index("--model-path") + 1]).name for cmd in servers] == [ + "checkpoint-50", + "checkpoint-100", + ] + for cmd in clients: + assert cmd[cmd.index("--env-name") + 1] == f"libero_sim/{task}" + assert cmd[cmd.index("--seed") + 1] == "42" + assert cmd[cmd.index("--n-episodes") + 1] == "10" + assert ( + clients[0][clients[0].index("--video-dir") + 1] + != clients[1][clients[1].index("--video-dir") + 1] + ) + + +def test_unknown_suite_is_rejected(tmp_path): + with pytest.raises(ValueError, match="Unknown LIBERO suite"): + comparison_context(tmp_path, "typo") + + +@pytest.mark.parametrize("steps", [2, 3, 100, 101]) +def test_training_command_saves_two_checkpoints(tmp_path, steps): + namespace = comparison_context(tmp_path) + namespace.update( + COMPLEX_MAX_STEPS=steps, + GLOBAL_BATCH_SIZE=32, + RUN_COMPLEX_FINETUNE=False, + FULL_GROOT_SUPPORTED=False, + ) + exec(cell_containing("complex_cmd = ["), namespace) + cmd = namespace["complex_cmd"] + interval = int(cmd[cmd.index("--save-steps") + 1]) + assert len(range(interval, steps + 1, interval)) >= 2 + assert cmd[cmd.index("--save-total-limit") + 1] == "2" + if steps == 100: + assert interval == 50 + + +@pytest.mark.parametrize("steps", [0, 1, -1, True, 2.5]) +def test_invalid_comparison_training_budget_is_rejected(steps): + with pytest.raises(ValueError, match="at least 2"): + notebook_functions()["checkpoint_save_interval"](steps) + + +@pytest.mark.parametrize("count", [0, 1]) +def test_no_checkpoints_keeps_data_path_usable_but_blocks_requested_comparison(tmp_path, count): + namespace = comparison_context(tmp_path) + if count: + (namespace["COMPLEX_OUT"] / "checkpoint-100").mkdir() + exec(cell_containing("COMPARISON_PAIR = None"), namespace) + assert namespace["COMPARISON_PAIR"] is None + assert namespace["LATE_CKPT"] is None + namespace["RUN_MUJOCO_BEFORE_AFTER"] = True + with pytest.raises(ValueError, match="Two checkpoints are required"): + exec(cell_containing("COMPARISON_PAIR = None"), namespace) + + +def test_duplicate_checkpoint_alias_blocks_commands_and_plot(tmp_path): + namespace = comparison_context(tmp_path, requested=True) + early = namespace["COMPLEX_OUT"] / "checkpoint-50" + early.mkdir() + late = namespace["COMPLEX_OUT"] / "checkpoint-100" + late.symlink_to(early, target_is_directory=True) + with pytest.raises(ValueError, match="different checkpoints"): + exec(cell_containing("COMPARISON_PAIR = None"), namespace) + with pytest.raises(ValueError, match="different checkpoints"): + namespace["compare_success_rates"]("success rate: 0.2", "success rate: 0.9") + + +def test_missing_checkpoint_directory_is_rejected(tmp_path): + with pytest.raises(ValueError, match="must exist"): + notebook_functions()["validate_checkpoint_pair"](tmp_path, tmp_path / "missing") + + +def test_checkpoint_listing_ignores_unfinished_and_non_numeric_entries(tmp_path): + for name in ("checkpoint-100", "checkpoint-50", "checkpoint-pending", "old-checkpoint-1"): + (tmp_path / name).mkdir() + (tmp_path / "checkpoint-75").write_text("not a directory") + assert [p.name for p in notebook_functions()["list_checkpoints"](tmp_path)] == [ + "checkpoint-50", + "checkpoint-100", + ]