diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 1505c2c..50fa7d6 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,27 +1,45 @@ -# Contributing +# Contributing to GR00T-H -Thanks for considering contributing to NVIDIA Isaac-GR00T! Please read this document to learn the various ways you can contribute to this project and how to go about doing it. +Thanks for considering contributing to GR00T-H! Please read this document to learn the various ways you can contribute to this project and how to go about doing it. + +> **Note:** GR00T-H is a healthcare-focused fork of [Isaac-GR00T](https://github.com/NVIDIA/Isaac-GR00T). For the latest on the base GR00T model, general robotics features, and non-healthcare contributions, please refer to the upstream [Isaac-GR00T repository](https://github.com/NVIDIA/Isaac-GR00T). Contributions specific to healthcare robotics, the Open-H dataset, or surgical embodiments belong here. + +## Signing Your Work + +We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. + +Any contribution which contains commits that are not Signed-Off will not be accepted. + +To sign off on a commit you simply use the `--signoff` (or `-s`) option when committing your changes: + +```bash +$ git commit -s -m "Add cool feature." +``` + +This will append the following to your commit message: + +``` +Signed-off-by: Your Name +``` ## Bug reports and feature requests ### Did you find a bug? -First, do [a quick search](https://github.com/NVIDIA/Isaac-GR00T/issues) to see whether your issue has already been reported. -If your issue has already been reported, please comment on the existing issue. +First, determine whether the issue is specific to GR00T-H (healthcare embodiments, Open-H dataset, surgical data pipeline) or affects the base GR00T model: + +- **GR00T-H specific issues**: Search [GR00T-H issues](https://github.com/NVIDIA-Medtech/GR00T-H/issues), then [open a new issue](https://github.com/NVIDIA-Medtech/GR00T-H/issues/new) if not already reported. +- **Base GR00T issues**: Report on the upstream [Isaac-GR00T issues](https://github.com/NVIDIA/Isaac-GR00T/issues). -Otherwise, open [a new GitHub issue](https://github.com/NVIDIA/Isaac-GR00T/issues/new). Be sure to include a clear title -and description. The description should include as much relevant information as possible. The description should -explain how to reproduce the erroneous behavior as well as the behavior you expect to see. Ideally you would include a -code sample or an executable test case demonstrating the expected behavior. +Be sure to include a clear title and description with as much relevant information as possible, including how to reproduce the issue and the behavior you expect to see. ### Do you have a suggestion for an enhancement or new feature? We use GitHub issues to track feature requests. Before you create a feature request: -* Make sure you have a clear idea of the enhancement you would like. If you have a vague idea, consider discussing -it first on a GitHub issue. +* Make sure you have a clear idea of the enhancement you would like. If you have a vague idea, consider discussing it first on a GitHub issue. * Check the documentation to make sure your feature does not already exist. -* Do [a quick search](https://github.com/NVIDIA/Isaac-GR00T/issues) to see whether your feature has already been suggested. +* Do [a quick search](https://github.com/NVIDIA-Medtech/GR00T-H/issues) to see whether your feature has already been suggested. When creating your request, please: @@ -41,38 +59,38 @@ When you're ready to contribute code to address an open issue, please follow the Then clone your fork locally with - git clone https://github.com/USERNAME/Isaac-GR00T.git + git clone https://github.com/USERNAME/GR00T-H.git - or + or - git clone git@github.com:USERNAME/Isaac-GR00T.git + git clone git@github.com:USERNAME/GR00T-H.git - At this point the local clone of your fork only knows that it came from *your* repo, github.com/USERNAME/Isaac-GR00T.git, but doesn't know anything the *main* repo, [https://github.com/NVIDIA/Isaac-GR00T.git](https://github.com/NVIDIA/Isaac-GR00T). You can see this by running + At this point the local clone of your fork only knows that it came from *your* repo, github.com/USERNAME/GR00T-H.git, but doesn't know anything the *main* repo, [https://github.com/NVIDIA-Medtech/GR00T-H.git](https://github.com/NVIDIA-Medtech/GR00T-H). You can see this by running git remote -v which will output something like this: - origin https://github.com/USERNAME/Isaac-GR00T.git (fetch) - origin https://github.com/USERNAME/Isaac-GR00T.git (push) + origin https://github.com/USERNAME/GR00T-H.git (fetch) + origin https://github.com/USERNAME/GR00T-H.git (push) - This means that your local clone can only track changes from your fork, but not from the main repo, and so you won't be able to keep your fork up-to-date with the main repo over time. Therefore you'll need to add another "remote" to your clone that points to [https://github.com/NVIDIA/Isaac-GR00T.git](https://github.com/NVIDIA/Isaac-GR00T). To do this, run the following: + This means that your local clone can only track changes from your fork, but not from the main repo, and so you won't be able to keep your fork up-to-date with the main repo over time. Therefore you'll need to add another "remote" to your clone that points to [https://github.com/NVIDIA-Medtech/GR00T-H.git](https://github.com/NVIDIA-Medtech/GR00T-H). To do this, run the following: - git remote add upstream https://github.com/NVIDIA/Isaac-GR00T.git + git remote add upstream https://github.com/NVIDIA-Medtech/GR00T-H.git Now if you do `git remote -v` again, you'll see - origin https://github.com/USERNAME/Isaac-GR00T.git (fetch) - origin https://github.com/USERNAME/Isaac-GR00T.git (push) - upstream https://github.com/NVIDIA/Isaac-GR00T.git (fetch) - upstream https://github.com/NVIDIA/Isaac-GR00T.git (push) + origin https://github.com/USERNAME/GR00T-H.git (fetch) + origin https://github.com/USERNAME/GR00T-H.git (push) + upstream https://github.com/NVIDIA-Medtech/GR00T-H.git (fetch) + upstream https://github.com/NVIDIA-Medtech/GR00T-H.git (push) - Finally, you'll need to create a Python 3 virtual environment suitable for working on this project. + Finally, you'll need to create a Python 3 virtual environment suitable for working on this project. ```bash uv pip install -e .[dev] ``` - The "editable mode" comes from the `-e` argument to `pip`, and essential just creates a symbolic link from the site-packages directory of your virtual environment to the source code in your local clone. That way any changes you make will be immediately reflected in your virtual environment. + The "editable mode" comes from the `-e` argument to `pip`, and essentially just creates a symbolic link from the site-packages directory of your virtual environment to the source code in your local clone. That way any changes you make will be immediately reflected in your virtual environment. @@ -80,7 +98,7 @@ When you're ready to contribute code to address an open issue, please follow the
Expand details πŸ‘‡
- Once you've added an "upstream" remote pointing to [https://github.com/NVIDIA/Isaac-GR00T.git](https://github.com/NVIDIA/Isaac-GR00T), keeping your fork up-to-date is easy: + Once you've added an "upstream" remote pointing to [https://github.com/NVIDIA-Medtech/GR00T-H.git](https://github.com/NVIDIA-Medtech/GR00T-H), keeping your fork up-to-date is easy: git checkout main # if not already on main git pull --rebase upstream main @@ -106,32 +124,23 @@ When you're ready to contribute code to address an open issue, please follow the
Expand details πŸ‘‡
- Our continuous integration (CI) testing runs [a number of checks](https://github.com/NVIDIA/Isaac-GR00T/actions) for each pull request on [GitHub Actions](https://github.com/features/actions). You can run most of these tests locally, which is something you should do *before* opening a PR to help speed up the review process and make it easier for us. - - First, you should run [`ruff`](https://docs.astral.sh/ruff/) to make sure you code is formatted consistently. - Many IDEs support code formatters as plugins, so you may be able to setup isort and black to run automatically everytime you save. - For example, [`black.vim`](https://github.com/psf/black/tree/master/plugin) will give you this functionality in Vim. But both `isort` and `black` are also easy to run directly from the command line. - Just run this from the root of your clone: + First, you should run [`ruff`](https://docs.astral.sh/ruff/) to make sure your code is formatted consistently: ```bash ruff format . ruff check --fix . ``` - We also strive to maintain high test coverage, so most contributions should include additions to [the unit tests](https://github.com/NVIDIA/Isaac-GR00T/tree/main/tests). These tests are run with [`pytest`](https://docs.pytest.org/en/latest/), which you can use to locally run any test modules that you've added or changed. - - For example, if you've fixed a bug in `isaac-gr00t/a/b.py`, you can run the tests specific to that module with - - pytest -v tests/a/b_test.py + We also strive to maintain high test coverage, so most contributions should include additions to [the unit tests](https://github.com/NVIDIA-Medtech/GR00T-H/tree/main/tests). These tests are run with [`pytest`](https://docs.pytest.org/en/latest/), which you can use to locally run any test modules that you've added or changed. - After all of the above checks have passed, you can now open [a new GitHub pull request](https://github.com/NVIDIA/Isaac-GR00T/pulls). + After all of the above checks have passed, you can now open [a new GitHub pull request](https://github.com/NVIDIA-Medtech/GR00T-H/pulls). Make sure you have a clear description of the problem and the solution, and include a link to relevant issues. We look forward to reviewing your PR!
-### Developer Certificate of Origin +## Full text of the DCO ``` Developer Certificate of Origin @@ -147,25 +156,11 @@ Developer's Certificate of Origin 1.1 By making a contribution to this project, I certify that: -(a) The contribution was created in whole or in part by me and I - have the right to submit it under the open source license - indicated in the file; or - -(b) The contribution is based upon previous work that, to the best - of my knowledge, is covered under an appropriate open source - license and I have the right under that license to submit that - work with modifications, whether created in whole or in part - by me, under the same open source license (unless I am - permitted to submit under a different license), as indicated - in the file; or - -(c) The contribution was provided directly to me by some other - person who certified (a), (b) or (c) and I have not modified - it. - -(d) I understand and agree that this project and the contribution - are public and that a record of the contribution (including all - personal information I submit with it, including my sign-off) is - maintained indefinitely and may be redistributed consistent with - this project or the open source license(s) involved. -``` \ No newline at end of file +(a) The contribution was created in whole or in part by me and I have the right to submit it under the open source license indicated in the file; or + +(b) The contribution is based upon previous work that, to the best of my knowledge, is covered under an appropriate open source license and I have the right under that license to submit that work with modifications, whether created in whole or in part by me, under the same open source license (unless I am permitted to submit under a different license), as indicated in the file; or + +(c) The contribution was provided directly to me by some other person who certified (a), (b) or (c) and I have not modified it. + +(d) I understand and agree that this project and the contribution are public and that a record of the contribution (including all personal information I submit with it, including my sign-off) is maintained indefinitely and may be redistributed consistent with this project or the open source license(s) involved. +``` diff --git a/README.md b/README.md index 6f1c4c2..19d5ba5 100644 --- a/README.md +++ b/README.md @@ -1,168 +1,61 @@ -
+# GR00T-H - NVIDIA Isaac GR00T N1.6 Header +[![License](https://img.shields.io/badge/License-NVIDIA--OneWay--Noncommercial--License-blue.svg)](LICENSE) +[![HuggingFace](https://img.shields.io/badge/%F0%9F%A4%97-Hugging%20Face-yellow)](https://huggingface.co/nvidia/GR00T-H) +[![Open-H Dataset](https://img.shields.io/badge/Dataset-Open--H-orange)](https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-Open-H-Embodiment) +[![Python](https://img.shields.io/badge/Python-3.10+-blue.svg)](https://python.org) - - -

- Website | - Model | - Dataset | - Paper | - Research Blog -

-
+A healthcare robotics variant of [GR00T N1.6](https://github.com/NVIDIA/Isaac-GR00T), post-trained on the [Open-H dataset](open_h/README.md) for multi-embodiment surgical and healthcare robot autonomy across 16 robot platforms and 10+ institutions. -## NVIDIA Isaac GR00T +

+ GR00T-H Header +

-
- GR00T Demo -
+## Overview -> We just released GR00T N1.6, an updated version of GR00T N1 with improved performance and new features. Check out the [release blog post](https://research.nvidia.com/labs/gear/gr00t-n1_6/) for more details. +GR00T-H post-trains the [GR00T N1.6](https://github.com/NVIDIA/Isaac-GR00T) vision-language-action (VLA) foundation model on surgical robot data from multiple institutions and robot platforms simultaneously. Each institution records data differently β€” different robots, coordinate conventions, frame rates, camera setups, and state/action representations. GR00T-H solves this by defining per-embodiment modality configs that convert each dataset into a common representation (REL_XYZ_ROT6D for EEF poses) while sharing the core VLA backbone. -> To use the older version, N1.5, please checkout the [n1.5-release](https://github.com/NVIDIA/Isaac-GR00T/tree/n1.5-release) branch. +The primary additions over upstream Isaac-GR00T live in [`open_h/`](open_h/README.md): +- Per-embodiment modality configs converting 16 healthcare robot datasets to a common action representation +- Multi-embodiment training config and dataset preparation tooling +- Extensions to the data pipeline (clutch-aware filtering, motion scaling, step filtering) -NVIDIA Isaac GR00T N1.6 is an open vision-language-action (VLA) model for generalized humanoid robot skills. This cross-embodiment model takes multimodal input, including language and images, to perform manipulation tasks in diverse environments. +For general robotics use cases, the upstream [Isaac-GR00T](https://github.com/NVIDIA/Isaac-GR00T) project is a better starting point. -GR00T N1.6 is trained on a diverse mixture of robot data including bimanual, semi-humanoid and an expansive humanoid dataset. It is adaptable through post-training for specific embodiments, tasks and environments. +## News -The neural network architecture of GR00T N1.6 is a combination of vision-language foundation model and diffusion transformer head that denoises continuous actions. Here is a schematic diagram of the architecture: +- **[March 2026]** β€” Released GR00T-H with pre-trained checkpoint and the [Open-H dataset](https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-Open-H-Embodiment) -
-model-architecture -
+## Model Variants -Here is the general procedure to use GR00T N1.6: +| Model | Base Model | Params | Capability | HuggingFace | License | +|-------|-----------|--------|------------|-------------|---------| +| GR00T-H | [GR00T-N1.6-3B](https://huggingface.co/nvidia/GR00T-N1.6-3B) | 3B | Multi-embodiment healthcare robotics (16 surgical platforms) | [Weights](https://huggingface.co/nvidia/GR00T-H) | [NVIDIA-OneWay-Noncommercial-License](https://developer.download.nvidia.com/licenses/NVIDIA-OneWay-Noncommercial-License-22Mar2022.pdf) | -1. We assume the user has already collected a dataset of robot demonstrations in the form of (video, state, action) triplets for a specific task. -2. The user will first convert the demonstration data into the LeRobot compatible data schema (more info in [`getting_started/data_preparation.md`](getting_started/data_preparation.md)), which is compatible with the upstream [Huggingface LeRobot Dataset V2](https://github.com/huggingface/lerobot). -3. Our repo provides convenient scripts to validate zero-shot performance of the pretrained model (see [Policy API Guide](getting_started/policy.md) and [RoboCasa Zero-Shot](examples/robocasa-gr1-tabletop-tasks/README.md)). -4. Our repo provides examples of different configurations for training with different robot embodiments (see [`examples/`](examples/) and [Fine-tuning Guide](getting_started/finetune_new_embodiment.md)). -5. Our repo provides convenient scripts for finetuning the pre-trained GR00T N1.6 model on user's data, and running inference, see [`examples`](examples). -6. Our repo provides convenient scripts to run academic simulation benchmarks with finetuned checkpoints (see [LIBERO](examples/LIBERO/README.md), [SimplerEnv](examples/SimplerEnv/README.md), [RoboCasa](examples/robocasa/README.md)). -7. The user will need to connect the `Gr00tPolicy` to the robot controller to execute actions on their target hardware. - -## What's New in GR00T N1.6 - -GR00T N1.6 represents a significant upgrade over GR00T N1.5, with improvements in both model architecture and data leading to better performance in many aspects. - -### Model and Data Improvements - -Architectural changes: -- Base VLM: We use an internal NVIDIA Cosmos-Reason-2B VLM variant. The VLM supports flexible resolution and can encode images in their native aspect ratio without padding. The VLM is trained both general vision-language tasks and embodied reasoning tasks like next action prediction. -- Uses 2x larger DiT (32 layers vs 16 layers in N1.5). -- Removes N1.5's post-VLM 4-layer transformer adapter. Instead, unfreezes top 4 layers of the VLM during pretraining. -- Predicts state-relative action chunks for most embodiments, rather than absolute joint angles or EEF positions. - -Beyond the N1.5 data mixture, the N1.6 pretraining data additionally includes several thousand hours of teleoperated data from: -- Bimanual YAM arms -- AGIBot Genie1 -- Simulated Galaxea R1 Pro on the BEHAVIOR suite -- Whole-Body Locomanipulation with Unitree G1 - -Other code-level improvements: -- Faster dataloader with sharded dataloader support. -- RTC and Async Policy Wrapper for inference (soon to release) -- Simplified data processing pipeline with `processing_gr00t_n1d6.py` -- Flexible Training configuration - -## Target Audience - -GR00T N1.6 is intended for researchers and professionals in robotics. This repository provides tools to: - -- Leverage a pre-trained foundation model for robot control -- Fine-tune on small, custom datasets -- Adapt the model to specific robotics tasks with minimal data -- Deploy the model for inference - -The focus is on enabling customization of robot behaviors through finetuning. - -## Installation Guide - -### Clone the Repository - -GR00T relies on submodules for certain dependencies. Include them when cloning: - -```sh -git clone --recurse-submodules https://github.com/NVIDIA/Isaac-GR00T -cd Isaac-GR00T -``` - -If you've already cloned without submodules, initialize them separately: - -```sh -git submodule update --init --recursive -``` - -### Set Up the Environment - -GR00T uses [uv](https://github.com/astral-sh/uv) for fast, reproducible dependency management. - -> **Requirement:** uv **v0.8.4+** is needed to parse `[tool.uv.extra-build-dependencies]` in `pyproject.toml` (required for building `flash-attn`). For RTX-5090, this was tested with CUDA 12.8, `flash-attn==2.8.0.post2`, `pytorch-cu128`. +## Quick Start -After installing uv, create the environment and install GR00T: +### Installation -```sh +```bash +git clone --recurse-submodules git@github.com:NVIDIA-Medtech/GR00T-H.git +cd GR00T-H uv sync --python 3.10 uv pip install -e . ``` -> Note: CUDA 12.4 is recommended and officially tested. However, CUDA 11.8 has also been verified to work. -> In such cases, make sure to install a compatible version of `flash-attn` manually (e.g., `flash-attn==2.8.2` was confirmed working with CUDA 11.8). - -For a containerized setup that avoids system-level dependency conflicts, see our [Docker Setup Guide](docker/README.md). - -For training and inference hardware recommendations (RTX PRO Servers, DGX, Jetson AGX Thor), see the [Hardware Recommendation Guide](getting_started/hardware_recommendation.md). - -## Model Checkpoints - -### Base Models -We provide pre-trained base VLA model checkpoints. These checkpoints have been pre-trained on 10k+ hours of robot data and can be used for finetuning on downstream tasks. +If `flash-attn` was not built during `uv sync`, install it manually: -| Model | Use Case | Description | Checkpoint Path | Branch | -| ----- | -------- | ----------- | --------------- | ------ | -| GR00T N1.5 | Finetuning | Base [GR00T N1.5 model](https://research.nvidia.com/labs/gear/gr00t-n1_5/) (3B parameters) | [nvidia/GR00T-N1.5-3B](https://huggingface.co/nvidia/GR00T-N1.5-3B) | [n1.5-release](https://github.com/NVIDIA/Isaac-GR00T/tree/n1.5-release) | -| GR00T N1.6 | Finetuning | Base [GR00T N1.6 model](https://research.nvidia.com/labs/gear/gr00t-n1_6/) (3B parameters) | [nvidia/GR00T-N1.6-3B](https://huggingface.co/nvidia/GR00T-N1.6-3B) | [main](https://github.com/NVIDIA/Isaac-GR00T) | - -### Finetuned Models -We also provide finetuned checkpoints for various robot platforms and benchmarks. These models are finetuned from the base models above and can be used directly for evaluation or as starting points for further finetuning. - -| Model | Base Model | Description | Checkpoint Path | Example | -| ----- | ---------- | ----------- | --------------- | ------- | -| GR00T-N1.6-bridge | [nvidia/GR00T-N1.6-3B](https://huggingface.co/nvidia/GR00T-N1.6-3B) | Fine-tuned on [Bridge dataset](https://rail-berkeley.github.io/bridgedata/) for WidowX robot on manipulation tasks | [nvidia/GR00T-N1.6-bridge](https://huggingface.co/nvidia/GR00T-N1.6-bridge) | [SimplerEnv](examples/SimplerEnv/README.md) | -| GR00T-N1.6-fractal | [nvidia/GR00T-N1.6-3B](https://huggingface.co/nvidia/GR00T-N1.6-3B) | Fine-tuned on [Fractal dataset](https://www.tensorflow.org/datasets/catalog/fractal20220817_data) for Google robot on manipulation tasks | [nvidia/GR00T-N1.6-fractal](https://huggingface.co/nvidia/GR00T-N1.6-fractal) | [SimplerEnv](examples/SimplerEnv/README.md) | -| GR00T-N1.6-BEHAVIOR1k | [nvidia/GR00T-N1.6-3B](https://huggingface.co/nvidia/GR00T-N1.6-3B) | Fine-tuned on [BEHAVIOR-1K](https://behavior.stanford.edu/) for Galaxea R1 Pro robot on loco-manipulation tasks | [nvidia/GR00T-N1.6-BEHAVIOR1k](https://huggingface.co/nvidia/GR00T-N1.6-BEHAVIOR1k) | [BEHAVIOR](examples/BEHAVIOR/README.md) | -| GR00T-N1.6-G1-PnPAppleToPlate | [nvidia/GR00T-N1.6-3B](https://huggingface.co/nvidia/GR00T-N1.6-3B) | Fine-tuned for Unitree G1 loco-manipulation pick-and-place tasks | [nvidia/GR00T-N1.6-G1-PnPAppleToPlate](https://huggingface.co/nvidia/GR00T-N1.6-G1-PnPAppleToPlate) | [G1 LocoManipulation](examples/GR00T-WholeBodyControl/README.md) | -| GR00T-N1.6-DROID | [nvidia/GR00T-N1.6-DROID](https://huggingface.co/nvidia/GR00T-N1.6-DROID) | Fine-tuned for DROID robot on manipulation tasks | [nvidia/GR00T-N1.6-DROID](https://huggingface.co/nvidia/GR00T-N1.6-DROID) | [DROID](examples/DROID/README.md) | - - - -## Quick Start - -We can quickly start by downloading a pre-trained checkpoint and starting the policy server for any pretrained embodiement, e.g. GR1 embodiment. ```bash -# On GPU server: Start the policy server -uv run --extra=gpu python gr00t/eval/run_gr00t_server.py --embodiment-tag GR1 --model-path nvidia/GR00T-N1.6-3B +uv pip install flash-attn==2.7.4.post1 --no-build-isolation ``` -Then, refer to the [robocasa-gr1-tabletop-tasks](examples/robocasa-gr1-tabletop-tasks/README.md) for more details on how to rollout the policy with `GR1` embodiment. +For containerized setup, see the [Docker Setup Guide](docker/README.md). -## Getting started with this repo - -We provide accessible Jupyter notebooks and detailed documentation in the [`./getting_started`](getting_started) folder. - -## 1. Data Preparation - -Please refer to the [data preparation guide](getting_started/data_preparation.md) for more details. - -## 2. Inference - -After data is prepared, the GR00T model can be used to generate output actions with the below simple inference script: +### Inference ```bash uv run python scripts/deployment/standalone_inference_script.py \ - --model-path nvidia/GR00T-N1.6-3B \ + --model-path nvidia/GR00T-H \ --dataset-path demo_data/gr1.PickNPlace \ --embodiment-tag GR1 \ --traj-ids 0 1 2 \ @@ -170,231 +63,82 @@ uv run python scripts/deployment/standalone_inference_script.py \ --action-horizon 8 ``` -GR00T-N1.6-3B inference timing (4 denoising steps, single view): - -| Device | Mode | Data Processing | Backbone | Action Head | E2E | Frequency | -|--------|------|-----------------|----------|-------------|-----|-----------| -| RTX 5090 | torch.compile | 2 ms | 18 ms | 16 ms | 37 ms | 27.3 Hz | -| H100 | torch.compile | 4 ms | 23 ms | 11 ms | 38 ms | 26.3 Hz | -| RTX 4090 | torch.compile | 2 ms | 25 ms | 17 ms | 44 ms | 22.8 Hz | -| Thor | torch.compile | 5 ms | 39 ms | 61 ms | 105 ms | 9.5 Hz | - -For more details, please check our full [inference guide](scripts/deployment/README.md) for more details including faster inference with `TensorRT` - -## 3. Finetuning - -### Fine-tune on Pre-registered Post-train Embodiment Tags - -GR00T provides several pre-registered embodiment tags with ready-to-use configurations: - -- `LIBERO_PANDA` -- `OXE_GOOGLE` -- `OXE_WIDOWX` -- `UNITREE_G1` -- `BEHAVIOR_R1_PRO` - -**Example:** To finetune Libero-Spatial on GR00T N1.6, follow the instructions in the [Libero finetuning guide](examples/LIBERO/README.md#finetune-libero-spatial-dataset). We also provide simulation environment setup for evaluation linked with post-train checkpoints and benchmark numbers. - -### Fine-tune on Custom Embodiments ("NEW_EMBODIMENT") - -To finetune GR00T on your own robot data and configuration, follow the detailed tutorial available at [`getting_started/finetune_new_embodiment.md`](getting_started/finetune_new_embodiment.md). +For full inference options including TensorRT, see the [inference guide](scripts/deployment/README.md). -#### Prerequisites +### Finetuning on Open-H Embodiments -Ensure your input data follows the **GR00T-flavored LeRobot v2 format**, and specify your modality configuration at `modality_config_path`. - -#### Run Fine-tuning Script ```bash -# Set number of GPUs -export NUM_GPUS=1 - -CUDA_VISIBLE_DEVICES=0 uv run python \ +uv run torchrun --nproc_per_node=8 --master_port=29500 \ gr00t/experiment/launch_finetune.py \ - --base-model-path nvidia/GR00T-N1.6-3B \ - --dataset-path \ - --embodiment-tag NEW_EMBODIMENT \ - --modality-config-path \ - --num-gpus $NUM_GPUS \ - --output-dir \ - --save-total-limit 5 \ - --save-steps 2000 \ - --max-steps 2000 \ - --use-wandb \ + --base-model-path nvidia/GR00T-H \ + --dataset-path /path/to/dataset \ + --embodiment-tag \ + --num-gpus 8 \ --global-batch-size 32 \ - --color-jitter-params brightness 0.3 contrast 0.4 saturation 0.5 hue 0.08 \ - --dataloader-num-workers 4 -``` - -> For more extensive finetuning configuration, use `gr00t/experiment/launch_train.py` instead to launch the training process. - -### Recommended Fine-tuning Configuration - -For optimal results, maximize your batch size based on available hardware and train for a few thousand steps. - -#### Hardware Performance Considerations - -**Fine-tuning Performance** -- We recommend using 1 H100 node or L40 node for optimal finetuning performance -- Other hardware configurations (e.g., A6000) will also work but may require longer training time -- Optimal batch size depends on your hardware and which model components are being tuned - -#### Training Variance - -Users may observe some variance in post-training results across runs, even when using the same configuration, seed, and dropout settings. In our experiments, we have observed performance differences as large as 5-6% between runs. This variance may be attributed to non-deterministic operations in image augmentations or other stochastic components. When comparing results to reported benchmarks, please keep this inherent variance in mind. - -## 4. Evaluation - -We recommend a two-stage evaluation approach: open-loop evaluation followed by simulation evaluation to comprehensively assess model quality. - -### 4.1 Open-Loop Evaluation - -Open-loop evaluation provides an offline assessment by comparing the model's predicted actions against ground truth data from your dataset. - -#### Running the Evaluation - -Execute the evaluation script with your newly trained model: -```bash -uv run python gr00t/eval/open_loop_eval.py \ - --dataset-path \ - --embodiment-tag NEW_EMBODIMENT \ - --model-path \ - --traj-ids 0 \ - --action-horizon 16 # ensure this is within the delta_indices of action's modality config. -``` - -#### Interpreting Results - -The evaluation generates a visualization saved at `/tmp/open_loop_eval/traj_{traj_id}.jpeg`, which includes: -- Ground truth actions vs. predicted actions -- Unnormalized mean squared error (MSE) metrics - -These plots provide a quick indicator of the policy's accuracy on the training dataset distribution. - -### 4.2 Closed-Loop Evaluation - -After validating performance through open-loop evaluation, test your model in closed-loop environments. - -#### Understanding the Policy API - -After training your model, you'll use the `Gr00tPolicy` class to load and run inference. The policy expects observations in a specific format (nested dictionaries with video, state, and language modalities) and returns actions ready for execution. - -**Quick Start with Server-Client Architecture:** - -```bash -# On GPU server: Start the policy server -uv run --extra=gpu python gr00t/eval/run_gr00t_server.py \ - --embodiment-tag NEW_EMBODIMENT \ - --model-path \ - --device cuda:0 \ - --host 0.0.0.0 \ - --port 5555 -``` - -```python -from gr00t.policy.server_client import PolicyClient - -policy = PolicyClient(host="localhost", port=5555) # Connect to the policy server -env = YourEnvironment() # Create an environment -obs, info = env.reset() # Reset the environment -if not policy.ping(): # Verify connection - raise RuntimeError("Cannot connect to policy server!") -action, info = policy.get_action(obs) # Run inference -obs, reward, done, truncated, info = env.step(action) # Execute the action -``` - -**Debugging with ReplayPolicy:** - -When developing a new environment integration or debugging your inference loop, you can use `ReplayPolicy` to replay recorded actions from an existing dataset. This helps verify that your environment setup, observation formatting, and action execution work correctlyβ€”without needing a trained model. - -```bash -# Start server with ReplayPolicy (replays actions from dataset) -uv run --extra=gpu python gr00t/eval/run_gr00t_server.py \ - --dataset-path \ - --embodiment-tag NEW_EMBODIMENT \ - --execution-horizon 8 # should match the executed action horizon in the environment + --max-steps 20000 \ + --output-dir /path/to/output ``` -The server will replay actions from the first episode of the dataset. Use `policy.reset(options={"episode_index": N})` on the client to switch to a different episode. - -**For detailed documentation on:** -- How to adapt the policy to your own environment -- Server-client architecture for remote inference -- Observation and action formats -- Querying modality configurations -- Batched inference -- Troubleshooting common errors - -See the complete [Policy API Guide](getting_started/policy.md). - -#### Evaluation Examples - -We support evaluation on available public benchmarks and our internal benchmarks. Our evaluation framework uses a server-client architecture that communicates via RESTful API. Both the policy server and simulation environment client use the same IP (usually localhost) and port to run simulation evaluation. - -For the policy server, we reuse the project root's uv environment (same as finetuning) to run `run_gr00t_server`. For simulation environment clients, we provide individual setup scripts to configure uv environments, as they typically conflict with each other when using a single shared environment. +See [open_h/README.md](open_h/README.md) for multi-embodiment training and dataset preparation. -You can use [the verification script](scripts/eval/check_sim_eval_ready.py) to verify that all dependencies and environments for simulation evaluation are properly configured. +## Open-H Dataset -Please refer to each benchmark link below for more details. +

+ Open-H Dataset +

-#### Adding a New Sim Benchmark +The [Open-H dataset](https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-Open-H-Embodiment) comprises 16 healthcare robot embodiments across 10+ institutions, stored in [LeRobot](https://github.com/huggingface/lerobot) format. See [open_h/embodiments/README.md](open_h/embodiments/README.md) for the full embodiment comparison table. -Each sim benchmark registers its environments under a gym env_name with the format `{prefix}/{task_name}` (e.g., `libero_sim/LIVING_ROOM_SCENE2_put_soup_in_basket`). The evaluation framework uses the prefix to look up the corresponding `EmbodimentTag` via a mapping in [`gr00t/eval/sim/env_utils.py`](gr00t/eval/sim/env_utils.py). +## Documentation -> **Important:** The env_name prefix and the `EmbodimentTag` value are often different. For example, `libero_sim` maps to `EmbodimentTag.LIBERO_PANDA` (`"libero_panda"`). Do not assume they match. +| Guide | Description | +|-------|-------------| +| [Open-H Overview](open_h/README.md) | GR00T-H additions, embodiment configs, dataset preparation, training | +| [Embodiment Comparison](open_h/embodiments/README.md) | All 16 embodiments β€” dimensions, cameras, action formats | +| [Action Configuration](open_h/docs/action_configuration.md) | REL_XYZ_ROT6D, rotation formats, adding new embodiments | +| [Data Preparation](open_h/docs/data_preparation.md) | Stats pipeline, temporal statistics, troubleshooting | +| [Inference Guide](scripts/deployment/README.md) | Inference options, TensorRT, server-client architecture | +| [Policy API](getting_started/policy.md) | Observation/action formats, batched inference, environment integration | +| [Finetuning Guide](getting_started/finetune_new_embodiment.md) | Custom embodiment finetuning tutorial | +| [Hardware Recommendations](getting_started/hardware_recommendation.md) | RTX PRO Servers, DGX, Jetson AGX Thor | +| [Docker Setup](docker/README.md) | Containerized environment setup | -To add a new benchmark: +## Base Model -1. Add an entry to `ENV_PREFIX_TO_EMBODIMENT_TAG` in `gr00t/eval/sim/env_utils.py`: - ```python - ENV_PREFIX_TO_EMBODIMENT_TAG = { - ... - "my_new_benchmark": EmbodimentTag.MY_ROBOT, - } - ``` -2. If the benchmark has multiple env_name prefixes (e.g., `my_benchmark_v1`, `my_benchmark_v2`), all related prefixes **must** map to the same `EmbodimentTag`. -3. Add corresponding test cases in `tests/gr00t/eval/sim/test_env_utils.py` and update the `test_all_known_prefixes_present` test. +GR00T-H builds on [GR00T N1.6](https://github.com/NVIDIA/Isaac-GR00T), a 3B-parameter vision-language-action model combining a Cosmos-Reason-2B VLM with a 32-layer diffusion transformer action head. The base model is pre-trained on 10k+ hours of robot data across bimanual, semi-humanoid, and humanoid embodiments. -**Zero-shot Evaluation** (evaluate without finetuning): -- **RoboCasa**: [Instructions](examples/robocasa/README.md) -- **RoboCasa GR1 Tabletop Tasks**: [Instructions](examples/robocasa-gr1-tabletop-tasks/README.md) +
+Base model architecture -**Finetuned Evaluation** (test after task-specific finetuning): -- **G1 LocoManipulation**: [Instructions](examples/GR00T-WholeBodyControl/README.md) -- **LIBERO**: [Instructions](examples/LIBERO/README.md) -- **SimplerEnv**: [Instructions](examples/SimplerEnv/README.md) -- **BEHAVIOR**: [Instructions](examples/BEHAVIOR/README.md) -- **PointNav**: [Instructions](examples/PointNav/README.md) -- **SO-100**: [Instructions](examples/SO100/README.md) +

+GR00T N1.6 Architecture +

+
-# Contributing +
+Base model inference timing (4 denoising steps, single view) -For more details, see [CONTRIBUTING.md](CONTRIBUTING.md) +| Device | Mode | Data Processing | Backbone | Action Head | E2E | Frequency | +|--------|------|-----------------|----------|-------------|-----|-----------| +| RTX 5090 | torch.compile | 2 ms | 18 ms | 16 ms | 37 ms | 27.3 Hz | +| H100 | torch.compile | 4 ms | 23 ms | 11 ms | 38 ms | 26.3 Hz | +| RTX 4090 | torch.compile | 2 ms | 25 ms | 17 ms | 44 ms | 22.8 Hz | +| Thor | torch.compile | 5 ms | 39 ms | 61 ms | 105 ms | 9.5 Hz | +
-## License +## License -``` -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -``` +| Component | License | +|-----------|---------| +| Source code | [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0) | +| GR00T-H model weights | [NVIDIA-OneWay-Noncommercial-License](https://developer.download.nvidia.com/licenses/NVIDIA-OneWay-Noncommercial-License-22Mar2022.pdf) | +This project will download and install additional third-party open source software projects. Review the license terms of these open source projects before use. ## Citation - -[Paper Site](https://research.nvidia.com/labs/lpr/publication/gr00tn1_2025/) + ```bibtex @inproceedings{gr00tn1_2025, archivePrefix = {arxiv}, @@ -406,3 +150,16 @@ For more details, see [CONTRIBUTING.md](CONTRIBUTING.md) booktitle = {ArXiv Preprint}, } ``` + +## Resources + +- [GR00T-H on HuggingFace](https://huggingface.co/nvidia/GR00T-H) β€” Model weights and checkpoints +- [Open-H Dataset](https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-Open-H-Embodiment) β€” Multi-embodiment healthcare robot benchmark +- [Isaac-GR00T](https://github.com/NVIDIA/Isaac-GR00T) β€” Upstream base model repository +- [GR00T N1.6 Blog Post](https://research.nvidia.com/labs/gear/gr00t-n1_6/) β€” Base model details +- [GR00T N1 Paper](https://research.nvidia.com/labs/lpr/publication/gr00tn1_2025/) β€” Research paper +- [NVIDIA MedTech Open Models](https://github.com/NVIDIA-Medtech) + +## Contributing + +See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines. diff --git a/gr00t/configs/base_config.py b/gr00t/configs/base_config.py index 7864cf5..db82c4e 100755 --- a/gr00t/configs/base_config.py +++ b/gr00t/configs/base_config.py @@ -48,6 +48,11 @@ def load(self, path: Path): def load_dict(self, data: dict): if "model" in data: self.model = self.model.__class__(**data["model"]) + # YAML loads sequences as lists, but some model fields expect tuples. + for attr in ("image_crop_size", "image_target_size"): + val = getattr(self.model, attr, None) + if isinstance(val, list): + setattr(self.model, attr, tuple(val)) if "data" in data: self.data = DataConfig(**data["data"]) # Ensure nested datasets are converted to dataclass instances diff --git a/gr00t/configs/data/data_config.py b/gr00t/configs/data/data_config.py index dbf3342..21b3a6d 100755 --- a/gr00t/configs/data/data_config.py +++ b/gr00t/configs/data/data_config.py @@ -31,6 +31,11 @@ class SingleDatasetConfig: # If not provided, falls back to dataset_paths for evaluation val_dataset_path: Optional[str] = None + # Optional split-based episode filtering (from meta/info.json splits) + # If include_splits is set, only those splits are used. Then exclude_splits is applied. + exclude_splits: List[str] | None = None + include_splits: List[str] | None = None + @dataclass class DataConfig: diff --git a/gr00t/configs/finetune_config.py b/gr00t/configs/finetune_config.py index 926c012..019227b 100644 --- a/gr00t/configs/finetune_config.py +++ b/gr00t/configs/finetune_config.py @@ -31,6 +31,18 @@ class FinetuneConfig: If None, use the pre-registered modality config in `gr00t/configs/data/embodiment_configs.py`. """ + include_splits: list[str] | None = None + """ + Optional allowlist of dataset splits (from meta/info.json) to include. + If provided, only these splits are used for training and stats. + """ + + exclude_splits: list[str] | None = None + """ + Optional denylist of dataset splits (from meta/info.json) to exclude. + Applied after include_splits (if set). Useful for skipping fail episodes. + """ + # --- Model Tuning Flags --- tune_llm: bool = False """If True, fine-tune the language model (LLM) backbone during training.""" @@ -49,6 +61,12 @@ class FinetuneConfig: Dropout probability applied to state inputs for regularization during training. """ + state_dropout_prob_per_embodiment: dict[str, float] | None = None + """ + Per-embodiment state dropout overrides. Keys are embodiment tag strings, + values are dropout probabilities in [0.0, 1.0]. + """ + # --- Data Augmentation --- random_rotation_angle: int | None = None """Maximum rotation angle (in degrees) for random rotation augmentation of input images.""" @@ -85,6 +103,23 @@ class FinetuneConfig: If None, no extra augmentations are applied. """ + image_size: tuple[int, int] | None = None + """ + Intermediate padded size as (height, width) for resize with padding. + Images are resized (preserving aspect ratio) and padded to this size. + Should be >= image_crop_size to allow for cropping augmentation. + Example: (540, 720) to pad all images to 540Γ—720 before cropping. + If None, uses shortest_image_edge with aspect-preserving crops (default albumentations). + """ + + image_crop_size: tuple[int, int] | None = None + """ + Final output size as (height, width) after cropping. Only used when image_size is set. + This is the resolution the model actually sees. + Example: (480, 640) means final images are 480Γ—640. + If None, defaults to image_size (no cropping, padded size is final size). + """ + # --- Training Configuration --- global_batch_size: int = 64 """Total effective batch size across all GPUs and accumulation steps.""" @@ -134,3 +169,25 @@ class FinetuneConfig: num_shards_per_epoch: int = int(1e5) """Number of shards to use for the dataset. reduce this number if vram is limited.""" + + # --- Statistics Calculation Flags --- + calculate_norm_stats: bool = False + """ + If True, only calculate normalization statistics and exit without training. + Uses skip_video=True for fast iteration. Statistics will be saved to + the norm_stats_output_path or the dataset's meta directory. + """ + + norm_stats_output_path: str | None = None + """ + Path to save calculated normalization statistics. If None, saves to + the dataset's meta/temporal_stats.json file. Unlike stats.json (raw + parquet data), temporal stats cover actions after REL_XYZ_ROT6D + conversion and include a temporal dimension for the action chunk. + """ + + stats_num_workers: int | None = None + """ + Number of parallel workers for statistics calculation. If None, uses CPU count. + Set to 1 to disable parallelism. Only used when calculate_norm_stats=True. + """ diff --git a/gr00t/configs/model/gr00t_n1d6.py b/gr00t/configs/model/gr00t_n1d6.py index 3c58f34..d67d98e 100755 --- a/gr00t/configs/model/gr00t_n1d6.py +++ b/gr00t/configs/model/gr00t_n1d6.py @@ -56,7 +56,7 @@ class Gr00tN1d6Config(PretrainedConfig): # Action head configuration parameters max_state_dim: int = 29 # Default from state_shape max_action_dim: int = 29 # Default from action_shape - action_horizon: int = 16 + action_horizon: int = 50 hidden_size: int = 1024 input_embedding_dim: int = 1536 @@ -98,6 +98,7 @@ class Gr00tN1d6Config(PretrainedConfig): # State Augmentation parameters state_dropout_prob: float = 0.0 # State dropout probability + state_dropout_prob_per_embodiment: dict[str, float] | None = None # Per-embodiment overrides state_additive_noise_scale: float = 0.0 # Scale for additive Gaussian noise on state features # Multi-embodiment parameters diff --git a/gr00t/data/dataset/factory.py b/gr00t/data/dataset/factory.py index 6f04aa8..4764f82 100644 --- a/gr00t/data/dataset/factory.py +++ b/gr00t/data/dataset/factory.py @@ -1,14 +1,19 @@ +import json +from pathlib import Path + import numpy as np import torch from tqdm import tqdm from gr00t.configs.base_config import Config +from gr00t.data.dataset.lerobot_episode_loader import LeRobotEpisodeLoader from gr00t.data.dataset.sharded_mixture_dataset import ShardedMixtureDataset from gr00t.data.dataset.sharded_single_step_dataset import ShardedSingleStepDataset from gr00t.data.embodiment_tags import EmbodimentTag from gr00t.data.interfaces import BaseProcessor -from gr00t.data.stats import generate_rel_stats, generate_stats -from gr00t.experiment.dist_utils import barrier +from gr00t.data.split_utils import load_info_json, resolve_episode_indices +from gr00t.data.stats import check_stats_validity, generate_rel_stats, generate_stats +from gr00t.data.types import ActionRepresentation class DatasetFactory: @@ -19,6 +24,178 @@ class DatasetFactory: def __init__(self, config: Config): self.config = config + def _get_allowed_episode_indices( + self, + dataset_path: str, + embodiment_tag: str, + include_splits: list[str] | None, + exclude_splits: list[str] | None, + ) -> np.ndarray | None: + """Resolve allowed episode indices based on include/exclude split settings. + + Args: + dataset_path: Path to the dataset root directory. + embodiment_tag: Embodiment tag string for loader fallback. + include_splits: Optional allowlist of split names. + exclude_splits: Optional denylist of split names. + + Returns: + Sorted numpy array of allowed episode indices, or None if no filtering + is requested. + """ + if not include_splits and not exclude_splits: + return None + + info = load_info_json(Path(dataset_path)) + total_episodes = info.get("total_episodes") + if total_episodes is None: + total_episodes = self._get_episode_count(dataset_path, embodiment_tag) + + allowed = resolve_episode_indices( + info, + include_splits=include_splits, + exclude_splits=exclude_splits, + total_episodes=int(total_episodes), + ) + return allowed + + def _get_episode_count(self, dataset_path: str, embodiment_tag: str) -> int: + """Get the number of episodes in a dataset without loading full data. + + Creates a lightweight episode loader that only reads metadata, + avoiding expensive video decoding. + + Args: + dataset_path: Path to the LeRobot format dataset + embodiment_tag: Embodiment tag for modality config lookup + + Returns: + Number of episodes in the dataset + """ + # Create lightweight loader just to get episode count + # skip_video=True avoids expensive video initialization + loader = LeRobotEpisodeLoader( + dataset_path=dataset_path, + modality_configs=self.config.data.modality_configs[embodiment_tag], + video_backend=self.config.data.video_backend, + skip_video=True, + ) + return len(loader) + + def _stats_exist(self, dataset_path: str) -> bool: + """Check whether a valid stats.json already exists for the dataset. + + Uses the same validation logic as gr00t.data.stats.check_stats_validity, + checking that the file exists and contains all required stat fields for + every float feature defined in info.json. + + Args: + dataset_path: Path to the dataset root directory. + + Returns: + True if stats.json exists and is valid; False otherwise. + """ + dp = Path(dataset_path) + info_path = dp / "meta" / "info.json" + if not info_path.exists(): + return False + + with open(info_path, "r") as f: + features = json.load(f).get("features", {}) + + lowdim_features = [k for k, v in features.items() if "float" in v.get("dtype", "")] + return check_stats_validity(dp, lowdim_features) + + def _get_relative_action_keys(self, embodiment_tag: str) -> list[str]: + """Resolve action keys that use RELATIVE representation for an embodiment. + + This inspects the configured modality configs rather than the dataset itself, + because the action representation is defined by the embodiment configuration. + + Args: + embodiment_tag: Embodiment tag string for modality config lookup. + + Returns: + Sorted list of action keys that require RELATIVE stats. + """ + modality_configs = self.config.data.modality_configs[embodiment_tag] + action_config = modality_configs.get("action") + if action_config is None or action_config.action_configs is None: + return [] + + relative_keys = [] + for key, config in zip(action_config.modality_keys, action_config.action_configs): + if config is not None and config.rep == ActionRepresentation.RELATIVE: + relative_keys.append(key) + + return sorted(relative_keys) + + def _load_percentile_stats( + self, + dataset_path: str, + consolidated_stats: dict | None, + ) -> dict | None: + """Load percentile stats for a dataset from consolidated or per-dataset files. + + The lookup order is controlled by whether a consolidated stats mapping is + provided. When consolidated stats are available, we only use those to avoid + mixing sources with potentially different normalization assumptions. + + Args: + dataset_path: Path to the dataset root directory. + consolidated_stats: Optional consolidated stats mapping keyed by repo_id. + + Returns: + Percentile stats dictionary for the dataset, or None if not found. + """ + if consolidated_stats is not None: + repo_id = Path(dataset_path).name + return consolidated_stats.get(repo_id) + + stats_path = Path(dataset_path) / "meta" / "temporal_stats.json" + if not stats_path.exists(): + return None + + with open(stats_path, "r") as f: + return json.load(f) + + def _percentile_stats_have_relative_action( + self, + dataset_path: str, + embodiment_tag: str, + consolidated_stats: dict | None, + ) -> bool: + """Check whether percentile stats already include RELATIVE action entries. + + This prevents recomputing relative_stats.json when percentile-based stats + already contain a "relative_action" section for the relevant action keys. + The check is intentionally lightweight (key existence only) to avoid heavy + schema validation during dataset initialization. + + Args: + dataset_path: Path to the dataset root directory. + embodiment_tag: Embodiment tag string for modality config lookup. + consolidated_stats: Optional consolidated stats mapping keyed by repo_id. + + Returns: + True if percentile stats exist and include relative action stats for all + RELATIVE action keys; False otherwise. + """ + relative_action_keys = self._get_relative_action_keys(embodiment_tag) + if not relative_action_keys: + # No RELATIVE actions configured; nothing to generate. + return True + + percentile_stats = self._load_percentile_stats(dataset_path, consolidated_stats) + if not percentile_stats: + return False + + relative_stats = percentile_stats.get("relative_action") + if not isinstance(relative_stats, dict): + return False + + return all(key in relative_stats for key in relative_action_keys) + def build( self, processor: BaseProcessor ) -> tuple[ShardedMixtureDataset, ShardedMixtureDataset | None]: @@ -29,24 +206,63 @@ def build( all_datasets = [] all_weights = [] + + consolidated_percentile_stats = None + for dataset_spec in tqdm( self.config.data.datasets, total=len(self.config.data.datasets), desc="Initializing datasets", ): - datasets = [] + datasets_for_spec = [] + for dataset_path in dataset_spec.dataset_paths: embodiment_tag = dataset_spec.embodiment_tag assert embodiment_tag is not None, "Embodiment tag is required" assert self.config.data.mode == "single_turn", "Only single turn mode is supported" + + # Generate statistics on rank 0, then barrier for distributed + allowed_episode_indices = self._get_allowed_episode_indices( + dataset_path=dataset_path, + embodiment_tag=embodiment_tag, + include_splits=dataset_spec.include_splits, + exclude_splits=dataset_spec.exclude_splits, + ) + if allowed_episode_indices is not None: + print( + f"Dataset {dataset_path}: using {len(allowed_episode_indices)} episodes " + f"after split filtering" + ) + if torch.distributed.is_initialized(): if torch.distributed.get_rank() == 0: - generate_stats(dataset_path) - generate_rel_stats(dataset_path, EmbodimentTag(embodiment_tag)) + if not self._stats_exist(dataset_path): + generate_stats(dataset_path, episode_indices=allowed_episode_indices) + if not self._percentile_stats_have_relative_action( + dataset_path=dataset_path, + embodiment_tag=embodiment_tag, + consolidated_stats=consolidated_percentile_stats, + ): + generate_rel_stats( + dataset_path, + EmbodimentTag(embodiment_tag), + episode_indices=allowed_episode_indices, + ) + torch.distributed.barrier() else: - generate_stats(dataset_path) - generate_rel_stats(dataset_path, EmbodimentTag(embodiment_tag)) - barrier() + if not self._stats_exist(dataset_path): + generate_stats(dataset_path, episode_indices=allowed_episode_indices) + if not self._percentile_stats_have_relative_action( + dataset_path=dataset_path, + embodiment_tag=embodiment_tag, + consolidated_stats=consolidated_percentile_stats, + ): + generate_rel_stats( + dataset_path, + EmbodimentTag(embodiment_tag), + episode_indices=allowed_episode_indices, + ) + dataset = ShardedSingleStepDataset( dataset_path=dataset_path, embodiment_tag=EmbodimentTag(embodiment_tag), @@ -56,11 +272,18 @@ def build( episode_sampling_rate=self.config.data.episode_sampling_rate, seed=self.config.data.seed, allow_padding=self.config.data.allow_padding, + episode_indices=allowed_episode_indices, ) - datasets.append(dataset) - dataset_lengths = np.array([len(dataset) for dataset in datasets]) - dataset_relative_lengths = dataset_lengths / dataset_lengths.sum() - for dataset, relative_length in zip(datasets, dataset_relative_lengths): + datasets_for_spec.append(dataset) + + # Calculate relative weights for training datasets in this spec + dataset_lengths = np.array([len(d) for d in datasets_for_spec]) + if dataset_lengths.sum() > 0: + dataset_relative_lengths = dataset_lengths / dataset_lengths.sum() + else: + dataset_relative_lengths = np.ones(len(datasets_for_spec)) / len(datasets_for_spec) + + for dataset, relative_length in zip(datasets_for_spec, dataset_relative_lengths): weight = relative_length * dataset_spec.mix_ratio all_datasets.append(dataset) all_weights.append(weight) diff --git a/gr00t/data/dataset/lerobot_episode_loader.py b/gr00t/data/dataset/lerobot_episode_loader.py index d6009d1..4046146 100755 --- a/gr00t/data/dataset/lerobot_episode_loader.py +++ b/gr00t/data/dataset/lerobot_episode_loader.py @@ -81,6 +81,9 @@ class LeRobotEpisodeLoader: that specify temporal sampling and data keys to load video_backend: Video decoding backend ('torchcodec', 'decord', etc.) video_backend_kwargs: Additional arguments for the video backend + skip_video: If True, skip loading video data entirely. Useful for fast + iteration during statistics calculation where only state/action + data is needed. Default False. Example: >>> loader = LeRobotEpisodeLoader( @@ -102,6 +105,8 @@ def __init__( modality_configs: dict[str, ModalityConfig], video_backend: str = "torchcodec", video_backend_kwargs: dict[str, Any] | None = None, + skip_video: bool = False, + require_stats: bool = True, ) -> None: """ Initialize LeRobot episode loader with dataset path and modality configurations. @@ -110,10 +115,21 @@ def __init__( 1. Loading all metadata files from the dataset 2. Parsing and validating modality configurations 3. Computing effective episode lengths based on action horizon + + Args: + dataset_path: Path to dataset root directory + modality_configs: Dictionary mapping modality names to ModalityConfig objects + video_backend: Video decoding backend ('torchcodec', 'decord', etc.) + video_backend_kwargs: Additional arguments for the video backend + skip_video: If True, skip loading video data (for fast stats calculation) + require_stats: If True, require stats.json to exist. Set to False when + calculating statistics for the first time. Default True. """ self.dataset_path = Path(dataset_path) self.video_backend = video_backend self.video_backend_kwargs = video_backend_kwargs + self.skip_video = skip_video + self.require_stats = require_stats if not self.dataset_path.is_dir(): raise FileNotFoundError(f"Dataset path does not exist: {self.dataset_path}") @@ -161,18 +177,22 @@ def _load_metadata(self) -> None: with open(modality_path, "r") as f: self.modality_meta = json.load(f) - # Load dataset statistics for normalization + # Load dataset statistics for normalization (optional when calculating stats) stats_path = meta_dir / LEROBOT_STATS_FILE_NAME - assert stats_path.exists(), ( - f"{stats_path} does not exist for {self.dataset_path}, please use gr00t/data/stats.py to generate it" - ) - with open(stats_path, "r") as f: - self.stats = json.load(f) + if self.require_stats: + assert stats_path.exists(), ( + f"{stats_path} does not exist for {self.dataset_path}, please use gr00t/data/stats.py to generate it" + ) + with open(stats_path, "r") as f: + self.stats = json.load(f) - relative_stats_path = meta_dir / LEROBOT_RELATIVE_STATS_FILE_NAME - if relative_stats_path.exists(): - with open(relative_stats_path, "r") as f: - self.stats["relative_action"] = json.load(f) + relative_stats_path = meta_dir / LEROBOT_RELATIVE_STATS_FILE_NAME + if relative_stats_path.exists(): + with open(relative_stats_path, "r") as f: + self.stats["relative_action"] = json.load(f) + else: + # Stats not required (e.g., when calculating them for the first time) + self.stats = {} # Extract key configuration parameters self.feature_config = self.info_meta.get("features", {}) @@ -284,6 +304,18 @@ def _extract_joint_groups( for group_name in joint_groups: if group_name in modality_info: group_info = modality_info[group_name] + if "original_keys" in group_info: + original_keys = group_info["original_keys"] + start_idx = group_info.get("start", 0) + end_idx = group_info.get("end", None) + joint_data[group_name] = self._concat_original_key_values( + df=df, + original_keys=original_keys, + start_idx=start_idx, + end_idx=end_idx, + ) + continue + start_idx = group_info["start"] end_idx = group_info["end"] original_key = group_info.get("original_key", DEFAULT_COLUMN_NAMES[modality_type]) @@ -299,13 +331,64 @@ def _extract_joint_groups( return joint_data + def _concat_original_key_values( + self, + df: pd.DataFrame, + original_keys: list[str], + start_idx: int = 0, + end_idx: int | None = None, + ) -> list[np.ndarray]: + """Concatenate multiple source columns into a single vector per row. + + This supports modality.json entries that define `original_keys` as a list + of columns (e.g., pose.x, pose.y, pose.z, quat.x, quat.y, quat.z, quat.w) + which should be combined into a single state/action vector. + + Args: + df: Episode DataFrame containing raw columns. + original_keys: Ordered list of column names to concatenate. + start_idx: Optional start index applied after concatenation. + end_idx: Optional end index applied after concatenation. + + Returns: + List of numpy arrays, one per row, containing the concatenated values. + + Raises: + KeyError: If any of the requested columns are missing from the DataFrame. + """ + missing = [key for key in original_keys if key not in df.columns] + if missing: + raise KeyError( + f"Missing original_keys columns in dataset: {missing}. " + f"Available columns: {list(df.columns)}" + ) + + columns = [df[key] for key in original_keys] + combined: list[np.ndarray] = [] + for row_idx in range(len(df)): + parts = [] + for col in columns: + value = col.iloc[row_idx] + if isinstance(value, np.ndarray): + array_value = value + elif isinstance(value, (list, tuple)): + array_value = np.asarray(value) + else: + array_value = np.asarray([value]) + parts.append(np.atleast_1d(array_value)) + concat_value = np.concatenate(parts, axis=0) + combined.append(concat_value[start_idx:end_idx]) + return combined + def _load_parquet_data(self, episode_index: int) -> pd.DataFrame: """ Load and process parquet data for a specific episode. Handles the complete data loading pipeline: 1. Load raw parquet file based on chunking structure - 2. Process language annotations (convert task indices to strings) + 2. Process language annotations: + - Convert task indices to strings using tasks.jsonl + - Or pass through raw text columns when annotation metadata sets is_text=True 3. Extract state and action joint groups Args: @@ -323,7 +406,7 @@ def _load_parquet_data(self, episode_index: int) -> pd.DataFrame: original_df = pd.read_parquet(parquet_path) loaded_df = pd.DataFrame() - # Process language annotations (convert task indices to task strings) + # Process language annotations (task index -> string, or raw text pass-through) if "language" in self.modality_configs: for key in self.modality_configs["language"].modality_keys: # these keys will be loaded separately from episodes.jsonl @@ -334,10 +417,44 @@ def _load_parquet_data(self, episode_index: int) -> pd.DataFrame: assert subkey in self.modality_meta["annotation"], ( f"Key {subkey} not found in language modality" ) - original_key = self.modality_meta["annotation"][subkey].get("original_key", key) - loaded_df[f"language.{key}"] = original_df[original_key].apply( - lambda x: self.tasks_map[x] - ) + annotation_meta = self.modality_meta["annotation"][subkey] + original_key = annotation_meta.get("original_key", key) + + if annotation_meta.get("is_text", False): + # Pass through raw text columns (e.g., instruction.text) + loaded_df[f"language.{key}"] = original_df[original_key].fillna("").astype(str) + else: + # Map task indices -> task strings (tasks.jsonl) + def _coerce_task_index(value: Any) -> int: + """Normalize task index values for tasks.jsonl lookup. + + Some datasets store task indices as length-1 arrays or lists. + This helper converts those into plain integers so they can be + used as keys into tasks.jsonl. + + Args: + value: Raw task index value from the parquet column. + + Returns: + Integer task index suitable for tasks.jsonl mapping. + + Raises: + ValueError: If the value cannot be coerced into a + single integer task index. + """ + if isinstance(value, np.ndarray): + if value.size == 1: + return int(value.item()) + raise ValueError(f"Task index array has size {value.size}") + if isinstance(value, (list, tuple)): + if len(value) == 1: + return int(value[0]) + raise ValueError(f"Task index list has length {len(value)}") + return int(value) + + loaded_df[f"language.{key}"] = original_df[original_key].apply( + lambda x: self.tasks_map[_coerce_task_index(x)] + ) # Extract joint groups for state and action modalities for modality_type in ["state", "action"]: @@ -468,27 +585,78 @@ def get_dataset_statistics(self) -> dict[str, Any]: for modality in mapping.keys(): # state, action for joint_key in self.modality_configs[modality].modality_keys: + modality_meta = self.modality_meta[modality][joint_key] # Determine which statistics key to use - if self.modality_meta[modality][joint_key].get("original_key", None) is not None: - stats_key = self.modality_meta[modality][joint_key]["original_key"] + if modality_meta.get("original_keys") is not None: + original_keys = modality_meta["original_keys"] + start_idx = modality_meta.get("start", 0) + end_idx = modality_meta.get("end", None) + stat_types = self.stats[original_keys[0]].keys() + for stat_type in stat_types: + combined_stats = self._concat_stats_for_original_keys( + original_keys=original_keys, + stat_type=stat_type, + ) + dataset_statistics[modality][joint_key][stat_type] = combined_stats[ + start_idx:end_idx + ] else: - stats_key = mapping[modality] - - # Extract the relevant slice of statistics - start_idx, end_idx = ( - self.modality_meta[modality][joint_key]["start"], - self.modality_meta[modality][joint_key]["end"], - ) - for stat_type in self.stats[stats_key].keys(): # mean, std, min, max, q01, q99 - dataset_statistics[modality][joint_key][stat_type] = self.stats[stats_key][ - stat_type - ][start_idx:end_idx] + if modality_meta.get("original_key", None) is not None: + stats_key = modality_meta["original_key"] + else: + stats_key = mapping[modality] + + # Extract the relevant slice of statistics + start_idx, end_idx = ( + modality_meta["start"], + modality_meta["end"], + ) + for stat_type in self.stats[stats_key].keys(): # mean, std, min, max, q01, q99 + dataset_statistics[modality][joint_key][stat_type] = self.stats[stats_key][ + stat_type + ][start_idx:end_idx] stats = _to_plain_dict(dataset_statistics) # Directly add relative action stats if "relative_action" in self.stats: stats["relative_action"] = self.stats["relative_action"] return stats + def _concat_stats_for_original_keys( + self, + original_keys: list[str], + stat_type: str, + ) -> np.ndarray: + """Concatenate per-key stats from stats.json into a single vector. + + Args: + original_keys: Ordered list of stats.json keys to concatenate. + stat_type: Stat type to extract (mean, std, min, max, q01, q99). + + Returns: + Concatenated numpy array for the requested stat type. + + Raises: + KeyError: If any of the stats keys or stat types are missing. + """ + parts = [] + missing = [key for key in original_keys if key not in self.stats] + if missing: + raise KeyError( + f"Missing stats entries for original_keys: {missing}. " + f"Available stats keys: {list(self.stats.keys())}" + ) + + for key in original_keys: + if stat_type not in self.stats[key]: + raise KeyError( + f"Missing stat '{stat_type}' for key '{key}'. " + f"Available stats: {list(self.stats[key].keys())}" + ) + stat_value = np.asarray(self.stats[key][stat_type]) + parts.append(np.atleast_1d(stat_value)) + + return np.concatenate(parts, axis=0) + def create_language_from_meta( self, episode_meta: dict, nframes: int, lang_key: str ) -> list[str]: @@ -555,15 +723,156 @@ def __getitem__(self, idx: int) -> pd.DataFrame: actual_length = min(len(df), nominal_length) df = df.iloc[:actual_length] - # Load synchronized video data - video_data = self._load_video_data(episode_id, np.arange(actual_length)) + # Load synchronized video data (skip if skip_video flag is set for fast stats calculation) + if not self.skip_video: + video_data = self._load_video_data(episode_id, np.arange(actual_length)) - # Add video frames to dataframe as PIL Images - for key in video_data.keys(): - assert len(video_data[key]) == len(df), ( - f"Video data for {key} has length {len(video_data[key])} but dataframe has length {len(df)}" - ) - df[f"video.{key}"] = [frame for frame in video_data[key]] + # Add video frames to dataframe as PIL Images + for key in video_data.keys(): + assert len(video_data[key]) == len(df), ( + f"Video data for {key} has length {len(video_data[key])} but dataframe has length {len(df)}" + ) + df[f"video.{key}"] = [frame for frame in video_data[key]] + + return df + + def get_sparse_episode( + self, + idx: int, + frame_indices: list[int] | np.ndarray, + video_indices: list[int] | np.ndarray | None = None, + allow_padding: bool = False, + ) -> pd.DataFrame: + """Load episode data for only specific frame indices (sparse loading). + + This is a performance optimization for sharded datasets. Instead of loading + ALL video frames for an episode (which can be hundreds or thousands of frames + per camera view), this method loads only the specific frames needed for the + timesteps assigned to a shard. + + The returned DataFrame preserves original frame indices in its index (not + reset to 0, 1, 2, ...), allowing extract_step_data() to correctly map + step_index + delta_indices to DataFrame positions via an index map. + + Video indices are separated from state/action indices because the video + modality typically uses only delta_indices=[0] (current frame), while + actions use a long horizon (e.g., 16-50 future steps). Without this + decoupling, loading one timestep's action labels would force decoding + 16-50 extra video frames for no reason. + + Args: + idx: Episode index to load. + frame_indices: Frame indices to load for non-video modalities + (state, action, language). These define the parquet rows to keep. + video_indices: Frame indices to load for video modalities. If None, + defaults to frame_indices. Typically much smaller than + frame_indices since video delta_indices are usually just [0]. + allow_padding: If True, clamp out-of-range indices to the valid + episode bounds. If False, raise on out-of-range indices to + surface data consistency issues. + + Returns: + DataFrame with columns for all modalities, but only rows for the + requested frame_indices. The DataFrame.index contains the original + frame indices (not 0, 1, 2, ...) so downstream code can correctly + map step_index + delta to positions. + + Raises: + IndexError: If episode index is out of bounds. + ValueError: If frame_indices is empty. + + Example: + >>> # Load only frames 10, 50, 100 (state/action) and frame 10 (video) + >>> sparse_df = loader.get_sparse_episode(0, [10, 50, 100], video_indices=[10]) + >>> sparse_df.index # Int64Index([10, 50, 100]), not RangeIndex(0, 3) + """ + if idx < 0 or idx >= len(self): + raise IndexError(f"Episode index {idx} out of bounds") + + frame_indices = np.asarray(frame_indices, dtype=np.int64) + if len(frame_indices) == 0: + raise ValueError("frame_indices cannot be empty") + if video_indices is None: + video_indices = frame_indices + else: + video_indices = np.asarray(video_indices, dtype=np.int64) + + episode_meta = self.episodes_metadata[idx] + episode_id = episode_meta["episode_index"] + nominal_length = episode_meta["length"] + + # Load and parse the parquet data (state/action/language β€” fast, no video) + df = self._load_parquet_data(episode_id) + + # Handle language from metadata (needs full frame count before subsetting) + if "language" in self.modality_configs: + lang_key = self.modality_configs["language"].modality_keys[0] + if lang_key in LANG_KEYS: + new_languages = self.create_language_from_meta(episode_meta, len(df), lang_key) + df["language." + lang_key] = new_languages + + # Clamp to actual DataFrame length + actual_length = min(len(df), nominal_length) + + # Validate and optionally clamp indices to valid range + def _validate_indices(name: str, indices: np.ndarray) -> None: + if len(indices) == 0: + return + if (indices < 0).any() or (indices >= actual_length).any(): + raise ValueError( + f"{name} out of range for episode {idx}: " + f"min={int(indices.min())}, max={int(indices.max())}, " + f"valid=[0, {actual_length - 1}]" + ) + + video_indices_same = video_indices is frame_indices + + if allow_padding: + frame_indices = np.clip(frame_indices, 0, actual_length - 1) + frame_indices = np.unique(frame_indices) # Remove duplicates and sort + if video_indices is not None and len(video_indices) > 0: + if video_indices_same: + video_indices = frame_indices + else: + video_indices = np.clip(video_indices, 0, actual_length - 1) + video_indices = np.unique(video_indices) + else: + _validate_indices("frame_indices", frame_indices) + frame_indices = np.unique(frame_indices) + if video_indices is not None and len(video_indices) > 0: + if video_indices_same: + video_indices = frame_indices + else: + _validate_indices("video_indices", video_indices) + video_indices = np.unique(video_indices) + + # Merge video indices into frame set so the DataFrame contains rows for + # both state/action AND video indices + if video_indices is not None and len(video_indices) > 0: + frame_indices = np.unique(np.concatenate([frame_indices, video_indices])) + + # Subset DataFrame to only requested frames, preserving original indices + # so that extract_step_data can map step_index + delta -> position + df = df.iloc[frame_indices].copy() + df.index = frame_indices # Preserve original frame indices in the index + + # Load ONLY the requested video frames (the key optimization). + # Instead of decoding all 700+ frames per camera view, we decode only + # the ~1-10 frames actually needed for the video modality's delta_indices. + if not self.skip_video and video_indices is not None and len(video_indices) > 0: + video_data = self._load_video_data(episode_id, video_indices) + + # Insert video frames only at the corresponding video_indices rows. + # Other rows (needed only for state/action) get None in video columns + # and are never accessed by the video modality. + for key in video_data.keys(): + df[f"video.{key}"] = [None] * len(df) + frames = [frame for frame in video_data[key]] + assert len(frames) == len(video_indices), ( + f"Video data for {key} has length {len(frames)} " + f"but expected {len(video_indices)}" + ) + df.loc[video_indices, f"video.{key}"] = frames # Load synchronized mask data mask_data = self._load_mask_data(episode_id, np.arange(actual_length)) diff --git a/gr00t/data/dataset/sharded_mixture_dataset.py b/gr00t/data/dataset/sharded_mixture_dataset.py index 7fce6dc..d809499 100755 --- a/gr00t/data/dataset/sharded_mixture_dataset.py +++ b/gr00t/data/dataset/sharded_mixture_dataset.py @@ -1,4 +1,5 @@ from concurrent.futures import Future, ThreadPoolExecutor +from pathlib import Path import time import numpy as np @@ -6,101 +7,130 @@ from torch.utils.data import IterableDataset, get_worker_info from gr00t.data.interfaces import BaseProcessor, ShardedDataset +from gr00t.data.percentile_merge import merge_piecewise_linear_quantiles + + +def _validate_required_stats( + stats: dict[str, list[float] | np.ndarray], + context: str, +) -> None: + """Validate that required statistic keys are present for percentile merging. + + Args: + stats: Statistic dictionary for a single joint group. + context: Context string for error reporting. + + Raises: + ValueError: If any required statistic keys are missing. + """ + required_keys = ("min", "max", "mean", "std", "q01", "q02", "q98", "q99") + missing = [key for key in required_keys if key not in stats] + if missing: + raise ValueError( + f"Missing required statistics for percentile merge. Missing {missing} in {context}." + ) def merge_statistics( per_dataset_stats: list[dict[str, dict[str, list[float] | np.ndarray]]], dataset_sampling_weights: list[float] | np.ndarray, is_relative_stats: bool = False, -) -> dict[str, dict[str, list[float]]]: - """ - Compute overall statistics from per-dataset statistics using weighted averaging. - - This function combines statistics from multiple datasets according to their sampling - weights, computing weighted means and variances while preserving min/max/quantile - information across all datasets. +) -> dict[str, dict[str, list[float] | np.ndarray]]: + """Merge per-dataset statistics for a single modality across joint groups. - The weighted variance computation uses the formula: - Var_combined = Ξ£(w_i * (Οƒ_iΒ² + ΞΌ_iΒ²)) - (Ξ£(w_i * ΞΌ_i))Β² + This function combines statistics from multiple datasets according to their + sampling weights. Means and variances are merged via weighted averaging, while + percentiles are merged via piecewise-linear CDF interpolation. Args: - per_dataset_stats: List of per-dataset statistics dictionaries. - Each element has structure: {modality: {joint_group: {stat_type: values}}} - Example: {"state": {"gripper": {"mean": [0.1, 0.2], "std": [0.5, 0.3]}}} - dataset_sampling_weights: Weights for combining dataset statistics. - Should sum to 1.0 or will be normalized. - is_relative_stats: Whether the statistics are relative (affects merging logic). + per_dataset_stats: List of per-dataset statistic dicts for a single modality. + Structure: {joint_group: {stat_type: values}}. + dataset_sampling_weights: Weights for combining dataset statistics. Will be + normalized to sum to 1.0. + is_relative_stats: Whether stats are for relative actions (used for context + in error messages). Returns: - Combined statistics dictionary with same structure as input, containing - weighted averages for mean/std and global min/max/quantiles across datasets. + Dictionary mapping joint_group to merged statistics for that group. """ - # Normalize sampling weights to sum to 1 - dataset_sampling_weights = np.array(dataset_sampling_weights) - normalized_weights = dataset_sampling_weights / dataset_sampling_weights.sum() - - # Initialize overall statistics dict - overall_stats: dict[str, dict[str, list[float]]] = {} - - # Process each modality (e.g., "state", "action") - for modality in per_dataset_stats[0]: - # Get dimensionality from first dataset (assumed consistent) - dim = ( - [len(per_dataset_stats[0][modality]["mean"])] - if not is_relative_stats - else np.array(per_dataset_stats[0][modality]["mean"]).shape - ) - - # Initialize accumulators for weighted mean and variance computation - weighted_means = np.zeros(dim) - weighted_squares = np.zeros(dim) - - # Collect min/max/quantiles from all datasets for global computation + if not per_dataset_stats: + raise ValueError("Cannot merge statistics: per_dataset_stats is empty.") + + normalized_weights = np.array(dataset_sampling_weights, dtype=np.float64) + weight_sum = normalized_weights.sum() + if weight_sum <= 0: + raise ValueError("Dataset sampling weights must sum to a positive value.") + normalized_weights = normalized_weights / weight_sum + + overall_stats: dict[str, dict[str, list[float] | np.ndarray]] = {} + joint_groups = per_dataset_stats[0].keys() + stats_type = "relative_action" if is_relative_stats else "modality" + + for joint_group in joint_groups: + weighted_means = None + weighted_squares = None min_list = [] max_list = [] - q01_list = [] - q99_list = [] + quantile_payloads = [] - # Accumulate weighted statistics across datasets for dataset_idx, dataset_stats in enumerate(per_dataset_stats): - w_i = normalized_weights[dataset_idx] - stats = dataset_stats[modality] - means = np.array(stats["mean"]) - stds = np.array(stats["std"]) - - # Update weighted sums for mean and variance calculation - weighted_means += w_i * means - weighted_squares += w_i * (stds**2 + means**2) - - # Collect extremes and quantiles for global computation - min_list.append(stats["min"]) - max_list.append(stats["max"]) - q01_list.append(stats["q01"]) - q99_list.append(stats["q99"]) - - # Compute final combined statistics - overall_mean = weighted_means.tolist() + if joint_group not in dataset_stats: + raise ValueError( + "Missing joint group statistics for merge. " + f"Joint group '{joint_group}' not found in dataset index {dataset_idx}." + ) + stats = dataset_stats[joint_group] + context = f"{stats_type} '{joint_group}', dataset index {dataset_idx}" + _validate_required_stats(stats, context=context) + means = np.array(stats["mean"], dtype=np.float64) + stds = np.array(stats["std"], dtype=np.float64) + + if weighted_means is None: + weighted_means = np.zeros_like(means, dtype=np.float64) + weighted_squares = np.zeros_like(means, dtype=np.float64) + + weight = normalized_weights[dataset_idx] + weighted_means += weight * means + weighted_squares += weight * (stds**2 + means**2) + + min_list.append(np.array(stats["min"], dtype=np.float64)) + max_list.append(np.array(stats["max"], dtype=np.float64)) + quantile_payloads.append( + { + "min": np.array(stats["min"], dtype=np.float64), + "q01": np.array(stats["q01"], dtype=np.float64), + "q02": np.array(stats["q02"], dtype=np.float64), + "q98": np.array(stats["q98"], dtype=np.float64), + "q99": np.array(stats["q99"], dtype=np.float64), + "max": np.array(stats["max"], dtype=np.float64), + } + ) + + if weighted_means is None or weighted_squares is None: + raise ValueError( + f"Cannot merge statistics: no valid datasets for joint group '{joint_group}'." + ) + + overall_mean = weighted_means overall_variance = weighted_squares - weighted_means**2 - overall_std = np.sqrt(overall_variance).tolist() - - # Global min/max across all datasets - overall_min = np.min(np.array(min_list), axis=0).tolist() - overall_max = np.max(np.array(max_list), axis=0).tolist() - - # Global quantiles (conservative bounds across datasets) - q01_array = np.array(q01_list) - q99_array = np.array(q99_list) - weighted_q01 = np.min(q01_array, axis=0).tolist() - weighted_q99 = np.max(q99_array, axis=0).tolist() - - # Store combined statistics for this modality - overall_stats[modality] = { - "min": overall_min, - "max": overall_max, - "mean": overall_mean, - "std": overall_std, - "q01": weighted_q01, - "q99": weighted_q99, + overall_std = np.sqrt(overall_variance) + + overall_min = np.min(np.stack(min_list, axis=0), axis=0) + overall_max = np.max(np.stack(max_list, axis=0), axis=0) + merged_quantiles = merge_piecewise_linear_quantiles( + per_dataset_quantiles=quantile_payloads, + weights=normalized_weights, + ) + + overall_stats[joint_group] = { + "min": overall_min.tolist(), + "max": overall_max.tolist(), + "mean": overall_mean.tolist(), + "std": overall_std.tolist(), + "q01": merged_quantiles["q01"].tolist(), + "q02": merged_quantiles["q02"].tolist(), + "q98": merged_quantiles["q98"].tolist(), + "q99": merged_quantiles["q99"].tolist(), } return overall_stats @@ -135,6 +165,7 @@ class ShardedMixtureDataset(IterableDataset): seed: Random seed for reproducible sampling training: Whether in training mode (affects sampling strategy) num_shards_per_epoch: Number of shards to sample per epoch during training + override_pretraining_statistics: Whether to override pretrained model statistics Example: >>> mixture = ShardedMixtureDataset( @@ -158,7 +189,17 @@ def __init__( num_shards_per_epoch: int = int(1e5), override_pretraining_statistics: bool = False, ): - """Initialize mixture dataset with datasets, weights, and configuration.""" + """Initialize mixture dataset with datasets, weights, and configuration. + + Args: + datasets: List of ShardedDataset instances to combine + weights: Mixing weights for each dataset (will be normalized) + processor: Data processor to apply to all datasets + seed: Random seed for reproducible sampling + training: Whether in training mode (affects sampling strategy) + num_shards_per_epoch: Number of shards to sample per epoch during training + override_pretraining_statistics: Whether to override pretrained model statistics + """ self.datasets = datasets self.weights = weights self.seed = seed @@ -171,7 +212,7 @@ def __init__( # Generate initial shard sampling schedule self.shard_sampling_schedule = self.generate_shard_sampling_schedule() - # Merge statistics across datasets and configure processor + # Merge statistics and configure processor self.merge_statistics() # Initialize distributed training parameters @@ -200,6 +241,7 @@ def merge_statistics(self): # Group datasets and weights by embodiment all_stats_by_emb: dict[str, list] = {} weights_by_emb: dict[str, list[float]] = {} + datasets_missing_percentiles: list[str] = [] for ds, w in zip(self.datasets, self.weights): emb = getattr(ds, "embodiment_tag", None) if emb is None: @@ -208,15 +250,39 @@ def merge_statistics(self): if emb not in all_stats_by_emb: all_stats_by_emb[emb] = [] weights_by_emb[emb] = [] - stats = ds.get_dataset_statistics() # type: ignore - all_stats_by_emb[emb].append(stats) + if not hasattr(ds, "get_percentile_statistics"): + dataset_name = getattr(ds, "repo_id", Path(ds.dataset_path).name) + raise ValueError( + "Per-embodiment percentile merge requires datasets that expose " + f"percentile statistics. Dataset '{dataset_name}' does not support it." + ) + percentile_stats = ds.get_percentile_statistics(None) # type: ignore + if percentile_stats is None: + dataset_name = getattr(ds, "repo_id", Path(ds.dataset_path).name) + datasets_missing_percentiles.append(dataset_name) + continue + all_stats_by_emb[emb].append(percentile_stats) weights_by_emb[emb].append(w) + if datasets_missing_percentiles: + raise ValueError( + "Per-embodiment percentile merge requires per-dataset percentile stats " + "for every dataset. Missing stats for: " + f"{sorted(datasets_missing_percentiles)}" + ) + # Merge statistics within each embodiment group stats_by_emb = {} for emb, stats in all_stats_by_emb.items(): stats_by_emb[emb] = {} for modality in ["state", "action", "relative_action"]: + if modality == "relative_action": + relative_action_presence = [modality in s for s in stats] + if any(relative_action_presence) and not all(relative_action_presence): + raise ValueError( + "Relative-action statistics must be present for all datasets or none " + f"within embodiment '{emb}'. Found presence flags: {relative_action_presence}" + ) if modality in stats[0]: modality_stats = [s[modality] for s in stats] stats_by_emb[emb][modality] = merge_statistics( @@ -351,6 +417,12 @@ def __iter__(self): # Initialize worker-specific shard schedule self.worker_shard_sampling_schedule = self.filter_shard_sample_schedule() self.curr_shard_index = -1 + + # Seed processor RNG streams for this (epoch, rank, worker) context + # before scheduling the first background cache job. + if self.processor is not None and hasattr(self.processor, "seed_vqa_rng"): + self.processor.seed_vqa_rng(self.seed, self.epoch, self.rank, self.worker_id or 0) + self.cache_next_shard() rng = np.random.default_rng(self.seed + self.epoch) @@ -396,6 +468,10 @@ def cache_next_shard(self): self.worker_shard_sampling_schedule = self.filter_shard_sample_schedule() self.curr_shard_index = -1 + # Reseed processor RNG streams for the new epoch context. + if self.processor is not None and hasattr(self.processor, "seed_vqa_rng"): + self.processor.seed_vqa_rng(self.seed, self.epoch, self.rank, self.worker_id or 0) + print(f"Rank {self.rank}, Worker {self.worker_id}: Caching shard...") next_dataset_idx, next_shard_idx = self.worker_shard_sampling_schedule[ self.curr_shard_index + 1 diff --git a/gr00t/data/dataset/sharded_single_step_dataset.py b/gr00t/data/dataset/sharded_single_step_dataset.py index 1dc40cc..b0a140b 100644 --- a/gr00t/data/dataset/sharded_single_step_dataset.py +++ b/gr00t/data/dataset/sharded_single_step_dataset.py @@ -1,3 +1,4 @@ +import json from pathlib import Path from typing import Any @@ -5,11 +6,20 @@ import pandas as pd from gr00t.data.interfaces import ShardedDataset +from gr00t.data.step_filtering import compute_valid_step_indices_parallel from gr00t.data.types import EmbodimentTag, MessageType, ModalityConfig, VLAStepData from .lerobot_episode_loader import LeRobotEpisodeLoader +# Maximum number of workers for parallel filtering +MAX_FILTER_WORKERS = 128 + + +# Constants for percentile stats file handling +PERCENTILE_STATS_FILENAME = "meta/temporal_stats.json" + + def extract_step_data( episode_data: pd.DataFrame, step_index: int, @@ -17,18 +27,80 @@ def extract_step_data( embodiment_tag: EmbodimentTag, allow_padding: bool = False, ) -> VLAStepData: + """Extract a single training sample from episode data at a given step index. + + Handles both dense DataFrames (from full episode loading, with a standard + RangeIndex 0..N-1) and sparse DataFrames (from get_sparse_episode, where + the index contains the original frame indices). For sparse DataFrames, an + index map translates original frame indices to positional iloc offsets. + + Args: + episode_data: Episode DataFrame, either dense (full episode) or sparse + (subset of frames with original indices preserved in df.index). + step_index: The anchor timestep to extract data for. + modality_configs: Per-modality configuration (delta_indices, keys). + embodiment_tag: Embodiment identifier for the dataset. + allow_padding: If True, clamp out-of-range indices to the valid range + instead of raising an error. + + Returns: + VLAStepData containing video, state, action, and language data for the + requested step_index with all configured delta offsets applied. + """ step_data = {} + # Detect sparse DataFrames from get_sparse_episode. + # Sparse DataFrames have a non-standard index (original frame indices) + # instead of the default RangeIndex(start=0, step=1). + is_sparse = not ( + isinstance(episode_data.index, pd.RangeIndex) + and episode_data.index.start == 0 + and episode_data.index.step == 1 + ) + + # Build index map for sparse DataFrames: {original_frame_idx: iloc_position} + # This lets us convert step_index + delta -> iloc position for data access. + # For dense DataFrames, index_map is None and positions == indices_to_load. + index_map: dict[int, int] | None = None + if is_sparse: + index_map = {int(idx): pos for pos, idx in enumerate(episode_data.index)} + # Extract data for each configured modality for modality, config in modality_configs.items(): step_data[modality] = {} # Sample timesteps according to delta indices configuration indices_to_load = [step_index + delta_index for delta_index in config.delta_indices] + if allow_padding: - indices_to_load = [max(0, min(idx, len(episode_data) - 1)) for idx in indices_to_load] + if is_sparse: + # For sparse DataFrames, clamp to the range of available indices + min_idx = int(episode_data.index.min()) + max_idx = int(episode_data.index.max()) + indices_to_load = [max(min_idx, min(idx, max_idx)) for idx in indices_to_load] + else: + indices_to_load = [ + max(0, min(idx, len(episode_data) - 1)) for idx in indices_to_load + ] + + # Convert original frame indices to iloc positions for sparse DataFrames. + # When allow_padding=True, clamp to the sparse DF's index range so + # out-of-range indices map to boundary frames (same as dense-path padding). + # When allow_padding=False, do a direct lookup β€” a missing index raises + # KeyError, consistent with the dense path's IndexError on out-of-bounds iloc. + if is_sparse and index_map is not None: + if allow_padding: + min_idx = int(episode_data.index.min()) + max_idx = int(episode_data.index.max()) + clamped_indices = [max(min_idx, min(idx, max_idx)) for idx in indices_to_load] + positions = [index_map[idx] for idx in clamped_indices] + else: + positions = [index_map[idx] for idx in indices_to_load] + else: + positions = indices_to_load + for key in config.modality_keys: if f"{modality}.{key}" in episode_data.columns: - modality_data = episode_data[f"{modality}.{key}"].iloc[indices_to_load] + modality_data = episode_data[f"{modality}.{key}"].iloc[positions] else: raise KeyError( f"{modality}.{key} not found in episode data, available keys: {episode_data.columns}" @@ -96,6 +168,9 @@ class ShardedSingleStepDataset(ShardedDataset): episode_sampling_rate: Fraction of episode timesteps to use (for efficiency) seed: Random seed for reproducible sharding and sampling allow_padding: Whether to allow padding of indices to valid range [0, max_length - 1] + episode_indices: Optional subset of episode indices to use. If provided, only these + episodes will be included in the dataset. This enables train/val splitting at + the episode level to prevent data leakage. If None, all episodes are used. Example: >>> dataset = ShardedSingleStepDataset( @@ -125,8 +200,23 @@ def __init__( episode_sampling_rate: float = 0.1, seed: int = 42, allow_padding: bool = False, + episode_indices: np.ndarray | None = None, ): - """Initialize single-step dataset with sharding configuration.""" + """Initialize single-step dataset with sharding configuration. + + Args: + dataset_path: Path to LeRobot format dataset directory + embodiment_tag: Embodiment identifier for cross-embodiment training + modality_configs: Configuration for each modality (sampling, keys) + video_backend: Video decoding backend ('torchcodec', 'decord', etc.) + video_backend_kwargs: Additional arguments for video backend + shard_size: Target number of timesteps per shard + episode_sampling_rate: Fraction of episode timesteps to use (for efficiency) + seed: Random seed for reproducible sharding and sampling + allow_padding: Whether to allow padding of indices to valid range + episode_indices: Optional subset of episode indices to use for train/val split. + If None, all episodes are used. + """ super().__init__(dataset_path) self.embodiment_tag = embodiment_tag self.modality_configs = modality_configs @@ -136,10 +226,16 @@ def __init__( self.episode_sampling_rate = episode_sampling_rate self.seed = seed self.allow_padding = allow_padding + self.episode_indices = episode_indices # Store for shard_dataset() self.processor = None self.rng = np.random.default_rng(seed) action_delta_indices = modality_configs["action"].delta_indices - self.action_horizon = max(action_delta_indices) - min(action_delta_indices) + 1 + + # BUG: The below assumed contiguous delta indices, but fails for [0,2,4,...,32] + # self.action_horizon = max(action_delta_indices) - min(action_delta_indices) + 1 + # action_horizon must account for the maximum delta index, not just the range + # This ensures step_index + max(delta_indices) < episode_length + self.action_horizon = max(action_delta_indices) + 1 self.episode_loader = LeRobotEpisodeLoader( dataset_path=dataset_path, @@ -156,28 +252,78 @@ def shard_dataset(self): Create balanced shards by distributing episode timesteps across shards. The sharding process: - 1. Shuffle episode order for randomization - 2. Split each episode into multiple sub-sequences based on sampling rate - 3. Distribute sub-sequences across shards to balance shard sizes - 4. Use greedy assignment to minimize shard size variance + 1. Run parallel clutch-aware filtering (PyArrow + ProcessPoolExecutor) + 2. Shuffle episode order for randomization + 3. Split each episode into multiple sub-sequences based on sampling rate + 4. Distribute sub-sequences across shards to balance shard sizes + 5. Use greedy assignment to minimize shard size variance This approach ensures: - Balanced shard sizes for consistent training batches - Diversity within shards (mix of episodes and timesteps) - Reproducible sharding based on seed + - Fast startup with parallel filtering (~15-30 seconds for 4792 episodes) """ - shuffled_episode_indices = self.rng.permutation(len(self.episode_loader.episode_lengths)) + # Run parallel filtering once at startup (auto-detects CMR data) + # Uses PyArrow for fast column selection + ProcessPoolExecutor for parallelism + # Worker count auto-detects available CPUs, capped at MAX_FILTER_WORKERS (128) + filter_results = self._filter_all_episodes_parallel() + + # Use provided episode_indices subset or all episodes + # This enables train/val splitting at the episode level + if self.episode_indices is not None: + all_episode_indices = self.episode_indices + else: + all_episode_indices = np.arange(len(self.episode_loader.episode_lengths)) + + shuffled_episode_indices = self.rng.permutation(all_episode_indices) num_splits = int(1 / self.episode_sampling_rate) assert len(shuffled_episode_indices) > 0, ( f"No valid trajectories found for dataset {self.dataset_path}" ) - # Calculate total timesteps and required number of shards - total_steps = np.sum( - [self.get_effective_episode_length(idx) for idx in shuffled_episode_indices] - ).astype(int) - num_shards = np.ceil(total_steps / self.shard_size).astype(int) + # Calculate total timesteps accounting for clutch filtering when applied + # Use filtered step counts for CMR data, raw episode lengths otherwise + if filter_results is not None: + # CMR data: count actual valid steps after filtering + total_steps = 0 + for idx in shuffled_episode_indices: + if idx in filter_results: + total_steps += len(filter_results[idx]) + # Episodes not in filter_results had 0 valid indices, skip them + else: + # Non-CMR data: use raw episode lengths + total_steps = int( + np.sum([self.get_effective_episode_length(idx) for idx in shuffled_episode_indices]) + ) + + # Ensure at least 1 shard, handle edge case of all data filtered out + assert total_steps > 0, ( + f"No valid timesteps after filtering for dataset {self.dataset_path}. " + f"All {len(shuffled_episode_indices)} episodes were filtered out by clutch-aware filtering. " + f"filter_results is None: {filter_results is None}, " + f"filter_results keys: {list(filter_results.keys()) if filter_results else 'N/A'}" + ) + + # Count episodes that will actually contribute data + if filter_results is not None: + num_episodes_with_data = sum( + 1 for idx in shuffled_episode_indices if idx in filter_results + ) + else: + num_episodes_with_data = len(shuffled_episode_indices) + + # Calculate shards needed, but cap at episodes * num_splits since each episode + # contributes to at most num_splits shards in the distribution loop + num_shards_by_steps = max(1, int(np.ceil(total_steps / self.shard_size))) + max_shards_by_episodes = num_episodes_with_data * num_splits + num_shards = min(num_shards_by_steps, max_shards_by_episodes) + + print( + f"Shard plan: total_steps={total_steps}, num_shards={num_shards}, shard_size={self.shard_size}, " + f"num_episodes_with_data={num_episodes_with_data}, num_splits={num_splits}" + ) # Initialize shard containers sharded_episodes = [[] for _ in range(num_shards)] @@ -185,19 +331,37 @@ def shard_dataset(self): # Distribute episode sub-sequences across shards for ep_idx in shuffled_episode_indices: - # Split episode timesteps into multiple sub-sequences - step_indices = np.arange(0, self.get_effective_episode_length(ep_idx)) + # Get step indices - either from pre-computed filter results or full range + if filter_results is not None: + # CMR data: only use episodes that passed filtering + if ep_idx in filter_results: + step_indices = filter_results[ep_idx].copy() + else: + # Episode had 0 valid indices after filtering, skip it + continue + else: + # Non-CMR data: use full range + step_indices = np.arange(0, self.get_effective_episode_length(ep_idx)) + + if len(step_indices) == 0: + continue # Skip episodes with no valid indices + self.rng.shuffle(step_indices) for i in range(num_splits): split_step_indices = step_indices[i::num_splits] + if len(split_step_indices) == 0: + continue # Skip empty splits # Assign to shard with minimum current length (greedy balancing) shard_index = np.argmin(shard_lengths) sharded_episodes[shard_index].append((ep_idx, split_step_indices)) shard_lengths[shard_index] += len(split_step_indices) # Validate shard creation - assert all(shard_lengths[i] > 0 for i in range(num_shards)), ( - "All shards must have length greater than 0" + empty_shards = [i for i in range(num_shards) if shard_lengths[i] == 0] + assert len(empty_shards) == 0, ( + f"All shards must have length > 0. Empty shards: {empty_shards}, " + f"shard_lengths: {shard_lengths.tolist()}, total distributed: {shard_lengths.sum()}, " + f"expected total_steps: {total_steps}" ) print(f"Generated {num_shards} shards for dataset {self.dataset_path}") @@ -212,6 +376,45 @@ def get_effective_episode_length(self, episode_index: int) -> int: original_length = self.episode_loader.get_episode_length(episode_index) return max(0, original_length - self.action_horizon + 1) + def _filter_all_episodes_parallel( + self, num_workers: int | None = None + ) -> dict[int, np.ndarray] | None: + """Parallel filtering of all episodes using PyArrow column selection. + + This function applies one of two dataset-specific filters: + 1) CMR clutch-aware filtering (uses observation.state keys) + 2) Terminal-step filtering for datasets with `next.done` padding + + Uses ProcessPoolExecutor to distribute filtering across multiple CPUs. + Each episode's parquet file is read with PyArrow (only required columns), + avoiding video decoding and full DataFrame loads. + + Args: + num_workers: Number of parallel workers. If None, uses min(cpu_count, MAX_FILTER_WORKERS). + + Returns: + None if no filtering is required for this dataset. + Dict mapping episode_idx to valid step indices if filtering is applied. + Empty dict {} means all episodes were filtered out. + """ + if self.episode_indices is not None: + episode_indices = list(self.episode_indices) + else: + episode_indices = list(range(len(self.episode_loader))) + effective_lengths = [self.get_effective_episode_length(i) for i in episode_indices] + return compute_valid_step_indices_parallel( + dataset_path=self.dataset_path, + embodiment_tag=self.embodiment_tag, + chunk_size=self.episode_loader.chunk_size, + data_path_pattern=self.episode_loader.data_path_pattern, + action_delta_indices=self.modality_configs["action"].delta_indices, + episode_indices=episode_indices, + effective_lengths=effective_lengths, + num_workers=num_workers, + max_filter_workers=MAX_FILTER_WORKERS, + show_progress=True, + ) + def __len__(self): """Return the number of shards in the dataset.""" return len(self.shard_lengths) @@ -232,6 +435,7 @@ def get_datapoint(self, episode_data: pd.DataFrame, step_index: int) -> dict: Raises: AssertionError: If processor is not set before calling this method + KeyError: If the expected stats key is not present in processor normalization params """ assert self.processor is not None, "Processor must be set before getting datapoints" vla_step_data = extract_step_data( @@ -250,11 +454,19 @@ def get_shard_length(self, idx: int) -> int: return self.shard_lengths[idx] def get_shard(self, idx: int) -> list: - """ - Load and process all timesteps in a specific shard. + """Load and process all timesteps in a specific shard using sparse loading. - Loads the required episodes and extracts all timesteps assigned to this shard, - applying the configured processor to each timestep. + Uses sparse episode loading to decode only the video frames and parquet + rows actually needed for the timesteps assigned to this shard, rather than + loading entire episodes. This reduces video decode work by 5-11x for + typical configurations where video delta_indices=[0] but action horizons + span 16-50 steps. + + For each episode referenced in the shard: + 1. Compute the union of all frame indices needed across all modalities + 2. Compute the (smaller) union of frame indices needed for video only + 3. Call get_sparse_episode to load just those frames + 4. Extract individual timestep datapoints from the sparse DataFrame Args: idx: Shard index to load @@ -265,16 +477,158 @@ def get_shard(self, idx: int) -> list: episodes = self.sharded_episodes[idx] datapoints = [] for ep_idx, step_indices in episodes: - # Load episode data once per episode in shard - episode_data = self.episode_loader[ep_idx] + # Compute the minimal set of frame indices needed for this episode's + # step_indices across all modality delta offsets + required_indices = self._compute_required_indices(ep_idx, step_indices) + required_video_indices = self._compute_required_video_indices(ep_idx, step_indices) + + # Load only the needed frames (sparse parquet subset + targeted video decode) + episode_data = self.episode_loader.get_sparse_episode( + ep_idx, + frame_indices=required_indices, + video_indices=required_video_indices, + allow_padding=self.allow_padding, + ) for step_index in step_indices: datapoints.append(self.get_datapoint(episode_data, step_index)) return datapoints + def _compute_required_indices( + self, + ep_idx: int, + step_indices: np.ndarray, + ) -> np.ndarray: + """Compute the union of all frame indices needed for a set of step indices. + + For sparse parquet loading, we need to know exactly which timesteps are + required across all step_indices and all modality delta_indices. This method + computes that union, handling padding by clamping to valid episode bounds. + + Video frames are included here too (the sparse DataFrame must contain rows + for both state/action AND video indices), but the actual video decode uses + the smaller set from _compute_required_video_indices. + + Args: + ep_idx: Episode index (for getting episode length bounds). + step_indices: Array of step indices assigned to this shard for this episode. + + Returns: + Sorted array of unique frame indices needed across all modalities. + """ + episode_length = self.episode_loader.get_episode_length(ep_idx) + required_indices = set() + + # Collect delta_indices from ALL modalities (state, action, video, language) + all_deltas = set() + for config in self.modality_configs.values(): + all_deltas.update(config.delta_indices) + + # Compute the union of all needed frame indices + for step_idx in step_indices: + for delta in all_deltas: + idx = step_idx + delta + if self.allow_padding: + # Clamp to valid range when padding is allowed + idx = max(0, min(idx, episode_length - 1)) + if 0 <= idx < episode_length: + required_indices.add(idx) + + return np.array(sorted(required_indices), dtype=np.int64) + + def _compute_required_video_indices( + self, + ep_idx: int, + step_indices: np.ndarray, + ) -> np.ndarray: + """Compute the union of video frame indices needed for a set of step indices. + + This deliberately uses ONLY the video modality's delta_indices, which are + typically just [0] (the current frame). This is the core of the sparse + loading optimization: by separating video indices from state/action indices, + we avoid decoding frames that are only needed for action labels. + + For example, with video delta_indices=[0] and action delta_indices=range(16), + a step_index of 100 needs video frame 100 but action frames 100-115. Without + this separation, we'd decode 16 video frames instead of 1. + + Args: + ep_idx: Episode index (for getting episode length bounds). + step_indices: Array of step indices assigned to this shard for this episode. + + Returns: + Sorted array of unique video frame indices to decode. + """ + if "video" not in self.modality_configs: + return np.array([], dtype=np.int64) + + episode_length = self.episode_loader.get_episode_length(ep_idx) + required_indices = set() + video_deltas = self.modality_configs["video"].delta_indices + + for step_idx in step_indices: + for delta in video_deltas: + idx = step_idx + delta + if self.allow_padding: + idx = max(0, min(idx, episode_length - 1)) + if 0 <= idx < episode_length: + required_indices.add(idx) + + return np.array(sorted(required_indices), dtype=np.int64) + def get_dataset_statistics(self) -> dict: """Get dataset statistics from the underlying episode loader.""" return self.episode_loader.get_dataset_statistics() + def get_percentile_statistics( + self, + consolidated_stats_path: str | Path | None = None, + ) -> dict | None: + """ + Get percentile statistics for this dataset from consolidated or per-dataset file. + + Percentile statistics contain q01, q02, q98, q99, mean, and std for both + state and action modalities. Action statistics are temporal-aware with + shape (horizon, dim), while state statistics have shape (dim,). + + Args: + consolidated_stats_path: Optional path to a consolidated stats JSON file + keyed by repo_id. If not provided, looks for + per-dataset stats in meta/temporal_stats.json. + + Returns: + Dictionary with structure: + { + "state": {key: {stat_type: values}}, + "action": {key: {stat_type: values}} + } + Returns None if no stats file is found. + + Raises: + FileNotFoundError: If require_stats=True and no stats file is found + KeyError: If using consolidated stats but this dataset's repo_id is missing + """ + # Try consolidated stats file first (if provided) + if consolidated_stats_path is not None: + consolidated_stats_path = Path(consolidated_stats_path) + if consolidated_stats_path.exists(): + with open(consolidated_stats_path, "r") as f: + all_stats = json.load(f) + repo_id = Path(self.dataset_path).name + if repo_id not in all_stats: + raise KeyError( + f"Dataset '{repo_id}' not found in consolidated stats file " + f"{consolidated_stats_path}. Available datasets: {list(all_stats.keys())}" + ) + return all_stats[repo_id] + + # Fall back to per-dataset stats file + per_dataset_stats_path = Path(self.dataset_path) / PERCENTILE_STATS_FILENAME + if per_dataset_stats_path.exists(): + with open(per_dataset_stats_path, "r") as f: + return json.load(f) + + return None + def get_initial_actions(self): """Get initial actions from the underlying episode loader.""" return self.episode_loader.get_initial_actions() diff --git a/gr00t/data/embodiment_tags.py b/gr00t/data/embodiment_tags.py index 67c6559..ffa7847 100755 --- a/gr00t/data/embodiment_tags.py +++ b/gr00t/data/embodiment_tags.py @@ -54,6 +54,150 @@ class EmbodimentTag(Enum): The Behavior R1 Pro robot. """ + ##### Open-H embodiment tags ##### + JHU_IMERSE_DVRK = "jhu_imerse_dvrk" + """ + The da Vinci Research Kit (dVRK) surgical robot. + Dual-arm (PSM1/PSM2) with REL_XYZ_ROT6D EEF control. + """ + + JHU_IMERSE_DVRK_MONO = "jhu_imerse_dvrk_mono" + """ + Monocular JHU dVRK surgical robot variant. + Uses only the left endoscope video stream with the same dual-arm state/action + representation as the standard dVRK embodiment. + """ + + JHU_LSCR_DVRK_MIRACLE = "jhu_lscr_dvrk_miracle" + """ + JHU LSCR MIRACLE datasets. + Dual-arm joint-angle control with RELATIVE actions (14D action, 12D joints + grippers). + Uses stereo endoscope cameras (15 Hz). + """ + + JHU_LSCR_DVRK_SMARTS = "jhu_lscr_dvrk_smarts" + """ + JHU LSCR SMARTS offline datasets. + Dual-arm joint-angle control with RELATIVE actions (PSM1/PSM2 joints + grippers). + Uses endoscope left/right plus side-view camera (10 Hz). + """ + + JHU_IMERSE_DVRK_STAR_IL = "jhu_imerse_star_il" + """ + JHU IMERSE star_IL dataset. + Single-arm KUKA with 7D pose actions (xyz + quat) and 8D state (7 joints + endo360). + Uses endoscope left + wrist left video. + """ + + CMR_VERSIUS = "cmr_versius" + """ + The CMR Versius surgical robot. + Dual-arm with REL_XYZ_ROT6D EEF control and energy buttons. + """ + + UCB_DVRK = "ucb_dvrk" + """ + The UCBerkeley dVRK debridement dataset. + Dual-arm (PSM1/PSM2) with cartesian EEF control (16D). + Uses REL_XYZ_ROT6D for pose actions with quaternion inputs. + """ + + OBUDA_DVRK = "obuda_dvrk" + """ + The Obuda University Open-H dVRK datasets. + Dual-arm (PSM1/PSM2) with REL_XYZ_ROT6D EEF control (16D). + Uses endoscope left + wrist left/right cameras; ECM is excluded in v1. + """ + + MOON_MAESTRO = "moon_maestro" + """ + The Moon Surgical Maestro assistant dataset. + Dual-arm robot with 18D joint state and 6D delta translation actions (xyz per arm). + """ + + UCSD_DVRK = "ucsd_dvrk" + """ + The UCSD surgical learning dataset. + Dual-arm (retraction + cutter) with delta EEF pose actions (16D). + Uses REL_XYZ_ROT6D for EEF pose actions with wxyz quaternion ordering. + """ + + STANFORD_DVRK_REAL = "stanford_dvrk_real" + """ + Stanford real-robot dVRK datasets (Needle Transfer, Tissue Retraction, Peg Transfer). + Dual-arm (PSM1/PSM2) with absolute EEF pose actions in Euler RPY (camera/ECM frame). + Uses REL_XYZ_ROT6D with Euler input and reference rotations. + """ + + SANOSCIENCE_SIM = "sanoscience_sim" + """ + The SanoScience simulated surgical robot. + 4 instruments with REL_XYZ_ROT6D EEF control (32D state/action). + Each instrument: xyz + quaternion (7D) + gripper (1D). + """ + + TUM_SONATA_FRANKA = "tum_sonata_franka" + """ + The TUM SonATA ultrasound sonography dataset. + Franka Panda robot with ultrasound probe end-effector. + REL_XYZ_ROT6D EEF control with Euler angles (6D xyz + RPY -> 9D xyz_rel + rot6d). + Includes force/torque sensor data and 3 camera views. + """ + + USTC_TORIN_TUODAO = "ustc_torin_tuodao" + """ + The USTC Torin surgical dataset. + Stereo endoscope video with 14D joint-angle state and 14D Cartesian delta actions + (xyz + roll/pitch/yaw + gripper per arm). Actions are treated as pass-through deltas + until a reliable Cartesian EEF state reference is available. + """ + + TUD_TUNDRA_UR5E = "tud_tundra_ur5e" + """ + The TUD TUNDRA UR5e surgical assistance dataset. + Stereo laparoscope video with grasping/retraction configured for REL_XYZ_ROT6D + EEF actions using absolute pose targets sourced from observation.state. + (Tundra Endoscope guidance requires an additional config.) + """ + + TURIN_MITIC_EX_VIVO = "turin_mitic_ex_vivo" + """ + The Turin MITIC ex vivo surgical dataset. + Dual-arm dVRK (PSM1/PSM2) with joint-angle state and absolute EEF pose actions + (xyz + quaternion per arm). REL_XYZ_ROT6D uses action[t=0] as the pose reference. + """ + + ROB_SURGICAL_BITRACK = "rob_surgical_bitrack" + """ + The Rob Surgical (bitrack) dataset. + Single endoscope video with 4-arm Cartesian EEF state/action (24D). + Uses REL_XYZ_ROT6D EEF control with Euler (RPY) rotation inputs. + """ + + HAMLYN_DVRK_15HZ = "hamlyn_dvrk_15hz" + """ + Hamlyn Centre dVRK surgical robot dataset - 15Hz tasks. + Tasks: knot_tying, needle_grasp_and_handover, peg_transfer, Suturing-1, + Suturing-2, suturing_single_loop_2, tissue_lifting. + Dual-arm with REL_XYZ_ROT6D EEF control (16D state/action). + Uses wxyz quaternion ordering (scalar-first). + """ + + HAMLYN_DVRK_30HZ = "hamlyn_dvrk_30hz" + """ + Hamlyn Centre dVRK surgical robot dataset - 30Hz tasks. + Tasks: suturing_single_loop_1, tissue_retraction. + Dual-arm with REL_XYZ_ROT6D EEF control (16D state/action). + Uses wxyz quaternion ordering (scalar-first). + """ + + POLYU_SIM = "polyu_sim" + """ + The PolyU OpenH_Dataset_full simulated surgical dataset. + Single-arm surgical robot with joint-angle state (10D) and cartesian pose state (7D). + Actions use REL_XYZ_ROT6D pose targets plus a gripper channel (1D). + """ + # New embodiment during post-training NEW_EMBODIMENT = "new_embodiment" """ diff --git a/gr00t/data/percentile_merge.py b/gr00t/data/percentile_merge.py new file mode 100644 index 0000000..dc6a087 --- /dev/null +++ b/gr00t/data/percentile_merge.py @@ -0,0 +1,356 @@ +"""Utilities for merging percentile statistics with piecewise-linear CDFs.""" + +from __future__ import annotations + +from typing import Callable + +import numpy as np + + +QUANTILE_KEY_ORDER = ("min", "q01", "q02", "q98", "q99", "max") +QUANTILE_PROBS = np.array([0.0, 0.01, 0.02, 0.98, 0.99, 1.0]) +TARGET_QUANTILE_KEYS = ("q01", "q02", "q98", "q99") +TARGET_QUANTILE_PROBS = np.array([0.01, 0.02, 0.98, 0.99]) +QUANTILE_PROB_BY_KEY = { + "min": 0.0, + "q01": 0.01, + "q02": 0.02, + "q98": 0.98, + "q99": 0.99, + "max": 1.0, +} + + +def _normalize_weights(weights: list[float] | np.ndarray) -> np.ndarray: + """Normalize dataset weights to sum to 1. + + Args: + weights: Raw dataset sampling weights. + + Returns: + Normalized weights as a numpy array that sums to 1. + + Raises: + ValueError: If weights sum to zero or contain negative values. + """ + weights_array = np.array(weights, dtype=np.float64) + if np.any(weights_array < 0): + raise ValueError("Dataset sampling weights must be non-negative.") + total = weights_array.sum() + if total <= 0: + raise ValueError("Dataset sampling weights must sum to a positive value.") + return weights_array / total + + +def _flatten_quantiles( + quantiles: dict[str, np.ndarray], + key_order: tuple[str, ...] = QUANTILE_KEY_ORDER, +) -> tuple[np.ndarray, tuple[int, ...]]: + """Flatten quantile arrays into a (num_points, num_features) matrix. + + Args: + quantiles: Dictionary mapping quantile keys to numpy arrays. + key_order: Ordered quantile keys to stack along the first dimension. + + Returns: + A tuple containing: + - stacked quantiles with shape (num_points, num_features) + - original feature shape (used for reshaping later) + """ + first = np.asarray(quantiles[key_order[0]], dtype=np.float64) + feature_shape = first.shape + stacked = np.stack( + [np.asarray(quantiles[key], dtype=np.float64).reshape(-1) for key in key_order], + axis=0, + ) + return stacked, feature_shape + + +def _validate_no_nan_or_none(values: np.ndarray, context: str) -> None: + """Validate that an array does not contain NaN/None-derived values. + + None values become NaN when cast to float arrays, so a NaN check catches both + invalid cases requested by callers. + + Args: + values: Numeric numpy array to validate. + context: Context string for error reporting. + + Raises: + ValueError: If the array contains NaN or None-derived values. + """ + if np.any(np.isnan(values)): + raise ValueError( + "Quantile merge inputs cannot contain NaN or None values. " + f"Invalid value detected in {context}." + ) + + +def _validate_quantile_monotonicity( + quantile_values: np.ndarray, + context: str, +) -> None: + """Validate that quantile values are non-decreasing per feature. + + Args: + quantile_values: Array of shape (num_points, num_features) containing + ordered quantile values per feature. + context: Context string for error reporting. + + Raises: + ValueError: If any feature has decreasing quantile values. + """ + _validate_no_nan_or_none(quantile_values, context=context) + diffs = np.diff(quantile_values, axis=0) + if np.any(diffs < 0): + raise ValueError( + "Quantile values must be non-decreasing for piecewise CDF merge. " + f"Invalid ordering detected in {context}." + ) + + +def _dedupe_piecewise_knots( + quantile_values: np.ndarray, + quantile_probs: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Collapse duplicate x-knots and keep monotonic probabilities. + + `np.interp` behaves ambiguously when x-knots are duplicated. This helper + collapses repeated quantile values to unique x locations and uses the + maximum probability seen at each location, then applies cumulative-max to + preserve non-decreasing CDF behavior. + + Args: + quantile_values: Non-decreasing quantile values, shape (num_points,). + quantile_probs: Corresponding quantile probabilities, shape (num_points,). + + Returns: + Tuple of (unique_quantile_values, monotonic_quantile_probs). + """ + unique_values, inverse_indices = np.unique(quantile_values, return_inverse=True) + dedup_probs = np.full(unique_values.shape, -np.inf, dtype=np.float64) + np.maximum.at(dedup_probs, inverse_indices, quantile_probs) + dedup_probs = np.maximum.accumulate(dedup_probs) + return unique_values, dedup_probs + + +def build_piecewise_cdf( + quantile_values: np.ndarray, + quantile_probs: np.ndarray, +) -> Callable[[np.ndarray], np.ndarray]: + """Create a piecewise-linear CDF for a single feature. + + Args: + quantile_values: Array of shape (num_points,) containing quantile values. + quantile_probs: Array of shape (num_points,) containing corresponding CDF + probabilities in ascending order. + + Returns: + Callable that maps input values to CDF probabilities. + """ + + quantile_values = np.asarray(quantile_values, dtype=np.float64) + quantile_probs = np.asarray(quantile_probs, dtype=np.float64) + _validate_no_nan_or_none(quantile_values, context="piecewise CDF quantile values") + _validate_no_nan_or_none(quantile_probs, context="piecewise CDF quantile probabilities") + if quantile_values.ndim != 1 or quantile_probs.ndim != 1: + raise ValueError( + "Piecewise CDF inputs must be 1D arrays. " + f"Got shapes values={quantile_values.shape}, probs={quantile_probs.shape}." + ) + if quantile_values.size == 0: + raise ValueError("Piecewise CDF requires at least one quantile knot.") + if quantile_values.size != quantile_probs.size: + raise ValueError( + "Piecewise CDF requires matching lengths for quantile values and probabilities. " + f"Got {quantile_values.size} values and {quantile_probs.size} probabilities." + ) + if np.any(np.diff(quantile_values) < 0): + raise ValueError("Piecewise CDF quantile values must be non-decreasing.") + if np.any(np.diff(quantile_probs) < 0): + raise ValueError("Piecewise CDF probabilities must be non-decreasing.") + if np.any((quantile_probs < 0.0) | (quantile_probs > 1.0)): + raise ValueError("Piecewise CDF probabilities must be within [0, 1].") + + unique_values, unique_probs = _dedupe_piecewise_knots(quantile_values, quantile_probs) + + def _cdf(x: np.ndarray) -> np.ndarray: + return np.interp(x, unique_values, unique_probs, left=0.0, right=1.0) + + return _cdf + + +def merge_cdfs( + cdfs: list[Callable[[np.ndarray], np.ndarray]], + weights: list[float] | np.ndarray, +) -> Callable[[np.ndarray], np.ndarray]: + """Create a weighted mixture of multiple CDFs. + + Args: + cdfs: List of CDF callables, each mapping values to probabilities. + weights: Sampling weights for each CDF. + + Returns: + Callable representing the weighted mixture CDF. + """ + if len(cdfs) != len(weights): + raise ValueError( + "CDF merge requires the same number of CDFs and weights. " + f"Got {len(cdfs)} CDFs and {len(weights)} weights." + ) + normalized_weights = _normalize_weights(weights) + + def _merged_cdf(x: np.ndarray) -> np.ndarray: + cdf_values = np.zeros_like(x, dtype=np.float64) + for cdf, weight in zip(cdfs, normalized_weights): + cdf_values += weight * cdf(x) + return cdf_values + + return _merged_cdf + + +def invert_cdf( + cdf: Callable[[np.ndarray], np.ndarray], + x_grid: np.ndarray, + target_probs: np.ndarray, +) -> np.ndarray: + """Invert a CDF over a predefined grid to obtain quantile values. + + Args: + cdf: Callable CDF to invert. + x_grid: Monotonic grid of x values to evaluate the CDF on. + target_probs: Target probabilities to invert (e.g., [0.01, 0.02, 0.98, 0.99]). + + Returns: + Array of quantile values corresponding to target_probs. + """ + if x_grid.size == 1: + return np.full_like(target_probs, x_grid[0], dtype=np.float64) + + cdf_values = np.maximum.accumulate(cdf(x_grid)) + cdf_values[0] = 0.0 + cdf_values[-1] = 1.0 + + quantiles = np.zeros_like(target_probs, dtype=np.float64) + for idx, prob in enumerate(target_probs): + if prob <= cdf_values[0]: + quantiles[idx] = x_grid[0] + continue + if prob >= cdf_values[-1]: + quantiles[idx] = x_grid[-1] + continue + right = np.searchsorted(cdf_values, prob, side="left") + left = max(right - 1, 0) + if cdf_values[right] == cdf_values[left]: + quantiles[idx] = x_grid[right] + else: + ratio = (prob - cdf_values[left]) / (cdf_values[right] - cdf_values[left]) + quantiles[idx] = x_grid[left] + ratio * (x_grid[right] - x_grid[left]) + return quantiles + + +def merge_piecewise_linear_quantiles( + per_dataset_quantiles: list[dict[str, np.ndarray]], + weights: list[float] | np.ndarray, + key_order: tuple[str, ...] = QUANTILE_KEY_ORDER, + target_keys: tuple[str, ...] = TARGET_QUANTILE_KEYS, + target_probs: np.ndarray = TARGET_QUANTILE_PROBS, +) -> dict[str, np.ndarray]: + """Merge per-dataset percentile statistics using piecewise-linear CDFs. + + Args: + per_dataset_quantiles: List of quantile dicts, each containing keys in key_order + and arrays of identical shape (e.g., (dim,) or (horizon, dim)). + weights: Sampling weights for each dataset. + key_order: Ordered quantile keys defining the CDF breakpoints. + target_keys: Output quantile keys to return. + target_probs: Target probabilities corresponding to target_keys. + + Returns: + Dictionary mapping target_keys to merged quantile arrays with the same shape + as the input quantile arrays. + """ + if not per_dataset_quantiles: + raise ValueError("Quantile merge requires at least one dataset quantile payload.") + if len(per_dataset_quantiles) != len(weights): + raise ValueError( + "Quantile merge requires the same number of quantile payloads and weights. " + f"Got {len(per_dataset_quantiles)} payloads and {len(weights)} weights." + ) + if len(key_order) == 0: + raise ValueError("Quantile merge requires a non-empty key_order.") + if len(set(key_order)) != len(key_order): + raise ValueError(f"Quantile merge key_order contains duplicate keys: {key_order}.") + if len(target_keys) == 0: + raise ValueError("Quantile merge requires at least one target quantile key.") + if len(set(target_keys)) != len(target_keys): + raise ValueError(f"Quantile merge target_keys contains duplicate keys: {target_keys}.") + + target_probs = np.asarray(target_probs, dtype=np.float64) + _validate_no_nan_or_none(target_probs, context="quantile merge target probabilities") + if target_probs.ndim != 1: + raise ValueError( + f"Quantile merge target_probs must be a 1D array. Got shape {target_probs.shape}." + ) + if len(target_keys) != target_probs.size: + raise ValueError( + "Quantile merge requires target_keys and target_probs to have equal lengths. " + f"Got {len(target_keys)} target keys and {target_probs.size} target probabilities." + ) + if np.any(np.diff(target_probs) < 0): + raise ValueError(f"Quantile merge target_probs must be non-decreasing. Got {target_probs}.") + if np.any((target_probs < 0.0) | (target_probs > 1.0)): + raise ValueError(f"Quantile merge target_probs must be within [0, 1]. Got {target_probs}.") + + try: + quantile_probs = np.array( + [QUANTILE_PROB_BY_KEY[key] for key in key_order], dtype=np.float64 + ) + except KeyError as exc: + raise ValueError(f"Unsupported quantile key in key_order: {exc}.") from exc + if np.any(np.diff(quantile_probs) < 0): + raise ValueError( + "Quantile probabilities must be non-decreasing. " + f"Got {quantile_probs} for key_order={key_order}." + ) + stacked_quantiles = [] + feature_shape = None + + for idx, quantiles in enumerate(per_dataset_quantiles): + missing = [key for key in key_order if key not in quantiles] + if missing: + raise ValueError( + f"Missing required quantiles for piecewise merge: {missing} in dataset index {idx}." + ) + stacked, shape = _flatten_quantiles(quantiles, key_order=key_order) + if feature_shape is None: + feature_shape = shape + elif feature_shape != shape: + raise ValueError( + "Quantile shapes must match across datasets. " + f"Expected {feature_shape}, got {shape} in dataset index {idx}." + ) + _validate_quantile_monotonicity( + stacked, + context=f"dataset index {idx}", + ) + stacked_quantiles.append(stacked) + + _, num_features = stacked_quantiles[0].shape + merged_quantiles = np.zeros((len(target_probs), num_features), dtype=np.float64) + + for feature_idx in range(num_features): + feature_quantiles = [dataset[:, feature_idx] for dataset in stacked_quantiles] + x_grid = np.unique(np.concatenate(feature_quantiles)) + if x_grid.size == 0: + raise ValueError( + f"Quantile merge failed: empty grid encountered for feature index {feature_idx}." + ) + cdfs = [build_piecewise_cdf(values, quantile_probs) for values in feature_quantiles] + merged_cdf = merge_cdfs(cdfs, weights) + merged_quantiles[:, feature_idx] = invert_cdf(merged_cdf, x_grid, target_probs) + + reshaped = { + key: merged_quantiles[idx].reshape(feature_shape) for idx, key in enumerate(target_keys) + } + return reshaped diff --git a/gr00t/data/split_utils.py b/gr00t/data/split_utils.py new file mode 100644 index 0000000..a6cad32 --- /dev/null +++ b/gr00t/data/split_utils.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import numpy as np + + +def load_info_json(dataset_path: Path | str) -> dict[str, Any]: + """ + Load a LeRobot dataset's `meta/info.json`. + + Args: + dataset_path: Path to the dataset root directory. + + Returns: + Parsed info.json contents as a dictionary. + + Raises: + FileNotFoundError: If info.json does not exist. + json.JSONDecodeError: If info.json is not valid JSON. + """ + dataset_path = Path(dataset_path) + info_path = dataset_path / "meta" / "info.json" + if not info_path.exists(): + raise FileNotFoundError(f"Missing info.json at {info_path}") + with open(info_path, "r") as f: + return json.load(f) + + +def parse_split_ranges(split_spec: str | list[str]) -> list[tuple[int, int]]: + """ + Parse split range specifications into inclusive (start, end) tuples. + + Supported formats: + - "0:499" (inclusive end) + - "0-499" (inclusive end) + - "42" (single episode index) + - Comma/whitespace-separated lists, e.g. "0:99, 200:299" + + Args: + split_spec: Split range specification string or list of strings. + + Returns: + List of inclusive (start, end) tuples. + + Raises: + ValueError: If a range is malformed or has end < start. + """ + if isinstance(split_spec, list): + raw_specs = split_spec + else: + raw_specs = [split_spec] + + ranges: list[tuple[int, int]] = [] + for raw in raw_specs: + parts = [p for p in raw.replace(",", " ").split() if p] + for part in parts: + if ":" in part: + start_str, end_str = part.split(":", 1) + elif "-" in part: + start_str, end_str = part.split("-", 1) + else: + start_str, end_str = part, part + + start = int(start_str) + end = int(end_str) + if end < start: + raise ValueError(f"Invalid split range '{part}': end < start") + ranges.append((start, end)) + + return ranges + + +def build_split_index_map(info: dict[str, Any]) -> dict[str, np.ndarray]: + """ + Build a mapping from split name to sorted episode indices. + + Args: + info: Parsed info.json dictionary containing a `splits` field. + + Returns: + Dict mapping lowercased split names to sorted numpy arrays of episode indices. + """ + splits = info.get("splits", {}) or {} + split_indices: dict[str, np.ndarray] = {} + for split_name, split_spec in splits.items(): + ranges = parse_split_ranges(split_spec) + indices: list[int] = [] + for start, end in ranges: + indices.extend(range(start, end + 1)) + split_indices[split_name.lower()] = np.array(sorted(set(indices)), dtype=int) + return split_indices + + +def resolve_episode_indices( + info: dict[str, Any], + include_splits: list[str] | None = None, + exclude_splits: list[str] | None = None, + *, + total_episodes: int | None = None, +) -> np.ndarray | None: + """ + Resolve an allowed episode index set from split allow/deny lists. + + Args: + info: Parsed info.json dictionary. + include_splits: Optional allowlist of split names to include. If provided, + only these splits are used as the base set. + exclude_splits: Optional denylist of split names to exclude. + total_episodes: Total episode count fallback if not present in info.json. + + Returns: + Sorted numpy array of allowed episode indices, or None if no filtering + is requested. + + Raises: + ValueError: If include_splits contains a split name not in info.json. + ValueError: If filtering is requested but no splits are defined. + """ + if not include_splits and not exclude_splits: + return None + + split_map = build_split_index_map(info) + if not split_map: + raise ValueError("Split filtering requested but info.json has no splits.") + + include_norm = [s.lower() for s in include_splits] if include_splits else [] + exclude_norm = [s.lower() for s in exclude_splits] if exclude_splits else [] + + if include_norm: + missing = [s for s in include_norm if s not in split_map] + if missing: + raise ValueError(f"Requested include_splits not found in info.json: {missing}") + allowed = np.concatenate([split_map[s] for s in include_norm], axis=0) + allowed_set = set(allowed.tolist()) + else: + total = info.get("total_episodes", total_episodes) + if total is None: + raise ValueError("total_episodes missing from info.json and not provided.") + allowed_set = set(range(int(total))) + + for split_name in exclude_norm: + if split_name not in split_map: + continue + allowed_set -= set(split_map[split_name].tolist()) + + if total_episodes is not None: + total = int(total_episodes) + allowed_set = {i for i in allowed_set if 0 <= i < total} + + if not allowed_set: + raise ValueError("Split filtering removed all episodes.") + + return np.array(sorted(allowed_set), dtype=int) diff --git a/gr00t/data/state_action/pose.py b/gr00t/data/state_action/pose.py index 8acba06..7487f36 100644 --- a/gr00t/data/state_action/pose.py +++ b/gr00t/data/state_action/pose.py @@ -13,6 +13,795 @@ PoseT = TypeVar("PoseT", bound="Pose") +# ============================================================================= +# Standalone rotation helper functions +# ============================================================================= + + +def rot6d_to_rotation_matrix(rot6d: np.ndarray) -> np.ndarray: + """ + Convert 6D rotation representation to 3x3 rotation matrix. + + The 6D representation consists of the first two columns of the rotation matrix, + flattened. This representation is continuous and avoids gimbal lock issues. + + Reference: Zhou et al., "On the Continuity of Rotation Representations in Neural Networks" + + Args: + rot6d: 6D rotation vector of shape (6,), representing first two columns + of rotation matrix flattened as [r00, r10, r20, r01, r11, r21] + + Returns: + 3x3 rotation matrix + + Example: + >>> rot6d = np.array([1, 0, 0, 0, 1, 0]) # Identity + >>> R = rot6d_to_rotation_matrix(rot6d) + >>> R.shape + (3, 3) + """ + # Reshape to (2, 3) then transpose to get columns as (3, 2) + rot6d = np.asarray(rot6d).reshape(2, 3).T + + # First two columns of the rotation matrix + col1 = rot6d[:, 0] + col2 = rot6d[:, 1] + + # Normalize first column + col1 = col1 / np.linalg.norm(col1) + + # Gram-Schmidt orthogonalization for second column + col2 = col2 - np.dot(col1, col2) * col1 + col2 = col2 / np.linalg.norm(col2) + + # Third column is cross product (ensures right-handed coordinate system) + col3 = np.cross(col1, col2) + + # Construct rotation matrix by stacking columns + rotation_matrix = np.column_stack([col1, col2, col3]) + + return rotation_matrix + + +def rotation_matrix_to_rot6d(rotation_matrix: np.ndarray) -> np.ndarray: + """ + Convert 3x3 rotation matrix to 6D rotation representation. + + Extracts the first two columns of the rotation matrix and flattens them. + This follows the convention from Zhou et al., "On the Continuity of + Rotation Representations in Neural Networks". + + Args: + rotation_matrix: 3x3 rotation matrix + + Returns: + 6D rotation vector of shape (6,), representing first two columns + flattened as [r00, r10, r20, r01, r11, r21] + + Example: + >>> R = np.eye(3) + >>> rot6d = rotation_matrix_to_rot6d(R) + >>> rot6d + array([1., 0., 0., 0., 1., 0.]) + """ + # Extract first two columns and flatten: [:, :2] gives (3, 2), .T gives (2, 3), flatten gives (6,) + return rotation_matrix[:, :2].T.flatten() + + +def rotation_matrices_to_rot6d(rotation_matrices: np.ndarray) -> np.ndarray: + """ + Batch convert rotation matrices to 6D rotation representations. + + Extracts the first two columns of each rotation matrix and flattens them. + This follows the convention from Zhou et al., "On the Continuity of + Rotation Representations in Neural Networks". + + Args: + rotation_matrices: Rotation matrices of shape (N, 3, 3) + + Returns: + 6D rotation vectors of shape (N, 6), where each vector contains + first two columns flattened as [r00, r10, r20, r01, r11, r21] + """ + # Extract first two columns: (N, 3, 3) -> (N, 3, 2) + # Transpose to (N, 2, 3) then reshape to (N, 6) + return rotation_matrices[:, :, :2].transpose(0, 2, 1).reshape(-1, 6) + + +def rot6ds_to_rotation_matrices(rot6ds: np.ndarray) -> np.ndarray: + """ + Batch convert 6D rotation representations to rotation matrices. + + Uses Gram-Schmidt orthogonalization to ensure valid rotation matrices. + The rot6d format stores the first two columns of the rotation matrix, + following the convention from Zhou et al., "On the Continuity of + Rotation Representations in Neural Networks". + + Args: + rot6ds: 6D rotation vectors of shape (N, 6), where each vector contains + the first two columns of a rotation matrix flattened as + [r00, r10, r20, r01, r11, r21] + + Returns: + Rotation matrices of shape (N, 3, 3) + """ + N = rot6ds.shape[0] + + # Reshape to (N, 2, 3) then transpose to get columns as (N, 3, 2) + cols = rot6ds.reshape(N, 2, 3).transpose(0, 2, 1) + + col1 = cols[:, :, 0] # (N, 3) + col2 = cols[:, :, 1] # (N, 3) + + # Gram-Schmidt orthogonalization (vectorized) + col1_norm = np.linalg.norm(col1, axis=1, keepdims=True) + col1 = col1 / np.maximum(col1_norm, 1e-8) + + # col2 = col2 - (col1 Β· col2) * col1 + dot = np.sum(col1 * col2, axis=1, keepdims=True) + col2 = col2 - dot * col1 + col2_norm = np.linalg.norm(col2, axis=1, keepdims=True) + col2 = col2 / np.maximum(col2_norm, 1e-8) + + # col3 = col1 Γ— col2 + col3 = np.cross(col1, col2) + + # Stack columns to form rotation matrices: (N, 3, 3) + return np.stack([col1, col2, col3], axis=2) + + +def quats_to_rotation_matrices(quats: np.ndarray, order: str = "xyzw") -> np.ndarray: + """ + Batch convert quaternions to rotation matrices. + + Args: + quats: Quaternions of shape (N, 4) + order: Quaternion ordering - "xyzw" (scipy default) or "wxyz" (scalar-first) + + Returns: + Rotation matrices of shape (N, 3, 3) + """ + if order.lower() == "wxyz": + # Convert from wxyz to xyzw: [w, x, y, z] -> [x, y, z, w] + quats = quats[:, [1, 2, 3, 0]] + return Rotation.from_quat(quats).as_matrix() + + +def rotation_matrices_to_quats(rotation_matrices: np.ndarray, order: str = "xyzw") -> np.ndarray: + """ + Batch convert rotation matrices to quaternions. + + Args: + rotation_matrices: Rotation matrices of shape (N, 3, 3) + order: Desired quaternion ordering - "xyzw" or "wxyz" + + Returns: + Quaternions of shape (N, 4) + """ + quats_xyzw = Rotation.from_matrix(rotation_matrices).as_quat() + if order.lower() == "wxyz": + # Convert from xyzw to wxyz: [x, y, z, w] -> [w, x, y, z] + return quats_xyzw[:, [3, 0, 1, 2]] + return quats_xyzw + + +def eulers_to_rotation_matrices(eulers: np.ndarray, seq: str = "xyz") -> np.ndarray: + """ + Batch convert Euler angles to rotation matrices. + + Args: + eulers: Euler angles of shape (N, 3) in radians (roll, pitch, yaw) + seq: Euler angle convention - "xyz" for extrinsic rotations (default) + This means: rotate about fixed X axis (roll), then Y (pitch), then Z (yaw) + + Returns: + Rotation matrices of shape (N, 3, 3) + + Note: + Euler angles can have discontinuities at Β±Ο€ (wraparound), but this conversion + handles them correctly because rotation matrices are continuous representations. + This is particularly important for REL_XYZ_ROT6D action conversion where + consecutive frames may have Euler angle jumps due to wraparound. + """ + return Rotation.from_euler(seq, eulers).as_matrix() + + +def rotation_matrices_to_eulers(rotation_matrices: np.ndarray, seq: str = "xyz") -> np.ndarray: + """ + Batch convert rotation matrices to Euler angles. + + Args: + rotation_matrices: Rotation matrices of shape (N, 3, 3) + seq: Euler angle convention - "xyz" for extrinsic rotations (default) + + Returns: + Euler angles of shape (N, 3) in radians (roll, pitch, yaw) + + Warning: + Euler angles have discontinuities at Β±Ο€. For continuous action representation, + prefer rot6d format which is always continuous. + """ + return Rotation.from_matrix(rotation_matrices).as_euler(seq) + + +def quat_to_rotation_matrix(quat: np.ndarray, order: str = "xyzw") -> np.ndarray: + """ + Convert quaternion to 3x3 rotation matrix. + + Args: + quat: Quaternion array of shape (4,) + order: Quaternion ordering - "xyzw" (scipy default) or "wxyz" (scalar-first) + + Returns: + 3x3 rotation matrix + + Example: + >>> quat = np.array([0, 0, 0, 1]) # Identity in xyzw + >>> R = quat_to_rotation_matrix(quat, order="xyzw") + >>> np.allclose(R, np.eye(3)) + True + """ + quat = np.asarray(quat) + if order.lower() == "wxyz": + # Convert from wxyz to xyzw (scipy uses xyzw) + quat = np.array([quat[1], quat[2], quat[3], quat[0]]) + return Rotation.from_quat(quat).as_matrix() + + +def rotation_matrix_to_quat(rotation_matrix: np.ndarray, order: str = "xyzw") -> np.ndarray: + """ + Convert 3x3 rotation matrix to quaternion. + + Args: + rotation_matrix: 3x3 rotation matrix + order: Desired quaternion ordering - "xyzw" or "wxyz" + + Returns: + Quaternion array of shape (4,) + + Example: + >>> R = np.eye(3) + >>> quat = rotation_matrix_to_quat(R, order="wxyz") + >>> quat + array([1., 0., 0., 0.]) + """ + quat_xyzw = Rotation.from_matrix(rotation_matrix).as_quat() + if order.lower() == "wxyz": + return np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]]) + return quat_xyzw + + +def convert_to_rel_xyz_rot6d( + action_data: np.ndarray, + eef_pose: np.ndarray, + input_rotation_format: str = "quat", + reference_rotation_format: str = "rot6d", + input_quat_order: str = "xyzw", + reference_quat_order: str = "xyzw", +) -> np.ndarray: + """ + Convert absolute action data to rel-xyz-rot6d representation. + + REL_XYZ_ROT6D means: + - Translation: relative to reference EEF position (delta from reference) + - Rotation: relative to reference orientation, expressed in 6D format + - Gripper: absolute (unchanged, handle gripper separately) + + This representation is useful for manipulation tasks where: + - Actions are predicted relative to a single reference state (the current observation) + - Position is a delta from the reference EEF position + - Rotation is relative to the reference EEF orientation + + Args: + action_data: Absolute action data of shape (H, D) where: + - H is the action horizon + - D = 3 (xyz) + 4 (quat) or D = 3 (xyz) + 6 (rot6d) + - Does NOT include gripper; handle gripper separately + eef_pose: Reference end-effector pose: + - Shape (9,) for rot6d format: xyz + rot6d + - Shape (7,) for quat format: xyz + quaternion + This is typically the current EEF pose from the observation. + input_rotation_format: Format of rotation in action_data: + - "quat": quaternion format + - "rot6d": 6D rotation representation + reference_rotation_format: Format of rotation in eef_pose: + - "quat": quaternion format (7D pose: xyz + quat) + - "rot6d": 6D rotation (9D pose: xyz + rot6d) + input_quat_order: Quaternion ordering for action_data when input_rotation_format="quat": + - "xyzw": Scalar-last (scipy convention, default) + - "wxyz": Scalar-first (e.g., Hamlyn dataset) + reference_quat_order: Quaternion ordering for eef_pose when reference_rotation_format="quat": + - "xyzw": Scalar-last (default) + - "wxyz": Scalar-first + + Returns: + REL_XYZ_ROT6D actions of shape (H, 9) with xyz (relative) + rot6d (relative) + + Example: + >>> action_data = np.random.randn(16, 7) # 16 steps, xyz + quat + >>> eef_pose = np.random.randn(7) # xyz + quat (xyzw) + >>> rel_xyz_rot6d = convert_to_rel_xyz_rot6d( + ... action_data, + ... eef_pose, + ... input_rotation_format="quat", + ... reference_rotation_format="quat", + ... ) + >>> rel_xyz_rot6d.shape + (16, 9) + """ + H, D = action_data.shape + + # Validate input dimensions + expected_dims = {"quat": 7, "rot6d": 9, "euler": 6} # xyz + rotation + if D not in expected_dims.values(): + raise ValueError( + f"Unexpected action dimension {D}. Expected 6 (xyz+euler), 7 (xyz+quat), or 9 (xyz+rot6d)" + ) + + # Extract reference position and rotation + ref_xyz = eef_pose[:3] + + # Parse reference rotation based on format + if reference_rotation_format == "quat": + ref_quat = eef_pose[3:7] + ref_R = quat_to_rotation_matrix(ref_quat, order=reference_quat_order) + elif reference_rotation_format == "rot6d": + ref_rot6d = eef_pose[3:9] + ref_R = rot6d_to_rotation_matrix(ref_rot6d) + elif reference_rotation_format == "euler": + # Euler angles in RPY order (roll, pitch, yaw) - radians + # Using 'xyz' extrinsic convention (standard robotics convention) + ref_euler = eef_pose[3:6] # (3,) + ref_R = Rotation.from_euler("xyz", ref_euler).as_matrix() + else: + raise ValueError(f"Unknown reference_rotation_format: {reference_rotation_format}") + + result = np.zeros((H, 9), dtype=np.float32) # Always output xyz + rot6d + + # Translation: vectorized subtraction (H, 3) - (3,) broadcasts correctly + result[:, :3] = action_data[:, :3] - ref_xyz + + # Rotation: batch convert to rotation matrices + if input_rotation_format == "quat": + action_quats = action_data[:, 3:7] # (H, 4) + action_Rs = quats_to_rotation_matrices(action_quats, order=input_quat_order) # (H, 3, 3) + elif input_rotation_format == "rot6d": + action_rot6ds = action_data[:, 3:9] # (H, 6) + action_Rs = rot6ds_to_rotation_matrices(action_rot6ds) # (H, 3, 3) + elif input_rotation_format == "euler": + # Euler angles in RPY order (roll, pitch, yaw) - radians + # Using 'xyz' extrinsic convention (standard robotics convention) + # This handles Euler wraparound correctly by going through rotation matrices + action_eulers = action_data[:, 3:6] # (H, 3) + action_Rs = eulers_to_rotation_matrices(action_eulers, seq="xyz") # (H, 3, 3) + else: + raise ValueError(f"Unknown input_rotation_format: {input_rotation_format}") + + # Relative rotation: R_ref^T @ R_action for all H matrices + # ref_R.T is (3, 3), action_Rs is (H, 3, 3) + # Use einsum for batch matrix multiply: (3, 3) @ (H, 3, 3) -> (H, 3, 3) + relative_Rs = np.einsum("ij,hjk->hik", ref_R.T, action_Rs) + + # Convert back to rot6d (batch) + result[:, 3:9] = rotation_matrices_to_rot6d(relative_Rs) + + return result + + +def convert_to_rel_xyz_rot6d_with_engagement( + action_data: np.ndarray, + eef_pose: np.ndarray, + engaged: np.ndarray, + input_rotation_format: str = "quat", + reference_rotation_format: str = "quat", + ref_engaged: bool = True, + input_quat_order: str = "xyzw", + reference_quat_order: str = "xyzw", +) -> np.ndarray: + """ + Compute rel-xyz-rot6d actions with engagement-aware delta re-integration. + + Instead of computing: action[t] = pose[t] - pose[ref] + This function computes: action[t] = sum(delta[i] * engaged[i] for i in ref+1..t) + + This correctly handles clutch scenarios in CMR Versius surgical robot data: + - Reference disengaged β†’ later engaged (no phantom jump from repositioning) + - Mid-horizon clutch events (disengaged deltas zeroed) + - Repositioning during clutch-out (not counted as arm motion) + + The key insight is that controller movement during disengaged periods doesn't + represent actual arm motion, so we should not include those deltas in the + cumulative relative action. + + Args: + action_data: Absolute action data of shape (T, D) where: + - T is the action horizon + - D = 3 (xyz) + 4 (quat) or D = 3 (xyz) + 6 (rot6d) + - Does NOT include gripper; handle gripper separately + eef_pose: Reference end-effector pose: + - Shape (7,) for quat format: xyz + quaternion + - Shape (9,) for rot6d format: xyz + rot6d + This is typically the current EEF pose from the observation (t=0). + engaged: Boolean engagement mask of shape (T,) + True where the surgeon is engaged (controlling the arm) + False where disengaged (clutched out / menu navigation) + input_rotation_format: Format of rotation in action_data: + - "quat": quaternion format + - "rot6d": 6D rotation representation + reference_rotation_format: Format of rotation in eef_pose: + - "quat": quaternion format (7D pose: xyz + quat) + - "rot6d": 6D rotation (9D pose: xyz + rot6d) + ref_engaged: Whether the reference frame (t=0 state) is engaged. + If False, the first delta (action[0] - eef_pose) is invalid and + will be masked out. This prevents "phantom jumps" from controller + repositioning while clutched out. Default is True for backward + compatibility, but should be set explicitly for CMR data. + input_quat_order: Quaternion ordering for action_data when input_rotation_format="quat": + - "xyzw": Scalar-last (scipy convention, default) + - "wxyz": Scalar-first (e.g., Hamlyn dataset) + reference_quat_order: Quaternion ordering for eef_pose when reference_rotation_format="quat": + - "xyzw": Scalar-last (default) + - "wxyz": Scalar-first + + Returns: + REL_XYZ_ROT6D actions of shape (T, 9) with xyz (relative) + rot6d (relative) + + Example: + >>> action_data = np.random.randn(16, 7) # 16 steps, xyz + quat + >>> eef_pose = np.random.randn(7) # xyz + quat (xyzw) + >>> engaged = np.array([1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=bool) + >>> rel_xyz_rot6d = convert_to_rel_xyz_rot6d_with_engagement( + ... action_data, + ... eef_pose, + ... engaged, + ... input_rotation_format="quat", + ... reference_rotation_format="quat", + ... ref_engaged=True, + ... ) + >>> rel_xyz_rot6d.shape + (16, 9) + """ + T = action_data.shape[0] + result = np.zeros((T, 9), dtype=np.float32) # xyz(3) + rot6d(6) + + # ========================================================================= + # STEP 1: Parse reference pose (single operation, not in loop) + # ========================================================================= + ref_xyz = eef_pose[:3] + + if reference_rotation_format == "quat": + ref_R = quat_to_rotation_matrix(eef_pose[3:7], order=reference_quat_order) + elif reference_rotation_format == "rot6d": + ref_R = rot6d_to_rotation_matrix(eef_pose[3:9]) + elif reference_rotation_format == "euler": + # Euler angles in RPY order (roll, pitch, yaw) - radians + ref_R = Rotation.from_euler("xyz", eef_pose[3:6]).as_matrix() + else: + raise ValueError(f"Unknown reference_rotation_format: {reference_rotation_format}") + + # ========================================================================= + # STEP 2: Batch convert all action rotations to matrices (vectorized) + # ========================================================================= + if input_rotation_format == "quat": + action_Rs = quats_to_rotation_matrices( + action_data[:, 3:7], order=input_quat_order + ) # (T, 3, 3) + elif input_rotation_format == "rot6d": + action_Rs = rot6ds_to_rotation_matrices(action_data[:, 3:9]) # (T, 3, 3) + elif input_rotation_format == "euler": + # Euler angles in RPY order (roll, pitch, yaw) - radians + # Handles wraparound correctly by going through rotation matrices + action_Rs = eulers_to_rotation_matrices(action_data[:, 3:6], seq="xyz") # (T, 3, 3) + else: + raise ValueError(f"Unknown input_rotation_format: {input_rotation_format}") + + # ========================================================================= + # STEP 3: Build validity mask (vectorized) + # Delta[t] is valid iff both endpoints are engaged: + # - t=0: ref_engaged AND engaged[0] + # - t>0: engaged[t-1] AND engaged[t] + # ========================================================================= + engaged_bool = engaged.astype(bool) + prev_engaged = np.concatenate([[ref_engaged], engaged_bool[:-1]]) # (T,) + delta_valid = prev_engaged & engaged_bool # (T,) + + # ========================================================================= + # STEP 4: Compute translation deltas and cumulative sum (fully vectorized) + # ========================================================================= + action_xyz = action_data[:, :3] # (T, 3) + + # Prepend reference xyz to compute deltas: delta[t] = xyz[t] - xyz[t-1] + all_xyz = np.vstack([ref_xyz[np.newaxis, :], action_xyz]) # (T+1, 3) + delta_xyz = np.diff(all_xyz, axis=0) # (T, 3) + + # Mask invalid deltas (set to zero) + masked_delta_xyz = delta_xyz * delta_valid[:, np.newaxis] # (T, 3) + + # Cumulative sum gives relative translation at each timestep + result[:, :3] = np.cumsum(masked_delta_xyz, axis=0) + + # ========================================================================= + # STEP 5: Compute rotation deltas (vectorized batch matrix multiply) + # delta_R[t] = prev_R[t].T @ curr_R[t] + # ========================================================================= + # Build array of previous rotations: [ref_R, action_Rs[0], ..., action_Rs[T-2]] + prev_Rs = np.concatenate([ref_R[np.newaxis, :, :], action_Rs[:-1]], axis=0) # (T, 3, 3) + + # Batch compute: delta_R[t] = prev_Rs[t].T @ action_Rs[t] + # Using einsum: 'tji,tjk->tik' transposes first matrix and multiplies + delta_Rs = np.einsum("tji,tjk->tik", prev_Rs, action_Rs) # (T, 3, 3) + + # For invalid deltas, set delta_R to identity (no rotation change) + identity = np.eye(3, dtype=np.float32) + delta_Rs = np.where(delta_valid[:, np.newaxis, np.newaxis], delta_Rs, identity) + + # ========================================================================= + # STEP 6: Cumulative rotation product (sequential - inherently not parallelizable) + # cumulative_R[t] = cumulative_R[t-1] @ delta_R[t] + # This is a prefix product of matrices, which must be computed sequentially. + # ========================================================================= + cumulative_R = np.eye(3, dtype=np.float32) + cumulative_Rs = np.zeros((T, 3, 3), dtype=np.float32) + + for t in range(T): + cumulative_R = cumulative_R @ delta_Rs[t] + cumulative_Rs[t] = cumulative_R + + # ========================================================================= + # STEP 7: Batch convert cumulative rotations to rot6d (vectorized) + # ========================================================================= + result[:, 3:9] = rotation_matrices_to_rot6d(cumulative_Rs) + + return result + + +def convert_from_rel_xyz_rot6d( + rel_xyz_rot6d_data: np.ndarray, + eef_pose: np.ndarray, + output_rotation_format: str = "rot6d", + reference_rotation_format: str = "rot6d", + output_quat_order: str = "xyzw", + reference_quat_order: str = "xyzw", +) -> np.ndarray: + """ + Convert rel-xyz-rot6d actions back to absolute representation. + + This is the inverse of convert_to_rel_xyz_rot6d. + + Args: + rel_xyz_rot6d_data: REL_XYZ_ROT6D actions of shape (H, 9) - xyz_rel + rot6d_rel + eef_pose: Reference end-effector pose: + - Shape (9,) for rot6d format: xyz + rot6d + - Shape (7,) for quat format: xyz + quaternion + This is typically the current EEF pose from the observation. + output_rotation_format: Desired output rotation format: + - "rot6d": 6D rotation (output shape: H, 9) + - "quat": quaternion format (output shape: H, 7) + reference_rotation_format: Format of rotation in eef_pose: + - "quat": quaternion format (7D pose: xyz + quat) + - "rot6d": 6D rotation (9D pose: xyz + rot6d) + output_quat_order: Quaternion ordering for output when output_rotation_format="quat": + - "xyzw": Scalar-last (scipy convention, default) + - "wxyz": Scalar-first + reference_quat_order: Quaternion ordering for eef_pose when reference_rotation_format="quat": + - "xyzw": Scalar-last (default) + - "wxyz": Scalar-first + + Returns: + Absolute actions with shape (H, 9) for rot6d or (H, 7) for quat + + Example: + >>> rel_xyz_rot6d = np.random.randn(16, 9) # 16 steps, xyz_rel + rot6d_rel + >>> eef_pose = np.random.randn(7) # xyz + quat (xyzw) + >>> absolute = convert_from_rel_xyz_rot6d( + ... rel_xyz_rot6d, + ... eef_pose, + ... output_rotation_format="rot6d", + ... reference_rotation_format="quat", + ... ) + >>> absolute.shape + (16, 9) + """ + H = rel_xyz_rot6d_data.shape[0] + + # Extract reference position and rotation + ref_xyz = eef_pose[:3] + + # Parse reference rotation based on format + if reference_rotation_format == "quat": + ref_quat = eef_pose[3:7] + ref_R = quat_to_rotation_matrix(ref_quat, order=reference_quat_order) + elif reference_rotation_format == "rot6d": + ref_rot6d = eef_pose[3:9] + ref_R = rot6d_to_rotation_matrix(ref_rot6d) + elif reference_rotation_format == "euler": + # Euler angles in RPY order (roll, pitch, yaw) - radians + ref_euler = eef_pose[3:6] + ref_R = Rotation.from_euler("xyz", ref_euler).as_matrix() + else: + raise ValueError(f"Unknown reference_rotation_format: {reference_rotation_format}") + + if output_rotation_format == "rot6d": + result = np.zeros((H, 9), dtype=np.float32) + elif output_rotation_format == "quat": + result = np.zeros((H, 7), dtype=np.float32) + elif output_rotation_format == "euler": + result = np.zeros((H, 6), dtype=np.float32) + else: + raise ValueError(f"Unknown output_rotation_format: {output_rotation_format}") + + # Translation: vectorized addition (H, 3) + (3,) broadcasts correctly + result[:, :3] = rel_xyz_rot6d_data[:, :3] + ref_xyz + + # Rotation: batch convert relative rot6d to rotation matrices + relative_rot6ds = rel_xyz_rot6d_data[:, 3:9] # (H, 6) + relative_Rs = rot6ds_to_rotation_matrices(relative_rot6ds) # (H, 3, 3) + + # Absolute rotation: R_action = R_ref @ R_relative for all H matrices + # ref_R is (3, 3), relative_Rs is (H, 3, 3) + # Use einsum for batch matrix multiply: (3, 3) @ (H, 3, 3) -> (H, 3, 3) + action_Rs = np.einsum("ij,hjk->hik", ref_R, relative_Rs) + + # Convert to output format (batch) + if output_rotation_format == "rot6d": + result[:, 3:9] = rotation_matrices_to_rot6d(action_Rs) + elif output_rotation_format == "quat": + result[:, 3:7] = rotation_matrices_to_quats(action_Rs, order=output_quat_order) + else: + # Note: Euler output has discontinuities at Β±Ο€. For continuous action + # representation during inference, prefer rot6d output. + result[:, 3:6] = rotation_matrices_to_eulers(action_Rs, seq="xyz") + + return result + + +# ============================================================================= +# Motion scaling functions for CMR Versius +# ============================================================================= + + +def scale_rot6d_by_angle(rot6d: np.ndarray, scale_factor: float) -> np.ndarray: + """ + Scale a rot6d representation by scaling its axis-angle magnitude. + + For relative rotations, this effectively scales the magnitude of the rotation + while preserving the axis of rotation. This is used for CMR Versius motion + scaling normalization. + + Args: + rot6d: 6D rotation representation of shape (6,) or (N, 6) + scale_factor: Factor to multiply the rotation angle by + + Returns: + Scaled rot6d representation with same shape as input + + Example: + >>> # Scale a 90-degree rotation by 0.5 to get 45-degree rotation + >>> rot6d = np.array([1, 0, 0, 0, 1, 0]) # ~identity + >>> scaled = scale_rot6d_by_angle(rot6d, 0.5) + """ + # Handle batch dimension + single_input = rot6d.ndim == 1 + if single_input: + rot6d = rot6d[np.newaxis, :] + + # Convert rot6d to rotation matrices + rot_matrices = rot6ds_to_rotation_matrices(rot6d) # (N, 3, 3) + + # Convert to axis-angle via scipy Rotation + rotations = Rotation.from_matrix(rot_matrices) + rotvecs = rotations.as_rotvec() # (N, 3) - axis * angle + + # Handle near-zero rotations to avoid numerical instability + angle_magnitudes = np.linalg.norm(rotvecs, axis=-1, keepdims=True) + epsilon = 1e-8 + + # Scale the rotation vector (this scales the angle while preserving axis) + # For very small rotations, keep them unchanged to avoid division issues + scaled_rotvecs = np.where( + angle_magnitudes > epsilon, + rotvecs * scale_factor, + rotvecs, # Keep unchanged for near-identity rotations + ) + + # Convert back to rotation matrices + scaled_rotations = Rotation.from_rotvec(scaled_rotvecs) + scaled_matrices = scaled_rotations.as_matrix() # (N, 3, 3) + + # Convert back to rot6d + scaled_rot6d = rotation_matrices_to_rot6d(scaled_matrices) # (N, 6) + + if single_input: + scaled_rot6d = scaled_rot6d[0] + + return scaled_rot6d + + +def apply_motion_scaling_to_rel_xyz_rot6d( + rel_xyz_rot6d_data: np.ndarray, + translation_scaling: float, + rotation_scaling: float, +) -> np.ndarray: + """ + Apply motion scaling normalization to rel-xyz-rot6d actions. + + This converts from "hand-controller-space" to "instrument-space" by + multiplying by the scaling factors. This ensures that the same visual + outcome (instrument movement) produces the same normalized action + regardless of the motion scaling settings used by the surgeon. + + Scaling relationship: instrument_movement = hand_movement * scaling + + Args: + rel_xyz_rot6d_data: REL_XYZ_ROT6D actions of shape (H, 9) - xyz_rel + rot6d_rel + These are in hand-controller-space (raw relative movements) + translation_scaling: Translation scaling factor (e.g., 0.333, 0.5, 1.0) + rotation_scaling: Rotation scaling factor (e.g., 1.0, 1.5, 2.0) + + Returns: + Motion-scaled REL_XYZ_ROT6D actions of shape (H, 9) in instrument-space + + Example: + >>> # With translationScaling=0.333, hand moves 3cm, instrument moves 1cm + >>> # After scaling: action represents 1cm instrument movement + >>> scaled = apply_motion_scaling_to_rel_xyz_rot6d( + ... rel_xyz_rot6d_data, translation_scaling=0.333, rotation_scaling=2.0 + ... ) + """ + result = np.zeros_like(rel_xyz_rot6d_data) + + # Scale translation: multiply by translationScaling to get instrument-space + # (larger controller movement -> smaller instrument movement when scaling < 1) + result[:, :3] = rel_xyz_rot6d_data[:, :3] * translation_scaling + + # Scale rotation: multiply by rotationScaling to get instrument-space + # Use axis-angle scaling for proper rotation scaling + rot6d = rel_xyz_rot6d_data[:, 3:9] + result[:, 3:9] = scale_rot6d_by_angle(rot6d, rotation_scaling) + + return result + + +def unapply_motion_scaling_from_rel_xyz_rot6d( + scaled_rel_xyz_rot6d_data: np.ndarray, + translation_scaling: float, + rotation_scaling: float, +) -> np.ndarray: + """ + Reverse motion scaling to convert back to hand-controller-space. + + Used during inference to convert model predictions (in instrument-space) + back to actual hand controller commands that the robot system expects. + + Args: + scaled_rel_xyz_rot6d_data: Motion-scaled REL_XYZ_ROT6D actions of shape (H, 9) + These are in instrument-space (model predictions) + translation_scaling: Translation scaling factor + rotation_scaling: Rotation scaling factor + + Returns: + Original REL_XYZ_ROT6D actions of shape (H, 9) in hand-controller-space + + Example: + >>> # Convert instrument-space prediction back to hand controller commands + >>> hand_controller_action = unapply_motion_scaling_from_rel_xyz_rot6d( + ... model_prediction, translation_scaling=0.333, rotation_scaling=2.0 + ... ) + """ + result = np.zeros_like(scaled_rel_xyz_rot6d_data) + + # Divide translation by scaling factor to get hand-controller-space + result[:, :3] = scaled_rel_xyz_rot6d_data[:, :3] / translation_scaling + + # Divide rotation by scaling factor (scale by 1/rotation_scaling) + rot6d = scaled_rel_xyz_rot6d_data[:, 3:9] + result[:, 3:9] = scale_rot6d_by_angle(rot6d, 1.0 / rotation_scaling) + + print("WARNING: Batch size > 1 may have problems due to global use of scaling") + + return result + + def invert_transformation(T: NDArray[np.float64]) -> NDArray[np.float64]: """ Invert a homogeneous transformation matrix. @@ -412,45 +1201,30 @@ def _rot6d_to_matrix(rot6d: np.ndarray) -> np.ndarray: """ Convert 6D rotation representation to rotation matrix. + Delegates to module-level rot6d_to_rotation_matrix function. + Args: rot6d: 6D rotation as (6,) array - first two rows of rotation matrix flattened Returns: Rotation matrix (3, 3) """ - rot6d = rot6d.reshape(2, 3) - - # First two rows - row1 = rot6d[0] - row2 = rot6d[1] - - # Normalize first row - row1 = row1 / np.linalg.norm(row1) - - # Gram-Schmidt orthogonalization for second row - row2 = row2 - np.dot(row1, row2) * row1 - row2 = row2 / np.linalg.norm(row2) - - # Third row is cross product - row3 = np.cross(row1, row2) - - # Construct rotation matrix - rotation_matrix = np.vstack([row1, row2, row3]) - - return rotation_matrix + return rot6d_to_rotation_matrix(rot6d) @staticmethod def _matrix_to_rot6d(rotation_matrix: np.ndarray) -> np.ndarray: """ Convert rotation matrix to 6D rotation representation. + Delegates to module-level rotation_matrix_to_rot6d function. + Args: rotation_matrix: Rotation matrix (3, 3) Returns: 6D rotation - (6,) array (first two rows flattened) """ - return rotation_matrix[:2, :].flatten() + return rotation_matrix_to_rot6d(rotation_matrix) def _set_rotation( self, diff --git a/gr00t/data/state_action/state_action_processor.py b/gr00t/data/state_action/state_action_processor.py index f8adec4..9724b68 100644 --- a/gr00t/data/state_action/state_action_processor.py +++ b/gr00t/data/state_action/state_action_processor.py @@ -4,20 +4,33 @@ Handles: - State normalization (min/max, mean/std, sin/cos encoding) - Action normalization -- Absolute <-> Relative action representation conversion +- Absolute <-> Relative <-> REL_XYZ_ROT6D action representation conversion - Action processing with state dependency """ from copy import deepcopy -from gr00t.configs.data.embodiment_configs import ( +from gr00t.data.state_action.action_chunking import EndEffectorActionChunk, JointActionChunk +from gr00t.data.state_action.pose import ( + EndEffectorPose, + JointPose, + apply_motion_scaling_to_rel_xyz_rot6d, + convert_from_rel_xyz_rot6d, + convert_to_rel_xyz_rot6d, + convert_to_rel_xyz_rot6d_with_engagement, + unapply_motion_scaling_from_rel_xyz_rot6d, +) +from gr00t.data.types import ( + CMR_ENGAGED_LEFT_KEY, + CMR_ENGAGED_RIGHT_KEY, + EEF_XYZ_ROT6D_DIM, + ROT6D_IDENTITY, + XYZ_DIM, ActionFormat, ActionRepresentation, ActionType, ModalityConfig, ) -from gr00t.data.state_action.action_chunking import EndEffectorActionChunk, JointActionChunk -from gr00t.data.state_action.pose import EndEffectorPose, JointPose from gr00t.data.utils import ( apply_sin_cos_encoding, nested_dict_to_numpy, @@ -37,14 +50,15 @@ class StateActionProcessor: Handles: - State normalization (min/max, mean/std, sin/cos encoding) - Action normalization - - Absolute <-> Relative action representation conversion + - Absolute <-> Relative <-> REL_XYZ_ROT6D action representation conversion - Action processing with state dependency + - CMR clutch-aware zeroing of disengaged actions """ def __init__( self, modality_configs: dict[str, dict[str, ModalityConfig]], - statistics: (dict[str, dict[str, dict[str, dict[str, list[float]]]]] | None) = None, + statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]] | None = None, use_percentiles: bool = False, clip_outliers: bool = True, apply_sincos_state_encoding: bool = False, @@ -61,11 +75,12 @@ def __init__( statistics: Optional nested dict with structure: {embodiment_tag: {modality: {joint_group: {stat_type: values}}}} where modality in ["state", "action", "relative_action"] - and stat_type in ["min", "max", "mean", "std", "q01", "q99"] + and stat_type in ["min", "max", "mean", "std", "q01", "q02", "q98", "q99"] Example: {"gr1": {"state": {"left_arm": {"min": [...], "max": [...], ...}}}} use_percentiles: Whether to use percentiles (q01/q99) instead of min/max clip_outliers: Whether to clip normalized values to [-1, 1] apply_sincos_state_encoding: Global flag to enable sin/cos encoding for states + use_relative_action: Whether to convert actions to relative representation """ self.modality_configs = parse_modality_configs(modality_configs) self.statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]] = {} @@ -77,11 +92,14 @@ def __init__( # Normalization parameters computed from statistics self.norm_params: dict[str, dict[str, dict[str, dict[str, np.ndarray]]]] = {} # Format: norm_params[embodiment_tag][modality][joint_group][stat_type] - # where stat_type in ["min", "max", "mean", "std", "dim"] + # where stat_type in ["min", "max", "mean", "std", "q01", "q02", "q98", "q99", "dim"] if statistics is not None: self.set_statistics(statistics) + # Initialize CMR detection for clutch-aware zeroing + self._init_cmr_action_indices() + self.train() def train(self): @@ -90,6 +108,191 @@ def train(self): def eval(self): self.training = False + def _init_cmr_action_indices(self) -> None: + """Pre-compute whether each embodiment has CMR clutch engagement keys.""" + self._cmr_action_indices: dict[str, bool] = {} + + for embodiment_tag, modality_configs in self.modality_configs.items(): + state_config = modality_configs.get("state") + if state_config is None: + self._cmr_action_indices[embodiment_tag] = False + continue + + state_keys = state_config.modality_keys or [] + self._cmr_action_indices[embodiment_tag] = ( + CMR_ENGAGED_LEFT_KEY in state_keys and CMR_ENGAGED_RIGHT_KEY in state_keys + ) + + def _zero_disengaged_actions( + self, + action: dict[str, np.ndarray], + state: dict[str, np.ndarray], + embodiment_tag: str, + ) -> dict[str, np.ndarray]: + """Zero out action targets for disengaged timesteps (CMR clutch-aware zeroing). + + For each timestep where hapticengaged=False for a particular arm, zero out all + action components (pose, energy) for that arm, however absolute actions with + hold_through_clutch=True will hold the last engaged value (Like the Gripper). + This teaches the model to predict "no movement" when the surgeon is clutched out. + + This happens PRE-normalization because zero is meaningful in raw action space. + + Temporal Engagement: + When hapticengaged_left/right are in the ACTION modality_keys, they provide + per-timestep engagement data [T, 1]. This enables proper zeroing for mid-horizon + clutch events. The hapticengaged keys are removed from the action dict after + zeroing (before normalization). + + Error (No Temporal Engagement): + If hapticengaged is only in STATE (not action), a ValueError is raised. + Temporal engagement in the action dict is required for proper clutch handling. + + Args: + action: Dict mapping joint_group -> raw action values [T, D] + state: Dict mapping joint_group -> raw state values + embodiment_tag: Embodiment identifier + + Returns: + Action dict with disengaged timesteps zeroed out (hapticengaged keys removed) + """ + # Check if this embodiment has CMR clutch engagement keys + if not self._cmr_action_indices.get(embodiment_tag, False): + return action # Not CMR data, pass through unchanged + + # Use centralized CMR engagement key constants from types.py + engaged_left_key = CMR_ENGAGED_LEFT_KEY + engaged_right_key = CMR_ENGAGED_RIGHT_KEY + + # Check for temporal engagement in action dict (preferred) + has_temporal_engagement = engaged_left_key in action and engaged_right_key in action + + if has_temporal_engagement: + # Temporal Zeroing/Sample-and-Hold: Per-timestep engagement + # + # Behavior determined by action representation and hold_through_clutch flag: + # - RELATIVE/REL_XYZ_ROT6D actions: always zero (no movement) + # - ABSOLUTE actions with hold_through_clutch=True: hold last engaged value + # (e.g., gripper holding a needle - don't drop it) + # - ABSOLUTE actions with hold_through_clutch=False: zero (fail-safe) + # (e.g., energy button - stop firing for safety) + # + # Default is fail-safe (zero) unless explicitly marked to hold. + engaged_left = action[engaged_left_key].astype(bool).flatten() # [T] + engaged_right = action[engaged_right_key].astype(bool).flatten() # [T] + + # Get modality keys and action configs for checking representation type + modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys + action_configs = self.modality_configs[embodiment_tag]["action"].action_configs + action_keys = [ + k for k in modality_keys if k not in [engaged_left_key, engaged_right_key] + ] + + # Build mapping from key to action config for checking hold_through_clutch + action_configs_by_key = {} + if action_configs: + for k, cfg in zip(modality_keys, action_configs): + action_configs_by_key[k] = cfg + + # Process each action key individually to handle sample-and-hold correctly + for key in action_keys: + action_data = action[key] # [T, D] + T = action_data.shape[0] + + # Determine behavior based on action config + cfg = action_configs_by_key.get(key) + is_relative = cfg is not None and cfg.rep in [ + ActionRepresentation.RELATIVE, + ActionRepresentation.REL_XYZ_ROT6D, + ] + + # For ABSOLUTE actions, check hold_through_clutch flag (defined in ActionConfig with default=False) + # - True: hold last engaged value (e.g., gripper holding a needle) + # - False: zero during clutch-out (fail-safe default, e.g., energy button) + should_hold = cfg is not None and not is_relative and cfg.hold_through_clutch + + # Determine which arm this key belongs to + is_left = "left_" in key + is_right = "right_" in key + + for t in range(T): + # Get engagement for this arm at this timestep + engaged = True # Default: engaged + if is_left: + engaged = engaged_left[t] + elif is_right: + engaged = engaged_right[t] + + if not engaged: + if should_hold: + # Sample-and-hold for ABSOLUTE actions with hold_through_clutch=True + if t > 0: + action_data[t] = action_data[t - 1] # Hold previous value + else: + # t=0 fallback - use robot state instead of controller value + # When disengaged, controller may have snapped to a different position + # (e.g., gripper opened) but robot is still holding (e.g., needle) + if cfg.state_key and cfg.state_key in state: + # Use robot state (e.g., 1.0 for closed gripper) + state_val = state[cfg.state_key] + if state_val.ndim > 1: + state_val = state_val.flatten() + # Handle dimension mismatch (state may be different shape) + if state_val.shape[0] >= action_data.shape[1]: + action_data[t] = state_val[: action_data.shape[1]] + else: + action_data[t] = state_val[0] + else: + # Error: hold_through_clutch=True requires state_key for t=0 fallback + # + # When t=0 is disengaged and hold_through_clutch=True, we need a fallback + # value from the robot's actual state. Without state_key, we cannot look + # up this value, and using the controller value is risky (it may have + # snapped to a different position, e.g., gripper opened to 0.0). + raise ValueError( + f"Action key '{key}' has hold_through_clutch=True but no state_key defined. " + f"This is required when t=0 is disengaged to provide a fallback value. " + f"Add state_key to the ActionConfig for '{key}'." + ) + else: + # Zero for RELATIVE actions OR ABSOLUTE with hold_through_clutch=False + # Use format-aware zeroing for rot6d components + # rot6d=[0,0,0,0,0,0] is invalid; identity=[1,0,0,0,1,0] means "no rotation" + if cfg is not None and cfg.format == ActionFormat.XYZ_ROT6D: + # XYZ_ROT6D format: xyz (3D) + rot6d (6D) = 9D + # Zero translation, identity rotation + action_data[t, :XYZ_DIM] = 0.0 + action_data[t, XYZ_DIM:EEF_XYZ_ROT6D_DIM] = ROT6D_IDENTITY + elif cfg is not None and cfg.format == ActionFormat.ROT6D: + # Pure rotation format: use identity + action_data[t] = ROT6D_IDENTITY + else: + # Default: zero (XYZ, DEFAULT, or other formats) + action_data[t] = 0.0 + + action[key] = action_data + + # Remove hapticengaged keys from action dict (not needed for model) + del action[engaged_left_key] + del action[engaged_right_key] + + # CMR data requires per-timestep engagement for proper clutch handling. + # Reference-frame-only engagement is insufficient because it cannot detect + # mid-horizon clutch events, doesn't support sample-and-hold for ABSOLUTE + # actions, and doesn't use format-aware zeroing (identity rotation for rot6d). + # + # Fix: Include hapticengaged_left/right in action.modality_keys and + # action.pass_through_keys in your modality config. + elif engaged_left_key in state and engaged_right_key in state: + raise ValueError( + f"CMR data detected (hapticengaged_left/right in state) but temporal engagement " + f"not found in action dict for embodiment '{embodiment_tag}'. " + f"Add '{engaged_left_key}' and '{engaged_right_key}' to action.modality_keys " + f"and action.pass_through_keys in your modality config for proper clutch handling." + ) + + return action + def set_statistics( self, statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]], @@ -110,41 +313,84 @@ def set_statistics( self._compute_normalization_parameters() def _compute_normalization_parameters(self) -> None: - """Compute and cache normalization parameters from statistics for all embodiments and modalities.""" + """Compute and cache normalization parameters from statistics for all embodiments and modalities. + + Statistics are keyed by embodiment_tag. + + Parameters stored per joint group: + - min, max: For minmax normalization (state) + - mean, std: For meanstd normalization + - q02, q98: For percentile normalization (2nd/98th percentiles) + - q01, q99: For alternate percentile normalization (1st/99th percentiles) + - dim: Feature dimension + """ for embodiment_tag in self.statistics: + stats_data = self.statistics[embodiment_tag] + self.norm_params[embodiment_tag] = {} for modality in ["state", "action"]: - if modality not in self.statistics[embodiment_tag]: + if modality not in stats_data: continue self.norm_params[embodiment_tag][modality] = {} - for joint_group, stats in self.statistics[embodiment_tag][modality].items(): + for joint_group, stats in stats_data[modality].items(): + # Legacy min/max handling if self.use_percentiles: min_vals = np.array(stats["q01"]) max_vals = np.array(stats["q99"]) else: - min_vals = np.array(stats["min"]) - max_vals = np.array(stats["max"]) + min_vals = np.array(stats["min"]) if "min" in stats else None + max_vals = np.array(stats["max"]) if "max" in stats else None mean_vals = np.array(stats["mean"]) std_vals = np.array(stats["std"]) - # Compute range, ensuring it's not zero - range_vals = max_vals - min_vals - range_vals = np.maximum(range_vals, 1e-8) + # New percentile parameters (q02/q98) for percentile normalization + q02_vals = np.array(stats["q02"]) if "q02" in stats else None + q98_vals = np.array(stats["q98"]) if "q98" in stats else None + q01_vals = np.array(stats["q01"]) if "q01" in stats else None + q99_vals = np.array(stats["q99"]) if "q99" in stats else None + + # Determine dimension from available stats + # For temporal stats, shape is (horizon, dim), so use last axis + if q02_vals is not None: + dim = q02_vals.shape[-1] if q02_vals.ndim > 1 else q02_vals.shape[0] + elif q01_vals is not None: + dim = q01_vals.shape[-1] if q01_vals.ndim > 1 else q01_vals.shape[0] + elif min_vals is not None: + dim = min_vals.shape[-1] if min_vals.ndim > 1 else min_vals.shape[0] + else: + dim = mean_vals.shape[-1] if mean_vals.ndim > 1 else mean_vals.shape[0] self.norm_params[embodiment_tag][modality][joint_group] = { - "min": min_vals, - "max": max_vals, - "dim": np.array(range_vals.shape[0]), + "dim": np.array(dim), "mean": mean_vals, "std": std_vals, } + # Add min/max if available + if min_vals is not None: + self.norm_params[embodiment_tag][modality][joint_group]["min"] = min_vals + if max_vals is not None: + self.norm_params[embodiment_tag][modality][joint_group]["max"] = max_vals + + # Add q02/q98 if available + if q02_vals is not None: + self.norm_params[embodiment_tag][modality][joint_group]["q02"] = q02_vals + if q98_vals is not None: + self.norm_params[embodiment_tag][modality][joint_group]["q98"] = q98_vals + if q01_vals is not None: + self.norm_params[embodiment_tag][modality][joint_group]["q01"] = q01_vals + if q99_vals is not None: + self.norm_params[embodiment_tag][modality][joint_group]["q99"] = q99_vals + # Override absolute action stats with relative stats where specified - if "action" in self.modality_configs[embodiment_tag]: + if ( + embodiment_tag in self.modality_configs + and "action" in self.modality_configs[embodiment_tag] + ): modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys action_configs = self.modality_configs[embodiment_tag]["action"].action_configs @@ -154,22 +400,36 @@ def _compute_normalization_parameters(self) -> None: action_config.rep == ActionRepresentation.RELATIVE and self.use_relative_action ): - if "relative_action" not in self.statistics[embodiment_tag]: + if "relative_action" not in stats_data: raise ValueError( f"Relative action statistics required for embodiment '{embodiment_tag}' " f"but 'relative_action' not found in statistics" ) - if key not in self.statistics[embodiment_tag]["relative_action"]: + if key not in stats_data["relative_action"]: raise ValueError( f"Relative action statistics required for key '{key}' " f"in embodiment '{embodiment_tag}' but not found" ) action_dim = self.norm_params[embodiment_tag]["action"][key]["dim"] self.norm_params[embodiment_tag]["action"][key] = nested_dict_to_numpy( - self.statistics[embodiment_tag]["relative_action"][key] + stats_data["relative_action"][key] ) self.norm_params[embodiment_tag]["action"][key]["dim"] = action_dim + # For EEF XYZ_ROT6D actions, stats are xyz-only (3D) but model output is 9D + # (xyz + rot6d). This dimension is used to SPLIT the concatenated model output, + # so it must match the model's internal format, not the final output format. + # The rot6dβ†’quat conversion (if needed) happens later in unapply_action. + if ( + action_config.type == ActionType.EEF + and action_config.format == ActionFormat.XYZ_ROT6D + and key in self.norm_params[embodiment_tag].get("action", {}) + ): + # Model internal format is always XYZ_ROT6D (9D) regardless of output format + self.norm_params[embodiment_tag]["action"][key]["dim"] = np.array( + EEF_XYZ_ROT6D_DIM + ) + def apply_state( self, state: dict[str, np.ndarray], @@ -181,7 +441,7 @@ def apply_state( Args: state: Dict mapping joint_group -> raw state values Shape per group: (..., D) where D is state dimension - embodiment_tag: Embodiment identifier (e.g., "gr1") + embodiment_tag: Embodiment identifier (e.g., "gr1") for modality config lookup Returns: Dict mapping joint_group -> processed state values @@ -198,6 +458,12 @@ def apply_state( if state_config and hasattr(state_config, "sin_cos_embedding_keys"): sin_cos_keys = state_config.sin_cos_embedding_keys + # Get pass-through keys (keys that should not be normalized) + pass_through_keys = None + state_config = self.modality_configs[embodiment_tag].get("state") + if state_config and hasattr(state_config, "pass_through_keys"): + pass_through_keys = state_config.pass_through_keys + for joint_group in self.modality_configs[embodiment_tag]["state"].modality_keys: if joint_group not in state: raise KeyError( @@ -222,8 +488,12 @@ def apply_state( normalized = normalize_values_meanstd(state[joint_group], params) normalized_values[joint_group] = normalized - # Strategy 3: Min/max normalization to [-1, 1] - else: + # Strategy 3: Explicit pass-through for auxiliary keys (e.g., motion scaling factors) + elif pass_through_keys and joint_group in pass_through_keys: + normalized_values[joint_group] = state[joint_group] + + # Strategy 4: Min/max normalization to [-1, 1] + elif joint_group in self.norm_params.get(embodiment_tag, {}).get("state", {}): params = self.norm_params[embodiment_tag]["state"][joint_group] normalized = normalize_values_minmax(state[joint_group], params) @@ -232,6 +502,14 @@ def apply_state( normalized_values[joint_group] = normalized + # Error: key not in any normalization strategy and not in pass_through_keys + else: + raise KeyError( + f"No normalization stats found for state key '{joint_group}' " + f"in embodiment '{embodiment_tag}'. Either add this key to pass_through_keys in the state " + f"ModalityConfig, or regenerate statistics to include this key." + ) + return normalized_values def unapply_state( @@ -261,6 +539,12 @@ def unapply_state( if state_config and hasattr(state_config, "sin_cos_embedding_keys"): sin_cos_keys = state_config.sin_cos_embedding_keys + # Get pass-through keys (keys that should not be normalized) + pass_through_keys = None + state_config = self.modality_configs[embodiment_tag].get("state") + if state_config and hasattr(state_config, "pass_through_keys"): + pass_through_keys = state_config.pass_through_keys + for joint_group in self.modality_configs[embodiment_tag]["state"].modality_keys: if joint_group not in state: raise KeyError( @@ -288,13 +572,25 @@ def unapply_state( unnormalized = unnormalize_values_meanstd(state[joint_group], params) unnormalized_values[joint_group] = unnormalized + # Explicit pass-through for auxiliary keys (e.g., motion scaling factors) + elif pass_through_keys and joint_group in pass_through_keys: + unnormalized_values[joint_group] = state[joint_group] + # Reverse min/max normalization - else: + elif joint_group in self.norm_params.get(embodiment_tag, {}).get("state", {}): params = self.norm_params[embodiment_tag]["state"][joint_group] unnormalized_values[joint_group] = unnormalize_values_minmax( state[joint_group], params ) + # Error: key not in any normalization strategy and not in pass_through_keys + else: + raise KeyError( + f"No normalization stats found for state key '{joint_group}' " + f"in embodiment '{embodiment_tag}'. Either add this key to pass_through_keys in the state " + f"ModalityConfig, or regenerate statistics to include this key." + ) + return unnormalized_values def apply_action( @@ -360,26 +656,161 @@ def apply_action( action_format=action_config.format, ) + elif ( + action_config.rep == ActionRepresentation.REL_XYZ_ROT6D + and self.use_relative_action + ): + # REL_XYZ_ROT6D: translation relative to current EEF, rotation relative to initial + if state is None: + raise ValueError( + f"State dict required for REL_XYZ_ROT6D action processing of key '{key}' " + f"in embodiment '{embodiment_tag}'" + ) + + state_key = action_config.state_key if action_config.state_key else key + + if state_key not in state: + raise KeyError( + f"Reference state key '{state_key}' not found in state dict " + f"for embodiment '{embodiment_tag}'" + ) + + # Reference EEF pose - single state with delta_indices=[0] + # Format depends on reference_rotation_format: + # - "quat": xyz (3) + quat (4) = 7D + # - "rot6d": xyz (3) + rot6d (6) = 9D + # State shape is (1, dim), so use [0] to get the single reference + eef_pose = state[state_key][0] + + # Check if per-timestep engagement data is available in action dict (CMR data) + # This enables engagement-aware delta re-integration to handle clutch events + # Use centralized CMR engagement key constants from types.py + if "left" in key: + engaged_key = CMR_ENGAGED_LEFT_KEY + elif "right" in key: + engaged_key = CMR_ENGAGED_RIGHT_KEY + else: + engaged_key = None + engaged = action.get(engaged_key) if engaged_key else None + + if engaged is not None: + # Get reference frame engagement from STATE dict. + # State has delta_indices=[0], so it contains the reference frame's engagement. + # If ref is disengaged, first delta (action[0] - eef_pose) is invalid. + ref_engaged = True # Default for non-CMR data + if engaged_key and engaged_key in state: + # State engagement is shape [1, 1] or [1,], extract scalar + ref_eng_val = state[engaged_key] + if ref_eng_val.ndim > 0: + ref_eng_val = ref_eng_val.flatten()[0] + ref_engaged = bool(ref_eng_val > 0.5) + + # Engagement-aware delta re-integration correctly handles: + # ref disengaged β†’ later engaged (no phantom jump), + # mid-horizon clutch events (disengaged deltas zeroed), and repositioning + action[key] = convert_to_rel_xyz_rot6d_with_engagement( + action_data=action[key], + eef_pose=eef_pose, + engaged=engaged.flatten().astype(bool), + input_rotation_format=action_config.input_rotation_format, + reference_rotation_format=action_config.reference_rotation_format, + ref_engaged=ref_engaged, + input_quat_order=action_config.input_quat_order, + reference_quat_order=action_config.reference_quat_order, + ) + else: + # Standard conversion for non-CMR data or when engagement not available + # Input: action with xyz + quat (7D) or xyz + rot6d (9D) + # Output: xyz_rel + rot6d_rel (9D) + action[key] = convert_to_rel_xyz_rot6d( + action_data=action[key], + eef_pose=eef_pose, + input_rotation_format=action_config.input_rotation_format, + reference_rotation_format=action_config.reference_rotation_format, + input_quat_order=action_config.input_quat_order, + reference_quat_order=action_config.reference_quat_order, + ) + + # Apply motion scaling if configured (CMR Versius specific) + # This converts from hand-controller-space to instrument-space + if action_config.translation_scaling_key or action_config.rotation_scaling_key: + trans_scale = 1.0 + rot_scale = 1.0 + + if action_config.translation_scaling_key: + if action_config.translation_scaling_key not in state: + raise KeyError( + f"Translation scaling key '{action_config.translation_scaling_key}' " + f"not found in state dict for embodiment '{embodiment_tag}'" + ) + trans_scale = float(state[action_config.translation_scaling_key][0]) + + if action_config.rotation_scaling_key: + if action_config.rotation_scaling_key not in state: + raise KeyError( + f"Rotation scaling key '{action_config.rotation_scaling_key}' " + f"not found in state dict for embodiment '{embodiment_tag}'" + ) + rot_scale = float(state[action_config.rotation_scaling_key][0]) + + action[key] = apply_motion_scaling_to_rel_xyz_rot6d( + action[key], trans_scale, rot_scale + ) + + # Step 1.5: Zero out actions for disengaged timesteps (CMR clutch-aware zeroing) + # This happens PRE-normalization because zero is meaningful in raw action space + if state is not None: + action = self._zero_disengaged_actions(action, state, embodiment_tag) + # Step 2: Normalize actions + # Skip pass_through_keys - they're used for processing but not sent to model normalized_values = {} - for joint_group in modality_keys: + action_config_obj = self.modality_configs[embodiment_tag]["action"] + pass_through_keys = set(action_config_obj.pass_through_keys or []) + + for idx, joint_group in enumerate(modality_keys): + # Skip pass_through_keys - they've been used for processing (e.g., clutch-aware zeroing) + # and should not be normalized or included in the output + if joint_group in pass_through_keys: + continue + if joint_group not in action: raise KeyError( f"Joint group '{joint_group}' not found in action dict for embodiment '{embodiment_tag}'" ) params = self.norm_params[embodiment_tag]["action"][joint_group] - if ( - self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys is not None - and joint_group - in self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys - ): - normalized = normalize_values_meanstd(action[joint_group], params) - else: - normalized = normalize_values_minmax(action[joint_group], params) - if self.clip_outliers: - normalized = np.clip(normalized, -1.0, 1.0) + # For EEF actions with XYZ_ROT6D format, only normalize xyz (rot6d is already bounded) + action_config = ( + action_configs[idx] if action_configs and idx < len(action_configs) else None + ) + is_eef_xyz_rot6d = ( + action_config is not None + and action_config.type == ActionType.EEF + and action_config.format == ActionFormat.XYZ_ROT6D + ) + + normalization_type = ( + action_config.normalization_type if action_config else "temporal_meanstd" + ) + + if normalization_type == "skip": + # Pass through without normalization (e.g., for pass-through action keys) + normalized = action[joint_group] + elif is_eef_xyz_rot6d: + # XYZ_ROT6D: normalize xyz with meanstd, pass through rot6d + action_data = action[joint_group] + xyz = action_data[..., :3] + rot6d = action_data[..., 3:] + normalized_xyz = normalize_values_meanstd(xyz, params) + normalized = np.concatenate([normalized_xyz, rot6d], axis=-1) + elif normalization_type == "minmax": + normalized = normalize_values_minmax(action[joint_group], params) + else: + # "temporal_meanstd" and "meanstd" both use meanstd normalization + # When stats have shape (horizon, dim), normalization is per-timestep + normalized = normalize_values_meanstd(action[joint_group], params) normalized_values[joint_group] = normalized @@ -414,10 +845,20 @@ def unapply_action( ValueError: If state is None but required for relative->absolute conversion """ # Step 1: Unnormalize actions + # Skip pass_through_keys - they were used for processing (e.g., clutch-aware zeroing) + # during apply_action and are not present in the normalized action dict. unnormalized_values = {} modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys + action_configs = self.modality_configs[embodiment_tag]["action"].action_configs + action_config_obj = self.modality_configs[embodiment_tag]["action"] + pass_through_keys = set(action_config_obj.pass_through_keys or []) + + for idx, joint_group in enumerate(modality_keys): + # Skip pass_through_keys - they were stripped during apply_action and are not + # present in the normalized output. Mirrors the skip in apply_action. + if joint_group in pass_through_keys: + continue - for joint_group in modality_keys: if joint_group not in action: raise KeyError( f"Joint group '{joint_group}' not found in action dict for embodiment '{embodiment_tag}'" @@ -426,22 +867,45 @@ def unapply_action( params = self.norm_params[embodiment_tag]["action"][joint_group] group_values = action[joint_group] - if ( - self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys is not None - and joint_group - in self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys - ): - unnormalized = unnormalize_values_meanstd(group_values, params) - else: + # For EEF actions with XYZ_ROT6D format, only unnormalize xyz (rot6d was not normalized) + action_config = ( + action_configs[idx] if action_configs and idx < len(action_configs) else None + ) + is_eef_xyz_rot6d = ( + action_config is not None + and action_config.type == ActionType.EEF + and action_config.format == ActionFormat.XYZ_ROT6D + ) + + normalization_type = ( + action_config.normalization_type if action_config else "temporal_meanstd" + ) + + if normalization_type == "skip": + unnormalized = group_values + elif is_eef_xyz_rot6d: + # Only unnormalize xyz (first 3 dims), pass through rot6d (last 6 dims) + xyz = group_values[..., :3] + rot6d = group_values[..., 3:] + unnormalized_xyz = unnormalize_values_meanstd(xyz, params) + unnormalized = np.concatenate([unnormalized_xyz, rot6d], axis=-1) + elif normalization_type == "minmax": unnormalized = unnormalize_values_minmax(group_values, params) + else: + # "temporal_meanstd" and "meanstd" both use meanstd + unnormalized = unnormalize_values_meanstd(group_values, params) unnormalized_values[joint_group] = unnormalized # Step 2: Convert relative actions to absolute (if needed) - action_configs = self.modality_configs[embodiment_tag]["action"].action_configs if action_configs is not None: for key, action_config in zip(modality_keys, action_configs): + # Skip pass_through_keys - not present in unnormalized_values + # (they were stripped during apply_action normalization) + if key in pass_through_keys: + continue + if action_config.rep == ActionRepresentation.RELATIVE and self.use_relative_action: if state is None: raise ValueError( @@ -490,6 +954,97 @@ def unapply_action( else: unnormalized_values[key] = absolute_actions[0] + elif ( + action_config.rep == ActionRepresentation.REL_XYZ_ROT6D + and self.use_relative_action + ): + # Convert REL_XYZ_ROT6D back to absolute + if state is None: + raise ValueError( + f"State dict required for REL_XYZ_ROT6D->absolute conversion of key '{key}' " + f"in embodiment '{embodiment_tag}'" + ) + + state_key = action_config.state_key if action_config.state_key else key + + if state_key not in state: + raise KeyError( + f"Reference state key '{state_key}' not found in state dict " + f"for embodiment '{embodiment_tag}'" + ) + + rel_xyz_rot6d_action = unnormalized_values[key] + + # Handle batched and unbatched cases + is_batched = rel_xyz_rot6d_action.ndim == 3 + if not is_batched: + assert rel_xyz_rot6d_action.ndim == 2 + reference_state = state[state_key] + if reference_state.ndim == 2: + reference_state = reference_state[None, :] + rel_xyz_rot6d_action = rel_xyz_rot6d_action[None, :] + else: + reference_state = state[state_key] + if reference_state.ndim == 2: + reference_state = reference_state[None, :] + + # Get motion scaling values if configured (CMR Versius specific) + has_motion_scaling = ( + action_config.translation_scaling_key or action_config.rotation_scaling_key + ) + num_samples = reference_state.shape[0] + trans_scales = np.ones(num_samples, dtype=np.float32) + rot_scales = np.ones(num_samples, dtype=np.float32) + + if has_motion_scaling: + if action_config.translation_scaling_key: + trans_scales = self._extract_scaling_values( + state, + action_config.translation_scaling_key, + num_samples, + is_batched, + embodiment_tag, + ) + if action_config.rotation_scaling_key: + rot_scales = self._extract_scaling_values( + state, + action_config.rotation_scaling_key, + num_samples, + is_batched, + embodiment_tag, + ) + + # Convert batched REL_XYZ_ROT6D actions to absolute + # State has delta_indices=[0], so s[0] is the single reference state + absolute_actions = [] + for sample_idx, (s, a) in enumerate(zip(reference_state, rel_xyz_rot6d_action)): + # Reference EEF pose - single state with delta_indices=[0] + eef_pose = s[0] + + # Unapply motion scaling first (convert instrument-space to hand-controller-space) + if has_motion_scaling: + a = unapply_motion_scaling_from_rel_xyz_rot6d( + a, + float(trans_scales[sample_idx]), + float(rot_scales[sample_idx]), + ) + + absolute_action = convert_from_rel_xyz_rot6d( + rel_xyz_rot6d_data=a, + eef_pose=eef_pose, + # Output in original input format (e.g., quat xyzw) to match GT data + output_rotation_format=action_config.input_rotation_format, + reference_rotation_format=action_config.reference_rotation_format, + output_quat_order=action_config.input_quat_order, + reference_quat_order=action_config.reference_quat_order, + ) + absolute_actions.append(absolute_action) + + if is_batched: + unnormalized_values[key] = np.stack(absolute_actions, axis=0) + else: + unnormalized_values[key] = absolute_actions[0] + return unnormalized_values def apply( @@ -567,10 +1122,11 @@ def get_state_dim(self, embodiment_tag: str, include_sincos_expansion: bool = Fa include_sincos_expansion: If True, accounts for sin/cos encoding doubling dimensions Returns: - Total state dimension across all joint groups + Total state dimension across model-input state joint groups. """ total_dim = 0 state_config = self.modality_configs[embodiment_tag]["state"] + pass_through_keys = set(state_config.pass_through_keys or []) # Get sin/cos embedding keys if enabled sin_cos_keys = set() @@ -578,6 +1134,8 @@ def get_state_dim(self, embodiment_tag: str, include_sincos_expansion: bool = Fa sin_cos_keys = set(state_config.sin_cos_embedding_keys) for joint_group in state_config.modality_keys: + if joint_group in pass_through_keys: + continue base_dim = self.norm_params[embodiment_tag]["state"][joint_group]["dim"].item() # Sin/cos encoding doubles the dimension @@ -590,19 +1148,66 @@ def get_state_dim(self, embodiment_tag: str, include_sincos_expansion: bool = Fa def get_action_dim(self, embodiment_tag: str) -> int: """ - Get total action dimension. + Get total action dimension (excluding pass_through_keys). + + Pass-through keys are used for data processing (e.g., clutch-aware zeroing) but + are not sent to the model. They are excluded from the dimension calculation. Args: embodiment_tag: Embodiment identifier Returns: - Total action dimension across all joint groups + Total action dimension across all joint groups (excluding pass_through_keys) """ total_dim = 0 - for joint_group in self.modality_configs[embodiment_tag]["action"].modality_keys: + action_config = self.modality_configs[embodiment_tag]["action"] + pass_through_keys = set(action_config.pass_through_keys or []) + + for joint_group in action_config.modality_keys: + # Skip pass_through_keys - they're not sent to the model + if joint_group in pass_through_keys: + continue total_dim += self.norm_params[embodiment_tag]["action"][joint_group]["dim"].item() return total_dim + @staticmethod + def _extract_scaling_values( + state: dict[str, np.ndarray], + scaling_key: str, + num_samples: int, + is_batched: bool, + embodiment_tag: str, + ) -> np.ndarray: + """Extract per-sample scaling values from state dict. + + Args: + state: Dict mapping joint_group -> state values + scaling_key: Key in state dict for the scaling factor + num_samples: Number of samples (batch size or 1) + is_batched: Whether the data is batched + embodiment_tag: Embodiment identifier (for error messages) + + Returns: + Array of shape (num_samples,) with per-sample scaling values + """ + if scaling_key not in state: + raise KeyError( + f"Scaling key '{scaling_key}' not found in state dict " + f"for embodiment '{embodiment_tag}'" + ) + scale_values = np.asarray(state[scaling_key], dtype=np.float32) + if is_batched and scale_values.ndim > 0 and scale_values.shape[0] == num_samples: + return np.array( + [float(np.asarray(scale_values[i]).reshape(-1)[0]) for i in range(num_samples)], + dtype=np.float32, + ) + else: + return np.full( + num_samples, + float(scale_values.reshape(-1)[0]), + dtype=np.float32, + ) + def _convert_to_relative_action( self, action: np.ndarray, diff --git a/gr00t/data/stats.py b/gr00t/data/stats.py index 8bcb285..0668100 100644 --- a/gr00t/data/stats.py +++ b/gr00t/data/stats.py @@ -1,28 +1,55 @@ #!/usr/bin/env python """ Calculate dataset statistics for LeRobot datasets. -Note: Please update the `gr00t/configs/data/embodiment_configs.py` file with the correct modality configurations for the dataset you are using before running this script. + +This module provides functions for computing normalization statistics including: +- Standard statistics (mean, std, min, max, q01, q99) for backward compatibility +- Temporal-aware percentile statistics for action normalization with shape (horizon, dim) +- Non-temporal percentile statistics for state normalization with shape (dim,) + +The temporal statistics are designed for percentile-based normalization using 2nd/98th +percentiles, which is more robust to outliers than min-max scaling. + +Note: Please update the `gr00t/configs/data/embodiment_configs.py` file with the correct +modality configurations for the dataset you are using before running this script. Usage: python gr00t/data/stats.py Args: dataset_path: Path to the dataset. - embodiment_tag: Embodiment tag to use to load modality configurations from `gr00t/configs/data/embodiment_configs.py`. + embodiment_tag: Embodiment tag to use to load modality configurations from + `gr00t/configs/data/embodiment_configs.py`. """ +from concurrent.futures import ProcessPoolExecutor import json from pathlib import Path +from typing import Any import numpy as np +import open_h.embodiments # noqa: F401 - registers Open-H embodiment configs import pandas as pd from tqdm import tqdm from gr00t.configs.data.embodiment_configs import MODALITY_CONFIGS from gr00t.data.dataset.lerobot_episode_loader import LeRobotEpisodeLoader +from gr00t.data.split_utils import load_info_json, resolve_episode_indices from gr00t.data.state_action.action_chunking import EndEffectorActionChunk, JointActionChunk -from gr00t.data.state_action.pose import EndEffectorPose, JointPose -from gr00t.data.types import ActionRepresentation, ActionType, EmbodimentTag, ModalityConfig +from gr00t.data.state_action.pose import ( + EndEffectorPose, + JointPose, + apply_motion_scaling_to_rel_xyz_rot6d, + convert_to_rel_xyz_rot6d, +) +from gr00t.data.step_filtering import compute_valid_step_indices_parallel +from gr00t.data.types import ( + ActionFormat, + ActionRepresentation, + ActionType, + EmbodimentTag, + ModalityConfig, +) from gr00t.data.utils import to_json_serializable @@ -30,6 +57,32 @@ LE_ROBOT_INFO_FILENAME = "meta/info.json" LE_ROBOT_STATS_FILENAME = "meta/stats.json" LE_ROBOT_REL_STATS_FILENAME = "meta/relative_stats.json" +LE_ROBOT_TEMPORAL_STATS_FILENAME = "meta/temporal_stats.json" + + +def _resolve_episode_ids( + total_episodes: int, + episode_indices: np.ndarray | list[int] | None = None, + max_episodes: int = -1, + context: str = "percentile stats", +) -> list[int]: + """Resolve episode IDs from optional indices and max episode cap.""" + if episode_indices is None: + episode_ids = list(range(total_episodes)) + else: + indices = np.unique(np.asarray(episode_indices, dtype=int)) + invalid = [int(i) for i in indices if i < 0 or i >= total_episodes] + if invalid: + raise ValueError( + f"Invalid episode indices for {context}: {invalid}. " + f"Valid range: 0-{total_episodes - 1}" + ) + episode_ids = indices.tolist() + + if max_episodes != -1: + episode_ids = episode_ids[:max_episodes] + + return episode_ids def calculate_dataset_statistics( @@ -62,6 +115,16 @@ def calculate_dataset_statistics( dataset_statistics = {} if features is None: features = list(all_low_dim_data.columns) + else: + # Some datasets list float features in info.json that never appear in parquet. + # Skip those to avoid KeyError during stats calculation. + missing_features = [feature for feature in features if feature not in all_low_dim_data] + if missing_features: + print( + "WARNING: skipping missing features during stats generation: " + f"{sorted(missing_features)}" + ) + features = [feature for feature in features if feature not in missing_features] for le_modality in features: print(f"Computing statistics for {le_modality}...") np_data = np.vstack( @@ -95,7 +158,14 @@ def check_stats_validity(dataset_path: Path | str, features: list[str]): return True -def generate_stats(dataset_path: Path | str): +def generate_stats(dataset_path: Path | str, episode_indices: np.ndarray | list[int] | None = None): + """Generate stats.json for low-dimensional float features. + + Args: + dataset_path: Path to the LeRobot dataset root. + episode_indices: Optional list/array of episode indices to include. If None, + compute stats over the full dataset. + """ dataset_path = Path(dataset_path) print(f"Generating stats for {str(dataset_path)}") lowdim_features = [] @@ -104,10 +174,34 @@ def generate_stats(dataset_path: Path | str): for feature in le_features: if "float" in le_features[feature]["dtype"]: lowdim_features.append(feature) - if check_stats_validity(dataset_path, lowdim_features): - return - parquet_files = list(dataset_path.glob(LE_ROBOT_DATA_FILENAME)) + if episode_indices is None: + parquet_files = list(dataset_path.glob(LE_ROBOT_DATA_FILENAME)) + else: + # Use loader to validate indices and resolve chunked parquet paths + loader = LeRobotEpisodeLoader( + dataset_path, + modality_configs={}, + skip_video=True, + require_stats=False, + ) + total_episodes = len(loader) + indices = np.unique(np.asarray(episode_indices, dtype=int)) + invalid = [int(i) for i in indices if i < 0 or i >= total_episodes] + if invalid: + raise ValueError( + f"Invalid episode indices for stats: {invalid}. Valid range: 0-{total_episodes - 1}" + ) + + parquet_files = [ + dataset_path + / loader.data_path_pattern.format( + episode_chunk=int(ep_idx) // loader.chunk_size, + episode_index=int(ep_idx), + ) + for ep_idx in indices + ] + stats = calculate_dataset_statistics(parquet_files, lowdim_features) stats_path = dataset_path / LE_ROBOT_STATS_FILENAME with open(stats_path, "w") as f: @@ -115,7 +209,21 @@ def generate_stats(dataset_path: Path | str): class RelativeActionLoader: - def __init__(self, dataset_path: Path | str, embodiment_tag: EmbodimentTag, action_key: str): + def __init__( + self, + dataset_path: Path | str, + embodiment_tag: EmbodimentTag, + action_key: str, + episode_indices: np.ndarray | list[int] | None = None, + ): + """Load episodes for RELATIVE action stats computation. + + Args: + dataset_path: Path to dataset root directory. + embodiment_tag: Embodiment tag for modality config lookup. + action_key: Action key to compute relative stats for. + episode_indices: Optional subset of episodes to include. + """ self.dataset_path = Path(dataset_path) self.modality_configs: dict[str, ModalityConfig] = {} self.action_key = action_key @@ -136,15 +244,36 @@ def __init__(self, dataset_path: Path | str, embodiment_tag: EmbodimentTag, acti delta_indices=MODALITY_CONFIGS[embodiment_tag.value]["state"].delta_indices, modality_keys=[state_key], ) - # Check state-action consistency - assert ( - self.modality_configs["state"].delta_indices[-1] - == self.modality_configs["action"].delta_indices[0] + # Check state-action consistency: + # Allow action horizons that start AFTER the reference state (e.g., [1..H]) + # so we can compute relative deltas to future targets. + state_delta = self.modality_configs["state"].delta_indices[-1] + action_start = self.modality_configs["action"].delta_indices[0] + assert state_delta <= action_start, ( + "State reference index must be <= first action delta index. " + f"Got state_delta={state_delta}, action_start={action_start}." ) self.loader = LeRobotEpisodeLoader(dataset_path, self.modality_configs) + self.episode_indices = None + if episode_indices is not None: + indices = np.unique(np.asarray(episode_indices, dtype=int)) + invalid = [int(i) for i in indices if i < 0 or i >= len(self.loader)] + if invalid: + raise ValueError( + f"Invalid episode indices for relative stats: {invalid}. " + f"Valid range: 0-{len(self.loader) - 1}" + ) + self.episode_indices = indices + + def _resolve_episode_index(self, local_index: int) -> int: + """Map a local episode index to a dataset episode index.""" + if self.episode_indices is None: + return local_index + return int(self.episode_indices[local_index]) def load_relative_actions(self, trajectory_id: int) -> list[np.ndarray]: - df = self.loader[trajectory_id] + episode_index = self._resolve_episode_index(trajectory_id) + df = self.loader[episode_index] # OPTIMIZATION: Extract columns once and convert to numpy arrays # This eliminates repeated DataFrame.__getitem__ and Series.__getitem__ calls @@ -183,7 +312,9 @@ def load_relative_actions(self, trajectory_id: int) -> list[np.ndarray]: return trajectories def __len__(self) -> int: - return len(self.loader) + if self.episode_indices is None: + return len(self.loader) + return len(self.episode_indices) def calculate_stats_for_key( @@ -191,8 +322,20 @@ def calculate_stats_for_key( embodiment_tag: EmbodimentTag, group_key: str, max_episodes: int = -1, + episode_indices: np.ndarray | list[int] | None = None, ) -> dict: - loader = RelativeActionLoader(dataset_path, embodiment_tag, group_key) + """Compute stats for a single RELATIVE action key. + + Args: + dataset_path: Path to dataset root directory. + embodiment_tag: Embodiment tag for modality config lookup. + group_key: Action key to compute stats for. + max_episodes: Optional cap on episodes to process (-1 means all). + episode_indices: Optional subset of episode indices to include. + """ + loader = RelativeActionLoader( + dataset_path, embodiment_tag, group_key, episode_indices=episode_indices + ) trajectories = [] for episode_id in tqdm(range(len(loader)), desc=f"Loading trajectories for key {group_key}"): if max_episodes != -1 and episode_id >= max_episodes: @@ -208,7 +351,19 @@ def calculate_stats_for_key( } -def generate_rel_stats(dataset_path: Path | str, embodiment_tag: EmbodimentTag) -> None: +def generate_rel_stats( + dataset_path: Path | str, + embodiment_tag: EmbodimentTag, + episode_indices: np.ndarray | list[int] | None = None, +) -> None: + """Generate relative_stats.json for RELATIVE action representations. + + Args: + dataset_path: Path to the LeRobot dataset root. + embodiment_tag: Embodiment tag for action config lookup. + episode_indices: Optional list/array of episode indices to include. If None, + compute stats over the full dataset. + """ dataset_path = Path(dataset_path) action_config = MODALITY_CONFIGS[embodiment_tag.value]["action"] if action_config.action_configs is None: @@ -219,23 +374,1058 @@ def generate_rel_stats(dataset_path: Path | str, embodiment_tag: EmbodimentTag) if action_config.rep == ActionRepresentation.RELATIVE ] stats_path = Path(dataset_path) / LE_ROBOT_REL_STATS_FILENAME - if stats_path.exists(): + if stats_path.exists() and episode_indices is None: with open(stats_path, "r") as f: stats = json.load(f) else: stats = {} for action_key in sorted(action_keys): - if action_key in stats: + if action_key in stats and episode_indices is None: continue print(f"Generating relative stats for {dataset_path} {embodiment_tag} {action_key}") - stats[action_key] = calculate_stats_for_key(dataset_path, embodiment_tag, action_key) + stats[action_key] = calculate_stats_for_key( + dataset_path, embodiment_tag, action_key, episode_indices=episode_indices + ) with open(stats_path, "w") as f: json.dump(to_json_serializable(dict(stats)), f, indent=4) -def main(dataset_path: Path | str, embodiment_tag: EmbodimentTag): - generate_stats(dataset_path) - generate_rel_stats(dataset_path, embodiment_tag) +def calculate_temporal_percentile_stats( + dataset_path: Path | str, + modality_configs: dict[str, ModalityConfig], + skip_video: bool = True, + max_episodes: int = -1, + episode_indices: np.ndarray | list[int] | None = None, + embodiment_tag: EmbodimentTag | None = None, +) -> dict[str, dict[str, Any]]: + """ + Calculate temporal-aware percentile statistics for actions and non-temporal stats for states. + + This function computes statistics needed for percentile-based normalization: + - For actions: temporal-aware stats with shape (horizon, dim) where each timestep + in the action horizon has its own statistics + - For states: non-temporal stats with shape (dim,) using single timestamp + + Statistics computed: + - q01, q02: 1st and 2nd percentiles (lower bounds for normalization) + - q98, q99: 98th and 99th percentiles (upper bounds for normalization) + - mean, std: Mean and standard deviation (for meanstd normalization) + - min, max: Absolute min/max values (for minmax normalization) + + Args: + dataset_path: Path to the LeRobot dataset directory + modality_configs: Dictionary mapping modality names ('state', 'action') to + ModalityConfig objects specifying keys and delta indices + skip_video: If True, skip loading video data for faster iteration. Default True. + max_episodes: Maximum number of episodes to process. -1 for all episodes. + episode_indices: Optional list/array of episode indices to include. If None, + uses all episodes in the dataset. + embodiment_tag: Optional embodiment tag used to gate dataset-specific behaviors + (e.g., skipping terminal `next.done` rows for UCSD). + + Returns: + Dictionary with structure: + { + "state": { + "": { + "q01": [dim], "q02": [dim], "q98": [dim], "q99": [dim], + "mean": [dim], "std": [dim], "min": [dim], "max": [dim] + } + }, + "action": { + "": { + "q01": [horizon, dim], "q02": [horizon, dim], + "q98": [horizon, dim], "q99": [horizon, dim], + "mean": [horizon, dim], "std": [horizon, dim], + "min": [horizon, dim], "max": [horizon, dim] + } + } + } + + Example: + >>> modality_configs = { + ... "state": ModalityConfig(delta_indices=[0], modality_keys=["arm_joints"]), + ... "action": ModalityConfig( + ... delta_indices=list(range(16)), modality_keys=["arm_joints"] + ... ), + ... } + >>> stats = calculate_temporal_percentile_stats("/path/to/dataset", modality_configs) + >>> print( + ... stats["action"]["arm_joints"]["q02"].shape + ... ) # (16, 6) for 16-step horizon, 6 joints + """ + dataset_path = Path(dataset_path) + loader = LeRobotEpisodeLoader( + dataset_path, + modality_configs, + skip_video=skip_video, + require_stats=False, # Stats don't exist yet - we're calculating them + ) + + # Determine action horizon from delta indices + action_delta_indices = np.array(modality_configs["action"].delta_indices) + max_delta = int(np.max(action_delta_indices)) + action_horizon = len(action_delta_indices) + + # Collect data for each modality key, organized by timestep for actions + # state_data[key] = list of (dim,) arrays + # action_data[key][timestep] = list of (dim,) arrays + state_data: dict[str, list[np.ndarray]] = { + key: [] for key in modality_configs.get("state", ModalityConfig([], [])).modality_keys + } + action_data: dict[str, dict[int, list[np.ndarray]]] = { + key: {t: [] for t in range(action_horizon)} + for key in modality_configs["action"].modality_keys + } + + # Collect relative action data for RELATIVE representation (joint-angle deltas) + # This is separate from action_data because we need both absolute and relative stats + # relative_action_data[key][timestep] = list of (dim,) arrays + action_configs = modality_configs["action"].action_configs + relative_action_keys = [] + if action_configs: + for key_idx, key in enumerate(modality_configs["action"].modality_keys): + if key_idx < len(action_configs): + ac = action_configs[key_idx] + if ac.rep == ActionRepresentation.RELATIVE: + relative_action_keys.append(key) + + relative_action_data: dict[str, dict[int, list[np.ndarray]]] = { + key: {t: [] for t in range(action_horizon)} for key in relative_action_keys + } + + episode_ids = _resolve_episode_ids( + total_episodes=len(loader), + episode_indices=episode_indices, + max_episodes=max_episodes, + ) + + # Compute per-episode effective lengths and train-time-valid step indices + effective_length_by_episode = { + ep_id: max(0, loader.get_episode_length(ep_id) - max_delta) for ep_id in episode_ids + } + valid_step_indices_by_episode = compute_valid_step_indices_parallel( + dataset_path=dataset_path, + embodiment_tag=embodiment_tag, + chunk_size=loader.chunk_size, + data_path_pattern=loader.data_path_pattern, + action_delta_indices=modality_configs["action"].delta_indices, + episode_indices=episode_ids, + effective_lengths=[effective_length_by_episode[ep_id] for ep_id in episode_ids], + show_progress=False, + ) + + for episode_id in tqdm(episode_ids, desc="Collecting data for percentile stats"): + effective_length = effective_length_by_episode[episode_id] + if effective_length <= 0: + continue + + if valid_step_indices_by_episode is None: + step_indices = np.arange(effective_length, dtype=np.int32) + else: + step_indices = valid_step_indices_by_episode.get(episode_id) + if step_indices is None or len(step_indices) == 0: + continue + + df = loader[episode_id] + + # Collect state data (non-temporal, single timestamp per sample) + if "state" in modality_configs: + state_delta_idx = modality_configs["state"].delta_indices[-1] + for state_key in modality_configs["state"].modality_keys: + col_name = f"state.{state_key}" + if col_name not in df.columns: + continue + state_col = df[col_name].values + for step_idx in step_indices: + state_idx = state_delta_idx + int(step_idx) + state_data[state_key].append(np.asarray(state_col[state_idx], dtype=np.float32)) + + # Collect action data (temporal, separate by timestep in horizon) + # Apply action transformations (rel-xyz-rot6d) if configured + action_configs = modality_configs["action"].action_configs + state_delta_idx = ( + modality_configs["state"].delta_indices[-1] if "state" in modality_configs else 0 + ) + + for key_idx, action_key in enumerate(modality_configs["action"].modality_keys): + col_name = f"action.{action_key}" + if col_name not in df.columns: + continue + action_col = df[col_name].values + + # Check if this action needs rel-xyz-rot6d conversion + action_config = ( + action_configs[key_idx] + if action_configs and key_idx < len(action_configs) + else None + ) + needs_rel_xyz_rot6d = ( + action_config is not None + and action_config.rep == ActionRepresentation.REL_XYZ_ROT6D + ) + + # Check if this action needs RELATIVE stats (joint-angle deltas) + needs_relative = ( + action_config is not None and action_config.rep == ActionRepresentation.RELATIVE + ) + + # Get state column for reference pose if needed (for REL_XYZ_ROT6D or RELATIVE) + if needs_rel_xyz_rot6d or needs_relative: + state_key = action_config.state_key or action_key + state_col_name = f"state.{state_key}" + if state_col_name not in df.columns: + raise ValueError( + f"State key '{state_key}' not found in data for REL_XYZ_ROT6D conversion. " + f"Available columns: {df.columns.tolist()}" + ) + state_col = df[state_col_name].values + + # Check if motion scaling is configured (CMR Versius specific) + has_motion_scaling = action_config is not None and ( + action_config.translation_scaling_key or action_config.rotation_scaling_key + ) + trans_scale_col = None + rot_scale_col = None + + if has_motion_scaling: + if action_config.translation_scaling_key: + trans_col_name = f"state.{action_config.translation_scaling_key}" + if trans_col_name in df.columns: + trans_scale_col = df[trans_col_name].values + if action_config.rotation_scaling_key: + rot_col_name = f"state.{action_config.rotation_scaling_key}" + if rot_col_name in df.columns: + rot_scale_col = df[rot_col_name].values + + for step_idx in step_indices: + step_idx = int(step_idx) + # Get reference state for rel-xyz-rot6d conversion + if needs_rel_xyz_rot6d: + ref_state_idx = state_delta_idx + step_idx + eef_pose = np.asarray(state_col[ref_state_idx], dtype=np.float32) + + for t, delta in enumerate(action_delta_indices): + action_idx = step_idx + int(delta) + action_value = np.asarray(action_col[action_idx], dtype=np.float32) + + # Apply rel-xyz-rot6d conversion if configured + if needs_rel_xyz_rot6d: + action_value = convert_to_rel_xyz_rot6d( + action_data=action_value[np.newaxis, :], # Add timestep dim + eef_pose=eef_pose, + input_rotation_format=action_config.input_rotation_format, + reference_rotation_format=action_config.reference_rotation_format, + input_quat_order=action_config.input_quat_order, + reference_quat_order=action_config.reference_quat_order, + )[0] # Remove timestep dim + + # Apply motion scaling if configured (CMR Versius specific) + # This converts from hand-controller-space to instrument-space + if has_motion_scaling: + trans_scale = 1.0 + rot_scale = 1.0 + if trans_scale_col is not None: + trans_scale = float(trans_scale_col[ref_state_idx]) + if rot_scale_col is not None: + rot_scale = float(rot_scale_col[ref_state_idx]) + action_value = apply_motion_scaling_to_rel_xyz_rot6d( + action_value[np.newaxis, :], trans_scale, rot_scale + )[0] + + # For EEF XYZ_ROT6D, only keep xyz (first 3 dims) for stats + # rot6d is already bounded [-1, 1] and doesn't need normalization + if ( + action_config.type == ActionType.EEF + and action_config.format == ActionFormat.XYZ_ROT6D + ): + action_value = action_value[:3] # Only xyz + + action_data[action_key][t].append(action_value) + + # Compute relative action (delta from reference state) for RELATIVE representation + # This is separate from REL_XYZ_ROT6D which handles EEF poses with rotation math + if needs_relative: + ref_state_idx = state_delta_idx + step_idx + ref_state = np.asarray(state_col[ref_state_idx], dtype=np.float32) + # Simple subtraction for joint angles: relative = action - reference_state + relative_value = action_value - ref_state + relative_action_data[action_key][t].append(relative_value) + + # Calculate statistics + result: dict[str, dict[str, Any]] = {"state": {}, "action": {}} + + # State statistics: non-temporal, shape (dim,) + for state_key, data_list in state_data.items(): + if not data_list: + continue + stacked = np.stack(data_list, axis=0) # (num_samples, dim) + result["state"][state_key] = { + "q01": np.percentile(stacked, 1, axis=0), + "q02": np.percentile(stacked, 2, axis=0), + "q98": np.percentile(stacked, 98, axis=0), + "q99": np.percentile(stacked, 99, axis=0), + "mean": np.mean(stacked, axis=0), + "std": np.std(stacked, axis=0), + "min": np.min(stacked, axis=0), + "max": np.max(stacked, axis=0), + } + + # Action statistics: temporal-aware, shape (horizon, dim) + for action_key, timestep_data in action_data.items(): + if not timestep_data[0]: # Check if any data was collected + continue + + # Build stats for each timestep + horizon_stats = { + "q01": [], + "q02": [], + "q98": [], + "q99": [], + "mean": [], + "std": [], + "min": [], + "max": [], + } + + for t in range(action_horizon): + stacked = np.stack(timestep_data[t], axis=0) # (num_samples, dim) + horizon_stats["q01"].append(np.percentile(stacked, 1, axis=0)) + horizon_stats["q02"].append(np.percentile(stacked, 2, axis=0)) + horizon_stats["q98"].append(np.percentile(stacked, 98, axis=0)) + horizon_stats["q99"].append(np.percentile(stacked, 99, axis=0)) + horizon_stats["mean"].append(np.mean(stacked, axis=0)) + horizon_stats["std"].append(np.std(stacked, axis=0)) + horizon_stats["min"].append(np.min(stacked, axis=0)) + horizon_stats["max"].append(np.max(stacked, axis=0)) + + # Stack to get shape (horizon, dim) + result["action"][action_key] = { + stat_name: np.stack(stat_values, axis=0) + for stat_name, stat_values in horizon_stats.items() + } + + # Relative action statistics: temporal-aware, shape (horizon, dim) + # For RELATIVE representation (joint-angle deltas), stored separately from absolute action stats + if relative_action_data: + result["relative_action"] = {} + for action_key, timestep_data in relative_action_data.items(): + if not timestep_data[0]: # Check if any data was collected + continue + + # Build stats for each timestep (same structure as action stats) + horizon_stats = { + "q01": [], + "q02": [], + "q98": [], + "q99": [], + "mean": [], + "std": [], + "min": [], + "max": [], + } + + for t in range(action_horizon): + stacked = np.stack(timestep_data[t], axis=0) # (num_samples, dim) + horizon_stats["q01"].append(np.percentile(stacked, 1, axis=0)) + horizon_stats["q02"].append(np.percentile(stacked, 2, axis=0)) + horizon_stats["q98"].append(np.percentile(stacked, 98, axis=0)) + horizon_stats["q99"].append(np.percentile(stacked, 99, axis=0)) + horizon_stats["mean"].append(np.mean(stacked, axis=0)) + horizon_stats["std"].append(np.std(stacked, axis=0)) + horizon_stats["min"].append(np.min(stacked, axis=0)) + horizon_stats["max"].append(np.max(stacked, axis=0)) + + # Stack to get shape (horizon, dim) + result["relative_action"][action_key] = { + stat_name: np.stack(stat_values, axis=0) + for stat_name, stat_values in horizon_stats.items() + } + + return result + + +def _collect_episodes_worker(args: tuple) -> tuple[dict, dict, dict]: + """ + Worker function for parallel episode processing in stats calculation. + + This function is designed to run in a separate process via ProcessPoolExecutor. + Each worker processes a chunk of episodes independently and returns collected + data for later aggregation. + + Architecture: + - Each worker creates its own LeRobotEpisodeLoader (required because loaders + can't be pickled across process boundaries) + - Workers process episodes in parallel, with no communication between them + - Results are merged in the main process after all workers complete + + Memory Optimization (Episode-at-a-Time Stacking): + Instead of appending individual frames one-by-one: + for i in range(usable_length): + state_data[key].append(np.asarray(state_col[i], dtype=np.float32)) + + We stack all frames for an episode into a single contiguous array: + stacked = np.stack(df[col_name].values, axis=0).astype(np.float32) + state_data[key].append(stacked[start:end]) + + This reduces numpy object overhead significantly: + - Before: ~3600 small arrays per episode Γ— 4792 episodes = ~17M array objects + - After: ~1 array per episode Γ— 4792 episodes = ~5K array objects + - Each numpy array has ~96 bytes of header overhead, so this saves ~1.6GB + + The total DATA volume is identical - only the number of Python objects differs. + This is 100% statistically equivalent to frame-by-frame collection. + + Data Structures Returned: + state_data: {state_key: [array(frames, dim), array(frames, dim), ...]} + - Each list item is a 2D array containing all frames from one episode + - Shape per item: (usable_frames_in_episode, state_dim) + + action_data: {action_key: {timestep: [array(frames, dim), ...]}} + - Organized by action key, then by timestep in the action horizon + - Each list item is a 2D array for one episode at one timestep + - Shape per item: (usable_frames_in_episode, action_dim) + + Args: + args: Tuple of + (dataset_path, modality_configs, skip_video, episode_ids, worker_id, + valid_step_indices_by_episode, target_keys) + - dataset_path: Path to LeRobot dataset (as string for pickling) + - modality_configs: Dict of ModalityConfig objects + - skip_video: Whether to skip video loading + - episode_ids: List of episode indices this worker should process + - worker_id: Integer ID for progress bar positioning + - valid_step_indices_by_episode: Optional map episode_idx -> valid step indices + - target_keys: Optional dict with "state" and/or "action" lists. If provided, + only those keys are collected to reduce memory. Missing entries default + to all keys for that modality. + + Returns: + Tuple of (state_data, action_data) dictionaries containing collected data + for all episodes processed by this worker. + """ + if len(args) == 5: + dataset_path, modality_configs, skip_video, episode_ids, worker_id = args + valid_step_indices_by_episode = None + target_keys = None + elif len(args) == 6: + dataset_path, modality_configs, skip_video, episode_ids, worker_id, maybe_valid = args + valid_step_indices_by_episode = maybe_valid if isinstance(maybe_valid, dict) else None + target_keys = None + else: + ( + dataset_path, + modality_configs, + skip_video, + episode_ids, + worker_id, + valid_step_indices_by_episode, + target_keys, + ) = args + dataset_path = Path(dataset_path) + + # Each worker creates its own loader + loader = LeRobotEpisodeLoader( + dataset_path, + modality_configs, + skip_video=skip_video, + require_stats=False, + ) + + # Determine action horizon from delta indices + action_delta_indices = np.array(modality_configs["action"].delta_indices) + max_delta = int(np.max(action_delta_indices)) + action_horizon = len(action_delta_indices) + + # Initialize data structures (optionally filtered by target_keys) + all_state_keys = modality_configs.get("state", ModalityConfig([], [])).modality_keys + if target_keys is None or "state" not in target_keys: + state_keys_to_collect = list(all_state_keys) + else: + target_state_keys = set(target_keys.get("state", [])) + state_keys_to_collect = [key for key in all_state_keys if key in target_state_keys] + + all_action_keys = modality_configs["action"].modality_keys + if target_keys is None or "action" not in target_keys: + action_keys_to_collect = list(all_action_keys) + else: + target_action_keys = set(target_keys.get("action", [])) + action_keys_to_collect = [key for key in all_action_keys if key in target_action_keys] + action_keys_to_collect_set = set(action_keys_to_collect) + + state_data: dict[str, list[np.ndarray]] = {key: [] for key in state_keys_to_collect} + action_data: dict[str, dict[int, list[np.ndarray]]] = { + key: {t: [] for t in range(action_horizon)} for key in action_keys_to_collect + } + + # Collect relative action data for RELATIVE representation (joint-angle deltas) + action_configs = modality_configs["action"].action_configs + relative_action_keys = [] + if action_configs: + for key_idx, key in enumerate(all_action_keys): + if key not in action_keys_to_collect_set: + continue + if key_idx < len(action_configs): + ac = action_configs[key_idx] + if ac.rep == ActionRepresentation.RELATIVE: + relative_action_keys.append(key) + + relative_action_data: dict[str, dict[int, list[np.ndarray]]] = { + key: {t: [] for t in range(action_horizon)} for key in relative_action_keys + } + + state_delta_idx = ( + modality_configs["state"].delta_indices[-1] if "state" in modality_configs else 0 + ) + + # Progress bar for this worker + pbar = tqdm( + episode_ids, + desc=f"Worker {worker_id:2d}", + position=worker_id, + leave=False, + ncols=80, + ) + + for episode_id in pbar: + df = loader[episode_id] + + # Determine valid step indices for this episode. + if valid_step_indices_by_episode is None: + usable_length = max(0, len(df) - max_delta) + if usable_length <= 0: + continue + step_indices = np.arange(usable_length, dtype=np.int64) + else: + step_indices = valid_step_indices_by_episode.get(episode_id) + if step_indices is None or len(step_indices) == 0: + continue + step_indices = np.asarray(step_indices, dtype=np.int64) + + # ----------------------------------------------------------------- + # Collect state data using VECTORIZED approach (episode-at-a-time) + # ----------------------------------------------------------------- + # This is the key memory optimization: instead of appending 3600 + # individual (dim,) arrays, we create ONE (3600, dim) array per episode. + # Same data, far fewer Python objects. + if "state" in modality_configs: + for state_key in state_keys_to_collect: + col_name = f"state.{state_key}" + if col_name not in df.columns: + continue + # Stack all frames: df[col].values is list of arrays β†’ (frames, dim) + stacked = np.stack(df[col_name].values, axis=0).astype(np.float32) + # Gather only valid step indices (accounting for state delta index) + state_indices = state_delta_idx + step_indices + sliced = stacked[state_indices] + # Append single 2D array instead of per-step 1D arrays + state_data[state_key].append(sliced) + + # ----------------------------------------------------------------- + # Collect action data using VECTORIZED approach (per timestep) + # ----------------------------------------------------------------- + # Actions are organized by timestep in the horizon because each + # timestep has its own normalization statistics (temporal-aware). + for key_idx, action_key in enumerate(all_action_keys): + if action_key not in action_keys_to_collect_set: + continue + col_name = f"action.{action_key}" + if col_name not in df.columns: + continue + + # Stack all action frames for this episode: (total_frames, action_dim) + action_stacked = np.stack(df[col_name].values, axis=0).astype(np.float32) + + # Check if this action key requires rel-xyz-rot6d transformation + # (converts absolute poses to relative-to-reference-frame) + action_config = ( + action_configs[key_idx] + if action_configs and key_idx < len(action_configs) + else None + ) + needs_rel_xyz_rot6d = ( + action_config is not None + and action_config.rep == ActionRepresentation.REL_XYZ_ROT6D + ) + + # Check if this action needs RELATIVE stats (joint-angle deltas) + needs_relative = ( + action_config is not None and action_config.rep == ActionRepresentation.RELATIVE + ) + + # Load reference state data for rel-xyz-rot6d or relative conversion + if needs_rel_xyz_rot6d or needs_relative: + state_key = action_config.state_key or action_key + state_col_name = f"state.{state_key}" + if state_col_name not in df.columns: + raise ValueError(f"State key '{state_key}' not found for relative conversion.") + ref_state_stacked = np.stack(df[state_col_name].values, axis=0).astype(np.float32) + + # Check if motion scaling is configured (CMR Versius specific) + has_motion_scaling = action_config is not None and ( + action_config.translation_scaling_key or action_config.rotation_scaling_key + ) + trans_scale_stacked = None + rot_scale_stacked = None + + if has_motion_scaling: + if action_config.translation_scaling_key: + trans_col_name = f"state.{action_config.translation_scaling_key}" + if trans_col_name in df.columns: + trans_scale_stacked = np.stack(df[trans_col_name].values, axis=0).astype( + np.float32 + ) + if action_config.rotation_scaling_key: + rot_col_name = f"state.{action_config.rotation_scaling_key}" + if rot_col_name in df.columns: + rot_scale_stacked = np.stack(df[rot_col_name].values, axis=0).astype( + np.float32 + ) + + # Process each timestep in the action horizon separately + # This is necessary because stats are computed per-timestep + for t, delta in enumerate(action_delta_indices): + # Gather actions for this timestep across valid frames only. + action_indices = step_indices + int(delta) + action_slice = action_stacked[action_indices] # (num_valid_steps, dim) + + if needs_rel_xyz_rot6d: + # Get reference states (current pose) for all valid frames. + ref_states = ref_state_stacked[state_delta_idx + step_indices] + + # NOTE: REL_XYZ_ROT6D requires frame-by-frame processing due to + # rotation math (can't be easily vectorized). This is the main + # computational bottleneck for EEF actions. + converted = [] + for i in range(len(step_indices)): + action_value = convert_to_rel_xyz_rot6d( + action_data=action_slice[i : i + 1], + eef_pose=ref_states[i], + input_rotation_format=action_config.input_rotation_format, + reference_rotation_format=action_config.reference_rotation_format, + input_quat_order=action_config.input_quat_order, + reference_quat_order=action_config.reference_quat_order, + )[0] + + # Apply motion scaling if configured (CMR Versius specific) + # This converts from hand-controller-space to instrument-space + if has_motion_scaling: + trans_scale = 1.0 + rot_scale = 1.0 + ref_idx = state_delta_idx + int(step_indices[i]) + if trans_scale_stacked is not None: + trans_scale = float(trans_scale_stacked[ref_idx]) + if rot_scale_stacked is not None: + rot_scale = float(rot_scale_stacked[ref_idx]) + action_value = apply_motion_scaling_to_rel_xyz_rot6d( + action_value[np.newaxis, :], trans_scale, rot_scale + )[0] + + # For EEF XYZ_ROT6D format, only keep xyz translation for stats + # rot6d is already bounded [-1, 1] and doesn't need normalization + if ( + action_config.type == ActionType.EEF + and action_config.format == ActionFormat.XYZ_ROT6D + ): + action_value = action_value[:3] + + converted.append(action_value) + action_slice = np.stack(converted, axis=0) + + # Append single 2D array: (num_valid_steps, dim) for this timestep. + action_data[action_key][t].append(action_slice) + + # Compute relative action (delta from reference state) for RELATIVE representation + # This is vectorized for efficiency since it's just subtraction (no rotation math) + if needs_relative: + ref_states = ref_state_stacked[state_delta_idx + step_indices] + # Simple subtraction for joint angles: relative = action - reference_state + relative_slice = action_slice - ref_states + relative_action_data[action_key][t].append(relative_slice) + + pbar.close() + return state_data, action_data, relative_action_data + + +def calculate_temporal_percentile_stats_parallel( + dataset_path: Path | str, + modality_configs: dict[str, ModalityConfig], + skip_video: bool = True, + max_episodes: int = -1, + num_workers: int | None = None, + episode_indices: np.ndarray | list[int] | None = None, + embodiment_tag: EmbodimentTag | None = None, +) -> dict[str, dict[str, Any]]: + """ + Parallel version of calculate_temporal_percentile_stats with episode-level parallelism. + + This implementation processes one modality key at a time (state or action) to + reduce peak memory while preserving exact numerical results. + + This function distributes episode processing across multiple CPU cores for faster + statistics computation on large datasets. + + Architecture Overview: + 1. SPLIT: Episodes are divided into chunks, one per worker + 2. KEY-BY-KEY: For each modality key, workers collect ONLY that key's data + 3. MERGE: Results are merged for that key and stats are computed immediately + 4. REPEAT: Memory is freed before moving to the next key + + Data Flow: + Episodes β†’ [Worker 1] β†’ state_data_1, action_data_1 ─┐ + Episodes β†’ [Worker 2] β†’ state_data_2, action_data_2 ─┼→ Merge β†’ Compute Stats + Episodes β†’ [Worker N] β†’ state_data_N, action_data_N β”€β”˜ + + Statistical Equivalence: + This parallel implementation produces IDENTICAL results to the sequential + version. All data is collected before computing percentiles - no approximation + or streaming algorithms are used. The only difference is the order in which + episodes are processed, which doesn't affect percentile computation. + + Note: True streaming/incremental percentile algorithms (like t-digest) would + NOT be statistically equivalent. We explicitly avoid such approaches. + + Memory Considerations: + - Each worker holds data for its chunk of episodes in memory + - Peak memory is reduced because only ONE key is merged at a time + - For memory-constrained systems, reduce num_workers (e.g., --stats-num-workers 8) + - The episode-at-a-time stacking optimization reduces numpy object overhead + but does not reduce total data volume + - Trade-off: the dataset is scanned once per key (higher runtime) + + Progress Monitoring: + - Each worker displays its own tqdm progress bar showing episodes processed + - A main "Merging results" progress bar shows workers completing + - Progress bars use position parameter for clean multi-line display + + Args: + dataset_path: Path to the LeRobot dataset directory + modality_configs: Dictionary mapping modality names ('state', 'action') to + ModalityConfig objects specifying keys, delta indices, and + action configs (for rel-xyz-rot6d conversion) + skip_video: If True, skip loading video data for faster iteration. Default True. + Video data is not needed for state/action statistics. + max_episodes: Maximum number of episodes to process. -1 for all episodes. + Useful for quick testing on a subset of data. + num_workers: Number of parallel workers. None uses os.cpu_count(). + Set to 1 to disable parallelism (falls back to sequential in + launch_finetune.py). Recommended: 8-16 for large datasets. + episode_indices: Optional list/array of episode indices to include. If None, + uses all episodes in the dataset. + embodiment_tag: Optional embodiment tag used to gate dataset-specific behaviors + (e.g., skipping terminal `next.done` rows for UCSD). + + Returns: + Dictionary with structure identical to calculate_temporal_percentile_stats: + { + "state": { + "": { + "q01": [dim], "q02": [dim], "q98": [dim], "q99": [dim], + "mean": [dim], "std": [dim], "min": [dim], "max": [dim] + } + }, + "action": { + "": { + "q01": [horizon, dim], "q02": [horizon, dim], ... + } + } + } + + Example: + >>> # Calculate stats with 8 workers + >>> stats = calculate_temporal_percentile_stats_parallel( + ... dataset_path="/path/to/large_dataset", + ... modality_configs=modality_configs, + ... num_workers=8, + ... ) + + See Also: + - calculate_temporal_percentile_stats: Sequential version + - _collect_episodes_worker: Worker function with memory optimization details + """ + import gc + import os + + dataset_path = Path(dataset_path) + action_delta_indices = np.array(modality_configs["action"].delta_indices) + max_delta = int(np.max(action_delta_indices)) + + # Get total episodes by creating a temporary loader + temp_loader = LeRobotEpisodeLoader( + dataset_path, + modality_configs, + skip_video=True, + require_stats=False, + ) + total_episodes = len(temp_loader) + episode_ids = _resolve_episode_ids( + total_episodes=total_episodes, + episode_indices=episode_indices, + max_episodes=max_episodes, + ) + effective_lengths = [ + max(0, temp_loader.get_episode_length(ep_id) - max_delta) for ep_id in episode_ids + ] + valid_step_indices_by_episode = compute_valid_step_indices_parallel( + dataset_path=dataset_path, + embodiment_tag=embodiment_tag, + chunk_size=temp_loader.chunk_size, + data_path_pattern=temp_loader.data_path_pattern, + action_delta_indices=modality_configs["action"].delta_indices, + episode_indices=episode_ids, + effective_lengths=effective_lengths, + show_progress=False, + ) + num_episodes = len(episode_ids) + del temp_loader + + # Determine number of workers + if num_workers is None: + num_workers = min(os.cpu_count() or 4, num_episodes) + num_workers = min(num_workers, num_episodes) + + if num_episodes == 0: + return {"state": {}, "action": {}} + + print(f"Processing {num_episodes} episodes with {num_workers} workers (key-by-key)") + + # ========================================================================= + # PHASE 1: SPLIT - Divide episodes into chunks for parallel processing + # ========================================================================= + chunk_size = (num_episodes + num_workers - 1) // num_workers # Ceiling division + episode_chunks = [episode_ids[i : i + chunk_size] for i in range(0, num_episodes, chunk_size)] + + # ========================================================================= + # PHASE 2: PROCESS - Run workers in parallel (key-by-key) + # ========================================================================= + action_horizon = len(action_delta_indices) + state_keys = modality_configs.get("state", ModalityConfig([], [])).modality_keys + action_keys = modality_configs["action"].modality_keys + + # Identify action keys that need RELATIVE stats + action_configs = modality_configs["action"].action_configs + relative_action_keys = [] + if action_configs: + for key_idx, key in enumerate(action_keys): + if key_idx < len(action_configs): + ac = action_configs[key_idx] + if ac.rep == ActionRepresentation.RELATIVE: + relative_action_keys.append(key) + + # Calculate statistics from merged data + result: dict[str, dict[str, Any]] = {"state": {}, "action": {}} + relative_action_result: dict[str, dict[str, Any]] = {} + + def _compute_stats_from_concatenated(concatenated: np.ndarray) -> dict[str, np.ndarray]: + """Compute scalar stats for a single (num_samples, dim) array.""" + return { + "q01": np.percentile(concatenated, 1, axis=0), + "q02": np.percentile(concatenated, 2, axis=0), + "q98": np.percentile(concatenated, 98, axis=0), + "q99": np.percentile(concatenated, 99, axis=0), + "mean": np.mean(concatenated, axis=0), + "std": np.std(concatenated, axis=0), + "min": np.min(concatenated, axis=0), + "max": np.max(concatenated, axis=0), + } + + def _compute_temporal_stats( + timestep_data: dict[int, list[np.ndarray]], + ) -> dict[str, np.ndarray]: + """Compute temporal stats for per-timestep lists of arrays.""" + if not timestep_data or not timestep_data[0]: + return {} + + horizon_stats = { + "q01": [], + "q02": [], + "q98": [], + "q99": [], + "mean": [], + "std": [], + "min": [], + "max": [], + } + + for t in range(action_horizon): + if not timestep_data[t]: + return {} + concatenated = np.concatenate(timestep_data[t], axis=0) + horizon_stats["q01"].append(np.percentile(concatenated, 1, axis=0)) + horizon_stats["q02"].append(np.percentile(concatenated, 2, axis=0)) + horizon_stats["q98"].append(np.percentile(concatenated, 98, axis=0)) + horizon_stats["q99"].append(np.percentile(concatenated, 99, axis=0)) + horizon_stats["mean"].append(np.mean(concatenated, axis=0)) + horizon_stats["std"].append(np.std(concatenated, axis=0)) + horizon_stats["min"].append(np.min(concatenated, axis=0)) + horizon_stats["max"].append(np.max(concatenated, axis=0)) + del concatenated + + return { + stat_name: np.stack(stat_values, axis=0) + for stat_name, stat_values in horizon_stats.items() + } + + from concurrent.futures import as_completed + + with ProcessPoolExecutor(max_workers=num_workers) as executor: + # --------------------------------------------------------------------- + # State statistics: non-temporal, processed key-by-key + # --------------------------------------------------------------------- + for state_key in state_keys: + target_keys = {"state": [state_key], "action": []} + args_list = [ + ( + str(dataset_path), + modality_configs, + skip_video, + chunk, + worker_id, + ( + None + if valid_step_indices_by_episode is None + else { + ep_id: valid_step_indices_by_episode[ep_id] + for ep_id in chunk + if ep_id in valid_step_indices_by_episode + } + ), + target_keys, + ) + for worker_id, chunk in enumerate(episode_chunks) + ] + + merged_state_data: list[np.ndarray] = [] + + # Print newlines to make room for worker progress bars + print("\n" * num_workers) + futures = { + executor.submit(_collect_episodes_worker, args): i + for i, args in enumerate(args_list) + } + main_pbar = tqdm( + total=len(futures), + desc=f"Merging state:{state_key}", + position=num_workers, + leave=True, + ncols=80, + ) + + for future in as_completed(futures): + state_data, _, _ = future.result() + merged_state_data.extend(state_data.get(state_key, [])) + main_pbar.update(1) + + main_pbar.close() + print("\n") + + if merged_state_data: + concatenated = np.concatenate(merged_state_data, axis=0) + result["state"][state_key] = _compute_stats_from_concatenated(concatenated) + del concatenated + + del merged_state_data + gc.collect() + + # --------------------------------------------------------------------- + # Action statistics: temporal-aware, processed key-by-key + # --------------------------------------------------------------------- + for action_key in action_keys: + target_keys = {"state": [], "action": [action_key]} + args_list = [ + ( + str(dataset_path), + modality_configs, + skip_video, + chunk, + worker_id, + ( + None + if valid_step_indices_by_episode is None + else { + ep_id: valid_step_indices_by_episode[ep_id] + for ep_id in chunk + if ep_id in valid_step_indices_by_episode + } + ), + target_keys, + ) + for worker_id, chunk in enumerate(episode_chunks) + ] + + merged_action_data: dict[int, list[np.ndarray]] = {t: [] for t in range(action_horizon)} + merged_relative_action_data: dict[int, list[np.ndarray]] | None = None + if action_key in relative_action_keys: + merged_relative_action_data = {t: [] for t in range(action_horizon)} + + print("\n" * num_workers) + futures = { + executor.submit(_collect_episodes_worker, args): i + for i, args in enumerate(args_list) + } + main_pbar = tqdm( + total=len(futures), + desc=f"Merging action:{action_key}", + position=num_workers, + leave=True, + ncols=80, + ) + + for future in as_completed(futures): + _, action_data, relative_action_data = future.result() + + for t in range(action_horizon): + merged_action_data[t].extend(action_data.get(action_key, {}).get(t, [])) + + if merged_relative_action_data is not None: + for t in range(action_horizon): + merged_relative_action_data[t].extend( + relative_action_data.get(action_key, {}).get(t, []) + ) + + main_pbar.update(1) + + main_pbar.close() + print("\n") + + action_stats = _compute_temporal_stats(merged_action_data) + if action_stats: + result["action"][action_key] = action_stats + + if merged_relative_action_data is not None: + relative_stats = _compute_temporal_stats(merged_relative_action_data) + if relative_stats: + relative_action_result[action_key] = relative_stats + + del merged_action_data + del merged_relative_action_data + gc.collect() + + if relative_action_result: + result["relative_action"] = relative_action_result + + return result + + +def main( + dataset_path: Path | str, + embodiment_tag: EmbodimentTag, + include_splits: list[str] | None = None, + exclude_splits: list[str] | None = None, +): + """CLI entrypoint for stats generation with optional split filtering.""" + episode_indices = None + if include_splits or exclude_splits: + info = load_info_json(dataset_path) + total_episodes = info.get("total_episodes") + episode_indices = resolve_episode_indices( + info, + include_splits=include_splits, + exclude_splits=exclude_splits, + total_episodes=total_episodes, + ) + generate_stats(dataset_path, episode_indices=episode_indices) + generate_rel_stats(dataset_path, embodiment_tag, episode_indices=episode_indices) if __name__ == "__main__": diff --git a/gr00t/data/step_filtering.py b/gr00t/data/step_filtering.py new file mode 100644 index 0000000..55e2dfa --- /dev/null +++ b/gr00t/data/step_filtering.py @@ -0,0 +1,300 @@ +""" +Shared step-index filtering utilities for dataset loading and stats generation. + +This module centralizes dataset-specific filtering rules so train-time sampling +and offline percentile-stat generation use the same valid step indices. +""" + +from __future__ import annotations + +from concurrent.futures import ProcessPoolExecutor, as_completed +from functools import partial +import json +import os +from pathlib import Path + +import numpy as np +import pyarrow.parquet as pq +from tqdm import tqdm + +from gr00t.data.types import ( + CMR_ENGAGED_LEFT_KEY, + CMR_ENGAGED_RIGHT_KEY, + CMR_RAW_INDEX_ARM_LINKED_LEFT, + CMR_RAW_INDEX_ARM_LINKED_RIGHT, + CMR_RAW_INDEX_HAPTIC_ENGAGED_LEFT, + CMR_RAW_INDEX_HAPTIC_ENGAGED_RIGHT, + EMBODIMENTS_SKIP_NEXT_DONE, + EmbodimentTag, +) + + +MAX_FILTER_WORKERS = 128 + + +def _filter_episode_fast( + episode_idx: int, + dataset_path: Path, + chunk_size: int, + data_path_pattern: str, + action_delta_indices: list[int], + effective_length: int, +) -> tuple[int, np.ndarray]: + """Filter one episode using CMR clutch-aware rules. + + Reads only `observation.state` via PyArrow and applies the same constraints + used at train time: + 1. Arm linkage cannot change within the action horizon. + 2. At least one side must be engaged within the action horizon. + + Args: + episode_idx: Episode index to process. + dataset_path: Dataset root directory. + chunk_size: Episodes per parquet chunk. + data_path_pattern: Parquet path pattern from LeRobotEpisodeLoader. + action_delta_indices: Action horizon offsets. + effective_length: Candidate anchor-step count for this episode. + + Returns: + Tuple of `(episode_idx, valid_step_indices)`. + """ + chunk_idx = episode_idx // chunk_size + parquet_path = dataset_path / data_path_pattern.format( + episode_chunk=chunk_idx, episode_index=episode_idx + ) + + table = pq.read_table(parquet_path, columns=["observation.state"]) + state_data = table.column("observation.state").to_pylist() + + eng_left = np.array([s[CMR_RAW_INDEX_HAPTIC_ENGAGED_LEFT] for s in state_data], dtype=bool) + eng_right = np.array([s[CMR_RAW_INDEX_HAPTIC_ENGAGED_RIGHT] for s in state_data], dtype=bool) + al_left = np.array([s[CMR_RAW_INDEX_ARM_LINKED_LEFT] for s in state_data], dtype=np.float32) + al_right = np.array([s[CMR_RAW_INDEX_ARM_LINKED_RIGHT] for s in state_data], dtype=np.float32) + + valid_indices: list[int] = [] + for step_idx in range(effective_length): + horizon_indices = np.array([step_idx + d for d in action_delta_indices]) + if horizon_indices[-1] >= len(state_data): + continue + if len(np.unique(al_left[horizon_indices])) > 1: + continue + if len(np.unique(al_right[horizon_indices])) > 1: + continue + if not eng_left[horizon_indices].any() and not eng_right[horizon_indices].any(): + continue + valid_indices.append(step_idx) + + return episode_idx, np.array(valid_indices, dtype=np.int32) + + +def _check_is_cmr_data(dataset_path: Path) -> bool: + """Check whether dataset metadata indicates CMR clutch-aware fields. + + Args: + dataset_path: Dataset root directory. + + Returns: + True if `meta/modality.json` contains both CMR engagement keys. + """ + try: + modality_path = dataset_path / "meta" / "modality.json" + if not modality_path.exists(): + return False + with open(modality_path, "r") as f: + modality_config = json.load(f) + state_keys = set(modality_config.get("state", {}).keys()) + action_keys = set(modality_config.get("action", {}).keys()) + all_keys = state_keys | action_keys + return CMR_ENGAGED_LEFT_KEY in all_keys and CMR_ENGAGED_RIGHT_KEY in all_keys + except Exception: + return False + + +def _check_has_next_done(dataset_path: Path) -> bool: + """Check whether dataset metadata declares a `next.done` feature. + + Args: + dataset_path: Dataset root directory. + + Returns: + True if `meta/info.json` lists `next.done` in features. + """ + info_path = dataset_path / "meta" / "info.json" + if not info_path.exists(): + return False + try: + with open(info_path, "r") as f: + info = json.load(f) + return "next.done" in info.get("features", {}) + except Exception: + return False + + +def _should_skip_next_done(embodiment_tag: EmbodimentTag | None, dataset_path: Path) -> bool: + """Determine whether terminal `next.done` rows should be skipped. + + Args: + embodiment_tag: Dataset embodiment tag. + dataset_path: Dataset root directory. + + Returns: + True when the embodiment is allowlisted and dataset metadata includes + `next.done`. + """ + if embodiment_tag is None: + return False + if embodiment_tag not in EMBODIMENTS_SKIP_NEXT_DONE: + return False + return _check_has_next_done(dataset_path) + + +def _filter_episode_done_fast( + episode_idx: int, + dataset_path: Path, + chunk_size: int, + data_path_pattern: str, + action_delta_indices: list[int], + effective_length: int, +) -> tuple[int, np.ndarray]: + """Filter one episode to exclude anchor steps crossing terminal `next.done`. + + For allowlisted embodiments with terminal padding, this trims candidate + anchor steps so the action horizon never includes rows at or after the first + `next.done=True`. + + Args: + episode_idx: Episode index to process. + dataset_path: Dataset root directory. + chunk_size: Episodes per parquet chunk. + data_path_pattern: Parquet path pattern from LeRobotEpisodeLoader. + action_delta_indices: Action horizon offsets. + effective_length: Candidate anchor-step count for this episode. + + Returns: + Tuple of `(episode_idx, valid_step_indices)`. + """ + chunk_idx = episode_idx // chunk_size + parquet_path = dataset_path / data_path_pattern.format( + episode_chunk=chunk_idx, episode_index=episode_idx + ) + + table = pq.read_table(parquet_path, columns=["next.done"]) + done = np.array(table.column("next.done").to_pylist(), dtype=bool) + + max_delta = int(max(action_delta_indices)) + usable_length = effective_length + if done.any(): + first_done = int(np.argmax(done)) + usable_length = min(usable_length, max(0, first_done - max_delta)) + + if usable_length <= 0: + return episode_idx, np.array([], dtype=np.int32) + return episode_idx, np.arange(usable_length, dtype=np.int32) + + +def compute_valid_step_indices_parallel( + dataset_path: Path | str, + embodiment_tag: EmbodimentTag | None, + chunk_size: int, + data_path_pattern: str, + action_delta_indices: list[int] | np.ndarray, + episode_indices: list[int] | np.ndarray, + effective_lengths: list[int] | np.ndarray, + num_workers: int | None = None, + max_filter_workers: int = MAX_FILTER_WORKERS, + show_progress: bool = True, +) -> dict[int, np.ndarray] | None: + """Compute valid step indices per episode using shared filtering semantics. + + Filtering mode is selected in this order: + 1. CMR clutch-aware filtering (if CMR keys are present in modality metadata). + 2. Terminal-step filtering using `next.done` (if embodiment is allowlisted). + 3. No special filtering (returns `None`). + + Args: + dataset_path: Dataset root directory. + embodiment_tag: Embodiment tag used for allowlisted behavior gates. + chunk_size: Episodes per parquet chunk. + data_path_pattern: Parquet path pattern from LeRobotEpisodeLoader. + action_delta_indices: Action horizon offsets. + episode_indices: Episode indices to process. + effective_lengths: Candidate anchor-step counts aligned to `episode_indices`. + num_workers: Optional worker count for ProcessPoolExecutor. + max_filter_workers: Upper bound on worker count. + show_progress: Whether to print filtering progress and summary. + + Returns: + - `None` when no special filtering is required. + - `dict[episode_idx, np.ndarray]` when filtering is applied. + - Empty dict when filtering is applied but no valid indices remain. + + Raises: + ValueError: If `episode_indices` and `effective_lengths` lengths differ. + """ + episode_indices = [int(i) for i in np.asarray(episode_indices, dtype=int).tolist()] + effective_lengths = [int(max(0, e)) for e in np.asarray(effective_lengths, dtype=int).tolist()] + if len(episode_indices) != len(effective_lengths): + raise ValueError( + "episode_indices and effective_lengths must have same length. " + f"Got {len(episode_indices)} and {len(effective_lengths)}." + ) + if not episode_indices: + return {} + + dataset_path = Path(dataset_path) + action_delta_indices = [int(d) for d in np.asarray(action_delta_indices, dtype=int).tolist()] + + if _check_is_cmr_data(dataset_path): + filter_fn = _filter_episode_fast + filter_desc = "Clutch-aware filtering" + elif _should_skip_next_done(embodiment_tag, dataset_path): + filter_fn = _filter_episode_done_fast + filter_desc = "Terminal-step filtering (next.done)" + else: + return None + + if num_workers is None: + num_workers = min(os.cpu_count() or 32, max_filter_workers) + num_workers = max(1, min(num_workers, max_filter_workers, len(episode_indices))) + + filter_fn = partial( + filter_fn, + dataset_path=dataset_path, + chunk_size=chunk_size, + data_path_pattern=data_path_pattern, + action_delta_indices=action_delta_indices, + ) + + if show_progress: + print(f"{filter_desc}: {len(episode_indices)} episodes with {num_workers} workers...") + + results: dict[int, np.ndarray] = {} + total_original = 0 + total_valid = 0 + + with ProcessPoolExecutor(max_workers=num_workers) as executor: + futures = { + executor.submit(filter_fn, ep_idx, effective_length=eff_len): (ep_idx, eff_len) + for ep_idx, eff_len in zip(episode_indices, effective_lengths) + } + + iter_futures = as_completed(futures) + if show_progress: + iter_futures = tqdm(iter_futures, total=len(futures), desc="Filtering episodes") + + for future in iter_futures: + ep_idx, valid_indices = future.result() + _, eff_len = futures[future] + total_original += eff_len + total_valid += len(valid_indices) + if len(valid_indices) > 0: + results[ep_idx] = valid_indices + + total_filtered = total_original - total_valid + if show_progress and total_filtered > 0 and total_original > 0: + print( + f"{filter_desc} complete: {total_filtered}/{total_original} indices filtered " + f"({100 * total_filtered / total_original:.1f}%)" + ) + + return results diff --git a/gr00t/data/types.py b/gr00t/data/types.py index 23a65db..72cd0b1 100644 --- a/gr00t/data/types.py +++ b/gr00t/data/types.py @@ -16,9 +16,20 @@ class MessageType(Enum): class ActionRepresentation(Enum): + """ + Defines how action values relate to the current robot state. + + - RELATIVE: Actions are deltas from current state (computed at training time) + - DELTA: Incremental changes per timestep + - ABSOLUTE: Target positions with no state dependency + - REL_XYZ_ROT6D: Translation and rotation is relative to current EEF, + gripper is absolute. Used for healthcare/manipulation tasks. + """ + RELATIVE = "relative" DELTA = "delta" ABSOLUTE = "absolute" + REL_XYZ_ROT6D = "rel_xyz_rot6d" class ActionType(Enum): @@ -27,7 +38,19 @@ class ActionType(Enum): class ActionFormat(Enum): + """ + Defines the format of action data components. + + - DEFAULT: Default format, no specific interpretation + - XYZ: Translation only (3D position) + - ROT6D: 6D rotation representation (first two columns of rotation matrix) + - XYZ_ROT6D: Combined translation and 6D rotation + - XYZ_ROTVEC: Combined translation and rotation vector + """ + DEFAULT = "default" + XYZ = "xyz" + ROT6D = "rot6d" XYZ_ROT6D = "xyz+rot6d" XYZ_ROTVEC = "xyz+rotvec" @@ -60,10 +83,66 @@ class VLAStepData: @dataclass class ActionConfig: + """ + Configuration for an action modality defining representation, type, format, and normalization. + + This config controls how action data is processed during training and inference, + including conversion to relative/rel-xyz-rot6d representations and normalization. + + Attributes: + rep: How action values relate to robot state (RELATIVE, ABSOLUTE, REL_XYZ_ROT6D, DELTA) + type: Whether this is end-effector control (EEF) or joint control (NON_EEF) + format: The format of the action data (XYZ, ROT6D, XYZ_ROT6D, etc.) + state_key: Which state key to use as reference for relative actions (e.g., "eef_pose") + input_rotation_format: Format of incoming rotation data (actions): + - "quat": Quaternion format - will be converted to rot6d for REL_XYZ_ROT6D + - "rot6d": Already in 6D rotation format - no conversion needed + input_quat_order: Quaternion component ordering when input_rotation_format="quat": + - "xyzw": Scalar-last order (x, y, z, w) - scipy convention, default + - "wxyz": Scalar-first order (w, x, y, z) - used by some datasets (e.g., Hamlyn) + reference_rotation_format: Format of rotation in the reference state (for REL_XYZ_ROT6D): + - "quat": Quaternion format - state is 7D: xyz + quat + - "rot6d": 6D rotation format - state is 9D: xyz + rot6d + reference_quat_order: Quaternion component ordering when reference_rotation_format="quat": + - "xyzw": Scalar-last order (default) + - "wxyz": Scalar-first order + translation_scaling_key: Optional state key containing translation scaling factor. + If provided, rel-xyz-rot6d translation is multiplied by this scaling factor + to convert from hand-controller-space to instrument-space. Used for CMR Versius. + rotation_scaling_key: Optional state key containing rotation scaling factor. + If provided, rel-xyz-rot6d rotation angle is multiplied by this scaling factor. + Uses axis-angle representation for proper rotation scaling. + hold_through_clutch: Whether this ABSOLUTE action should hold its value during clutch-out + (sample-and-hold) instead of being zeroed (fail-safe). Only applies to ABSOLUTE actions. + - True: Hold last engaged value (e.g., gripper holding a needle) + - False: Zero during clutch-out (e.g., energy button - fail-safe) + Default is False for safety - actions zero unless explicitly marked to hold. + For RELATIVE/REL_XYZ_ROT6D actions, this flag is ignored (always zeroed). + normalization_type: Which normalization strategy to use for this action group. + - "temporal_meanstd": Use meanstd normalization with temporal-aware stats (default). + Stats have shape (horizon, dim) so each timestep in the action chunk + is normalized independently, accounting for cumulative magnitude growth + in relative action representations like REL_XYZ_ROT6D. + - "meanstd": Use meanstd normalization (same as temporal_meanstd) + - "minmax": Use min/max normalization + - "skip": Skip normalization entirely (pass through raw values) + """ + rep: ActionRepresentation type: ActionType format: ActionFormat state_key: str | None = None + input_rotation_format: str = "quat" + input_quat_order: str = "xyzw" # "xyzw" (scipy default) or "wxyz" (scalar-first) + reference_rotation_format: str = "rot6d" + reference_quat_order: str = "xyzw" # "xyzw" (scipy default) or "wxyz" (scalar-first) + # Motion scaling keys for CMR Versius (optional, default None = no scaling) + translation_scaling_key: str | None = None + rotation_scaling_key: str | None = None + # Clutch-aware behavior for ABSOLUTE actions (default False = fail-safe zeroing) + hold_through_clutch: bool = False + # Normalization strategy for this action group + normalization_type: str = "temporal_meanstd" @dataclass @@ -82,10 +161,18 @@ class ModalityConfig: """Optional list of keys to apply sin/cos encoding. If None or empty, use min/max normalization for all keys.""" mean_std_embedding_keys: list[str] | None = None """Optional list of keys to apply mean/std normalization. If None or empty, use min/max normalization for all keys.""" + min_max_embedding_keys: list[str] | None = None + """Optional list of keys to apply min/max normalization. If None or empty, keys not in sin_cos or mean_std will use min/max.""" + pass_through_keys: list[str] | None = None + """Optional list of keys that are used for intermediate calculations, but ARE NOT sent to the model. Used for auxiliary data like motion scaling factors that are needed for action processing but should not be normalized.""" action_configs: list[ActionConfig] | None = None def __post_init__(self): - """Set default values for action-related fields if not specified.""" + """Parse action configs from dictionaries if provided as dicts. + + Converts dictionary-based action configs (e.g., from JSON) into ActionConfig + dataclass instances, handling enum parsing and default values. + """ if self.action_configs is not None: assert len(self.action_configs) == len(self.modality_keys), ( f"Number of action configs ({len(self.action_configs)}) must match number of modality keys ({len(self.modality_keys)})" @@ -98,6 +185,110 @@ def __post_init__(self): type=ActionType[action_config["type"]], format=ActionFormat[action_config["format"]], state_key=action_config.get("state_key", None), + input_rotation_format=action_config.get("input_rotation_format", "quat"), + input_quat_order=action_config.get("input_quat_order", "xyzw"), + reference_rotation_format=action_config.get( + "reference_rotation_format", "rot6d" + ), + reference_quat_order=action_config.get("reference_quat_order", "xyzw"), + translation_scaling_key=action_config.get("translation_scaling_key", None), + rotation_scaling_key=action_config.get("rotation_scaling_key", None), + hold_through_clutch=action_config.get("hold_through_clutch", False), + normalization_type=action_config.get( + "normalization_type", "temporal_meanstd" + ), ) parsed_action_configs.append(action_config) self.action_configs = parsed_action_configs + + +# ============================================================================= +# Dataset-specific runtime behavior flags +# ============================================================================= +# These flags gate behavior that should only apply to specific embodiments. +EMBODIMENTS_SKIP_NEXT_DONE: set[EmbodimentTag] = {EmbodimentTag.UCSD_DVRK} +""" +Embodiments that should skip terminal steps where `next.done == True`. + +These datasets include terminal padding where the final action is zeroed out +(including zero-norm quaternions). We exclude those steps from training and +stats to avoid invalid rotations and terminal snap-to-zero artifacts. +""" + + +# ============================================================================= +# CMR (Clutch-Mechanical-Robot) Data Keys +# ============================================================================= +# These keys are used for CMR Versius surgical robot data to track haptic engagement. +# When the surgeon "clutches out" (disengages the haptic controller), the robot should +# not move, and action targets should be zeroed (for RELATIVE/REL_XYZ_ROT6D) or +# held at the last engaged value (for ABSOLUTE actions with hold_through_clutch=True). + +CMR_ENGAGED_LEFT_KEY: str = "hapticengaged_left" +"""State/action key indicating if the left haptic controller is engaged (bool). +When False, the surgeon's left hand is "clutched out" and left arm should not move.""" + +CMR_ENGAGED_RIGHT_KEY: str = "hapticengaged_right" +"""State/action key indicating if the right haptic controller is engaged (bool). +When False, the surgeon's right hand is "clutched out" and right arm should not move.""" + +CMR_ARM_LINKED_LEFT_KEY: str = "armlinkedtohaptic_left" +"""State key indicating which robot arm (0-3) is linked to the left haptic controller. +Used for arm swapping detection and deriving arm_left_color from arm_X_color.""" + +CMR_ARM_LINKED_RIGHT_KEY: str = "armlinkedtohaptic_right" +"""State key indicating which robot arm (0-3) is linked to the right haptic controller. +Used for arm swapping detection and deriving arm_right_color from arm_X_color.""" + +# Arm color keys (indices 0-3 correspond to physical robot arms) +CMR_ARM_COLOR_KEYS: tuple[str, ...] = ("arm_0_color", "arm_1_color", "arm_2_color", "arm_3_color") +"""State keys for arm colors by physical arm index (0-3). +Used with armlinkedtohaptic_* to derive arm_left_color/arm_right_color.""" + +# ----------------------------------------------------------------------------- +# CMR Raw Parquet Indices +# ----------------------------------------------------------------------------- +# These indices refer to positions in the raw observation.state array in parquet files. +# Used for fast clutch-aware filtering with PyArrow (bypasses modality.json mapping). +# WARNING: These indices are specific to CMR Versius data format and must be updated +# if the observation.state layout changes. + +CMR_RAW_INDEX_HAPTIC_ENGAGED_LEFT: int = 16 +"""Index of hapticengaged_left in raw observation.state array.""" + +CMR_RAW_INDEX_HAPTIC_ENGAGED_RIGHT: int = 17 +"""Index of hapticengaged_right in raw observation.state array.""" + +CMR_RAW_INDEX_ARM_LINKED_LEFT: int = 20 +"""Index of armlinkedtohaptic_left in raw observation.state array.""" + +CMR_RAW_INDEX_ARM_LINKED_RIGHT: int = 21 +"""Index of armlinkedtohaptic_right in raw observation.state array.""" + +CMR_RAW_MIN_STATE_LENGTH: int = 22 +"""Minimum observation.state array length for CMR data (must include index 21).""" + + +# ============================================================================= +# Action Dimension Constants +# ============================================================================= +# These constants define the dimensions of different action format representations. + +XYZ_DIM: int = 3 +"""Dimension of XYZ translation component (3D position).""" + +ROT6D_DIM: int = 6 +"""Dimension of 6D rotation representation (first two columns of rotation matrix).""" + +EEF_XYZ_ROT6D_DIM: int = 9 +"""Total dimension for end-effector XYZ_ROT6D format: xyz (3) + rot6d (6) = 9.""" + +QUAT_DIM: int = 4 +"""Dimension of quaternion rotation representation (xyzw order).""" + +EEF_XYZ_QUAT_DIM: int = 7 +"""Total dimension for end-effector XYZ_QUAT format: xyz (3) + quat (4) = 7.""" + +ROT6D_IDENTITY: tuple[float, ...] = (1.0, 0.0, 0.0, 0.0, 1.0, 0.0) +"""Identity rotation in rot6d format: first two columns of the 3x3 identity matrix. +Used when zeroing out REL_XYZ_ROT6D actions for disengaged timesteps.""" diff --git a/gr00t/data/utils.py b/gr00t/data/utils.py index 69eab61..d77bf38 100644 --- a/gr00t/data/utils.py +++ b/gr00t/data/utils.py @@ -154,7 +154,7 @@ def normalize_values_meanstd(values, params): * Case 2 - 2D params: Shape (T, D) - different std per step Returns: - Normalized values using z-score normalization + Normalized values using z-score normalization, clipped to [-5, 5] - Same shape as input values: (T, D) or (B, T, D) - Values are transformed as: (x - mean) / std - For features where std == 0, normalized value equals original value @@ -181,7 +181,7 @@ def normalize_values_meanstd(values, params): # Keep original values for zero-std features normalized[..., ~mask] = values[..., ~mask] - return normalized + return np.clip(normalized, -5.0, 5.0) def unnormalize_values_meanstd(normalized_values, params): diff --git a/gr00t/experiment/experiment.py b/gr00t/experiment/experiment.py index eba3bec..3d88d08 100755 --- a/gr00t/experiment/experiment.py +++ b/gr00t/experiment/experiment.py @@ -84,7 +84,9 @@ def warn_configs(config: Config): assert ( config.model.shortest_image_edge is None and config.model.crop_fraction is None ), ( - "Do not set shortest_image_edge and crop_fraction together with image_crop_size and image_target_size" + "Do not set shortest_image_edge and crop_fraction together with image_crop_size and image_target_size. " + "If you are using a YAML config with only image_target_size/image_crop_size, set " + "shortest_image_edge: null and crop_fraction: null explicitly." ) if ( @@ -121,8 +123,10 @@ def run(config: Config): # Validate config config.validate() - # Create output directory - if config.training.experiment_name is None: + if os.environ.get("WANDB_RUN_ID") is not None: + experiment_name = os.environ.get("WANDB_RUN_ID") + output_dir = Path(config.training.output_dir) / experiment_name + elif config.training.experiment_name is None: output_dir = Path(config.training.output_dir) experiment_name = output_dir.name else: @@ -162,8 +166,10 @@ def run(config: Config): wandb.init( project=config.training.wandb_project, name=experiment_name, + id=experiment_name, config=config_dict, tags=[config.data.mode], + resume="allow", ) # Setup model training pipeline. diff --git a/gr00t/experiment/launch_finetune.py b/gr00t/experiment/launch_finetune.py index 54d8c44..4a2c563 100644 --- a/gr00t/experiment/launch_finetune.py +++ b/gr00t/experiment/launch_finetune.py @@ -5,6 +5,7 @@ import os from pathlib import Path +import open_h.embodiments # noqa: F401 β€” registers Open-H embodiment configs import tyro from gr00t.configs.base_config import get_default_config @@ -26,6 +27,118 @@ def load_modality_config(modality_config_path: str): raise FileNotFoundError(f"Modality config path does not exist: {modality_config_path}") +def calculate_norm_stats_only(ft_config: FinetuneConfig) -> None: + """ + Calculate normalization statistics for the dataset and exit without training. + + This function loads the dataset with skip_video=True for fast iteration, + calculates temporal percentile statistics (q01, q02, q98, q99, mean, std), + and saves them to the specified output path. + + Uses parallel processing by default for faster computation on large datasets. + + Args: + ft_config: FinetuneConfig containing dataset path, embodiment tag, and output settings + """ + import json + + from gr00t.configs.data.embodiment_configs import MODALITY_CONFIGS + from gr00t.data.split_utils import load_info_json, resolve_episode_indices + from gr00t.data.stats import ( + calculate_temporal_percentile_stats, + calculate_temporal_percentile_stats_parallel, + ) + from gr00t.data.utils import to_json_serializable + + # Note: For multi-dataset training with different embodiments, run stats calculation + # separately for each dataset. Stats are keyed by repo_id, so consolidate them + # into a single JSON file using --norm-stats-output-path. + embodiment_tag = ft_config.embodiment_tag.value + dataset_path = Path(ft_config.dataset_path) + + print(f"Calculating normalization statistics for {dataset_path}") + print(f"Embodiment: {embodiment_tag}") + + # Get modality configs for this embodiment + if embodiment_tag not in MODALITY_CONFIGS: + raise ValueError( + f"Embodiment '{embodiment_tag}' not found in MODALITY_CONFIGS. " + f"Available: {list(MODALITY_CONFIGS.keys())}" + ) + + modality_configs = MODALITY_CONFIGS[embodiment_tag] + + episode_indices = None + if ft_config.include_splits or ft_config.exclude_splits: + info = load_info_json(dataset_path) + total_episodes = info.get("total_episodes") + episode_indices = resolve_episode_indices( + info, + include_splits=ft_config.include_splits, + exclude_splits=ft_config.exclude_splits, + total_episodes=total_episodes, + ) + assert episode_indices is not None + print(f"Using {len(episode_indices)} episodes after split filtering") + + # Calculate temporal percentile statistics with skip_video=True + # Use parallel version by default (num_workers=None uses CPU count) + num_workers = ft_config.stats_num_workers + if num_workers == 1: + # Single worker - use sequential version + stats = calculate_temporal_percentile_stats( + dataset_path=dataset_path, + modality_configs=modality_configs, + skip_video=True, + max_episodes=-1, + episode_indices=episode_indices, + embodiment_tag=ft_config.embodiment_tag, + ) + else: + # Parallel version (default) + stats = calculate_temporal_percentile_stats_parallel( + dataset_path=dataset_path, + modality_configs=modality_configs, + skip_video=True, + max_episodes=-1, + num_workers=num_workers, + episode_indices=episode_indices, + embodiment_tag=ft_config.embodiment_tag, + ) + + # Determine output path + if ft_config.norm_stats_output_path: + output_path = Path(ft_config.norm_stats_output_path) + # If it's a consolidated file, key by dataset name + if output_path.suffix == ".json": + # Check if file exists and load existing stats + if output_path.exists(): + with open(output_path, "r") as f: + all_stats = json.load(f) + else: + all_stats = {} + + # Add this dataset's stats keyed by repo_id + repo_id = dataset_path.name + all_stats[repo_id] = to_json_serializable(stats) + + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as f: + json.dump(all_stats, f, indent=2) + print(f"Saved statistics to {output_path} (keyed by '{repo_id}')") + else: + raise ValueError(f"Output path must be a .json file: {output_path}") + else: + # Save to dataset's meta directory + output_path = dataset_path / "meta" / "temporal_stats.json" + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as f: + json.dump(to_json_serializable(stats), f, indent=2) + print(f"Saved statistics to {output_path}") + + print("Statistics calculation complete. Exiting without training.") + + if __name__ == "__main__": # Set LOGURU_LEVEL environment variable if not already set (default: INFO) if "LOGURU_LEVEL" not in os.environ: @@ -38,6 +151,11 @@ def load_modality_config(modality_config_path: str): if ft_config.modality_config_path is not None: load_modality_config(ft_config.modality_config_path) + # Handle stats-only calculation mode + if ft_config.calculate_norm_stats: + calculate_norm_stats_only(ft_config) + exit(0) + config = get_default_config().load_dict( { "data": { @@ -47,6 +165,8 @@ def load_modality_config(modality_config_path: str): "dataset_paths": [ft_config.dataset_path], "mix_ratio": 1.0, "embodiment_tag": embodiment_tag, + "include_splits": ft_config.include_splits, + "exclude_splits": ft_config.exclude_splits, } ], } @@ -60,6 +180,7 @@ def load_modality_config(modality_config_path: str): config.model.tune_projector = ft_config.tune_projector config.model.tune_diffusion_model = ft_config.tune_diffusion_model config.model.state_dropout_prob = ft_config.state_dropout_prob + config.model.state_dropout_prob_per_embodiment = ft_config.state_dropout_prob_per_embodiment config.model.random_rotation_angle = ft_config.random_rotation_angle config.model.color_jitter_params = ft_config.color_jitter_params if ft_config.extra_augmentation_config: @@ -67,6 +188,16 @@ def load_modality_config(modality_config_path: str): else: config.model.extra_augmentation_config = None + # Image size configuration - when set, uses letterbox + resize + crop pipeline + if ft_config.image_size is not None: + config.model.image_target_size = ft_config.image_size + # Use specified crop size or default to target size (no crop augmentation) + config.model.image_crop_size = ft_config.image_crop_size or ft_config.image_size + # Use torchvision pipeline which has letterbox padding + config.model.use_albumentations_transforms = False + config.model.shortest_image_edge = None + config.model.crop_fraction = None + config.model.load_bf16 = False config.model.reproject_vision = False config.model.eagle_collator = True diff --git a/gr00t/experiment/launch_train.py b/gr00t/experiment/launch_train.py index 00e4c88..d1c1668 100644 --- a/gr00t/experiment/launch_train.py +++ b/gr00t/experiment/launch_train.py @@ -1,7 +1,10 @@ +import argparse import logging import os from pathlib import Path +import sys +import open_h.embodiments # noqa: F401 β€” registers Open-H embodiment configs import tyro from gr00t.configs.base_config import Config, get_default_config @@ -12,17 +15,24 @@ # Set LOGURU_LEVEL environment variable if not already set (default: INFO) if "LOGURU_LEVEL" not in os.environ: os.environ["LOGURU_LEVEL"] = "INFO" - # Use tyro for clean CLI - config = tyro.cli(Config, default=get_default_config(), description=__doc__) - # Load config from path if provided - if config.load_config_path: - assert Path(config.load_config_path).exists(), ( - f"Config path does not exist: {config.load_config_path}" - ) - config = config.load(Path(config.load_config_path)) # inplace loading - config.load_config_path = None - logging.info(f"Loaded config from {config.load_config_path}") + argv = sys.argv[1:] + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--load-config-path") + parsed_args, remaining_args = parser.parse_known_args(argv) - # Override with command-line. - config = tyro.cli(Config, default=config, description=__doc__) + config_default = get_default_config() + if parsed_args.load_config_path: + config_path = Path(parsed_args.load_config_path) + assert config_path.exists(), f"Config path does not exist: {config_path}" + if remaining_args: + parser.error( + "`--load-config-path` cannot be combined with additional CLI overrides. " + "Put overrides in the YAML config instead." + ) + config = config_default.load(config_path) + config.load_config_path = None + logging.info(f"Loaded config from {config_path}") + else: + # Use tyro for clean CLI + config = tyro.cli(Config, default=config_default, args=remaining_args, description=__doc__) run(config) diff --git a/gr00t/model/gr00t_n1d6/gr00t_n1d6.py b/gr00t/model/gr00t_n1d6/gr00t_n1d6.py index db5c5af..aafc77c 100755 --- a/gr00t/model/gr00t_n1d6/gr00t_n1d6.py +++ b/gr00t/model/gr00t_n1d6/gr00t_n1d6.py @@ -72,9 +72,24 @@ def __init__(self, config: Gr00tN1d6Config): # State dropout parameters self.state_dropout_prob = config.state_dropout_prob + self.state_dropout_prob_per_embodiment = getattr( + config, "state_dropout_prob_per_embodiment", None + ) + + # Build per-embodiment dropout lookup buffer + if self.state_dropout_prob_per_embodiment: + from .processing_gr00t_n1d6 import EMBODIMENT_TAG_TO_PROJECTOR_INDEX + + dropout_buf = torch.zeros(config.max_num_embodiments) + for tag, prob in self.state_dropout_prob_per_embodiment.items(): + if tag in EMBODIMENT_TAG_TO_PROJECTOR_INDEX: + dropout_buf[EMBODIMENT_TAG_TO_PROJECTOR_INDEX[tag]] = prob + self.register_buffer("dropout_prob_by_embodiment", dropout_buf) + + has_any_dropout = self.state_dropout_prob > 0 or self.state_dropout_prob_per_embodiment self.mask_token = ( nn.Parameter(0.02 * torch.randn(1, 1, self.input_embedding_dim)) - if self.state_dropout_prob > 0 + if has_any_dropout else None ) @@ -101,7 +116,7 @@ def set_trainable_parameters( self.action_decoder.requires_grad_(False) if self.config.add_pos_embed: self.position_embedding.requires_grad_(False) - if self.state_dropout_prob > 0: + if self.mask_token is not None: self.mask_token.requires_grad_(False) if not tune_diffusion_model: self.model.requires_grad_(False) @@ -175,17 +190,31 @@ def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> # Get embodiment ID. embodiment_id = action_input.embodiment_id + # Per-embodiment state dropout: zero state BEFORE encoding. + if self.state_dropout_prob_per_embodiment and hasattr(self, "dropout_prob_by_embodiment"): + dropout_probs = self.dropout_prob_by_embodiment[embodiment_id] # (B,) + if self.training: + do_dropout = ( + torch.rand(action_input.state.shape[0], device=action_input.state.device) + < dropout_probs + ) + else: + do_dropout = dropout_probs > 0.999 # deterministic at inference + do_dropout = do_dropout[:, None, None].to(dtype=action_input.state.dtype) + action_input.state = action_input.state * (1 - do_dropout) + # Embed state. state_features = self.state_encoder(action_input.state, embodiment_id) - # Dropout state features. - if self.state_dropout_prob > 0: - do_dropout = ( - torch.rand(state_features.shape[0], device=state_features.device) - < self.state_dropout_prob - ) - do_dropout = do_dropout[:, None, None].to(dtype=state_features.dtype) - state_features = state_features * (1 - do_dropout) + self.mask_token * do_dropout + # Global state dropout: replace encoded features with learned mask_token. + if self.mask_token is not None and self.state_dropout_prob > 0: + if self.training: + do_dropout = ( + torch.rand(state_features.shape[0], device=state_features.device) + < self.state_dropout_prob + ) + do_dropout = do_dropout[:, None, None].to(dtype=state_features.dtype) + state_features = state_features * (1 - do_dropout) + self.mask_token * do_dropout # Add Gaussian noise to state features. if self.training and self.state_additive_noise_scale > 0: @@ -280,8 +309,16 @@ def _encode_features( vl_embeds = backbone_output.backbone_features embodiment_id = action_input.embodiment_id + state = action_input.state + + # Per-embodiment state dropout: zero state before encoding (deterministic at inference) + if self.state_dropout_prob_per_embodiment and hasattr(self, "dropout_prob_by_embodiment"): + dropout_probs = self.dropout_prob_by_embodiment[embodiment_id] + do_dropout = (dropout_probs > 0.999)[:, None, None].to(dtype=state.dtype) + state = state * (1 - do_dropout) + # Embed state. - state_features = self.state_encoder(action_input.state, embodiment_id) + state_features = self.state_encoder(state, embodiment_id) return BatchFeature(data={"backbone_features": vl_embeds, "state_features": state_features}) diff --git a/gr00t/model/gr00t_n1d6/image_augmentations.py b/gr00t/model/gr00t_n1d6/image_augmentations.py index 07b8716..5001663 100755 --- a/gr00t/model/gr00t_n1d6/image_augmentations.py +++ b/gr00t/model/gr00t_n1d6/image_augmentations.py @@ -527,6 +527,74 @@ def __call__(self, img: torch.Tensor) -> torch.Tensor: return padded_img +class ResizeWithPadding: + """Resize image to target size while preserving aspect ratio, padding with black bars. + + This is different from LetterBoxTransform which pads to square. This transform + resizes and pads to an exact target size (height, width). + + Pipeline: + 1. Resize image so it fits within target size (preserving aspect ratio) + 2. Pad with black bars to reach exact target dimensions + """ + + def __init__(self, target_size: tuple[int, int]): + """ + Args: + target_size: Target (height, width) for the output image + """ + self.target_h, self.target_w = target_size + + def __call__(self, img: torch.Tensor) -> torch.Tensor: + """ + Resize and pad image to target size. + + Args: + img: Image tensor of shape (..., C, H, W) + + Returns: + Image tensor of shape (..., C, target_h, target_w) + """ + *leading_dims, c, h, w = img.shape + + # Calculate scale to fit within target while preserving aspect ratio + scale = min(self.target_h / h, self.target_w / w) + new_h = int(h * scale) + new_w = int(w * scale) + + # Handle leading dimensions by reshaping + if leading_dims: + batch_size = torch.tensor(leading_dims).prod().item() + img_reshaped = img.reshape(batch_size, c, h, w) + else: + img_reshaped = img.unsqueeze(0) + + # Resize preserving aspect ratio + resized = transforms.functional.resize(img_reshaped, [new_h, new_w], antialias=True) + + # Calculate padding needed to reach target size + pad_h = self.target_h - new_h + pad_w = self.target_w - new_w + pad_top = pad_h // 2 + pad_bottom = pad_h - pad_top + pad_left = pad_w // 2 + pad_right = pad_w - pad_left + + # Apply padding + padded = transforms.functional.pad( + resized, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0 + ) + + # Reshape back + if leading_dims: + output_shape = leading_dims + [c, self.target_h, self.target_w] + padded = padded.reshape(output_shape) + else: + padded = padded.squeeze(0) + + return padded + + def build_image_transformations( image_target_size, image_crop_size, random_rotation_angle, color_jitter_params ): @@ -544,11 +612,8 @@ def build_image_transformations( """ transform_list = [ transforms.ToImage(), - LetterBoxTransform(), - # transforms.ToDtype(torch.get_default_dtype(), scale=True), - transforms.Resize(size=image_target_size), + ResizeWithPadding(target_size=image_target_size), transforms.RandomCrop(size=image_crop_size), - transforms.Resize(size=image_target_size), ] if random_rotation_angle is not None and random_rotation_angle != 0: transform_list.append( @@ -559,12 +624,9 @@ def build_image_transformations( train_image_transform = transforms.Compose(transform_list) eval_image_transform = transforms.Compose( [ - transforms.ToImage(), - # transforms.ToDtype(torch.get_default_dtype(), scale=True), - LetterBoxTransform(), - transforms.Resize(size=image_target_size), + transforms.ToImage(), # Convert numpy/PIL to tensor + ResizeWithPadding(target_size=image_target_size), transforms.CenterCrop(size=image_crop_size), - transforms.Resize(size=image_target_size), ] ) return train_image_transform, eval_image_transform diff --git a/gr00t/model/gr00t_n1d6/processing_gr00t_n1d6.py b/gr00t/model/gr00t_n1d6/processing_gr00t_n1d6.py index d4063ed..ba4600f 100755 --- a/gr00t/model/gr00t_n1d6/processing_gr00t_n1d6.py +++ b/gr00t/model/gr00t_n1d6/processing_gr00t_n1d6.py @@ -40,8 +40,29 @@ "libero_panda": 2, "oxe_google": 0, "oxe_widowx": 1, - "oxe_droid": 16, + "oxe_droid": 29, "new_embodiment": 10, + ##### Open-H embodiment ids ##### + "jhu_imerse_dvrk": 3, + "cmr_versius": 4, + "ucb_dvrk": 5, + "sanoscience_sim": 6, + "tum_sonata_franka": 7, + "hamlyn_dvrk_15hz": 9, + "hamlyn_dvrk_30hz": 11, + "ustc_torin_tuodao": 12, + "ucsd_dvrk": 14, + "jhu_imerse_dvrk_mono": 15, + "rob_surgical_bitrack": 16, + "stanford_dvrk_real": 17, + "obuda_dvrk": 18, + "polyu_sim": 19, + "moon_maestro": 21, + "jhu_lscr_dvrk_miracle": 22, + "jhu_lscr_dvrk_smarts": 23, + "jhu_imerse_star_il": 27, + "tud_tundra_ur5e": 25, + "turin_mitic_ex_vivo": 26, } @@ -235,12 +256,21 @@ def decode_action( embodiment_tag: EmbodimentTag, state: dict[str, np.ndarray] | None = None, ): - """Undo action normalization and convert relative actions to absolute.""" - # Split concatenated action into joint groups + """Undo action normalization and convert relative actions to absolute. + + Args: + action: Normalized action array from model + embodiment_tag: Embodiment tag for modality config lookup + state: Optional current state for relative-to-absolute conversion + """ + # Split concatenated action into joint groups (excluding pass_through_keys) out_dict = {} start_idx = 0 - joint_groups = self.modality_configs[embodiment_tag.value]["action"].modality_keys - action_horizon = len(self.modality_configs[embodiment_tag.value]["action"].delta_indices) + action_config = self.modality_configs[embodiment_tag.value]["action"] + pass_through_keys = set(action_config.pass_through_keys or []) + joint_groups = [k for k in action_config.modality_keys if k not in pass_through_keys] + + action_horizon = len(action_config.delta_indices) for key in joint_groups: joint_dim = self.state_action_processor.norm_params[embodiment_tag.value]["action"][ key @@ -306,8 +336,10 @@ def __call__( ) if normalized_actions: - # Concatenate actions - action_keys = self.modality_configs[embodiment_tag.value]["action"].modality_keys + # Concatenate actions (excluding pass_through_keys which were removed by StateActionProcessor) + action_config = self.modality_configs[embodiment_tag.value]["action"] + pass_through_keys = set(action_config.pass_through_keys or []) + action_keys = [k for k in action_config.modality_keys if k not in pass_through_keys] normalized_actions = torch.cat( [torch.from_numpy(normalized_actions[key]) for key in action_keys], dim=-1 ) # (t, d) @@ -344,20 +376,28 @@ def __call__( normalized_actions = None action_mask = None - # Concatenate states + # Concatenate states (excluding pass-through keys) + state_config = self.modality_configs[embodiment_tag.value].get("state") + pass_through_keys = getattr(state_config, "pass_through_keys", None) or [] + pass_through_key_set = set(pass_through_keys) state_keys = self.modality_configs[embodiment_tag.value]["state"].modality_keys - normalized_states = torch.cat( + if pass_through_key_set: + state_keys = [key for key in state_keys if key not in pass_through_key_set] + normalized_states_concat = torch.cat( [torch.from_numpy(normalized_states[key]) for key in state_keys], dim=-1 ) - normalized_states = torch.cat( + normalized_states_concat = torch.cat( [ - normalized_states, + normalized_states_concat, torch.zeros( - normalized_states.shape[0], self.max_state_dim - normalized_states.shape[1] + normalized_states_concat.shape[0], + self.max_state_dim - normalized_states_concat.shape[1], ), ], dim=-1, ) + # Rename for consistency with rest of code + normalized_states = normalized_states_concat # Crop and resize images. if self.training: @@ -529,12 +569,19 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | Path, **kwargs): "color_jitter_params", "use_relative_action", "extra_augmentation_config", + "use_percentiles", + "image_crop_size", + "image_target_size", + "use_albumentations", + "shortest_image_edge", + "crop_fraction", + "max_action_horizon", ] for key in override_keys: if key in kwargs: override = kwargs.pop(key) - if override is not None: - processor_kwargs[key] = override + # Allow None as a valid override value for image config + processor_kwargs[key] = override return cls(**processor_kwargs, transformers_loading_kwargs=transformers_loading_kwargs) diff --git a/gr00t/model/gr00t_n1d6/setup.py b/gr00t/model/gr00t_n1d6/setup.py index b112a99..20359e5 100755 --- a/gr00t/model/gr00t_n1d6/setup.py +++ b/gr00t/model/gr00t_n1d6/setup.py @@ -75,6 +75,7 @@ def _create_model(self): tune_vlln=self.config.model.tune_vlln, state_dropout_prob=self.config.model.state_dropout_prob, backbone_trainable_params_fp32=self.config.model.backbone_trainable_params_fp32, + action_horizon=self.config.model.action_horizon, transformers_loading_kwargs=self.transformers_loading_kwargs, output_loading_info=True, **self.transformers_loading_kwargs, diff --git a/gr00t/policy/gr00t_policy.py b/gr00t/policy/gr00t_policy.py index f6457ae..ded04a2 100644 --- a/gr00t/policy/gr00t_policy.py +++ b/gr00t/policy/gr00t_policy.py @@ -375,8 +375,14 @@ def check_action(self, action: dict[str, Any]) -> None: Raises: AssertionError: If any validation check fails """ - # Validate each action key defined in the modality config - for action_key in self.modality_configs["action"].modality_keys: + # Validate each action key defined in the modality config (excluding pass_through_keys) + action_config = self.modality_configs["action"] + + # Skip keys used for intermediate action calculations + pass_through_keys = set(action_config.pass_through_keys or []) + for action_key in action_config.modality_keys: + if action_key in pass_through_keys: + continue # Pass-through keys are removed during processing # Check that the expected action key exists assert action_key in action, f"Action key '{action_key}' must be in action" @@ -635,9 +641,13 @@ def check_action(self, action: dict[str, Any]) -> None: AssertionError: If any validation check fails """ modality_configs = self.get_modality_config() + action_config = modality_configs["action"] + pass_through_keys = set(action_config.pass_through_keys or []) - # Validate each action key defined in the modality config - for action_key in modality_configs["action"].modality_keys: + # Validate each action key defined in the modality config (excluding pass_through_keys) + for action_key in action_config.modality_keys: + if action_key in pass_through_keys: + continue # Pass-through keys are removed during processing # Construct flat key expected in Gr00t sim environment (e.g., 'action.joints') parsed_key = f"action.{action_key}" assert parsed_key in action, f"Action key '{parsed_key}' must be in action" diff --git a/gr00t/policy/replay_policy.py b/gr00t/policy/replay_policy.py index 13005e3..02a5317 100644 --- a/gr00t/policy/replay_policy.py +++ b/gr00t/policy/replay_policy.py @@ -99,7 +99,11 @@ def __init__( def _preload_actions(self) -> None: """Preload all actions from the current episode for efficient replay.""" - action_keys = self.modality_configs["action"].modality_keys + action_config = self.modality_configs["action"] + + # Skip keys used for intermediate action calculations + pass_through_keys = set(action_config.pass_through_keys or []) + action_keys = [k for k in action_config.modality_keys if k not in pass_through_keys] self.actions: dict[str, np.ndarray] = {} for key in action_keys: @@ -267,7 +271,11 @@ def check_action(self, action: dict[str, Any]) -> None: Raises: AssertionError: If any validation check fails """ - for action_key in self.modality_configs["action"].modality_keys: + action_config = self.modality_configs["action"] + pass_through_keys = set(action_config.pass_through_keys or []) + for action_key in action_config.modality_keys: + if action_key in pass_through_keys: + continue # Pass-through keys are removed during processing assert action_key in action, f"Action key '{action_key}' must be in action" action_arr = action[action_key] diff --git a/media/gr00t-h-header.png b/media/gr00t-h-header.png new file mode 100644 index 0000000..ecdbd8b Binary files /dev/null and b/media/gr00t-h-header.png differ diff --git a/media/open-h-collage.jpg b/media/open-h-collage.jpg new file mode 100644 index 0000000..6060ded Binary files /dev/null and b/media/open-h-collage.jpg differ diff --git a/open_h/README.md b/open_h/README.md new file mode 100644 index 0000000..de9d5af --- /dev/null +++ b/open_h/README.md @@ -0,0 +1,126 @@ +
+ +# Open-H: Multi-Embodiment Healthcare Robot Training + +Open-H Dataset Collage + +*Screen captures from various datasets included in the Open-H dataset.* + +
+ +GR00T-H post-trains GR00T N1.6 on surgical robot data from multiple institutions and robot platforms simultaneously. The core challenge is that each institution records data differently β€” different robots, coordinate conventions, frame rates, camera setups, and state/action representations. GR00T-H solves this by defining per-embodiment modality configs that convert each dataset into a common representation (REL_XYZ_ROT6D for EEF poses) while preserving robot-specific details like clutch handling and motion scaling. + +Each embodiment gets its own projector index in the model, enabling embodiment-specific learned projections while sharing the core vision-language-action backbone. + +## Documentation + +| Guide | Description | +|-------|-------------| +| [Overview](docs/overview.md) | What's different from core GR00T, auto-registration, quick start workflows | +| [Action Configuration](docs/action_configuration.md) | REL_XYZ_ROT6D, rotation formats, the copy-EEF pattern, adding new embodiments | +| [Data Preparation](docs/data_preparation.md) | Stats pipeline, temporal statistics, troubleshooting | +| [Embodiment Comparison](embodiments/README.md) | All 16 embodiments at a glance β€” dimensions, cameras, action formats | + +For core GR00T concepts (LeRobot format, base ModalityConfig, inference API), see [`getting_started/`](../getting_started/). + +## Embodiments + +All supported embodiments live under `open_h/embodiments/`. Each subdirectory contains: +- A `*_config.py` that defines the modality configuration and registers it with GR00T +- A `modality.json` that maps raw dataset columns/indices to named keys +- A `README.md` with embodiment-specific details (data format, preparation steps, etc.) + +See [`open_h/embodiments/README.md`](embodiments/README.md) for a comparison table of all embodiments β€” covering dataset type, state/action formats, final action dimensions, number of arms, and number of cameras. + +Configs are auto-registered on import: `open_h/embodiments/__init__.py` discovers and executes every `*_config.py` file, which calls `register_modality_config()` to populate the global registry. Both `gr00t/experiment/launch_train.py` and `gr00t/experiment/launch_finetune.py` import `open_h.embodiments`, so built-in Open-H embodiments are available automatically. + +For built-in Open-H embodiment tags, do not pass `--modality-config-path` during finetuning or stats generation. None of the embodiments under `open_h/embodiments/` require it. The only time you should pass `--modality-config-path` from an Open-H workflow is when you are using the `NEW_EMBODIMENT` tag for a brand-new embodiment that is not already included in the Open-H registry. + +## Dataset Preparation + +Before training, each dataset needs normalization statistics computed. This is a prerequisite β€” training will fail without them. + +`prepare_datasets.sh` handles the full preparation pipeline for a set of datasets: + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag \ + --modality-json \ + /path/to/dataset_a /path/to/dataset_b ... +``` + +For each dataset, this script: +1. Copies the `modality.json` into the dataset's `meta/` directory +2. Generates `stats.json` (normalization stats over the raw parquet data) via `gr00t/data/stats.py` +3. Generates `temporal_stats.json` (normalization stats for actions after REL_XYZ_ROT6D conversion, with a temporal dimension for the action chunk) via `launch_finetune.py --calculate-norm-stats` + +Each dataset gets its own `temporal_stats.json`, but at training time the stats from all datasets sharing the same embodiment tag are merged (weighted by mix ratio) into a single set of normalization statistics per embodiment. + +## Training + +The primary training config is [`open_h/gr00t_h_config.yaml`](gr00t_h_config.yaml). It specifies all dataset paths, mix ratios, embodiment tags, model settings, and training hyperparameters. + +Before launching, replace every `REPLACE_WITH_OPEN_H_DATA_PATH` entry in that YAML with the absolute path to your local Open-H dataset root. + +### Multi-Embodiment Training + +The released GR00T-H checkpoint was trained on 4 nodes with 8 GPUs each: + +```bash +uv run torchrun --nnodes=4 --nproc_per_node=8 \ + --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \ + gr00t/experiment/launch_train.py \ + --load-config-path open_h/gr00t_h_config.yaml +``` + +If using fewer nodes or GPUs, adjust both `global_batch_size` and `num_gpus` in the config accordingly. For example, on a single node with 8 GPUs, set `num_gpus: 8` and `global_batch_size: 256`. + +Edit `gr00t_h_config.yaml` to add/remove datasets, adjust mix ratios, or change hyperparameters. Each dataset entry specifies an `embodiment_tag` that maps to its registered modality config. + +### Single-Embodiment Finetuning + +```bash +uv run torchrun --nproc_per_node=8 --master_port=29500 \ + gr00t/experiment/launch_finetune.py \ + --base-model-path nvidia/GR00T-H \ + --dataset-path /path/to/dataset \ + --embodiment-tag \ + --num-gpus 8 \ + --global-batch-size 32 \ + --max-steps 20000 \ + --output-dir /path/to/output +``` + +Key flags: +- `--embodiment-tag`: Must match the tag registered in the config file (e.g., `CMR_VERSIUS`, `JHU_IMERSE_DVRK`) +- `--modality-config-path`: Leave unset for all built-in Open-H embodiments; use it only when finetuning with `NEW_EMBODIMENT` for a brand-new embodiment config +- `--calculate-norm-stats`: Compute normalization statistics and exit (no training) + +## File Structure + +``` +open_h/ +β”œβ”€β”€ README.md # This file +β”œβ”€β”€ __init__.py # Package marker +β”œβ”€β”€ gr00t_h_config.yaml # Multi-embodiment GR00T-H training configuration +β”œβ”€β”€ prepare_datasets.sh # Dataset preparation script (stats generation) +└── embodiments/ # All embodiment definitions + β”œβ”€β”€ __init__.py # Auto-discovers and registers all *_config.py files + β”œβ”€β”€ README.md # Comparison table of all embodiments + β”œβ”€β”€ cmr_versius/ + β”œβ”€β”€ hamlyn_dvrk/ + β”œβ”€β”€ jhu_imerse_dvrk/ + β”œβ”€β”€ jhu_lscr_dvrk/ + β”œβ”€β”€ moon_maestro/ + β”œβ”€β”€ obuda_dvrk/ + β”œβ”€β”€ polyu_sim/ + β”œβ”€β”€ rob_surgical_bitrack/ + β”œβ”€β”€ sanoscience_sim/ + β”œβ”€β”€ stanford_dvrk_real/ + β”œβ”€β”€ tud_tundra_ur5e/ + β”œβ”€β”€ tum_sonata_franka/ + β”œβ”€β”€ turin_mitic_ex_vivo/ + β”œβ”€β”€ ucb_dvrk/ + β”œβ”€β”€ ucsd_dvrk/ + └── ustc_torin_tuodao/ +``` diff --git a/open_h/__init__.py b/open_h/__init__.py new file mode 100644 index 0000000..af901df --- /dev/null +++ b/open_h/__init__.py @@ -0,0 +1 @@ +# Open-H: Multi-Embodiment Healthcare Robot Training diff --git a/open_h/docs/action_configuration.md b/open_h/docs/action_configuration.md new file mode 100644 index 0000000..48aa580 --- /dev/null +++ b/open_h/docs/action_configuration.md @@ -0,0 +1,249 @@ +# Action Configuration for GR00T-H + +This document explains the REL_XYZ_ROT6D action representation, how to configure it for your embodiment, and how to handle the common case where absolute EEF poses live in the state column rather than the action column. For background on `ModalityConfig` basics (delta_indices, modality_keys, normalization keys), see [`getting_started/data_config.md`](../../getting_started/data_config.md). + + +## 1. Why REL_XYZ_ROT6D + +Surgical robot datasets store actions in many formats. Some use quaternions (xyzw or wxyz), others use Euler angles or axis-angle vectors. Some record absolute EEF poses, others record delta commands or joint velocities. Training a single model across all of them requires a common representation. + +GR00T-H unifies EEF actions into **REL_XYZ_ROT6D**: + +- **Translation**: relative XYZ displacement from the current EEF position (3D) +- **Rotation**: relative rotation in 6D representation, specifically the first two columns of the rotation matrix (6D) +- **Total per end-effector key**: 9D (3D xyz_rel + 6D rot6d_rel) + +Non-EEF actions (gripper open/close, energy buttons, jaw angles) stay **ABSOLUTE** and keep their original dimensionality. They aren't converted. + +The conversion happens dynamically inside `StateActionProcessor`. You configure what your data provides, the processor handles the rest. The core math lives in `gr00t/data/state_action/pose.py`, specifically `convert_to_rel_xyz_rot6d()` for training and `convert_from_rel_xyz_rot6d()` for inference. + + +## 2. The Core Requirement: Absolute EEF Measurement + +REL_XYZ_ROT6D computes the cumulative displacement from the current EEF pose (at the reference timestep) to each future timestep in the action chunk. You need an absolute EEF pose (xyz + rotation) accessible at each timestep. + +This can come from: +- **Your action column directly**, if actions are absolute EEF setpoints (e.g., JHU IMERSE dVRK, Hamlyn dVRK) +- **Your state column**, if only state has the absolute EEF pose. This uses the "Copy EEF" pattern described in Section 5 (e.g., TUD TUNDRA UR5e, PolyU Sim) + +The rotation component can be quaternion (4D), Euler angles (3D), or already rot6d (6D). Set `input_rotation_format` accordingly. If you have no absolute EEF poses anywhere in your data, you cannot use REL_XYZ_ROT6D. Consider joint-space RELATIVE actions instead. + + +## 3. ActionConfig Fields Reference + +Every action modality key needs an `ActionConfig`. The fields are defined in `gr00t/data/types.py`: + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `rep` | `ActionRepresentation` | (required) | `REL_XYZ_ROT6D` for EEF pose keys, `ABSOLUTE` for gripper/non-EEF | +| `type` | `ActionType` | (required) | `EEF` for end-effector, `NON_EEF` for everything else | +| `format` | `ActionFormat` | (required) | `XYZ_ROT6D` for EEF after conversion, `DEFAULT` for non-EEF | +| `state_key` | `str` or `None` | `None` | Which state key provides the EEF reference frame for relative conversion | +| `input_rotation_format` | `str` | `"quat"` | Rotation format in your action data: `"quat"`, `"rot6d"`, or `"euler"` | +| `input_quat_order` | `str` | `"xyzw"` | Quaternion ordering if input is quat: `"xyzw"` (scipy) or `"wxyz"` (scalar-first) | +| `reference_rotation_format` | `str` | `"rot6d"` | Rotation format in the reference state | +| `reference_quat_order` | `str` | `"xyzw"` | Quaternion ordering if reference is quat | +| `normalization_type` | `str` | `"temporal_meanstd"` | `"temporal_meanstd"` (recommended), `"meanstd"`, `"minmax"`, or `"skip"` | +| `hold_through_clutch` | `bool` | `False` | For ABSOLUTE actions only: hold last engaged value during clutch-out instead of zeroing. Used by CMR Versius; applicable to any embodiment with a controller clutching mechanism | +| `translation_scaling_key` | `str` or `None` | `None` | State key whose value scales relative translation. Used by CMR Versius where hand controller kinematics differ from instrument kinematics; applicable to any embodiment with such a scaling mismatch | +| `rotation_scaling_key` | `str` or `None` | `None` | State key whose value scales relative rotation angle. Same use case as `translation_scaling_key` | + +A typical EEF config and a typical non-EEF config, side by side: + +```python +# EEF pose (converted to REL_XYZ_ROT6D) +ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="eef_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + reference_quat_order="xyzw", +) + +# Gripper (stays absolute) +ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="meanstd", +) +``` + + +## 4. Rotation Format Variations Across Open-H + +Different embodiments store rotations differently. Here's what the actual configs use: + +| Embodiment | `input_rotation_format` | `input_quat_order` | `reference_rotation_format` | `reference_quat_order` | Notes | +|---|---|---|---|---|---| +| TUD TUNDRA UR5e | `quat` | `xyzw` | `quat` | `xyzw` | Standard scipy convention | +| PolyU Sim | `quat` | `xyzw` | `quat` | `xyzw` | Standard scipy convention | +| JHU IMERSE dVRK | `quat` | `xyzw` | `quat` | `xyzw` | Standard scipy convention | +| Hamlyn dVRK (15Hz + 30Hz) | `quat` | `wxyz` | `quat` | `wxyz` | Scalar-first quaternion | +| Stanford dVRK Real | `euler` | (n/a) | `euler` | (n/a) | Euler RPY in radians, `xyz` extrinsic convention | +| CMR Versius | `quat` | `xyzw` | `quat` | `xyzw` | Quaternion with motion scaling keys | + +If your dataset uses quaternions, check whether they're `xyzw` (scipy convention, most common in Open-H) or `wxyz` (scalar-first, used by Hamlyn and some other robotics frameworks). Getting this wrong causes silent corruption: rotations will be nonsensical, and the model will train on garbage without raising any errors. + +If your dataset uses Euler angles, set `input_rotation_format="euler"`. The processor assumes `xyz` extrinsic convention (roll, pitch, yaw) in radians. + + +## 5. The Copy-EEF Pattern: Sourcing Actions from State + +Some datasets store delta commands, joint velocities, or joint angles in the action column rather than absolute EEF poses. REL_XYZ_ROT6D needs absolute EEF poses at future timesteps. If the absolute EEF pose only exists in the state column, you can redirect the action loader to read from state instead. + +### The Solution + +Use the `original_key` field in `modality.json` to tell the data loader where to actually read from. Combined with `delta_indices=list(range(1, H+1))`, this reads future states as action targets. + +### Example 1: TUD TUNDRA UR5e + +The TUD dataset's action column contains delta commands (`dx, dy, dz, droll`), which can't be used for REL_XYZ_ROT6D directly. The absolute EEF pose lives in `observation.state` at indices 26-33. + +The modality JSON (`modality_grasping_retraction.json`): + +```json +{ + "state": { + "joint_position": {"start": 8, "end": 14, "original_key": "observation.state"}, + "eef_pose": {"start": 26, "end": 33, "original_key": "observation.state"} + }, + "action": { + "eef_pose": {"start": 26, "end": 33, "original_key": "observation.state"}, + "gripper": {"start": 4, "end": 5, "original_key": "action"} + } +} +``` + +How it works: + +- `action.eef_pose` has `"original_key": "observation.state"`. This tells the data loader to read from the state column, not the action column. +- The indices `[26:33]` select the 7D EEF pose (xyz + quaternion) from the 33D state vector. +- The config sets `delta_indices=list(range(1, 51))`. This reads `state[t+1], state[t+2], ..., state[t+50]`, creating a 50-step trajectory of absolute EEF poses from future timesteps. +- `action.gripper` uses `"original_key": "action"` to read from the actual action column. Ensure the gripper values at each timestep are aligned with the corresponding action timestep (e.g., `gripper[t]` corresponds to `state[t+1]` when using `delta_indices` starting from 1). +- The `+1` offset in delta_indices (starting from 1, not 0) means the first action target is the next timestep's pose. + +### Example 2: PolyU Sim + +The PolyU simulation stores joint deltas in its action column, but absolute Cartesian pose lives in a separate `observation.cartesian_state` column. + +The modality JSON: + +```json +{ + "state": { + "psm_joints": {"start": 0, "end": 10}, + "psm_cartesian_pose": {"start": 0, "end": 7, "original_key": "observation.cartesian_state"} + }, + "action": { + "psm_cartesian_pose": {"start": 0, "end": 7, "original_key": "observation.cartesian_state"}, + "psm_gripper": {"start": 10, "end": 11, "original_key": "action"} + } +} +``` + +Same pattern, but reading from a different source column (`observation.cartesian_state` instead of `observation.state`). The `original_key` can point to any column in your parquet file. This flexibility lets you source action data from whichever column actually contains the absolute EEF pose. + + +## 6. The state_key and pass_through_keys + +`ActionConfig.state_key = "eef_pose"` tells the processor: use this state key as the reference frame for relative conversion. The reference pose at the current timestep becomes the origin, and all future EEF poses in the action chunk are expressed as cumulative displacements from it. + +`pass_through_keys` is an optional field on `ModalityConfig` that lists state keys which are loaded from the dataset for intermediate calculations (such as providing the REL_XYZ_ROT6D reference frame) but are **not** sent to the model's state encoder. Example: + +```python +"state": ModalityConfig( + delta_indices=[0], + modality_keys=["joint_position", "eef_pose"], + mean_std_embedding_keys=["joint_position"], + pass_through_keys=["eef_pose"], # loaded for processing, NOT sent to model +) +``` + +In this example, `joint_position` is embedded by the state encoder and fed to the model. `eef_pose` is used only for the REL_XYZ_ROT6D conversion and then discarded. The `state_key` referenced by `ActionConfig` does not have to be in `pass_through_keys` β€” if it isn't, it will be both embedded as model input and used as the conversion reference. + +Note that GR00T-H was trained with `state_dropout_prob_per_embodiment: 1.0` for all embodiments in `gr00t_h_config.yaml`. This zeros out all state inputs before the encoder, making the model vision-only at inference. The `pass_through_keys` mechanism is separate from state dropout β€” it controls which keys reach the encoder, while state dropout controls whether the encoder output is used. + + +## 7. Dual-Arm Configurations + +For multi-arm robots, each arm gets its own modality key and `ActionConfig`. Here's the structure from JHU IMERSE dVRK (dual-arm da Vinci): + +```python +"action": ModalityConfig( + delta_indices=list(range(50)), + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + action_configs=[ + ActionConfig(rep=ActionRepresentation.REL_XYZ_ROT6D, type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, state_key="psm1_pose", + input_rotation_format="quat", reference_rotation_format="quat", + normalization_type="temporal_meanstd"), + ActionConfig(rep=ActionRepresentation.ABSOLUTE, type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, normalization_type="temporal_meanstd"), + ActionConfig(rep=ActionRepresentation.REL_XYZ_ROT6D, type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, state_key="psm2_pose", + input_rotation_format="quat", reference_rotation_format="quat", + normalization_type="temporal_meanstd"), + ActionConfig(rep=ActionRepresentation.ABSOLUTE, type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, normalization_type="temporal_meanstd"), + ], +) +``` + +Key rules: +- `action_configs` must have the same length as `modality_keys`. One `ActionConfig` per key, in order. +- Each EEF `ActionConfig` points to its own `state_key` for the reference frame (`"psm1_pose"` for PSM1, `"psm2_pose"` for PSM2). +- Non-EEF keys (grippers) don't need a `state_key`. + +The SanoScience Sim embodiment extends this to 4 arms (8 action configs). The pattern scales naturally. + + +## 8. Adding Your Own Embodiment: Decision Guide + +### Step 1: Do you have absolute EEF poses? + +- **In your action column** (e.g., absolute Cartesian setpoints): map the action keys directly in modality.json using `start`/`end` indices. No `original_key` is needed β€” the loader reads from the `action` column by default. See `jhu_imerse_dvrk/modality.json` for an example. +- **Only in your state column**: use the Copy-EEF pattern from Section 5. Set `"original_key": "observation.state"` (or whatever column contains the EEF pose). Then, determine the proper `delta_indices` offset to use, such that `action[t] = state[t+1]` +- **Nowhere** (only joint angles or delta commands with no FK available): you cannot use REL_XYZ_ROT6D for EEF actions. Consider joint-space RELATIVE representation instead. + +### Step 2: What rotation format does your data use? + +- Quaternion xyzw: `input_rotation_format="quat"`, `input_quat_order="xyzw"` +- Quaternion wxyz: `input_rotation_format="quat"`, `input_quat_order="wxyz"` +- Already 6D rotation: `input_rotation_format="rot6d"` +- Euler angles (radians, xyz extrinsic): `input_rotation_format="euler"` +- Euler angles in some other convention: convert to one of the above during dataset creation + +### Step 3: Single arm or multi-arm? + +- **Single arm**: one EEF key + one gripper key (if applicable). See TUD TUNDRA UR5e or PolyU Sim. +- **Dual arm**: separate keys per arm. See JHU IMERSE dVRK or Hamlyn dVRK. +- **More than two**: same pattern, more keys. See SanoScience Sim (4 instruments) or Rob Surgical BiTrack (3 arms). + +### Step 4: Create your config files + +1. **Add your embodiment tag** to `gr00t/data/embodiment_tags.py` (add a new entry to the `EmbodimentTag` enum). +2. **Add a projector index** in `gr00t/model/gr00t_n1d6/processing_gr00t_n1d6.py` (maps your tag to a unique projector slot). +3. **Create your config directory**: `open_h/embodiments/your_robot/` +4. **Write `your_robot_config.py`**: define the `ModalityConfig` dict with video, state, action, and language entries. Call `register_modality_config()` at module level. +5. **Write `modality.json`**: map your parquet columns and indices to named keys. Use `original_key` to redirect reads when needed. +6. **Run `prepare_datasets.sh`** to generate `stats.json` and `temporal_stats.json`. + +### Reference configs by pattern + +| Pattern | Example embodiment | Key features | +|---------|-------------------|--------------| +| Single arm, copy-EEF | `tud_tundra_ur5e` | Actions sourced from state column, quaternion xyzw | +| Dual arm, actions in action column | `jhu_imerse_dvrk` | Direct action column, quaternion xyzw | +| Dual arm, wxyz quaternions | `hamlyn_dvrk` | Scalar-first quaternion ordering | +| Dual arm, Euler angles | `stanford_dvrk_real` | Euler RPY input, xyz extrinsic | +| Dual arm, clutch-aware | `cmr_versius` | Motion scaling, hold-through-clutch, engagement filtering | +| Single arm, simulated | `polyu_sim` | Copy-EEF from separate cartesian_state column | diff --git a/open_h/docs/data_preparation.md b/open_h/docs/data_preparation.md new file mode 100644 index 0000000..baf2c6a --- /dev/null +++ b/open_h/docs/data_preparation.md @@ -0,0 +1,133 @@ +# Preparing Datasets for GR00T-H + +## Prerequisites + +Your data must be in LeRobot v2 format. If you're converting from v3, run `scripts/lerobot_conversion/convert_v3_to_v2.py` first. For the full format specification (parquet layout, episode structure, video conventions), see the [Data Preparation Guide](../../getting_started/data_preparation.md). + +Each dataset needs a `modality.json` that maps raw parquet columns and index ranges to named keys (e.g., `eef_pose`, `gripper`). If your dataset stores delta commands or joint angles in the action column rather than absolute EEF poses, see the [Copy-EEF pattern](action_configuration.md#5-the-copy-eef-pattern-sourcing-actions-from-state), which sources action data from `observation.state` at future timesteps. + +Your embodiment must have a registered config. Built-in Open-H embodiments are auto-registered on import, so no extra setup is needed. For new embodiments, follow [Adding Your Own Embodiment](action_configuration.md#8-adding-your-own-embodiment-decision-guide). + +## The Three Statistics Files + +GR00T-H training relies on normalization statistics stored alongside each dataset. Three files live under `meta/`, each serving a distinct purpose. + +### `meta/stats.json` + +Flat statistics computed directly from raw parquet data. + +- **Contains**: mean, std, min, max, q01, q99 per state and action key +- **Shape**: `(dim,)` ... one value per dimension, no temporal component +- **Used for**: state normalization (min-max or mean-std scaling) +- **Generated by**: `gr00t/data/stats.py` + +### `meta/temporal_stats.json` + +Per-timestep statistics computed *after* REL_XYZ_ROT6D action conversion. This is a GR00T-H addition for action normalization across multi-step horizons. + +- **Contains**: mean, std, q02, q98 per action key, per timestep in the action chunk +- **Shape**: `(horizon, dim)` ... separate statistics for every step in the chunk +- **Used for**: action normalization with `normalization_type="temporal_meanstd"` +- **Generated by**: `launch_finetune.py --calculate-norm-stats` + +#### Why temporal stats are beneficial + +REL_XYZ_ROT6D computes cumulative displacements from the current EEF pose to each future timestep in the action chunk. The magnitude of those displacements grows with the prediction horizon: + +- At step 1 (33ms into the future at 30Hz), displacement is small. +- At step 50 (1.7s into the future), cumulative displacement is much larger. + +Flat normalization (a single mean and std across all steps) applies the same scaling to both, which can under-normalize early steps and over-normalize late ones. Temporal stats address this by providing each timestep with its own mean and std, calibrated to the actual displacement magnitude at that horizon. The released GR00T-H checkpoint was trained with these per-timestep statistics. + +Here's what `temporal_stats.json` looks like in practice: + +```json +{ + "action": { + "eef_pose": { + "mean": [[0.001, 0.002, ...], [0.003, 0.005, ...], ...], + "std": [[0.01, 0.01, ...], [0.02, 0.03, ...], ...], + "q02": [[-0.02, -0.02, ...], [-0.05, -0.06, ...], ...], + "q98": [[0.02, 0.02, ...], [0.05, 0.06, ...], ...] + } + } +} +``` + +Each inner list has length equal to the action dimension (e.g., 9 for a single-arm EEF: 3 xyz + 6 rot6d). The outer list has `horizon` entries (e.g., 50 for a 50-step action chunk). + +### `meta/relative_stats.json` + +Generated alongside `stats.json`. Contains per-timestep stats for RELATIVE (non-EEF) actions, with the same `(horizon, dim)` shape concept as `temporal_stats.json` but computed differently. You generally don't need to worry about this file. It's produced automatically during stats generation. + +## The prepare_datasets.sh Pipeline + +Run the preparation script from the repository root: + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag \ + --modality-json \ + /path/to/dataset_a /path/to/dataset_b ... +``` + +The script processes each dataset through three steps. + +### Step 1: Copy modality mapping + +```bash +cp /meta/modality.json +``` + +Places the column-to-key mapping where the data loader expects it. If the `meta/` directory doesn't exist, the script creates it. + +### Step 2: Generate flat statistics + +```bash +uv run python gr00t/data/stats.py \ + --dataset-path \ + --embodiment-tag +``` + +Reads all parquet files, computes per-column statistics, and writes `meta/stats.json` (plus `meta/relative_stats.json`). + +### Step 3: Generate temporal statistics + +```bash +uv run python gr00t/experiment/launch_finetune.py \ + --base-model-path nvidia/GR00T-H \ + --dataset-path \ + --embodiment-tag \ + --calculate-norm-stats +``` + +This step loads the modality config for the embodiment (auto-registered, no `--modality-config-path` needed), iterates through episodes, performs the actual REL_XYZ_ROT6D conversion for each action chunk, computes per-timestep percentile statistics on the converted actions, and writes `meta/temporal_stats.json`. + +## Concrete Example: TUD TUNDRA UR5e + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag TUD_TUNDRA_UR5E \ + --modality-json open_h/embodiments/tud_tundra_ur5e/modality_grasping_retraction.json \ + /path/to/Surgical/TUD/260131_TUNDRA_dataset/grasping_retraction +``` + +After running, verify the output: + +```bash +ls /path/to/Surgical/TUD/260131_TUNDRA_dataset/grasping_retraction/meta/ +# Expected: episodes.jsonl info.json modality.json stats.json tasks.jsonl temporal_stats.json +``` + +## Stats Merging at Training Time + +When training on multiple datasets under the same embodiment tag (e.g., twelve dVRK datasets under `jhu_imerse_dvrk`), the training pipeline automatically merges their statistics. Each dataset keeps its own `temporal_stats.json`. At training time, `gr00t/data/percentile_merge.py` merges them, weighted by the mix ratios specified in the training config. + +## Troubleshooting + +| Problem | Cause | Fix | +|---------|-------|-----| +| `AssertionError: Embodiment tag already registered` | Passed `--modality-config-path` for a built-in tag | Remove `--modality-config-path`; the config is auto-registered | +| `IndexError: boolean index did not match indexed array along dimension 0` | `temporal_stats.json` shape doesn't match current `delta_indices` | Regenerate temporal stats with `prepare_datasets.sh` | +| `FileNotFoundError: meta/temporal_stats.json` | Stats not generated before training | Run `prepare_datasets.sh` before training | +| `KeyError: '' not in MODALITY_CONFIGS` | Embodiment tag misspelled or new embodiment not registered | Check spelling against `EmbodimentTag` enum; for new embodiments, ensure the config calls `register_modality_config()` | diff --git a/open_h/docs/overview.md b/open_h/docs/overview.md new file mode 100644 index 0000000..e1c943f --- /dev/null +++ b/open_h/docs/overview.md @@ -0,0 +1,122 @@ +# GR00T-H Overview + +## Intended Use + +GR00T-H is intended for use in robotics R&D, including exploration of surgical robotics and robotic ultrasound policies, benchmarking, and method development. It is not intended for clinical deployment, patient care, or medical decision-making. + +GR00T-H is not expected to work out of the box for arbitrary embodiments or tasks. It may produce reasonable behavior for the specific robots and tasks represented in the Open-H dataset, but the primary value of this release is as a pretrained VLA checkpoint for healthcare robotics. The expected workflow is to **finetune from GR00T-H** on your own robot's data β€” not to deploy it zero-shot. Zero-shot deployment on new embodiments or tasks is unlikely to produce usable results. + +## What Changed from Core GR00T + +This repository is a superset of upstream [Isaac-GR00T](https://github.com/NVIDIA/Isaac-GR00T). The table below summarizes every file added or modified. + +| File | Change | +|------|--------| +| `open_h/` (new package) | All embodiment configs, training YAML, dataset prep script | +| `gr00t/data/embodiment_tags.py` | 20 new surgical robot tags added to `EmbodimentTag` enum | +| `gr00t/data/types.py` | `REL_XYZ_ROT6D` action representation, clutch/motion-scaling fields, temporal normalization type | +| `gr00t/data/state_action/pose.py` | `convert_to_rel_xyz_rot6d()` / `convert_from_rel_xyz_rot6d()` conversion math | +| `gr00t/data/state_action/state_action_processor.py` | REL_XYZ_ROT6D conversion path, clutch-aware zeroing, motion scaling, temporal normalization | +| `gr00t/data/stats.py` | Temporal stats generation (per-timestep normalization for action chunks) | +| `gr00t/experiment/launch_finetune.py` | Auto-registers Open-H configs on import, adds `--calculate-norm-stats` mode | +| `gr00t/experiment/launch_train.py` | Auto-registers Open-H configs on import, refactored config loading | +| `gr00t/model/gr00t_n1d6/processing_gr00t_n1d6.py` | Projector index mappings for all Open-H embodiments | + +## Embodiment Auto-Registration + +All Open-H embodiment configs are auto-registered at import time. When any training or stats script runs, `open_h/embodiments/__init__.py` discovers and executes every `*_config.py` file, which calls `register_modality_config()` for each embodiment tag. This means: + +- **Built-in Open-H embodiment**: pass `--embodiment-tag ` only. Do not pass `--modality-config-path` β€” the config is already registered and doing so will raise an `AssertionError`. +- **Brand-new robot not in Open-H**: pass `--embodiment-tag NEW_EMBODIMENT --modality-config-path your_config.py`. + +Note: the `--modality-json` flag in `prepare_datasets.sh` is unrelated β€” it points to the JSON column mapping, not the Python config. + +## Quick Start + +### A. Finetune a built-in Open-H embodiment (single dataset) + +Four steps: prepare dataset statistics, verify data with replay, finetune, evaluate. + +```bash +# 1. Prepare dataset (copies modality.json, generates stats.json + temporal_stats.json) +bash open_h/prepare_datasets.sh \ + --embodiment-tag TUD_TUNDRA_UR5E \ + --modality-json open_h/embodiments/tud_tundra_ur5e/modality_grasping_retraction.json \ + /path/to/tud_dataset +``` + +```bash +# 2. Replay recorded actions to verify data processing and action conversion +uv run python gr00t/eval/run_gr00t_server.py \ + --dataset-path /path/to/tud_dataset \ + --embodiment-tag TUD_TUNDRA_UR5E \ + --execution-horizon 8 +``` + +Run this before any training to confirm your embodiment config loads correctly and actions are converted as expected. The server replays ground-truth actions from the dataset through the full processing pipeline (modality loading, REL_XYZ_ROT6D conversion, normalization). See the [Policy Guide](../../getting_started/policy.md#debugging-with-replaypolicy) for client-side usage and episode switching. + +```bash +# 3. Finetune from GR00T-H checkpoint +uv run torchrun --nproc_per_node=8 --master_port=29500 \ + gr00t/experiment/launch_finetune.py \ + --base-model-path nvidia/GR00T-H \ + --dataset-path /path/to/tud_dataset \ + --embodiment-tag TUD_TUNDRA_UR5E \ + --num-gpus 8 \ + --global-batch-size 32 \ + --max-steps 20000 \ + --output-dir /path/to/output +``` + +```bash +# 4. Open-loop evaluation +uv run python gr00t/eval/open_loop_eval.py \ + --dataset-path /path/to/tud_dataset \ + --embodiment-tag TUD_TUNDRA_UR5E \ + --model-path /path/to/output/checkpoint-20000 \ + --traj-ids 0 \ + --action-horizon 50 \ + --steps 400 +``` + +Note: no `--modality-config-path` anywhere. The config is already registered. + +The evaluation script writes trajectory visualizations to `/tmp/open_loop_eval/` showing predicted vs. ground-truth actions with per-dimension MSE. + +### B. Multi-embodiment training + +The released GR00T-H checkpoint was trained on 4 nodes with 8 GPUs each, using the config at `open_h/gr00t_h_config.yaml`. + +Before launching, replace every `REPLACE_WITH_OPEN_H_DATA_PATH` in the YAML with your local data root. + +```bash +uv run torchrun --nnodes=4 --nproc_per_node=8 \ + --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \ + gr00t/experiment/launch_train.py \ + --load-config-path open_h/gr00t_h_config.yaml +``` + +For fewer GPUs, adjust both `global_batch_size` and `num_gpus` in the YAML accordingly. On a single 8-GPU node, set `num_gpus: 8` and `global_batch_size: 256`. + +### C. Add your own healthcare embodiment + +If your robot isn't already in Open-H, you'll need to: + +1. Convert your data to LeRobot v2 format +2. Write a `modality.json` mapping your dataset columns to named keys +3. Define a `ModalityConfig` with `ActionConfig` entries for your end-effector(s) +4. Run the stats pipeline to generate normalization statistics + +See [Action Configuration](action_configuration.md) for how to write a config and [Data Preparation](data_preparation.md) for the stats pipeline. + +## Further Reading + +| Topic | Document | +|-------|----------| +| Action config deep-dive | [Action Configuration](action_configuration.md) | +| Data preparation pipeline | [Data Preparation](data_preparation.md) | +| Embodiment comparison table | [Embodiment Overview](../embodiments/README.md) | +| LeRobot v2 format | [Data Preparation Guide](../../getting_started/data_preparation.md) | +| Base ModalityConfig reference | [Data Config Guide](../../getting_started/data_config.md) | +| Inference / Policy API | [Policy Guide](../../getting_started/policy.md) | +| Core GR00T finetuning | [Finetune New Embodiment](../../getting_started/finetune_new_embodiment.md) | diff --git a/open_h/embodiments/README.md b/open_h/embodiments/README.md new file mode 100644 index 0000000..0650f42 --- /dev/null +++ b/open_h/embodiments/README.md @@ -0,0 +1,32 @@ +# Open-H Embodiment Overview + +Summary of all surgical robot embodiments currently supported by GR00T-H. + +## Embodiment Comparison + +| Embodiment | Dataset Type | Raw State Format | Raw Action Format | Final Action Format | Arms Used | Cameras Used | +|---|---|---|---|---|---|---| +| **CMR Versius** (`cmr_versius`) | Clinical | Cartesian EEF pose + gripper per arm (26D) | Cartesian EEF pose + gripper per arm (26D) | 2Γ—9D pose + 2Γ—1D gripper = **20D** | 2 | 1 | +| **JHU IMERSE dVRK** (`jhu_imerse_dvrk`) | Surgical tabletop | Cartesian EEF pose + gripper per arm (16D) | Cartesian EEF setpoint + gripper per arm (16D) | 2Γ—9D pose + 2Γ—1D gripper = **20D** | 2 | 3 | +| **JHU LSCR dVRK** (`jhu_lscr_dvrk`) | Surgical tabletop | Cartesian EEF pose + gripper per arm (16D) | Cartesian EEF setpoint + gripper per arm (16D) | 2Γ—9D pose + 2Γ—1D gripper = **20D** | 2 | 2-3 | +| **UCB dVRK** (`ucb_dvrk`) | Surgical tabletop | Cartesian EEF pose + gripper (16D) + joints (14D) | Cartesian EEF setpoint + gripper per arm (16D) | 2Γ—9D pose + 2Γ—1D gripper = **20D** | 2 | 2 | +| **Obuda dVRK** (`obuda_dvrk`) | Surgical tabletop | Cartesian EEF pose + gripper per arm (16D) | Cartesian EEF setpoint + gripper per arm (16D) | 2Γ—9D pose + 2Γ—1D gripper = **20D** | 2 | 3 | +| **Stanford dVRK Real** (`stanford_dvrk_real`) | Surgical tabletop | Cartesian EEF pose (Euler) + gripper per arm (14D) | Cartesian EEF pose (Euler) + gripper per arm (14D) | 2Γ—9D pose + 2Γ—1D gripper = **20D** | 2 | 2 | +| **UCSD dVRK** (`ucsd_dvrk`) | Ex-vivo | Cartesian EEF pose + gripper per arm (16D) | Cartesian EEF pose + gripper per arm (16D) | 2Γ—9D pose + 2Γ—1D gripper = **20D** | 2 | 2 | +| **Hamlyn dVRK** (`hamlyn_dvrk`) | Ex-vivo | Cartesian EEF pose + gripper per arm (16D) | Cartesian EEF pose + gripper per arm (16D) | 2Γ—9D pose + 2Γ—1D gripper = **20D** | 2 | 3 | +| **Turin MITIC** (`turin_mitic_ex_vivo`) | Ex-vivo | Joint angles per arm (12D) + EEF pass-through (14D) | Cartesian EEF pose per arm (14D) | 2Γ—9D pose = **18D** | 2 | 2 | +| **USTC Torin/Tuodao** (`ustc_torin_tuodao`) | Clinical + Ex-vivo | Joint angles per arm (14D) + EEF pass-through (16D) | Cartesian absolute pose + gripper per arm (16D) | 2Γ—9D pose + 2Γ—1D gripper = **20D** | 2 | 2 | +| **TUD TUNDRA UR5e** (`tud_tundra_ur5e`) | Clinical (porcine) | Joint position (6D) + EEF pass-through (7D) | Cartesian EEF pose (7D) + gripper (1D) | 1Γ—9D pose + 1Γ—1D gripper = **10D** | 1 | 2 | +| **TUM SonATA Franka** (`tum_sonata_franka`) | Ultrasound phantom | Joint angles (7D) + force/torque (6D) + EEF pass-through (6D) | Cartesian EEF pose, Euler (6D) | 1Γ—9D pose = **9D** | 1 | 3 | +| **Moon Maestro** (`moon_maestro`) | Surgical tabletop | Joint angles per arm (18D) | Delta translation per arm (6D) | 2Γ—3D delta xyz = **6D** | 2 | 2 | +| **Rob Surgical Bitrack** (`rob_surgical_bitrack`) | Surgical tabletop | Cartesian EEF pose (Euler) per arm (18D) | Cartesian EEF pose (Euler) per arm (18D) | 3Γ—9D pose = **27D** | 3 | 1 | +| **SanoScience Sim** (`sanoscience_sim`) | Simulation | Cartesian EEF pose + gripper per instrument (32D) | Cartesian EEF pose + gripper per instrument (32D) | 4Γ—9D pose + 4Γ—1D gripper = **40D** | 4 | 1 | +| **PolyU Sim** (`polyu_sim`) | Simulation | Joint angles (10D) + EEF pass-through (7D) | Cartesian EEF pose (7D) + gripper (1D) | 1Γ—9D pose + 1Γ—1D gripper = **10D** | 1 | 1 | + +## Column Definitions + +- **Raw State Format**: Observation state representation as stored in the dataset, before any model preprocessing. "EEF pass-through" indicates the state is not embedded by the model but is used as a reference frame for relative action conversion. +- **Raw Action Format**: Action representation as stored in the dataset, before conversion to the model's internal format. +- **Final Action Format**: The model's output action dimension after REL_XYZ_ROT6D conversion. Each pose key becomes 9D (3D relative xyz + 6D rotation). Gripper and other scalar keys stay at their original dimension. The breakdown shows exactly how the total is computed. +- **Arms Used**: Number of independently controlled robot arms or instruments used by the model. +- **Cameras Used**: Number of camera views used by the model. diff --git a/open_h/embodiments/__init__.py b/open_h/embodiments/__init__.py new file mode 100644 index 0000000..592accc --- /dev/null +++ b/open_h/embodiments/__init__.py @@ -0,0 +1,17 @@ +"""Auto-discover and register all Open-H embodiment configs. + +Importing this module finds every *_config.py file under each embodiment +subdirectory and executes it, which triggers register_modality_config() +calls that populate the global MODALITY_CONFIGS registry. +""" + +import importlib.util +from pathlib import Path + + +_embodiments_dir = Path(__file__).parent + +for _config_file in sorted(_embodiments_dir.glob("*/*_config.py")): + _spec = importlib.util.spec_from_file_location(_config_file.stem, _config_file) + _module = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_module) diff --git a/open_h/embodiments/cmr_versius/README.md b/open_h/embodiments/cmr_versius/README.md new file mode 100644 index 0000000..01cf7d8 --- /dev/null +++ b/open_h/embodiments/cmr_versius/README.md @@ -0,0 +1,183 @@ +# CMR Versius + +## Embodiment Configuration + +| Property | Value | +|----------|-------| +| **Embodiment Tag** | `cmr_versius` | +| **Projector Index** | 4 | +| **Config File** | `open_h/embodiments/cmr_versius/cmr_versius_config.py` | + +This enables combined training with other surgical embodiments, with embodiment-specific learned projections in the model. + +--- + +## Action & State Format + +CMR Versius records **hand controller (haptic) poses**, not robot end-effector poses. All kinematics represent the surgeon's controller position in camera frame. Actions and state are reformatted with left arm first, then right arm, pose components grouped together (xyz + rotation + gripper). State uses quaternion (xyzw); actions use 6D rotation for REL_XYZ_ROT6D. + +### Action Space (20D) + +| Idx | Name | Type | Representation | Description | +|-----|------|------|----------------|-------------| +| **Left Arm (0-9)** ||||| +| 0-2 | `xyz_left` | float | REL_XYZ_ROT6D | Translation relative to current hand controller pose | +| 3-8 | `rot6d_left` | float | REL_XYZ_ROT6D | 6D rotation relative to current hand controller pose | +| 9 | `pince_left` | float | ABSOLUTE | Gripper/pince [0-1] (sample-and-hold during clutch) | +| **Right Arm (10-19)** ||||| +| 10-12 | `xyz_right` | float | REL_XYZ_ROT6D | Translation relative to current hand controller pose | +| 13-18 | `rot6d_right` | float | REL_XYZ_ROT6D | 6D rotation relative to current hand controller pose | +| 19 | `pince_right` | float | ABSOLUTE | Gripper/pince [0-1] (sample-and-hold during clutch) | + +**Pass-Through Keys** (used for processing, removed before model): + +| Name | Type | Purpose | +|------|------|---------| +| `hapticengaged_left` | bool[T] | Per-timestep clutch state for left arm | +| `hapticengaged_right` | bool[T] | Per-timestep clutch state for right arm | + +### State Space (16D embedded) + +| Name | Dims | Description | +|------|------|-------------| +| `left_pose` | 7D | Hand controller xyz + quat_xyzw (from action column) | +| `left_gripper` | 1D | Pince [0-1] | +| `right_pose` | 7D | Hand controller xyz + quat_xyzw (from action column) | +| `right_gripper` | 1D | Pince [0-1] | + +**Pass-Through Keys** (used for processing, not sent to model): + +| Name | Dims | Purpose | +|------|------|---------| +| `translation_scaling` | 1D | Motion scaling factor for REL_XYZ_ROT6D conversion | +| `rotation_scaling` | 1D | Rotation scaling factor for REL_XYZ_ROT6D conversion | +| `hapticengaged_left` | 1D | Clutch state for left arm (filtering + delta re-integration) | +| `hapticengaged_right` | 1D | Clutch state for right arm (filtering + delta re-integration) | + +Instrument type, arm color, and arm-to-haptic linkage are encoded as per-timestep **language prompts** (`instruction.text_with_state`) rather than state embeddings. See [State Prompt Preprocessing](#state-prompt-preprocessing). + +### CMR-Specific Details + +- **Extra context**: Electrosurgery mode, arm linking, instrument type, arm color β€” sent via language prompts, not state embedding +- **Clutch handling**: Full clutch-aware processing pipeline (see below) +- **State source**: Extracted from `action` column via `original_key` in modality.json + +--- + +## Clutch-Aware Processing + +CMR Versius data presents a unique challenge: surgeons frequently **clutch out** (disengage) during procedures to reposition their hands without moving the robot arms. Computing relative actions naively across these clutch events produces invalid training data. + +### The Problem + +In teleoperation with clutch: +- **Controller (master):** Moves freely when disengaged +- **Robot (slave):** Holds position when disengaged +- **Data recorded:** Controller position (which diverges from robot during clutch) + +Standard REL_XYZ_ROT6D computation (`action[t] - state[ref]`) fails because: +1. **Phantom jumps:** Controller repositioning during clutch appears as large "movements" +2. **Invalid targets:** Model learns to predict controller motion, not robot motion +3. **Gripper drops:** Zeroing gripper during clutch teaches model to "drop the needle" + +### Dataset Statistics + +Analysis of cholecystectomy data (100 episodes, 289K samples): +- **76.3%** of episodes have clutch transitions +- **Mean 5.7** clutch events per episode +- **13.4%** of samples have engagement changes within 2s action horizon + +### The Solution: Multi-Stage Pipeline + +#### Stage 1: Load-Time Filtering + +Automatically discards samples that cannot produce valid training data: +- `armlinkedtohaptic` changes within action horizon (arm swap mid-sequence) +- Both arms fully disengaged for entire horizon (no valid signal) + +#### Stage 2: Engagement-Aware Delta Re-integration + +Instead of direct subtraction (`pose[t] - pose[ref]`), we: +1. Compute frame-to-frame deltas +2. Mask deltas where either endpoint is disengaged +3. Re-integrate to get cumulative motion + +This correctly handles: +- Reference disengaged β†’ later engaged (no phantom jump) +- Mid-horizon clutch events (disengaged deltas zeroed) +- Repositioning during clutch (not counted as arm motion) + +``` +Standard: action[t] = pose[t] - pose[ref] # WRONG: includes clutch repositioning +Ours: action[t] = Ξ£(delta[i] * engaged[i]) for i in ref+1..t +``` + +#### Stage 3: Sample-and-Hold for Absolute Actions + +Different action types require different clutch behavior: + +| Action Type | Clutch Behavior | Rationale | +|-------------|-----------------|-----------| +| REL_XYZ_ROT6D (pose) | Zero | No movement = zero delta | +| ABSOLUTE (gripper) | Sample-and-hold | Don't drop the needle | + +For grippers with `hold_through_clutch=True`: +- `t > 0`: Hold previous value (`action[t] = action[t-1]`) +- `t = 0`: Fall back to robot state (`action[0] = state[gripper]`) + +The t=0 fallback is critical because the controller may have snapped to a different position (e.g., gripper opened) but the robot is still holding (e.g., needle grasped). + +### Implementation Files + +| File | Component | +|------|-----------| +| `gr00t/data/dataset/sharded_single_step_dataset.py` | Load-time filtering | +| `gr00t/data/state_action/pose.py` | Delta re-integration | +| `gr00t/data/state_action/state_action_processor.py` | Sample-and-hold, action zeroing | +| `open_h/embodiments/cmr_versius/cmr_versius_config.py` | ActionConfig with `hold_through_clutch` | + +--- + +## State Prompt Preprocessing + +CMR datasets require a preprocessing step that reads per-frame `observation.state`, extracts instrument type / arm color / arm-to-haptic linkage, and writes an `instruction.text_with_state` column into each parquet in-place. Run this before training: + +```bash +uv run python open_h/embodiments/cmr_versius/utils/cmr_add_state_prompts.py # write in-place +uv run python open_h/embodiments/cmr_versius/utils/cmr_add_state_prompts.py --dry-run # preview only +``` + +--- + +## Dataset Preparation + +Use the shared `open_h/prepare_datasets.sh` script to copy the modality JSON into each dataset's `meta/` folder and generate normalization statistics (`stats.json` and `temporal_stats.json`). + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag CMR_VERSIUS \ + --modality-json open_h/embodiments/cmr_versius/modality.json \ + /path/to/cholecystectomy_50hz_480p /path/to/hysterectomy_50hz_480p ... +``` + +--- + +## Notes on State Extraction + +In CMR Versius, the `action` column contains the hand controller pose, which **is** the current state (no separate state column). We use `"original_key": "action"` in modality.json to extract pose state, and `delta_indices=[2, 4, ..., 100]` (shifted by 1*FRAME_STRIDE) because action[0] would be identical to state[0]. + +--- + +## File Structure + +``` +open_h/embodiments/cmr_versius/ +β”œβ”€β”€ README.md # This file +β”œβ”€β”€ cmr_versius_config.py # Modality config for GR00T training +β”œβ”€β”€ modality.json # Index mappings (uses original_key for runtime extraction) +└── utils/ + β”œβ”€β”€ cmr_add_state_prompts.py # Adds per-timestep instruction.text_with_state to parquets + └── cmr_state_prompt_prefix.py # Shared helper for constructing state prompt prefixes +``` + +--- diff --git a/open_h/embodiments/cmr_versius/cmr_versius_config.py b/open_h/embodiments/cmr_versius/cmr_versius_config.py new file mode 100644 index 0000000..8692824 --- /dev/null +++ b/open_h/embodiments/cmr_versius/cmr_versius_config.py @@ -0,0 +1,188 @@ +""" +CMR Versius modality configuration for GR00T N1.6. + +This configuration supports dual-arm surgical robot (Left and Right hand controllers) with: +- REL_XYZ_ROT6D action representation for EEF poses +- Temporal mean-std normalization for actions +- Clutch-aware filtering and action zeroing +- 1 camera view (endoscope) +- Per-timestep language prompts encoding instrument type, color, and arm linkage + +State Fields (2 categories): + Embedded Keys (mean-std normalized): + - left_pose (7D): xyz + quat_xyzw from action[0:7] + - left_gripper (1D): pince from action[10] + - right_pose (7D): xyz + quat_xyzw from action[13:20] + - right_gripper (1D): pince from action[23] + + Pass-Through Keys (not embedded, used for processing only): + - translation_scaling (1D): from observation.state[12] - motion scaling factor + - rotation_scaling (1D): from observation.state[13] - rotation scaling factor + - hapticengaged_left (1D): from observation.state[16] - clutch filtering + - hapticengaged_right (1D): from observation.state[17] - clutch filtering + +Language Prompts (per-timestep, from parquet column instruction.text_with_state): + Previously-embedded state info is now sent through the VLM backbone as language: + - armlinkedtohaptic_left/right, instrtype_left/right, arm color + - Format: "arm left: . left instrument: (). arm right: . + right instrument: (). do a " + See open_h/embodiments/cmr_versius/utils/cmr_add_state_prompts.py for the script that writes these prompts. + +Action: Extracted from action array + - left_pose (7D): xyz + quat_xyzw -> converted to xyz + rot6d (9D) + - left_gripper (1D): pince + - right_pose (7D): xyz + quat_xyzw -> converted to xyz + rot6d (9D) + - right_gripper (1D): pince +REL_XYZ_ROT6D Conversion: +- Pose actions are converted to REL_XYZ_ROT6D: translation and rotation relative to current EEF +- Output: 9D per arm (xyz_rel + rot6d_rel) + 1D gripper = 10D per arm + +Final Action Output: 20D = left(10) + right(10) + +Clutch-Aware Processing: +- Load-time filtering: Discards samples where armlinkedtohaptic changes within horizon + or where both arms are fully disengaged +- Action zeroing: Zeros action targets for arms where hapticengaged=False + +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +ACTION_HORIZON = 50 +# Frame stride for downsampling 60Hz -> 30Hz +FRAME_STRIDE = 2 + +cmr_versius_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "endoscope", + ], + ), + "state": ModalityConfig( + delta_indices=[0], # Single reference state for REL_XYZ_ROT6D + modality_keys=[ + # === Embedded keys (mean-std normalized) === + "left_pose", + "left_gripper", + "right_pose", + "right_gripper", + # === Pass-through keys (never embedded to model) === + # Used for action normalization only + "translation_scaling", + "rotation_scaling", + # Used for clutch-aware filtering only + "hapticengaged_left", + "hapticengaged_right", + ], + # Pass-through keys: NEVER sent to model, only used for data processing + pass_through_keys=[ + "translation_scaling", # For action normalization + "rotation_scaling", # For action normalization + "hapticengaged_left", # For clutch-aware filtering + "hapticengaged_right", # For clutch-aware filtering + ], + # Mean-std normalization for continuous values (pose, gripper) + mean_std_embedding_keys=[ + "left_pose", + "left_gripper", + "right_pose", + "right_gripper", + ], + ), + "action": ModalityConfig( + # [2, 4, 6, ..., 100] - stride=2 for 30Hz effective rate, starts at 2 to skip current state + delta_indices=list(range(FRAME_STRIDE, ACTION_HORIZON * FRAME_STRIDE + 1, FRAME_STRIDE)), + modality_keys=[ + "left_pose", + "left_gripper", + "right_pose", + "right_gripper", + # Haptic engagement for per-timestep action zeroing (removed before normalization) + "hapticengaged_left", + "hapticengaged_right", + ], + # Pass-through: hapticengaged used for action zeroing, then removed before model + pass_through_keys=[ + "hapticengaged_left", + "hapticengaged_right", + ], + action_configs=[ + # Left pose: REL_XYZ_ROT6D EEF action with motion scaling + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="left_pose", # Reference state for relative conversion + normalization_type="temporal_meanstd", + input_rotation_format="quat", # Input actions are xyz + quaternion + reference_rotation_format="quat", # Reference state is also xyz + quaternion + translation_scaling_key="translation_scaling", # CMR motion scaling + rotation_scaling_key="rotation_scaling", + ), + # Left gripper: absolute pince value (hold through clutch - don't drop needle) + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key="left_gripper", # Enables hold if clutch disengaged at t=1 + normalization_type="temporal_meanstd", + hold_through_clutch=True, # Gripper should hold position during clutch-out + ), + # Right pose: REL_XYZ_ROT6D EEF action with motion scaling + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="right_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + translation_scaling_key="translation_scaling", # CMR motion scaling + rotation_scaling_key="rotation_scaling", + ), + # Right gripper: absolute pince value (hold through clutch - don't drop needle) + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key="right_gripper", # For t=0 fallback when disengaged (Bug 6 fix) + normalization_type="temporal_meanstd", + hold_through_clutch=True, # Gripper should hold position during clutch-out + ), + # Haptic engaged left: pass-through for action zeroing (removed before normalization) + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="skip", + ), + # Haptic engaged right: pass-through for action zeroing (removed before normalization) + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="skip", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["annotation.human.task_description"], + ), +} + +# Register with CMR_VERSIUS tag for surgical robot finetuning +register_modality_config(cmr_versius_config, embodiment_tag=EmbodimentTag.CMR_VERSIUS) diff --git a/open_h/embodiments/cmr_versius/modality.json b/open_h/embodiments/cmr_versius/modality.json new file mode 100644 index 0000000..e7c339a --- /dev/null +++ b/open_h/embodiments/cmr_versius/modality.json @@ -0,0 +1,83 @@ +{ + "state": { + "left_pose": { + "start": 0, + "end": 7, + "original_key": "action" + }, + "left_gripper": { + "start": 10, + "end": 11, + "original_key": "action" + }, + "right_pose": { + "start": 13, + "end": 20, + "original_key": "action" + }, + "right_gripper": { + "start": 23, + "end": 24, + "original_key": "action" + }, + "translation_scaling": { + "start": 12, + "end": 13, + "original_key": "observation.state" + }, + "rotation_scaling": { + "start": 13, + "end": 14, + "original_key": "observation.state" + }, + "hapticengaged_left": { + "start": 16, + "end": 17, + "original_key": "observation.state" + }, + "hapticengaged_right": { + "start": 17, + "end": 18, + "original_key": "observation.state" + } + }, + "action": { + "left_pose": { + "start": 0, + "end": 7 + }, + "left_gripper": { + "start": 10, + "end": 11 + }, + "right_pose": { + "start": 13, + "end": 20 + }, + "right_gripper": { + "start": 23, + "end": 24 + }, + "hapticengaged_left": { + "start": 16, + "end": 17, + "original_key": "observation.state" + }, + "hapticengaged_right": { + "start": 17, + "end": 18, + "original_key": "observation.state" + } + }, + "video": { + "endoscope": { + "original_key": "observation.images.endoscope" + } + }, + "annotation": { + "human.task_description": { + "original_key": "instruction.text_with_state", + "is_text": true + } + } +} diff --git a/open_h/embodiments/cmr_versius/utils/cmr_add_state_prompts.py b/open_h/embodiments/cmr_versius/utils/cmr_add_state_prompts.py new file mode 100644 index 0000000..2b51bc5 --- /dev/null +++ b/open_h/embodiments/cmr_versius/utils/cmr_add_state_prompts.py @@ -0,0 +1,265 @@ +"""Add per-timestep state-derived language prompts to CMR parquet files. + +This script reads the per-frame observation.state column from each episode parquet, +extracts instrument type, arm color, and arm-to-haptic linkage, and constructs a +natural-language prompt that encodes this information. The prompt is written as a new +column `instruction.text_with_state` directly into each parquet file in-place. + +This replaces the previous approach of sending these values as always-visible state +inputs to the DiT. Instead, the information goes through the VLM backbone as language. + +Output format per row: + arm left: . left instrument: (). + arm right: . right instrument: (). + do a + +Observation.state index mapping: + [2:6] -> arm_0_color through arm_3_color (color enum per arm slot) + [20] -> armlinkedtohaptic_left (which arm slot the left controller uses, -1 = none) + [21] -> armlinkedtohaptic_right (which arm slot the right controller uses, -1 = none) + [22] -> instrtype_left (instrument type enum for left controller) + [23] -> instrtype_right (instrument type enum for right controller) + +Color derivation: + The color for each controller is derived from armlinkedtohaptic: + - If linked >= 0: color = arm_{linked}_color (observation.state[2 + linked]) + - If linked == -1: color = 0 (None / not connected) + +Usage: + # Dry-run (print first 5 episodes per dataset, don't write): + python open_h/embodiments/cmr_versius/utils/cmr_add_state_prompts.py --dry-run + + # Write in-place: + python open_h/embodiments/cmr_versius/utils/cmr_add_state_prompts.py +""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence +import json +import os +from pathlib import Path + +import pandas as pd + + +try: + from cmr_state_prompt_prefix import build_state_prefix_from_fields, extract_state_prompt_fields +except ModuleNotFoundError: + # Fallback for environments importing this script via repository-root paths. + from open_h.embodiments.cmr_versius.utils.cmr_state_prompt_prefix import ( + build_state_prefix_from_fields, + extract_state_prompt_fields, + ) + +# ────────────────────────────────────────────────────────────────────── +# Constants +# ────────────────────────────────────────────────────────────────────── + +# Root directory containing the 4 CMR datasets +CMR_ROOT = Path( + os.environ.get("OPEN_H_DATA_PATH", "."), + "cmr-surgical", +) + +# Only process the 4 variants +TARGET_DATASETS = [ + "cholecystectomy", + "hysterectomy", + "inguinal_hernia", + "prostatectomy", +] + +# Folder name -> procedure string for the language prompt suffix +PROCEDURE_MAP = { + "cholecystectomy": "do a cholecystectomy", + "hysterectomy": "do a hysterectomy", + "inguinal_hernia": "do an inguinal hernia repair", + "prostatectomy": "do a prostatectomy", +} + +# Destination column written into each parquet +DST_COL = "instruction.text_with_state" + +# ────────────────────────────────────────────────────────────────────── +# Prompt construction +# ────────────────────────────────────────────────────────────────────── + + +def build_prompt(state: Sequence[float], procedure: str) -> str: + """Build the full per-timestep language prompt from observation.state. + + Format: + arm left: . left instrument: (). + arm right: . right instrument: (). + do a + + Args: + state: The full observation.state array for this timestep (list or np array). + procedure: The procedure suffix string (e.g. "do a prostatectomy"). + + Returns: + Complete prompt string for this timestep. + """ + prefix_fields = extract_state_prompt_fields(state) + prefix = build_state_prefix_from_fields(**prefix_fields) + return f"{prefix}. {procedure}" + + +# ────────────────────────────────────────────────────────────────────── +# Episode processing +# ────────────────────────────────────────────────────────────────────── + + +def process_episode(parquet_path: Path, procedure: str, dry_run: bool = False) -> int: + """Add the instruction.text_with_state column to a single episode parquet. + + Reads the observation.state column, constructs a per-row prompt string, and + writes it back as a new column. The original parquet is overwritten in-place. + + Args: + parquet_path: Path to the episode parquet file. + procedure: Procedure suffix string (e.g. "do a prostatectomy"). + dry_run: If True, print sample output but do not write. + + Returns: + Number of rows processed. + """ + df = pd.read_parquet(parquet_path) + + # observation.state is stored as a list-of-floats per row + state_col = df["observation.state"] + + # Vectorized prompt construction: apply build_prompt to each row's state + prompts = state_col.apply(lambda s: build_prompt(s, procedure)) + df[DST_COL] = prompts + + if dry_run: + # Show first and last row as samples + print(f" {parquet_path.name} ({len(df)} rows):") + print(f' row 0: "{prompts.iloc[0]}"') + if len(df) > 1: + print(f' row {len(df) - 1}: "{prompts.iloc[-1]}"') + else: + df.to_parquet(parquet_path, index=False) + + return len(df) + + +def get_episode_parquets(dataset_root: Path) -> list[Path]: + """Discover all episode parquet files for a dataset using info.json metadata. + + Reads info.json to get the data_path pattern, total_episodes, and chunks_size, + then constructs the path to each episode parquet. + + Args: + dataset_root: Root directory of the dataset (e.g. .../cholecystectomy). + + Returns: + Sorted list of Path objects to episode parquet files. + + Raises: + FileNotFoundError: If info.json is missing. + """ + info_path = dataset_root / "meta" / "info.json" + if not info_path.exists(): + raise FileNotFoundError(f"info.json not found at {info_path}") + + info = json.loads(info_path.read_text()) + data_path_pattern = info["data_path"] + total_episodes = info["total_episodes"] + chunks_size = info["chunks_size"] + + parquets = [] + for ep_idx in range(total_episodes): + chunk = ep_idx // chunks_size + rel_path = data_path_pattern.format(episode_chunk=chunk, episode_index=ep_idx) + parquets.append(dataset_root / rel_path) + + return parquets + + +# ────────────────────────────────────────────────────────────────────── +# CLI +# ────────────────────────────────────────────────────────────────────── + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments. + + Returns: + Parsed arguments namespace. + """ + parser = argparse.ArgumentParser( + description="Add per-timestep state-derived language prompts to CMR parquets." + ) + parser.add_argument( + "--cmr-root", + type=Path, + default=CMR_ROOT, + help="Root directory containing the CMR dataset folders.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print sample outputs for the first 5 episodes per dataset without writing.", + ) + return parser.parse_args() + + +def main() -> None: + """Entry point: iterate all 4 CMR datasets and add state-derived prompts. + + For each dataset, discovers episode parquets via info.json, constructs per-row + prompts from observation.state fields, and writes them as a new parquet column. + + Raises: + FileNotFoundError: If dataset directories or info.json files are missing. + SystemExit: If no target datasets are found. + """ + args = parse_args() + mode_label = "DRY RUN" if args.dry_run else "WRITING" + + print(f"[{mode_label}] CMR state-to-language prompt writer") + print(f" Root: {args.cmr_root}") + print(f" Target column: {DST_COL}") + print() + + grand_total_rows = 0 + grand_total_episodes = 0 + + for ds_name in TARGET_DATASETS: + ds_root = args.cmr_root / ds_name + if not ds_root.exists(): + print(f"WARNING: Dataset directory not found, skipping: {ds_root}") + continue + + procedure = PROCEDURE_MAP[ds_name] + parquets = get_episode_parquets(ds_root) + total_eps = len(parquets) + limit = 5 if args.dry_run else total_eps + + print(f'--- {ds_name} ({total_eps} episodes, procedure="{procedure}") ---') + + ds_rows = 0 + for i, pq_path in enumerate(parquets[:limit]): + if not pq_path.exists(): + print(f" WARNING: Missing parquet: {pq_path}") + continue + ds_rows += process_episode(pq_path, procedure, dry_run=args.dry_run) + + # Progress reporting every 500 episodes (when not dry-run) + if not args.dry_run and (i + 1) % 500 == 0: + print(f" Processed {i + 1}/{total_eps} episodes ({ds_rows} rows so far)...") + + print(f" => {min(limit, total_eps)} episodes, {ds_rows} rows") + print() + grand_total_rows += ds_rows + grand_total_episodes += min(limit, total_eps) + + print(f"Done. Processed {grand_total_rows} total rows across {grand_total_episodes} episodes.") + + +if __name__ == "__main__": + main() diff --git a/open_h/embodiments/cmr_versius/utils/cmr_state_prompt_prefix.py b/open_h/embodiments/cmr_versius/utils/cmr_state_prompt_prefix.py new file mode 100644 index 0000000..c287ea6 --- /dev/null +++ b/open_h/embodiments/cmr_versius/utils/cmr_state_prompt_prefix.py @@ -0,0 +1,196 @@ +"""Shared helpers for constructing CMR state-derived prompt prefixes. + +This module centralizes the common "state prefix" logic used by multiple CMR +prompt-writing scripts. The prefix format is: + + arm left: . left instrument: (). + arm right: . right instrument: () + +The final task/procedure suffix (for example, "do a prostatectomy" or +"do the suturing task") should be appended by the calling script. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TypedDict + + +# arm_0_color through arm_3_color live in observation.state[2:6] +ARM_COLOR_START_IDX = 2 +ARM_COLOR_COUNT = 4 + +# armlinkedtohaptic_left / right +LINKED_LEFT_IDX = 20 +LINKED_RIGHT_IDX = 21 + +# instrtype_left / right +INSTRTYPE_LEFT_IDX = 22 +INSTRTYPE_RIGHT_IDX = 23 + + +class CMRStatePromptFields(TypedDict): + """Strongly-typed dictionary of CMR fields needed for prefix generation.""" + + armlinkedtohaptic_left: int + armlinkedtohaptic_right: int + instrtype_left: int + instrtype_right: int + arm_0_color: int + arm_1_color: int + arm_2_color: int + arm_3_color: int + + +# InstrType enum -> human-readable instrument name +INSTRTYPE_NAMES: dict[int, str] = { + 0: "Tool #0", + 1: "Tool #1", + 2: "Tool #2", + 3: "Tool #3", + 4: "Tool #4", + 5: "Tool #5", + 6: "Tool #6", + 7: "Tool #7", + 8: "Tool #8", + 9: "Tool #9", + 10: "Tool #10", + 11: "Tool #11", + 12: "Tool #12", + 13: "Tool #13", + 14: "Tool #14", + 15: "Tool #15", + 16: "Tool #16", + 17: "Tool #17", + 18: "Tool #18", +} + +# Color enum -> human-readable color name +COLOR_NAMES: dict[int, str] = { + 0: "None", + 1: "Green", + 2: "Blue", + 3: "Cyan", + 4: "Orange", + 5: "Purple", + 6: "White", + 7: "Pink", +} + + +def instrtype_to_name(val: int) -> str: + """Map an instrument enum value to a human-readable name. + + Args: + val: Integer value from CMR `instrtype_*`. + + Returns: + Human-readable instrument name. Unknown values fall back to + ``"Instrument"``. + """ + + return INSTRTYPE_NAMES.get(val, "Instrument") + + +def color_to_name(val: int) -> str: + """Map a color enum value to a human-readable color name. + + Args: + val: Integer value from CMR arm color fields. + + Returns: + Human-readable color name. Unknown values fall back to ``"None"``. + """ + + return COLOR_NAMES.get(val, "None") + + +def _derive_controller_color(linked_val: int, arm_colors: Sequence[int]) -> int: + """Derive controller color using linkage and the four arm colors. + + Args: + linked_val: Linkage arm slot index (0-3) or -1 for disconnected. + arm_colors: Sequence of exactly four arm color enum values. + + Returns: + The selected color enum value, or 0 (None) for disconnected/invalid + linkage values. + """ + + if linked_val < 0: + return 0 + if linked_val >= len(arm_colors): + return 0 + return int(arm_colors[linked_val]) + + +def extract_state_prompt_fields(state: Sequence[float]) -> CMRStatePromptFields: + """Extract all CMR fields needed for prefix construction from state. + + Args: + state: Full per-timestep ``observation.state`` sequence. + + Returns: + Dictionary containing the required linked/instrument/arm-color fields. + + Raises: + ValueError: If ``state`` does not contain enough entries for required + indices. + """ + + min_required_len = INSTRTYPE_RIGHT_IDX + 1 + if len(state) < min_required_len: + raise ValueError( + f"Expected observation.state length >= {min_required_len}, got {len(state)}" + ) + + return CMRStatePromptFields( + armlinkedtohaptic_left=int(round(state[LINKED_LEFT_IDX])), + armlinkedtohaptic_right=int(round(state[LINKED_RIGHT_IDX])), + instrtype_left=int(round(state[INSTRTYPE_LEFT_IDX])), + instrtype_right=int(round(state[INSTRTYPE_RIGHT_IDX])), + arm_0_color=int(round(state[ARM_COLOR_START_IDX + 0])), + arm_1_color=int(round(state[ARM_COLOR_START_IDX + 1])), + arm_2_color=int(round(state[ARM_COLOR_START_IDX + 2])), + arm_3_color=int(round(state[ARM_COLOR_START_IDX + 3])), + ) + + +def build_state_prefix_from_fields( + *, + armlinkedtohaptic_left: int, + armlinkedtohaptic_right: int, + instrtype_left: int, + instrtype_right: int, + arm_0_color: int, + arm_1_color: int, + arm_2_color: int, + arm_3_color: int, +) -> str: + """Build the shared CMR prompt prefix from explicit state fields. + + Args: + armlinkedtohaptic_left: Left controller linked arm slot index. + armlinkedtohaptic_right: Right controller linked arm slot index. + instrtype_left: Instrument enum on the left controller. + instrtype_right: Instrument enum on the right controller. + arm_0_color: Color enum for arm slot 0. + arm_1_color: Color enum for arm slot 1. + arm_2_color: Color enum for arm slot 2. + arm_3_color: Color enum for arm slot 3. + + Returns: + Prefix string containing left/right linkage and instrument/color details. + """ + + arm_colors = [arm_0_color, arm_1_color, arm_2_color, arm_3_color] + left_color = _derive_controller_color(armlinkedtohaptic_left, arm_colors) + right_color = _derive_controller_color(armlinkedtohaptic_right, arm_colors) + + parts = [ + f"arm left: {armlinkedtohaptic_left}", + f"left instrument: {instrtype_to_name(instrtype_left)} ({color_to_name(left_color)})", + f"arm right: {armlinkedtohaptic_right}", + f"right instrument: {instrtype_to_name(instrtype_right)} ({color_to_name(right_color)})", + ] + return ". ".join(parts) diff --git a/open_h/embodiments/hamlyn_dvrk/README.md b/open_h/embodiments/hamlyn_dvrk/README.md new file mode 100644 index 0000000..a7be4bb --- /dev/null +++ b/open_h/embodiments/hamlyn_dvrk/README.md @@ -0,0 +1,202 @@ +# Hamlyn Centre Surgical Robot Dataset + +This document describes the Hamlyn dataset, a collection of surgical robot demonstrations recorded on the da Vinci Research Kit (dVRK) at the Hamlyn Centre for Robotic Surgery at Imperial College London. + +## Embodiment Tags + +Due to different frame rates across tasks, **two embodiment tags** are provided: + +| Embodiment Tag | FPS | Action Horizon | Time Window | Tasks | +|----------------|-----|----------------|-------------|-------| +| `hamlyn_dvrk_15hz` | 15 | 25 steps | 1.67s | 7 tasks (knot_tying, needle_grasp_and_handover, peg_transfer, Suturing-1, Suturing-2, suturing_single_loop_2, tissue_lifting) | +| `hamlyn_dvrk_30hz` | 30 | 50 steps | 1.67s | 2 tasks (suturing_single_loop_1, tissue_retraction) | + +## Dataset Overview + +| Property | Value | +|----------|-------| +| **Robot Type** | dVRK (da Vinci Research Kit) | +| **Format** | LeRobot v2.1 | +| **Total Episodes (cleaned)** | 1,019 | +| **Total Frames (cleaned)** | 552,753 | +| **Total Size** | ~16 GB (cleaned folders only; duplicates archived separately) | +| **Frame Rate** | 15 fps (7 tasks) / 30 fps (2 tasks) | +| **Quaternion Order** | **wxyz** (scalar-first) | + +### Tasks (15 Hz) -- `hamlyn_dvrk_15hz` + +- knot_tying +- needle_grasp_and_handover +- peg_transfer +- Suturing-1 +- Suturing-2 +- suturing_single_loop_2 +- tissue_lifting + +### Tasks (30 Hz) -- `hamlyn_dvrk_30hz` + +- suturing_single_loop_1 +- tissue_retraction + +### Task Details + +**Multi-task datasets** (multiple task_index values): +- `peg_transfer`: 5 task variants (different colored pegs) +- `needle_grasp_and_handover`: 2 task variants (right-hand grasp, or left-grasp-then-transfer) + +**Single-task datasets** (task_index=0 only): +- `knot_tying`, `Suturing-1`, `Suturing-2`, `tissue_retraction` + +## Data Format + +### State Representation (16D) + +The state is composed of four modality keys: + +| Key | Dim | Description | +|-----|-----|-------------| +| `left_arm_pose` | 7D | Left arm xyz position (3D) + quaternion wxyz (4D) | +| `left_arm_gripper` | 1D | Left arm jaw angle (0 = closed, ~0.5-0.6 = open) | +| `right_arm_pose` | 7D | Right arm xyz position (3D) + quaternion wxyz (4D) | +| `right_arm_gripper` | 1D | Right arm jaw angle | + +Raw data source: `observation.state.left_arm_cartesian` (8D) and `observation.state.right_arm_cartesian` (8D), sliced via `modality.json`. + +### Action Representation (16D) + +Actions use `REL_XYZ_ROT6D` for EEF poses and `ABSOLUTE` for grippers: + +| Key | Rep | Description | +|-----|-----|-------------| +| `left_arm_pose` | REL_XYZ_ROT6D | Relative translation + 6D rotation for left arm | +| `left_arm_gripper` | ABSOLUTE | Left arm jaw angle | +| `right_arm_pose` | REL_XYZ_ROT6D | Relative translation + 6D rotation for right arm | +| `right_arm_gripper` | ABSOLUTE | Right arm jaw angle | + +Raw data source: `action.cartesian_absolute` (16D), sliced via `modality.json`. + +### Video Streams + +Three camera views are used for training: + +| Modality Key | Raw Key | Description | +|--------------|---------|-------------| +| `endoscope` | `observation.images.color` | Stereo endoscope color image | +| `wrist_left` | `observation.images.wrist_left` | Left wrist camera | +| `wrist_right` | `observation.images.wrist_right` | Right wrist camera | + +Depth images (`observation.images.depth`) are available in the dataset but are **not** used in training. + +Video codec: AV1 or H264 depending on task. + +### Language + +| Modality Key | Source | +|--------------|--------| +| `task` | `task_index` (from `tasks.jsonl`) | + +### Metadata + +Additional per-frame metadata: + +| Key | Description | +|-----|-------------| +| `observation.meta.left_arm_tool` | Tool attached to left arm (e.g., "debakey_forceps") | +| `observation.meta.right_arm_tool` | Tool attached to right arm (e.g., "needle_driver") | +| `observation.meta.left_arm_tpv_cali_mtx` | Left arm calibration matrix (7D) | +| `observation.meta.right_arm_tpv_cali_mtx` | Right arm calibration matrix (7D) | + +Tool types observed: +- `debakey_forceps` - Grasping forceps +- `needle_driver` - For needle manipulation +- `large_needle_driver` - Larger variant + +## Data Splits + +Each task has predefined train/val/test splits plus recovery and failure episodes: + +| Task | Train | Val | Test | Recovery | Failure | +|------|-------|-----|------|----------|---------| +| knot_tying | 0:50 | 50:57 | 57:72 | 72:73 | 73:77 | +| needle_grasp_and_handover | 0:91 | 91:104 | 104:130 | 130:137 | - | +| peg_transfer | 0:207 | 207:237 | 237:296 | 296:315 | 315:317 | +| Suturing-1 | 0:75 | 75:86 | 86:107 | 107:148 | 148:180 | +| Suturing-2 | 0:71 | 71:81 | 81:102 | 102:142 | 142:186 | +| tissue_lifting | 0:51 | 51:58 | 58:73 | - | 73:75 | +| tissue_retraction | 0:50 | 50:57 | 57:71 | 71:73 | 73:75 | + +**Note**: `suturing_single_loop_1` is a 30Hz task; `suturing_single_loop_2` and `tissue_lifting` are 15Hz tasks. + +## Hamlyn-Specific Details + +- **Quaternion order**: wxyz (scalar-first); handled via `input_quat_order`/`reference_quat_order` in config +- **Video naming**: `observation.images.color` (endoscope), `wrist_left`/`wrist_right` (wrist cameras) +- **Extra data**: Tool metadata and depth images available (depth not used in training) + +## Directory Structure + +``` +Hamlyn/ +β”œβ”€β”€ knot_tying/ +β”‚ β”œβ”€β”€ data/chunk-000/episode_*.parquet +β”‚ β”œβ”€β”€ videos/chunk-000/ +β”‚ β”‚ β”œβ”€β”€ observation.images.color/episode_*.mp4 +β”‚ β”‚ β”œβ”€β”€ observation.images.depth/episode_*.mp4 +β”‚ β”‚ β”œβ”€β”€ observation.images.wrist_left/episode_*.mp4 +β”‚ β”‚ └── observation.images.wrist_right/episode_*.mp4 +β”‚ └── meta/ +β”‚ β”œβ”€β”€ info.json +β”‚ └── tasks.jsonl +β”œβ”€β”€ needle_grasp_and_handover/ +β”œβ”€β”€ peg_transfer/ +β”œβ”€β”€ Suturing-1/ +β”œβ”€β”€ Suturing-2/ +β”œβ”€β”€ tissue_retraction/ +β”œβ”€β”€ tissue_lifting/ +β”œβ”€β”€ suturing_single_loop_1/ +└── suturing_single_loop_2/ +``` + +## Dataset Preparation + +Use the shared preparation script to copy the modality JSON into each dataset's `meta/` folder and generate normalization statistics: + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag hamlyn_dvrk_30hz \ + --modality-json open_h/embodiments/hamlyn_dvrk/modality.json \ + /path/to/dataset +``` + +For 15 Hz tasks, replace `hamlyn_dvrk_30hz` with `hamlyn_dvrk_15hz`. + +## Usage Notes + +### Loading with LeRobot + +```python +from lerobot.common.datasets.lerobot_dataset import LeRobotDataset + +# Load a single task +dataset = LeRobotDataset( + repo_id="hamlyn/tissue_retraction", + root="${OPEN_H_DATA_PATH}/Hamlyn/tissue_retraction" +) +``` + +### Embodiment Configuration + +Use the Hamlyn-specific embodiment tags to keep frame rates and quaternion ordering correct: +- `hamlyn_dvrk_15hz` for 15Hz tasks (knot_tying, needle_grasp_and_handover, peg_transfer, Suturing-1, Suturing-2, suturing_single_loop_2, tissue_lifting) +- `hamlyn_dvrk_30hz` for 30Hz tasks (suturing_single_loop_1, tissue_retraction) + +Depth images are available but intentionally **not** used. The modality config only uses +the color endoscope stream plus left/right wrist cameras. + +### Key Considerations for Training + +1. **Frame Rate Variation**: Use the matching Hamlyn tag (`hamlyn_dvrk_15hz` or `hamlyn_dvrk_30hz`) to keep the ~1.67s action horizon consistent +2. **Quaternion Order**: Hamlyn uses w,x,y,z and is handled via `input_quat_order`/`reference_quat_order` in the config +3. **Multi-Task Training**: The `task_index` field indicates which sub-task within a dataset +4. **Recovery/Failure Episodes**: May be useful for training robust policies but should be handled carefully + diff --git a/open_h/embodiments/hamlyn_dvrk/hamlyn_dvrk_config.py b/open_h/embodiments/hamlyn_dvrk/hamlyn_dvrk_config.py new file mode 100644 index 0000000..864ae49 --- /dev/null +++ b/open_h/embodiments/hamlyn_dvrk/hamlyn_dvrk_config.py @@ -0,0 +1,208 @@ +""" +Hamlyn Centre Surgical Robot Dataset Embodiment Configurations. + +This module defines modality configurations for the Hamlyn dataset, which contains +surgical robot demonstrations recorded on the dVRK (da Vinci Research Kit) at +Imperial College London. + +Two configurations are provided: +1. `hamlyn_15hz_config`: For 7 tasks recorded at 15Hz (knot_tying, needle_grasp_and_handover, + peg_transfer, Suturing-1, Suturing-2, suturing_single_loop_2, tissue_lifting) +2. `hamlyn_30hz_config`: For 2 tasks recorded at 30Hz (suturing_single_loop_1, tissue_retraction) + +Key differences from JHU dVRK dataset: +- Quaternion ordering: wxyz (scalar-first) vs xyzw (scalar-last) +- Video keys: observation.images.color vs observation.images.endoscope.left +Data format: +- State: 16D Cartesian EEF (left arm 8D + right arm 8D) + - Per arm: xyz (3D) + quaternion wxyz (4D) + gripper (1D) +- Action: 16D Cartesian absolute (same format as state) +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# ============================================================================= +# 15Hz Configuration (7 tasks) +# ============================================================================= +# Tasks: knot_tying, needle_grasp_and_handover, peg_transfer, Suturing-1, +# Suturing-2, suturing_single_loop_2, tissue_lifting +# +# Action horizon: 25 steps = 1.67 seconds at 15Hz +# This keeps ~85% of samples usable for episodes with 100+ frames. + +ACTION_HORIZON_15HZ = 25 + +hamlyn_15hz_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=["endoscope", "wrist_left", "wrist_right"], + ), + "state": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "left_arm_pose", # 7D: xyz + quat (wxyz order) + "left_arm_gripper", # 1D: jaw angle + "right_arm_pose", # 7D: xyz + quat (wxyz order) + "right_arm_gripper", # 1D: jaw angle + ], + mean_std_embedding_keys=[ + "left_arm_pose", + "left_arm_gripper", + "right_arm_pose", + "right_arm_gripper", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(ACTION_HORIZON_15HZ)), + modality_keys=[ + "left_arm_pose", + "left_arm_gripper", + "right_arm_pose", + "right_arm_gripper", + ], + action_configs=[ + # Left arm pose: REL_XYZ_ROT6D with wxyz quaternion order + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="left_arm_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="wxyz", # Hamlyn uses wxyz (scalar-first) + reference_rotation_format="quat", + reference_quat_order="wxyz", + ), + # Left arm gripper: ABSOLUTE (jaw angle doesn't need relative conversion) + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + # Right arm pose: REL_XYZ_ROT6D with wxyz quaternion order + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="right_arm_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="wxyz", + reference_rotation_format="quat", + reference_quat_order="wxyz", + ), + # Right arm gripper: ABSOLUTE + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["task"], # Uses tasks.jsonl + ), +} + + +# ============================================================================= +# 30Hz Configuration (2 tasks) +# ============================================================================= +# Tasks: suturing_single_loop_1, tissue_retraction +# +# Action horizon: 50 steps = 1.67 seconds at 30Hz (same time window as 15Hz config) +# This keeps ~85% of samples usable for episodes with 100+ frames. + +ACTION_HORIZON_30HZ = 50 + +hamlyn_30hz_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=["endoscope", "wrist_left", "wrist_right"], + ), + "state": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "left_arm_pose", + "left_arm_gripper", + "right_arm_pose", + "right_arm_gripper", + ], + mean_std_embedding_keys=[ + "left_arm_pose", + "left_arm_gripper", + "right_arm_pose", + "right_arm_gripper", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(ACTION_HORIZON_30HZ)), + modality_keys=[ + "left_arm_pose", + "left_arm_gripper", + "right_arm_pose", + "right_arm_gripper", + ], + action_configs=[ + # Left arm pose: REL_XYZ_ROT6D with wxyz quaternion order + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="left_arm_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="wxyz", + reference_rotation_format="quat", + reference_quat_order="wxyz", + ), + # Left arm gripper: ABSOLUTE + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + # Right arm pose: REL_XYZ_ROT6D with wxyz quaternion order + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="right_arm_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="wxyz", + reference_rotation_format="quat", + reference_quat_order="wxyz", + ), + # Right arm gripper: ABSOLUTE + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["task"], + ), +} + + +# Register both configurations +register_modality_config(hamlyn_15hz_config, embodiment_tag=EmbodimentTag.HAMLYN_DVRK_15HZ) +register_modality_config(hamlyn_30hz_config, embodiment_tag=EmbodimentTag.HAMLYN_DVRK_30HZ) diff --git a/open_h/embodiments/hamlyn_dvrk/modality.json b/open_h/embodiments/hamlyn_dvrk/modality.json new file mode 100644 index 0000000..8d2db42 --- /dev/null +++ b/open_h/embodiments/hamlyn_dvrk/modality.json @@ -0,0 +1,62 @@ +{ + "state": { + "left_arm_pose": { + "start": 0, + "end": 7, + "original_key": "observation.state.left_arm_cartesian" + }, + "left_arm_gripper": { + "start": 7, + "end": 8, + "original_key": "observation.state.left_arm_cartesian" + }, + "right_arm_pose": { + "start": 0, + "end": 7, + "original_key": "observation.state.right_arm_cartesian" + }, + "right_arm_gripper": { + "start": 7, + "end": 8, + "original_key": "observation.state.right_arm_cartesian" + } + }, + "action": { + "left_arm_pose": { + "start": 0, + "end": 7, + "original_key": "action.cartesian_absolute" + }, + "left_arm_gripper": { + "start": 7, + "end": 8, + "original_key": "action.cartesian_absolute" + }, + "right_arm_pose": { + "start": 8, + "end": 15, + "original_key": "action.cartesian_absolute" + }, + "right_arm_gripper": { + "start": 15, + "end": 16, + "original_key": "action.cartesian_absolute" + } + }, + "video": { + "endoscope": { + "original_key": "observation.images.color" + }, + "wrist_left": { + "original_key": "observation.images.wrist_left" + }, + "wrist_right": { + "original_key": "observation.images.wrist_right" + } + }, + "annotation": { + "task": { + "original_key": "task_index" + } + } +} diff --git a/open_h/embodiments/jhu_imerse_dvrk/README.md b/open_h/embodiments/jhu_imerse_dvrk/README.md new file mode 100644 index 0000000..6cb9d27 --- /dev/null +++ b/open_h/embodiments/jhu_imerse_dvrk/README.md @@ -0,0 +1,112 @@ +# JHU IMERSE dVRK + +Surgical robot data collected from the da Vinci Research Kit (dVRK) at JHU's IMERSE lab. + +## Embodiment Configuration + +| Property | Value | +|----------|-------| +| **Embodiment Tag** | `jhu_imerse_dvrk` | +| **Config File** | `open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_dvrk_config.py` | + +## IMERSE Datasets (JHU) + +- **NephFat**: Matches the standard dVRK 16D state/action layout; can reuse `open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_dvrk_config.py` and `open_h/embodiments/jhu_imerse_dvrk/modality.json`. The dataset includes stereo endoscope and both wrists; the default modality maps `endoscope_left` plus both wrists. +- **star_IL**: Single-arm KUKA + endo360; use `open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_star_il_config.py` with `open_h/embodiments/jhu_imerse_dvrk/modality_imerse_star_il.json`. + +## Mono Configuration + +The `jhu_imerse_dvrk_mono_config.py` provides a configuration variant for training +with a single endoscope view only (monocular), as opposed to the standard stereo +endoscope + wrist camera setup. It uses the `jhu_imerse_dvrk_mono` embodiment tag. +All other aspects (dual-arm REL_XYZ_ROT6D actions, 16D state/action format) remain +identical to the standard configuration. Use this when your dataset lacks wrist +camera views or when you want to evaluate endoscope-only performance. + +## Data Format + +### State (16D) + +| Key | Dim | Description | +|-----|-----|-------------| +| `psm1_pose` | 7D | PSM1 xyz position (3D) + quaternion xyzw (4D) | +| `psm1_gripper` | 1D | PSM1 jaw angle | +| `psm2_pose` | 7D | PSM2 xyz position (3D) + quaternion xyzw (4D) | +| `psm2_gripper` | 1D | PSM2 jaw angle | + +### Action (16D, horizon 50) + +Actions use `REL_XYZ_ROT6D` for EEF poses and `ABSOLUTE` for grippers: + +| Key | Rep | Description | +|-----|-----|-------------| +| `psm1_pose` | REL_XYZ_ROT6D | Relative translation + 6D rotation for PSM1 | +| `psm1_gripper` | ABSOLUTE | PSM1 jaw angle | +| `psm2_pose` | REL_XYZ_ROT6D | Relative translation + 6D rotation for PSM2 | +| `psm2_gripper` | ABSOLUTE | PSM2 jaw angle | + +### Video (3 camera views) + +| Modality Key | Raw Key | +|--------------|---------| +| `endoscope_left` | `observation.images.endoscope.left` | +| `wrist_left` | `observation.images.wrist.left` | +| `wrist_right` | `observation.images.wrist.right` | + +### Language + +| Modality Key | Source | +|--------------|--------| +| `annotation.human.task_description` | `task_index` | + +## Dataset Preparation + +Prepare datasets (copy modality JSON and compute normalization statistics) using the +centralized preparation script: + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag JHU_IMERSE_DVRK \ + --modality-json open_h/embodiments/jhu_imerse_dvrk/modality.json \ + /path/to/dataset1 /path/to/dataset2 +``` + +This copies `modality.json` into each dataset's `meta/` directory and computes normalization statistics (`stats.json` and `temporal_stats.json`). Each dataset gets its own statistics file, but at training time these are merged across all datasets sharing the same embodiment tag β€” so each embodiment trains with a single unified set of normalization statistics. + +## Open-Loop Evaluation + +```bash +uv run python gr00t/eval/open_loop_eval.py \ + --dataset-path /path/to/your/dvrk-dataset \ + --embodiment-tag JHU_IMERSE_DVRK \ + --model-path /path/to/checkpoint \ + --traj-ids 0 1 2 \ + --action-horizon 50 \ + --steps 500 +``` + +## Open-Loop Evaluation + +Evaluate the finetuned model against ground truth trajectories: + +```bash +uv run python gr00t/eval/open_loop_eval.py \ + --dataset-path /path/to/dvrk-dataset \ + --embodiment-tag JHU_IMERSE_DVRK \ + --model-path /path/to/checkpoint \ + --traj-ids 0 1 2 \ + --action-horizon 16 \ + --steps 500 +``` + +## File Structure + +``` +open_h/embodiments/jhu_imerse_dvrk/ +β”œβ”€β”€ README.md # This file +β”œβ”€β”€ jhu_imerse_dvrk_config.py # Built-in config module (auto-registered) +β”œβ”€β”€ jhu_imerse_dvrk_mono_config.py # Mono (endoscope-only) built-in config module +β”œβ”€β”€ jhu_imerse_star_il_config.py # IMERSE star_IL built-in config module +β”œβ”€β”€ modality.json # Data key mappings (copy to dataset/meta/) +└── modality_imerse_star_il.json # IMERSE star_IL data key mappings +``` diff --git a/open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_dvrk_config.py b/open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_dvrk_config.py new file mode 100644 index 0000000..88eadac --- /dev/null +++ b/open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_dvrk_config.py @@ -0,0 +1,111 @@ +""" +dVRK (da Vinci Research Kit) modality configuration for GR00T N1.6. + +This configuration supports dual-arm surgical robot (PSM1 and PSM2) with: +- REL_XYZ_ROT6D action representation for EEF poses +- Percentile-based normalization with per-dataset statistics +- 3 camera views (endoscope_left + wrist_left + wrist_right) + +Data Format: +- State: 16D = psm1_pose(7: xyz + quat_xyzw) + psm1_jaw(1) + psm2_pose(7: xyz + quat_xyzw) + psm2_jaw(1) +- Action: 16D (same format as state, representing setpoints) + +REL_XYZ_ROT6D Conversion: +- PSM pose actions are converted to REL_XYZ_ROT6D: translation and rotation relative to current EEF +- Output: 9D per arm (xyz_rel + rot6d_rel) +- Gripper actions remain absolute (jaw angle) +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +ACTION_HORIZON = 50 + +dvrk_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "endoscope_left", + "wrist_left", + "wrist_right", + ], + ), + "state": ModalityConfig( + delta_indices=[0], # Single reference state for REL_XYZ_ROT6D + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + mean_std_embedding_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(ACTION_HORIZON)), # [0, 1, 2, ..., 50] + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + action_configs=[ + # PSM1 pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm1_pose", # Reference state for relative conversion + normalization_type="temporal_meanstd", + input_rotation_format="quat", # Input actions are xyz + quaternion + reference_rotation_format="quat", # Reference state is also xyz + quaternion + ), + # PSM1 gripper: absolute jaw angle (no rotation conversion needed) + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + # PSM2 pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm2_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + ), + # PSM2 gripper: absolute jaw angle (no rotation conversion needed) + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["annotation.human.task_description"], + ), +} + +# Register with JHU_IMERSE_DVRK tag for surgical robot finetuning +register_modality_config(dvrk_config, embodiment_tag=EmbodimentTag.JHU_IMERSE_DVRK) diff --git a/open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_dvrk_mono_config.py b/open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_dvrk_mono_config.py new file mode 100644 index 0000000..369470e --- /dev/null +++ b/open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_dvrk_mono_config.py @@ -0,0 +1,106 @@ +""" +Monocular dVRK (JHU) modality configuration for GR00T N1.6. + +This configuration mirrors the standard dVRK embodiment but restricts the +video modality to a single camera view: + +- endoscope_left only (no wrist cameras) +- same dual-arm REL_XYZ_ROT6D action representation +- same 16D state/action format as the standard dVRK config + +Use this when training models that should rely solely on a monocular +endoscope feed while keeping the rest of the embodiment consistent. +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +ACTION_HORIZON = 50 + +jhu_dvrk_mono_config = { + "video": ModalityConfig( + delta_indices=[0], + # Monocular endoscope input only (no wrist cameras) + modality_keys=[ + "endoscope_left", + ], + ), + "state": ModalityConfig( + delta_indices=[0], # Single reference state for REL_XYZ_ROT6D + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + mean_std_embedding_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(ACTION_HORIZON)), + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + action_configs=[ + # PSM1 pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm1_pose", # Reference state for relative conversion + normalization_type="temporal_meanstd", + input_rotation_format="quat", # Input actions are xyz + quaternion + reference_rotation_format="quat", # Reference state is also xyz + quaternion + ), + # PSM1 gripper: absolute jaw angle (no rotation conversion needed) + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + # PSM2 pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm2_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + ), + # PSM2 gripper: absolute jaw angle (no rotation conversion needed) + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["annotation.human.task_description"], + ), +} + +# Register with the monocular dVRK tag for surgical robot finetuning +register_modality_config(jhu_dvrk_mono_config, embodiment_tag=EmbodimentTag.JHU_IMERSE_DVRK_MONO) diff --git a/open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_star_il_config.py b/open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_star_il_config.py new file mode 100644 index 0000000..dea5e22 --- /dev/null +++ b/open_h/embodiments/jhu_imerse_dvrk/jhu_imerse_star_il_config.py @@ -0,0 +1,79 @@ +"""IMERSE star_IL modality configuration for GR00T N1.6. + +This configuration supports the JHU IMERSE star_IL dataset with: +- Single-arm KUKA pose actions (7D: xyz + quat_xyzw) +- KUKA joint positions (7D) + endo360 joint (1D) in state +- Endoscope left + wrist left video streams +- Natural language commands from instruction.text + +Action Semantics: +- REL_XYZ_ROT6D conversion uses the action pose at t=0 as the reference. +- The reference pose is injected into state as a pass-through key via modality.json. +- In this dataset, the EEF pose in `action` represents the measured point (not a setpoint), + so using `action[t]` as the state reference is consistent with `action[t] == state[t]`. +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +ACTION_HORIZON = 50 + +imerse_star_il_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "endoscope_left", + "wrist_left", + ], + ), + "state": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "kuka_joint_pos", + "endo360_joint_pos", + "kuka_pose", + ], + mean_std_embedding_keys=[ + "kuka_joint_pos", + "endo360_joint_pos", + ], + pass_through_keys=[ + "kuka_pose", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(1, ACTION_HORIZON + 1)), + modality_keys=[ + "kuka_pose", + ], + action_configs=[ + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="kuka_pose", + normalization_type="percentile", + input_rotation_format="quat", + reference_rotation_format="quat", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["annotation.human.task_description"], + ), +} + +# Register with the JHU IMERSE star_IL embodiment tag. +register_modality_config( + imerse_star_il_config, embodiment_tag=EmbodimentTag.JHU_IMERSE_DVRK_STAR_IL +) diff --git a/open_h/embodiments/jhu_imerse_dvrk/modality.json b/open_h/embodiments/jhu_imerse_dvrk/modality.json new file mode 100644 index 0000000..b26ece4 --- /dev/null +++ b/open_h/embodiments/jhu_imerse_dvrk/modality.json @@ -0,0 +1,54 @@ +{ + "state": { + "psm1_pose": { + "start": 0, + "end": 7 + }, + "psm1_gripper": { + "start": 7, + "end": 8 + }, + "psm2_pose": { + "start": 8, + "end": 15 + }, + "psm2_gripper": { + "start": 15, + "end": 16 + } + }, + "action": { + "psm1_pose": { + "start": 0, + "end": 7 + }, + "psm1_gripper": { + "start": 7, + "end": 8 + }, + "psm2_pose": { + "start": 8, + "end": 15 + }, + "psm2_gripper": { + "start": 15, + "end": 16 + } + }, + "video": { + "endoscope_left": { + "original_key": "observation.images.endoscope.left" + }, + "wrist_left": { + "original_key": "observation.images.wrist.left" + }, + "wrist_right": { + "original_key": "observation.images.wrist.right" + } + }, + "annotation": { + "human.task_description": { + "original_key": "task_index" + } + } +} diff --git a/open_h/embodiments/jhu_imerse_dvrk/modality_imerse_star_il.json b/open_h/embodiments/jhu_imerse_dvrk/modality_imerse_star_il.json new file mode 100644 index 0000000..e591b9e --- /dev/null +++ b/open_h/embodiments/jhu_imerse_dvrk/modality_imerse_star_il.json @@ -0,0 +1,37 @@ +{ + "state": { + "kuka_joint_pos": { + "start": 0, + "end": 7 + }, + "endo360_joint_pos": { + "start": 7, + "end": 8 + }, + "kuka_pose": { + "start": 0, + "end": 7, + "original_key": "action" + } + }, + "action": { + "kuka_pose": { + "start": 0, + "end": 7 + } + }, + "video": { + "endoscope_left": { + "original_key": "observation.images.endoscope.left" + }, + "wrist_left": { + "original_key": "observation.images.wrist.left" + } + }, + "annotation": { + "human.task_description": { + "original_key": "instruction.text", + "is_text": true + } + } +} diff --git a/open_h/embodiments/jhu_lscr_dvrk/README.md b/open_h/embodiments/jhu_lscr_dvrk/README.md new file mode 100644 index 0000000..2276712 --- /dev/null +++ b/open_h/embodiments/jhu_lscr_dvrk/README.md @@ -0,0 +1,94 @@ +# JHU LSCR (Open-H Surgical) + +JHU LSCR is a collection of dVRK-based surgical datasets spanning three sub-benchmarks: +ARCADE (cartesian EEF setpoints), MIRACLE (cartesian EEF actions), and SMARTS +(cartesian EEF actions with additional cameras). + +## Embodiment Configuration + +| Property | Value | +|----------|-------| +| **Embodiment Tags** | `JHU_IMERSE_DVRK` (ARCADE), `JHU_LSCR_DVRK_MIRACLE`, `JHU_LSCR_DVRK_SMARTS` | +| **Config File** | `open_h/embodiments/jhu_lscr_dvrk/jhu_lscr_dvrk_config.py` | + +ARCADE uses the standard dVRK embodiment config; MIRACLE and SMARTS are registered +in the LSCR config file above. + +## Data Format + +### State + +| Sub-benchmark | Keys | Notes | +|---------------|------|-------| +| ARCADE | `observation.state` / `action` | EEF pose + gripper via dVRK schema | +| MIRACLE | `psm1_pose` (7), `psm1_gripper` (1), `psm2_pose` (7), `psm2_gripper` (1) | PSM cartesian pose per arm | +| SMARTS | `psm1_pose` (7), `psm1_gripper` (1), `psm2_pose` (7), `psm2_gripper` (1) | Per-key arrays (not a single `observation.state` vector) | + +### Action + +All sub-benchmarks use `REL_XYZ_ROT6D` for EEF pose actions and `ABSOLUTE` for +non-EEF (gripper). MIRACLE and SMARTS use `delta_indices` starting at `t+1` +(reference pose at `t`, actions at `t+1..t+H`), while ARCADE starts at `t+0` +(using the standard dVRK config). + +| Sub-benchmark | Action Horizon | FPS | +|---------------|----------------|-----| +| ARCADE | 50 | 30 Hz | +| MIRACLE | 25 | 15 Hz | +| SMARTS | 16 | 10 Hz | + +### Cameras + +| Sub-benchmark | Camera Streams (used in training) | +|---------------|----------------------------------| +| ARCADE | `endoscope_left`, `wrist_left`, `wrist_right` | +| MIRACLE | `camera_left` (mono endoscope) | +| SMARTS | `endoscope_left`, `camera_side_view` | + +Note: MIRACLE also has `camera_right` and SMARTS has `endoscope_right` in their +modality JSONs, but these are commented out in the config (mono only). + +### Language + +All configs use `tasks.jsonl` via `task_index`. ARCADE uses +`annotation.human.task_description`; MIRACLE/SMARTS use `annotation.task`. + +## Modality Mapping + +Each sub-benchmark has its own modality JSON: + +- **ARCADE**: `modality_arcade.json` -- EEF pose + gripper from `observation.state`/`action`, annotation key aligned to the dVRK schema. +- **MIRACLE**: `modality_miracle.json` -- PSM cartesian pose + grippers from `observation.state` slices. +- **SMARTS**: `modality_smarts.json` -- PSM cartesian pose + grippers from per-key arrays. + +## Dataset Preparation + +Run `prepare_datasets.sh` for each sub-benchmark with the matching modality file +(the script copies the modality JSON into `meta/` automatically): + +```bash +# ARCADE (30 Hz, action horizon 50) +bash open_h/prepare_datasets.sh \ + --embodiment-tag JHU_IMERSE_DVRK \ + --modality-json open_h/embodiments/jhu_lscr_dvrk/modality_arcade.json \ + /path/to/JHU_LSCR/ARCADE + +# MIRACLE (15 Hz, action horizon 25) +bash open_h/prepare_datasets.sh \ + --embodiment-tag JHU_LSCR_DVRK_MIRACLE \ + --modality-json open_h/embodiments/jhu_lscr_dvrk/modality_miracle.json \ + /path/to/JHU_LSCR/MIRACLE + +# SMARTS (10 Hz, action horizon 16) +bash open_h/prepare_datasets.sh \ + --embodiment-tag JHU_LSCR_DVRK_SMARTS \ + --modality-json open_h/embodiments/jhu_lscr_dvrk/modality_smarts.json \ + /path/to/JHU_LSCR/SMARTS +``` + +## Notes + +- MIRACLE publishes full PSM cartesian pose (XYZ + quaternion) per arm, so the modality mapping uses EEF pose keys and REL_XYZ_ROT6D actions to avoid losing pose information. +- ARCADE includes `instruction.text` in the dataset, but tasks are currently sourced from `tasks.jsonl` for consistency across LSCR datasets. +- SMARTS uses separate per-key arrays for cartesian pose and grippers, not a single `observation.state` vector. +- ARCADE/cautery is missing videos for episodes 12-21 across all camera streams; training should exclude a `missing_videos` split once it is defined in `meta/info.json`. diff --git a/open_h/embodiments/jhu_lscr_dvrk/jhu_lscr_dvrk_config.py b/open_h/embodiments/jhu_lscr_dvrk/jhu_lscr_dvrk_config.py new file mode 100644 index 0000000..cf4f526 --- /dev/null +++ b/open_h/embodiments/jhu_lscr_dvrk/jhu_lscr_dvrk_config.py @@ -0,0 +1,186 @@ +""" +JHU LSCR (Laboratory for Surgical & Computational Robotics) modality configs. + +This module defines built-in LSCR configurations that are auto-registered when +`open_h.embodiments` is imported. These cover the LSCR schemas that are not +already covered by the standard dVRK embodiment: +- MIRACLE (15 Hz): cartesian EEF actions (REL_XYZ_ROT6D) +- SMARTS (10 Hz): cartesian EEF actions (REL_XYZ_ROT6D) + +ARCADE (30 Hz) is handled by the standard dVRK configuration, with a modality +mapping that mirrors the dVRK annotation schema. +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizons tuned to dataset FPS. +ACTION_HORIZON_15HZ = 25 +ACTION_HORIZON_10HZ = 16 + + +jhu_lscr_miracle_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "camera_left", + # "camera_right", # Mono Only + ], + ), + "state": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + mean_std_embedding_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + ), + "action": ModalityConfig( + # Cartesian pose actions are derived from the same EEF pose state, so + # the action horizon starts at t+1 relative to the reference pose at t. + delta_indices=list(range(1, ACTION_HORIZON_15HZ + 1)), + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + action_configs=[ + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm1_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="xyzw", + reference_rotation_format="quat", + reference_quat_order="xyzw", + ), + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm2_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="xyzw", + reference_rotation_format="quat", + reference_quat_order="xyzw", + ), + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["annotation.task"], + ), +} + + +jhu_lscr_smarts_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "endoscope_left", + # "endoscope_right", # Mono Only + "camera_side_view", + ], + ), + "state": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + mean_std_embedding_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(1, ACTION_HORIZON_10HZ + 1)), + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + action_configs=[ + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm1_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="xyzw", + reference_rotation_format="quat", + reference_quat_order="xyzw", + ), + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm2_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="xyzw", + reference_rotation_format="quat", + reference_quat_order="xyzw", + ), + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["annotation.task"], + ), +} + + +register_modality_config( + jhu_lscr_miracle_config, embodiment_tag=EmbodimentTag.JHU_LSCR_DVRK_MIRACLE +) +register_modality_config(jhu_lscr_smarts_config, embodiment_tag=EmbodimentTag.JHU_LSCR_DVRK_SMARTS) diff --git a/open_h/embodiments/jhu_lscr_dvrk/modality_arcade.json b/open_h/embodiments/jhu_lscr_dvrk/modality_arcade.json new file mode 100644 index 0000000..c670b7e --- /dev/null +++ b/open_h/embodiments/jhu_lscr_dvrk/modality_arcade.json @@ -0,0 +1,62 @@ +{ + "state": { + "psm1_pose": { + "start": 0, + "end": 7, + "original_key": "observation.state" + }, + "psm1_gripper": { + "start": 7, + "end": 8, + "original_key": "observation.state" + }, + "psm2_pose": { + "start": 8, + "end": 15, + "original_key": "observation.state" + }, + "psm2_gripper": { + "start": 15, + "end": 16, + "original_key": "observation.state" + } + }, + "action": { + "psm1_pose": { + "start": 0, + "end": 7, + "original_key": "action" + }, + "psm1_gripper": { + "start": 7, + "end": 8, + "original_key": "action" + }, + "psm2_pose": { + "start": 8, + "end": 15, + "original_key": "action" + }, + "psm2_gripper": { + "start": 15, + "end": 16, + "original_key": "action" + } + }, + "video": { + "endoscope_left": { + "original_key": "observation.images.endoscope.left" + }, + "wrist_left": { + "original_key": "observation.images.wrist.left" + }, + "wrist_right": { + "original_key": "observation.images.wrist.right" + } + }, + "annotation": { + "human.task_description": { + "original_key": "task_index" + } + } +} diff --git a/open_h/embodiments/jhu_lscr_dvrk/modality_miracle.json b/open_h/embodiments/jhu_lscr_dvrk/modality_miracle.json new file mode 100644 index 0000000..54d8ed4 --- /dev/null +++ b/open_h/embodiments/jhu_lscr_dvrk/modality_miracle.json @@ -0,0 +1,59 @@ +{ + "state": { + "psm1_pose": { + "start": 18, + "end": 25, + "original_key": "observation.state" + }, + "psm1_gripper": { + "start": 17, + "end": 18, + "original_key": "observation.state" + }, + "psm2_pose": { + "start": 32, + "end": 39, + "original_key": "observation.state" + }, + "psm2_gripper": { + "start": 31, + "end": 32, + "original_key": "observation.state" + } + }, + "action": { + "psm1_pose": { + "start": 18, + "end": 25, + "original_key": "observation.state" + }, + "psm1_gripper": { + "start": 6, + "end": 7, + "original_key": "action" + }, + "psm2_pose": { + "start": 32, + "end": 39, + "original_key": "observation.state" + }, + "psm2_gripper": { + "start": 13, + "end": 14, + "original_key": "action" + } + }, + "video": { + "camera_left": { + "original_key": "observation.images.left" + }, + "camera_right": { + "original_key": "observation.images.right" + } + }, + "annotation": { + "task": { + "original_key": "task_index" + } + } +} diff --git a/open_h/embodiments/jhu_lscr_dvrk/modality_smarts.json b/open_h/embodiments/jhu_lscr_dvrk/modality_smarts.json new file mode 100644 index 0000000..1544b23 --- /dev/null +++ b/open_h/embodiments/jhu_lscr_dvrk/modality_smarts.json @@ -0,0 +1,94 @@ +{ + "state": { + "psm1_pose": { + "start": 0, + "end": 7, + "original_keys": [ + "observation.cartesian_state.psm1.pose.position.x", + "observation.cartesian_state.psm1.pose.position.y", + "observation.cartesian_state.psm1.pose.position.z", + "observation.cartesian_state.psm1.pose.orientation.x", + "observation.cartesian_state.psm1.pose.orientation.y", + "observation.cartesian_state.psm1.pose.orientation.z", + "observation.cartesian_state.psm1.pose.orientation.w" + ] + }, + "psm1_gripper": { + "start": 0, + "end": 1, + "original_key": "observation.state.psm1.gripper" + }, + "psm2_pose": { + "start": 0, + "end": 7, + "original_keys": [ + "observation.cartesian_state.psm2.pose.position.x", + "observation.cartesian_state.psm2.pose.position.y", + "observation.cartesian_state.psm2.pose.position.z", + "observation.cartesian_state.psm2.pose.orientation.x", + "observation.cartesian_state.psm2.pose.orientation.y", + "observation.cartesian_state.psm2.pose.orientation.z", + "observation.cartesian_state.psm2.pose.orientation.w" + ] + }, + "psm2_gripper": { + "start": 0, + "end": 1, + "original_key": "observation.state.psm2.gripper" + } + }, + "action": { + "psm1_pose": { + "start": 0, + "end": 7, + "original_keys": [ + "observation.cartesian_state.psm1.pose.position.x", + "observation.cartesian_state.psm1.pose.position.y", + "observation.cartesian_state.psm1.pose.position.z", + "observation.cartesian_state.psm1.pose.orientation.x", + "observation.cartesian_state.psm1.pose.orientation.y", + "observation.cartesian_state.psm1.pose.orientation.z", + "observation.cartesian_state.psm1.pose.orientation.w" + ] + }, + "psm1_gripper": { + "start": 0, + "end": 1, + "original_key": "action.psm1.gripper" + }, + "psm2_pose": { + "start": 0, + "end": 7, + "original_keys": [ + "observation.cartesian_state.psm2.pose.position.x", + "observation.cartesian_state.psm2.pose.position.y", + "observation.cartesian_state.psm2.pose.position.z", + "observation.cartesian_state.psm2.pose.orientation.x", + "observation.cartesian_state.psm2.pose.orientation.y", + "observation.cartesian_state.psm2.pose.orientation.z", + "observation.cartesian_state.psm2.pose.orientation.w" + ] + }, + "psm2_gripper": { + "start": 0, + "end": 1, + "original_key": "action.psm2.gripper" + } + }, + "video": { + "endoscope_left": { + "original_key": "observation.images.endoscope_left" + }, + "endoscope_right": { + "original_key": "observation.images.endoscope_right" + }, + "camera_side_view": { + "original_key": "observation.images.camera_side_view" + } + }, + "annotation": { + "task": { + "original_key": "task_index" + } + } +} diff --git a/open_h/embodiments/moon_maestro/README.md b/open_h/embodiments/moon_maestro/README.md new file mode 100644 index 0000000..e508bf9 --- /dev/null +++ b/open_h/embodiments/moon_maestro/README.md @@ -0,0 +1,59 @@ +# Moon Surgical (Maestro) Embodiment + +Dual-arm laparoscopic assistant robot with 20 task-level commands covering scope maintenance, anatomy navigation, instrument tracking, and zoom. + +## Embodiment Configuration + +| Property | Value | +|----------|-------| +| **Embodiment Tag** | `MOON_MAESTRO` | +| **Config File** | `open_h/embodiments/moon_maestro/moon_maestro_config.py` | +| **Modality Mapping** | `open_h/embodiments/moon_maestro/modality.json` | + +## Data Summary + +| Property | Value | +|----------|-------| +| Source | Moon Surgical (Maestro laparoscopic assistant) | +| Tasks | 20 assistant commands | +| Episodes | 65 | +| Frames | 12,020 | +| FPS | 30 | + +## Data Format + +### State +- 18D joint positions (9 joints per arm) + - `right_arm_joints`: joint_0_arm_0 .. joint_8_arm_0 + - `left_arm_joints`: joint_0_arm_1 .. joint_8_arm_1 + +### Action +- 6D delta translation in base frame + - `right_arm_delta_xyz`: d_tx_arm_0, d_ty_arm_0, d_tz_arm_0 + - `left_arm_delta_xyz`: d_tx_arm_1, d_ty_arm_1, d_tz_arm_1 +- Action representation in config: `DELTA` + `ActionFormat.XYZ` + +### Cameras +- `scope`: 960x540 +- `topcam`: 1280x720 + +### Language +- `tasks.jsonl` via `annotation.task` mapping in `modality.json` + +## Action Horizon +- 30 Hz, `ACTION_HORIZON = 50` (1.67 s prediction) + +## Dataset Preparation + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag MOON_MAESTRO \ + --modality-json open_h/embodiments/moon_maestro/modality.json \ + /path/to/moon_maestro_dataset +``` + +Note: Stats generation and the load test require `meta/modality.json` in the dataset path. + +## Notes +- No end-effector pose in the raw dataset (joint-only state). +- Actions are translation-only deltas (no rotation, no gripper signals). diff --git a/open_h/embodiments/moon_maestro/modality.json b/open_h/embodiments/moon_maestro/modality.json new file mode 100644 index 0000000..908bd6d --- /dev/null +++ b/open_h/embodiments/moon_maestro/modality.json @@ -0,0 +1,35 @@ +{ + "state": { + "right_arm_joints": { + "start": 0, + "end": 9 + }, + "left_arm_joints": { + "start": 9, + "end": 18 + } + }, + "action": { + "right_arm_delta_xyz": { + "start": 0, + "end": 3 + }, + "left_arm_delta_xyz": { + "start": 3, + "end": 6 + } + }, + "video": { + "scope": { + "original_key": "observation.images.scope" + }, + "topcam": { + "original_key": "observation.images.topcam" + } + }, + "annotation": { + "task": { + "original_key": "task_index" + } + } +} diff --git a/open_h/embodiments/moon_maestro/moon_maestro_config.py b/open_h/embodiments/moon_maestro/moon_maestro_config.py new file mode 100644 index 0000000..4dd892e --- /dev/null +++ b/open_h/embodiments/moon_maestro/moon_maestro_config.py @@ -0,0 +1,77 @@ +""" +Moon Surgical (Maestro) modality configuration for GR00T N1.6. + +This configuration targets the Moon Surgical assistant dataset: +- 65 episodes, 30 Hz +- Dual-arm robot with 9 joint positions per arm +- Delta Cartesian translation actions per arm (xyz only) + +Key design choices: +- Use joint positions as state +- Treat action deltas as ActionRepresentation.DELTA with ActionFormat.XYZ +- Use tasks.jsonl via annotation.task mapping +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +# 30 FPS -> 50 frames = 1.67 seconds of prediction +ACTION_HORIZON = 50 + +moon_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "scope", + "topcam", + ], + ), + "state": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "right_arm_joints", + "left_arm_joints", + ], + mean_std_embedding_keys=[ + "right_arm_joints", + "left_arm_joints", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(ACTION_HORIZON)), + modality_keys=[ + "right_arm_delta_xyz", + "left_arm_delta_xyz", + ], + action_configs=[ + ActionConfig( + rep=ActionRepresentation.DELTA, + type=ActionType.EEF, + format=ActionFormat.XYZ, + normalization_type="temporal_meanstd", + ), + ActionConfig( + rep=ActionRepresentation.DELTA, + type=ActionType.EEF, + format=ActionFormat.XYZ, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["annotation.task"], + ), +} + +# Register Moon Surgical embodiment config +register_modality_config(moon_config, embodiment_tag=EmbodimentTag.MOON_MAESTRO) diff --git a/open_h/embodiments/obuda_dvrk/README.md b/open_h/embodiments/obuda_dvrk/README.md new file mode 100644 index 0000000..f9fccc3 --- /dev/null +++ b/open_h/embodiments/obuda_dvrk/README.md @@ -0,0 +1,98 @@ +# Obuda dVRK + +Obuda University dVRK surgical datasets for Open-H tasks: +- FRS Dome (knot tying, suturing) +- Needle Threading +- Peg Transfer +- Rollercoaster +- Seaspike + +All raw datasets are LeRobot v2.1 with 30Hz video and 23D state/action. + +--- + +## Data Summary + +### Datasets + +| Dataset | Episodes | Tasks | +| --- | --- | --- | +| `FRS_Dome_1` | 102 | Knot tying, Suturing | +| `NeedleThreading_1` | 196 | Needle Threading | +| `PegTransfer_1` | 216 | Peg Transfer | +| `Rollercoaster_1` | 95 | Rollercoaster | +| `Seaspike_1` | 207 | Seaspike | + + +### Cameras + +| View | Original Key | +| --- | --- | +| `endoscope_left` | `observation.images.endoscope.left` | +| `wrist_left` | `observation.images.wrist.left` | +| `wrist_right` | `observation.images.wrist.right` | + +### State (16D) + +We use PSM1/PSM2 EEF poses only (ECM excluded): +- `psm1_pose` (7D): xyz + quaternion (xyzw) +- `psm1_gripper` (1D): jaw angle +- `psm2_pose` (7D): xyz + quaternion (xyzw) +- `psm2_gripper` (1D): jaw angle + +### Action (16D) + +Same layout as state, with absolute setpoints: +- EEF poses converted to `REL_XYZ_ROT6D` (xyz_rel + rot6d_rel) +- Grippers remain absolute + +--- + +## Modality Mapping + +Use `open_h/embodiments/obuda_dvrk/modality.json`: +- State: `psm1_pose`, `psm1_gripper`, `psm2_pose`, `psm2_gripper` +- Action: `psm1_pose`, `psm1_gripper`, `psm2_pose`, `psm2_gripper` +- Video: `endoscope_left`, `wrist_left`, `wrist_right` +- Language: `task` (from `task_index`) + +--- + +## Embodiment Configuration + +Embodiment tag: `OBUDA_DVRK` + +Action configuration: +- `REL_XYZ_ROT6D` for EEF poses (quaternion -> rot6d) +- `ABSOLUTE` for grippers +- `ACTION_HORIZON = 50` (30Hz -> 1.67s) + +See `open_h/embodiments/obuda_dvrk/obuda_dvrk_config.py`. + +--- + +## Dataset Preparation + +Generate stats files (copies modality JSON into `meta/`, generates `stats.json` and `temporal_stats.json`): + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag OBUDA_DVRK \ + --modality-json open_h/embodiments/obuda_dvrk/modality.json \ + /path/to/obuda_dataset +``` + +Repeat for each sub-dataset (FRS_Dome_1, NeedleThreading_1, etc.). + +--- + +--- + +## File Structure + +``` +open_h/embodiments/obuda_dvrk/ + modality.json # State/action/video field mapping + obuda_dvrk_config.py # Embodiment config (action configs, modality keys) + README.md # This file +``` diff --git a/open_h/embodiments/obuda_dvrk/modality.json b/open_h/embodiments/obuda_dvrk/modality.json new file mode 100644 index 0000000..1dd8fbc --- /dev/null +++ b/open_h/embodiments/obuda_dvrk/modality.json @@ -0,0 +1,49 @@ +{ + "state": { + "psm1_pose": { + "start": 0, + "end": 7 + }, + "psm1_gripper": { + "start": 7, + "end": 8 + }, + "psm2_pose": { + "start": 8, + "end": 15 + }, + "psm2_gripper": { + "start": 15, + "end": 16 + } + }, + "action": { + "psm1_pose": { + "start": 0, + "end": 7 + }, + "psm1_gripper": { + "start": 7, + "end": 8 + }, + "psm2_pose": { + "start": 8, + "end": 15 + }, + "psm2_gripper": { + "start": 15, + "end": 16 + } + }, + "video": { + "endoscope_left": { + "original_key": "observation.images.endoscope.left" + }, + "wrist_left": { + "original_key": "observation.images.wrist.left" + }, + "wrist_right": { + "original_key": "observation.images.wrist.right" + } + } +} diff --git a/open_h/embodiments/obuda_dvrk/obuda_dvrk_config.py b/open_h/embodiments/obuda_dvrk/obuda_dvrk_config.py new file mode 100644 index 0000000..8f301a5 --- /dev/null +++ b/open_h/embodiments/obuda_dvrk/obuda_dvrk_config.py @@ -0,0 +1,117 @@ +""" +Obuda dVRK modality configuration for GR00T N1.6. + +This configuration targets the Obuda University Open-H surgical datasets: +- FRS_Dome_1 (knot tying, suturing) +- NeedleThreading_1 +- PegTransfer_1 +- Rollercoaster_1 +- Seaspike_1 + +Design decisions: +- State/action use absolute Cartesian EEF poses for PSM1/PSM2 only (16D). +- ECM pose is excluded in v1 to align with standard dVRK state/action size. +- Cameras use endoscope.left + both wrist views (3 views). +- Actions use REL_XYZ_ROT6D for EEF poses and ABSOLUTE for grippers. +- Language uses tasks.jsonl via the "task" key. +- ACTION_HORIZON = 50 (all episodes are >= 50 frames). +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future steps to predict). +# We keep the 50-step horizon to match other Open-H surgical configs. +ACTION_HORIZON = 50 + +obuda_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "endoscope_left", + "wrist_left", + "wrist_right", + ], + ), + "state": ModalityConfig( + # Single reference state for REL_XYZ_ROT6D conversion. + delta_indices=[0], + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + # Use mean/std normalization for continuous pose and gripper values. + mean_std_embedding_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + ), + "action": ModalityConfig( + # Use consecutive indices [0..49] so horizon covers ~1.67s at 30 FPS. + delta_indices=list(range(ACTION_HORIZON)), + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + action_configs=[ + # PSM1 pose: REL_XYZ_ROT6D EEF action (xyz + rot6d output). + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm1_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + ), + # PSM1 gripper: absolute jaw angle. + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + # PSM2 pose: REL_XYZ_ROT6D EEF action (xyz + rot6d output). + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm2_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + ), + # PSM2 gripper: absolute jaw angle. + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + # "task" reads from tasks.jsonl via task_index. + modality_keys=["task"], + ), +} + +# Register with the Obuda dVRK embodiment tag. +register_modality_config(obuda_config, embodiment_tag=EmbodimentTag.OBUDA_DVRK) diff --git a/open_h/embodiments/polyu_sim/README.md b/open_h/embodiments/polyu_sim/README.md new file mode 100644 index 0000000..b10eedd --- /dev/null +++ b/open_h/embodiments/polyu_sim/README.md @@ -0,0 +1,57 @@ +# PolyU Simulated Surgical Embodiment + +Tissue retraction task in laparoscopic cholecystectomy using the PolyU Open-H surgical simulation dataset (msr_surgical robot). + +## Embodiment Configuration + +| Property | Value | +|----------|-------| +| **Embodiment Tag** | `POLYU_SIM` | +| **Config File** | `open_h/embodiments/polyu_sim/polyu_sim_config.py` | +| **Modality Mapping** | `open_h/embodiments/polyu_sim/modality.json` | + +## Data Summary + +| Property | Value | +|----------|-------| +| Source | PolyU (Open-H Surgical) | +| Task | Tissue Retraction in Laparoscopic Cholecystectomy | +| Episodes | 11,520 | +| Frames | 5,760,000 | +| FPS | 30 | + +## Data Format + +### State +- 10D model input + 7D pass-through: + - `psm_joints` (10D): joint angles from `observation.state` + - `psm_cartesian_pose` (7D, **pass-through**): xyz + quat_xyzw from `observation.cartesian_state` -- used as the reference state for REL_XYZ_ROT6D action conversion but not fed to the model as state input + +### Action +- 8D total: + - `psm_cartesian_pose` (7D): pose from `observation.cartesian_state` (aligned to t+1) + - `psm_gripper` (1D): gripper signal from `action` (likely all values are 1.0) + +### Cameras +- `endoscope`: `observation.images.main` (1280x720) + +### Language +- `task` from `meta/tasks.jsonl` + +## Action Horizon +- 50 steps @ 30 Hz (~1.7 s look-ahead) + +## Dataset Preparation + +Use the shared preparation script to copy the modality JSON into each dataset's `meta/` folder and generate normalization statistics: + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag polyu_sim \ + --modality-json open_h/embodiments/polyu_sim/modality.json \ + /path/to/dataset +``` + +## Notes +- Pose actions are converted to REL_XYZ_ROT6D using `psm_cartesian_pose` as the reference state. +- Joint deltas in `action[:10]` are present in the raw dataset, but the config uses pose actions. diff --git a/open_h/embodiments/polyu_sim/modality.json b/open_h/embodiments/polyu_sim/modality.json new file mode 100644 index 0000000..1a1ebf7 --- /dev/null +++ b/open_h/embodiments/polyu_sim/modality.json @@ -0,0 +1,13 @@ +{ + "video": { + "endoscope": {"original_key": "observation.images.main"} + }, + "state": { + "psm_joints": {"start": 0, "end": 10}, + "psm_cartesian_pose": {"start": 0, "end": 7, "original_key": "observation.cartesian_state"} + }, + "action": { + "psm_cartesian_pose": {"start": 0, "end": 7, "original_key": "observation.cartesian_state"}, + "psm_gripper": {"start": 10, "end": 11, "original_key": "action"} + } +} diff --git a/open_h/embodiments/polyu_sim/polyu_sim_config.py b/open_h/embodiments/polyu_sim/polyu_sim_config.py new file mode 100644 index 0000000..034a6f5 --- /dev/null +++ b/open_h/embodiments/polyu_sim/polyu_sim_config.py @@ -0,0 +1,84 @@ +"""PolyU OpenH_Dataset_full modality configuration for GR00T N1.6. + +This configuration supports the PolyU surgical dataset with: +- Joint-angle state (10D) that the policy sees +- Cartesian pose state (7D) used as a pass-through reference for REL_XYZ_ROT6D +- End-effector pose actions (7D: xyz + quat_xyzw) converted to REL_XYZ_ROT6D +- Gripper action (1D) sourced from the action column +- Single endoscope video stream + +Data Format: +- State: 17D = psm_joints(10) + psm_cartesian_pose(7: xyz + quat_xyzw) +- Action: 8D = psm_cartesian_pose(7) + psm_gripper(1) + +Action Semantics: +- Pose actions are sourced from observation.cartesian_state and aligned to t+1. +- REL_XYZ_ROT6D conversion uses psm_cartesian_pose as the reference state. +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict). +ACTION_HORIZON = 50 + +polyu_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "endoscope", + ], + ), + "state": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "psm_joints", + "psm_cartesian_pose", + ], + pass_through_keys=[ + "psm_cartesian_pose", + ], + mean_std_embedding_keys=[ + "psm_joints", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(1, ACTION_HORIZON + 1)), + modality_keys=[ + "psm_cartesian_pose", + "psm_gripper", + ], + action_configs=[ + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm_cartesian_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + ), + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["task"], + ), +} + +# Register with PolyU simulated surgical tag for Open-H data. +register_modality_config(polyu_config, embodiment_tag=EmbodimentTag.POLYU_SIM) diff --git a/open_h/embodiments/rob_surgical_bitrack/README.md b/open_h/embodiments/rob_surgical_bitrack/README.md new file mode 100644 index 0000000..f3249f5 --- /dev/null +++ b/open_h/embodiments/rob_surgical_bitrack/README.md @@ -0,0 +1,64 @@ +# Rob Surgical (bitrack) Dataset + +This directory contains the GR00T N1.6 integration assets for the Rob Surgical +bitrack dataset. The dataset uses LeRobot v2.1 format with a single endoscope +video stream and 3-arm Cartesian kinematics (left, right, aux). + +## Key characteristics + +- **FPS:** 30 Hz +- **Camera:** `observation.images.endoscope` (720 x 1280, AV1, yuv420p) +- **Language:** `instruction.text_with_tool` +- **Arms:** left, right, aux (3 active in model config; `lap_pose` is defined in `modality.json` for data loading but excluded from `rob_surgical_bitrack_config.py`) +- **State (EEF):** 18D = 3 arms * (xyz + roll + pitch + yaw) +- **Action (EEF):** 18D = 3 arms * (xyz + roll + pitch + yaw) +- **Action horizon:** 50 steps (~1.67s at 30 Hz) +- **Action representation:** REL_XYZ_ROT6D with Euler (RPY) input + +## Dataset preparation + +### 1. Add tool prompts to language instructions + +Before training, run the tool prompt script to generate `instruction.text_with_tool` from per-frame tool metadata. This prefixes each language instruction with the active tool names (e.g., `left tool: grasper. right tool: scissors. aux tool: none.`): + +```bash +uv run python open_h/embodiments/rob_surgical_bitrack/utils/rob_surgical_add_tool_prompts.py \ + --dataset-path /path/to/rob_surgical_dataset +``` + +### 2. Generate normalization statistics +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag ROB_SURGICAL_BITRACK \ + --modality-json open_h/embodiments/rob_surgical_bitrack/modality.json \ + /path/to/rob_surgical_dataset +``` + +## Modality mapping + +See `modality.json` for exact key mappings: +- `endoscope` -> `observation.images.endoscope` +- `left_pose/right_pose/aux_pose` -> slices of `observation.end_effector_state` (`lap_pose` is loaded via `modality.json` but excluded from model config) +- `action` uses the same slices from `action` +- `annotation.instruction` -> `instruction.text_with_tool` (raw text passthrough) + +## Base dataset fields (quick note) + +- The merged dataset does **not** include gripper/jaw values; EEF state/action are 18D poses only (3 active arms). +- Tool name strings are present (`observation.meta.left_tool/right_tool/aux_tool`). +- Arm visibility flags existed only in one source dataset and were dropped in the merge. + +## Action vs state alignment (Rob Surgical) + +- `state` is the current EEF pose at timestep `t` (`delta_indices=[0]`). +- `action` is an **absolute** target pose sequence at timesteps `t..t+49` + (`delta_indices=list(range(50))`), then converted to REL_XYZ_ROT6D w.r.t. + the current state during training. +- In the merged dataset, `action[t]` is **not** equal to `state[t]` or `state[t+1]`; + absolute differences can be large (max diff ~800–1200, mean diff ~6–15). + +## Notes on EEF NaNs + +> **Warning: Data Quality Note** β€” Some original episodes contain NaNs in the EEF `l_x` and `r_x` components. +> **A cleaning script was used to impute these values from the corresponding action x-values** so +> that REL_XYZ_ROT6D conversion can proceed without dropping frames. diff --git a/open_h/embodiments/rob_surgical_bitrack/modality.json b/open_h/embodiments/rob_surgical_bitrack/modality.json new file mode 100644 index 0000000..27a7d18 --- /dev/null +++ b/open_h/embodiments/rob_surgical_bitrack/modality.json @@ -0,0 +1,20 @@ +{ + "video": { + "endoscope": {"original_key": "observation.images.endoscope"} + }, + "state": { + "left_pose": {"start": 0, "end": 6, "original_key": "observation.end_effector_state"}, + "right_pose": {"start": 6, "end": 12, "original_key": "observation.end_effector_state"}, + "lap_pose": {"start": 12, "end": 18, "original_key": "observation.end_effector_state"}, + "aux_pose": {"start": 18, "end": 24, "original_key": "observation.end_effector_state"} + }, + "action": { + "left_pose": {"start": 0, "end": 6}, + "right_pose": {"start": 6, "end": 12}, + "lap_pose": {"start": 12, "end": 18}, + "aux_pose": {"start": 18, "end": 24} + }, + "annotation": { + "instruction": {"original_key": "instruction.text_with_tool", "is_text": true} + } +} diff --git a/open_h/embodiments/rob_surgical_bitrack/rob_surgical_bitrack_config.py b/open_h/embodiments/rob_surgical_bitrack/rob_surgical_bitrack_config.py new file mode 100644 index 0000000..79d0cee --- /dev/null +++ b/open_h/embodiments/rob_surgical_bitrack/rob_surgical_bitrack_config.py @@ -0,0 +1,117 @@ +""" +Rob Surgical (bitrack) modality configuration for GR00T N1.6. + +This configuration supports the Rob Surgical dataset with: +- Single endoscope video stream +- 3-arm Cartesian EEF state (left, right, aux) +- 3-arm absolute EEF action targets (xyz + roll/pitch/yaw) +- REL_XYZ_ROT6D action conversion using Euler angles + +Data representation: +- State: 18D = 3 * (xyz + roll + pitch + yaw) +- Action: 18D = 3 * (xyz + roll + pitch + yaw) +- Rotation format: Euler RPY (radians), converted to rot6d internally +- Note: lap_pose is defined in modality.json for data loading but is + excluded from this model config (commented out below). + +Notes: +- EEF x-value NaNs were imputed during the dataset merge step by copying + the action x-values for the affected arm. This config assumes merged data. +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +# 30 Hz -> 50 frames ~= 1.67 seconds +ACTION_HORIZON = 50 + +rob_surgical_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "endoscope", + ], + ), + "state": ModalityConfig( + delta_indices=[0], # Single reference state for REL_XYZ_ROT6D conversion + modality_keys=[ + "left_pose", + "right_pose", + # "lap_pose", + "aux_pose", + ], + mean_std_embedding_keys=[ + "left_pose", + "right_pose", + # "lap_pose", + "aux_pose", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(ACTION_HORIZON)), + modality_keys=[ + "left_pose", + "right_pose", + # "lap_pose", + "aux_pose", + ], + action_configs=[ + # Left arm pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="left_pose", + normalization_type="temporal_meanstd", + input_rotation_format="euler", + reference_rotation_format="euler", + ), + # Right arm pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="right_pose", + normalization_type="temporal_meanstd", + input_rotation_format="euler", + reference_rotation_format="euler", + ), + # Lap arm pose: REL_XYZ_ROT6D EEF action + # ActionConfig( + # rep=ActionRepresentation.REL_XYZ_ROT6D, + # type=ActionType.EEF, + # format=ActionFormat.XYZ_ROT6D, + # state_key="lap_pose", + # normalization_type="percentile", + # input_rotation_format="euler", + # reference_rotation_format="euler", + # ), + # Aux arm pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="aux_pose", + normalization_type="temporal_meanstd", + input_rotation_format="euler", + reference_rotation_format="euler", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["annotation.instruction"], # instruction.text (raw strings) + ), +} + +# Register with ROB_SURGICAL_BITRACK tag for Rob Surgical dataset integration +register_modality_config(rob_surgical_config, embodiment_tag=EmbodimentTag.ROB_SURGICAL_BITRACK) diff --git a/open_h/embodiments/rob_surgical_bitrack/utils/rob_surgical_add_tool_prompts.py b/open_h/embodiments/rob_surgical_bitrack/utils/rob_surgical_add_tool_prompts.py new file mode 100644 index 0000000..a150cd2 --- /dev/null +++ b/open_h/embodiments/rob_surgical_bitrack/utils/rob_surgical_add_tool_prompts.py @@ -0,0 +1,216 @@ +"""Add tool-augmented instruction text to Rob Surgical parquet files. + +This script reads the existing instruction.text prompts and the per-frame tool +metadata columns (observation.meta.left_tool, observation.meta.right_tool, +observation.meta.aux_tool), then constructs a new instruction column that +prefixes the original prompt with structured tool information. + +Output format per row: + left tool: . right tool: . aux tool: . + +If a tool column is missing or empty, the literal string "none" is used. + +The new column is written as `instruction.text_with_tool` into each episode +parquet file in-place. + +Usage: + uv run python open_h/embodiments/rob_surgical_bitrack/utils/rob_surgical_add_tool_prompts.py \ + --dataset-path /path/to/rob_surgical_dataset + + # Dry-run (print first 5 episodes, don't write): + uv run python open_h/embodiments/rob_surgical_bitrack/utils/rob_surgical_add_tool_prompts.py \ + --dataset-path /path/to/rob_surgical_dataset --dry-run +""" + +from __future__ import annotations + +import argparse +import glob +from pathlib import Path + +import pandas as pd + + +DEFAULT_DATASET_PATH = "/path/to/rob_surgical_dataset" + +# Parquet columns that contain the per-frame tool type string. +TOOL_COLUMNS = { + "left": "observation.meta.left_tool", + "right": "observation.meta.right_tool", + "aux": "observation.meta.aux_tool", +} + +# Source instruction column and the new output column. +SRC_INSTRUCTION_COL = "instruction.text" +DST_INSTRUCTION_COL = "instruction.text_with_tool" + + +def _safe_tool_name(value: object) -> str: + """Return a cleaned tool name, falling back to 'none' for missing values. + + Args: + value: Raw cell value from the tool metadata column. May be a string, + NaN, None, or empty string. + + Returns: + Lowercase tool name string, or the literal string 'none'. + + Examples: + >>> _safe_tool_name("bi_maryland_forceps") + 'bi_maryland_forceps' + >>> _safe_tool_name(None) + 'none' + >>> _safe_tool_name("") + 'none' + """ + if value is None: + return "none" + s = str(value).strip().lower() + if s in ("", "nan", "null"): + return "none" + return s + + +def build_tool_prefix(row: pd.Series) -> str: + """Construct the tool-description prefix for a single row. + + Args: + row: A pandas Series containing the tool metadata columns. + + Returns: + A string in the format: + 'left tool: . right tool: . aux tool: .' + + Examples: + >>> import pandas as pd + >>> row = pd.Series( + ... { + ... "observation.meta.left_tool": "mono_hook", + ... "observation.meta.right_tool": "forceps", + ... "observation.meta.aux_tool": "bi_maryland_forceps", + ... } + ... ) + >>> build_tool_prefix(row) + 'left tool: mono_hook. right tool: forceps. aux tool: bi_maryland_forceps.' + """ + parts = [] + for arm_label, col_name in TOOL_COLUMNS.items(): + tool_name = _safe_tool_name(row.get(col_name)) + parts.append(f"{arm_label} tool: {tool_name}") + return ". ".join(parts) + "." + + +def build_augmented_instruction(row: pd.Series) -> str: + """Build the full augmented instruction string for a single row. + + Combines the tool prefix with the original instruction text. If the + original instruction is missing or empty, only the tool prefix is returned. + + Args: + row: A pandas Series with tool metadata and instruction.text columns. + + Returns: + Augmented instruction string. + + Examples: + >>> import pandas as pd + >>> row = pd.Series( + ... { + ... "observation.meta.left_tool": "mono_hook", + ... "observation.meta.right_tool": "forceps", + ... "observation.meta.aux_tool": "bi_maryland_forceps", + ... "instruction.text": "The surgery process is a hemicolectomy.Suturing", + ... } + ... ) + >>> build_augmented_instruction(row) + 'left tool: mono_hook. right tool: forceps. aux tool: bi_maryland_forceps. The surgery process is a hemicolectomy.Suturing' + """ + prefix = build_tool_prefix(row) + original = str(row.get(SRC_INSTRUCTION_COL, "")).strip() + if not original or original.lower() in ("nan", "none"): + return prefix + return f"{prefix} {original}" + + +def process_episode(parquet_path: Path, dry_run: bool = False) -> int: + """Add the instruction.text_with_tool column to a single episode parquet. + + Args: + parquet_path: Path to the episode parquet file. + dry_run: If True, print sample output but do not write. + + Returns: + Number of rows processed. + + Raises: + KeyError: If the source instruction column is missing. + """ + df = pd.read_parquet(parquet_path) + + # Build augmented instruction for every row + df[DST_INSTRUCTION_COL] = df.apply(build_augmented_instruction, axis=1) + + if dry_run: + # Show first row as sample + sample = df[DST_INSTRUCTION_COL].iloc[0] + print(f' {parquet_path.name}: "{sample}"') + else: + df.to_parquet(parquet_path, index=False) + + return len(df) + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments. + + Returns: + Parsed arguments namespace. + """ + parser = argparse.ArgumentParser( + description="Add tool-augmented instruction text to Rob Surgical parquets." + ) + parser.add_argument( + "--dataset-path", + type=Path, + default=Path(DEFAULT_DATASET_PATH), + help="Path to the Rob Surgical LeRobot dataset root.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print sample outputs for the first 5 episodes without writing.", + ) + return parser.parse_args() + + +def main() -> None: + """Entry point: iterate all episode parquets and add augmented instructions. + + Raises: + FileNotFoundError: If no parquet files are found at the dataset path. + """ + args = parse_args() + pattern = str(args.dataset_path / "data" / "chunk-*" / "episode_*.parquet") + parquet_files = sorted(glob.glob(pattern)) + + if not parquet_files: + raise FileNotFoundError(f"No parquet files found matching: {pattern}") + + total_files = len(parquet_files) + limit = 5 if args.dry_run else total_files + mode_label = "DRY RUN" if args.dry_run else "WRITING" + + print(f"[{mode_label}] Processing {min(limit, total_files)} / {total_files} episodes...") + print(f" Source column: {SRC_INSTRUCTION_COL}") + print(f" Target column: {DST_INSTRUCTION_COL}") + print() + + total_rows = 0 + for i, f in enumerate(parquet_files[:limit]): + total_rows += process_episode(Path(f), dry_run=args.dry_run) + + print(f"\nDone. Processed {total_rows} total rows across {min(limit, total_files)} episodes.") + + +if __name__ == "__main__": + main() diff --git a/open_h/embodiments/sanoscience_sim/README.md b/open_h/embodiments/sanoscience_sim/README.md new file mode 100644 index 0000000..5038967 --- /dev/null +++ b/open_h/embodiments/sanoscience_sim/README.md @@ -0,0 +1,127 @@ +# SanoScience Sim + +Simulated surgical robot data from SanoScience with 4-instrument control. + +## Dataset Statistics + +> Dataset: SanoScience v1.2 merged (480p, cleaned) β€” 5 training groups + +| Property | Value | +|----------|-------| +| **Episodes** | 14,004 | +| **Total Frames** | 1,574,892 | +| **Avg Frames/Episode** | ~112 | + +Training groups: +- Expert_demonstrations: 5,454 episodes, 603,054 frames (25 fps) +- NonExpert_full_modalities_clean_final: 6,156 episodes, 713,070 frames (30 fps) +- NonExpert_partial_modalities_clean_final: 666 episodes, 58,752 frames (30 fps) +- NonExpert_recovery_clean_final: 126 episodes, 20,376 frames (30 fps) +- NonExpert_stereo_clean_final: 1,602 episodes, 179,640 frames (30 fps) + +### Episode Length Distribution + +**Key stats:** Min=16, Max=669, Median=102, Mean=112.5, Q25=90, Q75=122 + +### Action Horizon Selection + +Episodes shorter than the action horizon are **skipped during training** (not padded). We use `ACTION_HORIZON = 36`: + +- 13,968 / 14,004 episodes usable (99.7%) β€” only 36 episodes dropped (< 36 frames) +- Usable training steps: 1,085,094 + +## Embodiment Configuration + +| Property | Value | +|----------|-------| +| **Embodiment Tag** | `sanoscience_sim` | +| **Config File** | `open_h/embodiments/sanoscience_sim/sanoscience_sim_config.py` | + + + +## Data Format + +### State (32D) + +The state vector is organized by instrument, with each instrument having a 7D pose and 1D gripper: +State is sourced from `action.cartesian_absolute` via `modality.json` to keep state/action formats consistent. + +| Instrument | Indices | Dimensions | Description | +|------------|---------|------------|-------------| +| inst_0 | 0-7 | 7D + 1D | xyz (3D) + quaternion xyzw (4D) + gripper_angle_rad (1D) | +| inst_1 | 8-15 | 7D + 1D | xyz (3D) + quaternion xyzw (4D) + gripper_angle_rad (1D) | +| inst_2 | 16-23 | 7D + 1D | xyz (3D) + quaternion xyzw (4D) + gripper_angle_rad (1D) | +| inst_3 | 24-31 | 7D + 1D | xyz (3D) + quaternion xyzw (4D) + gripper_angle_rad (1D) | + +### Action (32D) + +Same format as state. Actions are extracted from the `action.cartesian_absolute` column (specified in `modality.json`). + +After REL_XYZ_ROT6D conversion, the action output becomes 40D: +- 4 instruments x (9D REL_XYZ_ROT6D pose + 1D gripper) +- Pose: xyz_rel (3D) + rot6d_rel (6D) = 9D +- Gripper: absolute angle (1D) + +Action horizon and indexing: +- `ACTION_HORIZON = 36` β€” 99.7% of episodes usable. +- `delta_indices = [1, 2, 3, ..., 36]` so index 0 is the state reference (CMR-style offset). + +### Video (1 view) + +| Camera | Original Key | +|--------|--------------| +| `camera_color` | `observation.images.color` | + +## Dataset Preparation + +Generate normalization statistics (from repo root): + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag SANOSCIENCE_SIM \ + --modality-json open_h/embodiments/sanoscience_sim/modality.json \ + /path/to/sanoscience_dataset +``` + +## Training + +### Single-Dataset Training + +```bash +CUDA_VISIBLE_DEVICES=0 uv run python \ + gr00t/experiment/launch_finetune.py \ + --base-model-path nvidia/GR00T-N1.6-3B \ + --dataset-path /path/to/your/sanoscience-dataset \ + --embodiment-tag SANOSCIENCE_SIM \ + --num-gpus 1 \ + --output-dir /path/to/output/checkpoints \ + --save-steps 2000 \ + --save-total-limit 5 \ + --max-steps 50000 \ + --warmup-ratio 0.05 \ + --weight-decay 1e-5 \ + --learning-rate 1e-4 \ + --use-wandb \ + --global-batch-size 32 +``` + +## Open-Loop Evaluation + +```bash +uv run python gr00t/eval/open_loop_eval.py \ + --dataset-path /path/to/sanoscience-dataset \ + --embodiment-tag SANOSCIENCE_SIM \ + --model-path /path/to/checkpoint \ + --traj-ids 0 1 2 \ + --action-horizon 36 \ + --steps 500 +``` + +## File Structure + +``` +open_h/embodiments/sanoscience_sim/ +β”œβ”€β”€ README.md # This file +β”œβ”€β”€ sanoscience_sim_config.py # Built-in config module (auto-registered) +└── modality.json # Data key mappings (copy to dataset/meta/) +``` diff --git a/open_h/embodiments/sanoscience_sim/modality.json b/open_h/embodiments/sanoscience_sim/modality.json new file mode 100644 index 0000000..04358b0 --- /dev/null +++ b/open_h/embodiments/sanoscience_sim/modality.json @@ -0,0 +1,91 @@ +{ + "video": { + "camera_color": { + "original_key": "observation.images.color" + } + }, + "state": { + "inst_0_pose": { + "start": 0, + "end": 7, + "original_key": "action.cartesian_absolute" + }, + "inst_0_gripper": { + "start": 7, + "end": 8, + "original_key": "action.cartesian_absolute" + }, + "inst_1_pose": { + "start": 8, + "end": 15, + "original_key": "action.cartesian_absolute" + }, + "inst_1_gripper": { + "start": 15, + "end": 16, + "original_key": "action.cartesian_absolute" + }, + "inst_2_pose": { + "start": 16, + "end": 23, + "original_key": "action.cartesian_absolute" + }, + "inst_2_gripper": { + "start": 23, + "end": 24, + "original_key": "action.cartesian_absolute" + }, + "inst_3_pose": { + "start": 24, + "end": 31, + "original_key": "action.cartesian_absolute" + }, + "inst_3_gripper": { + "start": 31, + "end": 32, + "original_key": "action.cartesian_absolute" + } + }, + "action": { + "inst_0_pose": { + "start": 0, + "end": 7, + "original_key": "action.cartesian_absolute" + }, + "inst_0_gripper": { + "start": 7, + "end": 8, + "original_key": "action.cartesian_absolute" + }, + "inst_1_pose": { + "start": 8, + "end": 15, + "original_key": "action.cartesian_absolute" + }, + "inst_1_gripper": { + "start": 15, + "end": 16, + "original_key": "action.cartesian_absolute" + }, + "inst_2_pose": { + "start": 16, + "end": 23, + "original_key": "action.cartesian_absolute" + }, + "inst_2_gripper": { + "start": 23, + "end": 24, + "original_key": "action.cartesian_absolute" + }, + "inst_3_pose": { + "start": 24, + "end": 31, + "original_key": "action.cartesian_absolute" + }, + "inst_3_gripper": { + "start": 31, + "end": 32, + "original_key": "action.cartesian_absolute" + } + } +} diff --git a/open_h/embodiments/sanoscience_sim/sanoscience_sim_config.py b/open_h/embodiments/sanoscience_sim/sanoscience_sim_config.py new file mode 100644 index 0000000..8b4b869 --- /dev/null +++ b/open_h/embodiments/sanoscience_sim/sanoscience_sim_config.py @@ -0,0 +1,215 @@ +""" +SanoScience modality configuration for GR00T N1.6. + +This configuration supports a simulated surgical robot with 4 instruments: +- REL_XYZ_ROT6D action representation for EEF poses +- Temporal mean-std normalization for actions +- Mean-std normalization for state +- Vision+language only training supported via state dropout +- 1 camera view (color) + +Dataset: SanoScience v1.2 merged + +Training groups: + - Expert_demonstrations: 5,454 episodes, 603,054 frames, 25 fps + - NonExpert_full_modalities_clean_final: 6,156 episodes, 713,070 frames, 30 fps + - NonExpert_partial_modalities_clean_final: 666 episodes, 58,752 frames, 30 fps + - NonExpert_recovery_clean_final: 126 episodes, 20,376 frames, 30 fps + - NonExpert_stereo_clean_final: 1,602 episodes, 179,640 frames, 30 fps + Total: 14,004 episodes, 1,574,892 frames + +Episode Length Statistics (across all 5 training groups): +- Min length: 16 frames +- Max length: 669 frames +- Mean length: 112.5 frames +- Median length: 102 frames +- Quartiles: Q25=90, Q75=122 +- 100% of episodes usable with ACTION_HORIZON=8 (all episodes >= 16 frames) +- 99.7% of episodes usable with ACTION_HORIZON=36 (only 36 dropped) + +Action Horizon Selection: +- Using ACTION_HORIZON = 36 +- 13,968 / 14,004 episodes usable (99.7%) β€” 36 episodes dropped (< 36 frames) +- Usable training steps: 1,085,094 + +Data Format: +- State: 32D total from action.cartesian_absolute, split by instrument: + - inst_0_pose (7D): xyz + quat_xyzw (indices 0-7) + - inst_0_gripper (1D): gripper_angle_rad (index 7) + - inst_1_pose (7D): xyz + quat_xyzw (indices 8-15) + - inst_1_gripper (1D): gripper_angle_rad (index 15) + - inst_2_pose (7D): xyz + quat_xyzw (indices 16-23) + - inst_2_gripper (1D): gripper_angle_rad (index 23) + - inst_3_pose (7D): xyz + quat_xyzw (indices 24-31) + - inst_3_gripper (1D): gripper_angle_rad (index 31) + +- Action: 32D (same format as state, representing setpoints) + - Extracted from action.cartesian_absolute column via modality.json original_key + +- Additional columns available but NOT used by this config: + - action (64D): full action vector including handle data + - action.camera_absolute (32D): camera-frame actions + - action.joint_positions (48D): joint positions from inverse kinematics + - observation.state (64D): full state vector + - observation.state.joint_positions (48D) + - observation.state.joint_velocities (48D) + - instruction.text: per-frame language (identical to tasks.jsonl task string) + +Final Action Output: 40D = 4 instruments x (9D REL_XYZ_ROT6D pose + 1D gripper) +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +# With horizon=36: 99.7% of episodes usable (13,968/14,004), 1,085,094 training steps +# Only 36 episodes dropped (those with < 36 frames) +ACTION_HORIZON = 36 + +sanoscience_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "camera_color", + ], + ), + "state": ModalityConfig( + delta_indices=[0], # Single reference state for REL_XYZ_ROT6D + modality_keys=[ + # Instrument 0 + "inst_0_pose", + "inst_0_gripper", + # Instrument 1 + "inst_1_pose", + "inst_1_gripper", + # Instrument 2 + "inst_2_pose", + "inst_2_gripper", + # Instrument 3 + "inst_3_pose", + "inst_3_gripper", + ], + # Mean-std normalization for all state keys + mean_std_embedding_keys=[ + "inst_0_pose", + "inst_0_gripper", + "inst_1_pose", + "inst_1_gripper", + "inst_2_pose", + "inst_2_gripper", + "inst_3_pose", + "inst_3_gripper", + ], + ), + "action": ModalityConfig( + # Start at +1 so index 0 is state reference (CMR-style offset) + delta_indices=list(range(1, ACTION_HORIZON + 1)), # [1, 2, 3, ..., 36] + modality_keys=[ + # Instrument 0 + "inst_0_pose", + "inst_0_gripper", + # Instrument 1 + "inst_1_pose", + "inst_1_gripper", + # Instrument 2 + "inst_2_pose", + "inst_2_gripper", + # Instrument 3 + "inst_3_pose", + "inst_3_gripper", + ], + action_configs=[ + # ===== Instrument 0 ===== + # Pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="inst_0_pose", # Reference state for relative conversion + normalization_type="temporal_meanstd", + input_rotation_format="quat", # Input actions are xyz + quaternion (xyzw) + reference_rotation_format="quat", # Reference state is also xyz + quaternion + ), + # Gripper: absolute angle (no rotation conversion needed) + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + # ===== Instrument 1 ===== + # Pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="inst_1_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + ), + # Gripper: absolute angle + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + # ===== Instrument 2 ===== + # Pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="inst_2_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + ), + # Gripper: absolute angle + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + # ===== Instrument 3 ===== + # Pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="inst_3_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + ), + # Gripper: absolute angle + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["task"], # Uses tasks.jsonl + ), +} + +# Register with SANOSCIENCE_SIM tag for simulated surgical robot finetuning +register_modality_config(sanoscience_config, embodiment_tag=EmbodimentTag.SANOSCIENCE_SIM) diff --git a/open_h/embodiments/stanford_dvrk_real/README.md b/open_h/embodiments/stanford_dvrk_real/README.md new file mode 100644 index 0000000..2e65b84 --- /dev/null +++ b/open_h/embodiments/stanford_dvrk_real/README.md @@ -0,0 +1,95 @@ +# Stanford Real Robot (dVRK) + +Stanford real-robot dVRK datasets for Open-H surgical tasks: +- Needle Transfer +- Tissue Retraction +- Peg Transfer + +All datasets are LeRobot v2.1 with 30Hz stereo endoscopic video. + +--- + +## Data Summary + +### Episodes and Frames (from `meta/info.json`) +- Needle Transfer: 700 episodes, 313,882 frames +- Tissue Retraction: 698 episodes, 291,826 frames +- Peg Transfer: 598 episodes, 268,729 frames + +### Cameras +- `observation.images.camera_left` (540x960, 30Hz) +- `observation.images.camera_right` (540x960, 30Hz) + +### State (`observation.state`, shape=26) +We only use the EEF pose + gripper for each arm: +- PSM1 gripper: index 6 +- PSM1 EEF pose (xyz + roll/pitch/yaw): indices 7-12 +- PSM2 gripper: index 19 +- PSM2 EEF pose (xyz + roll/pitch/yaw): indices 20-25 + +### Action (`action`, shape=14) +Absolute Cartesian EEF pose + gripper in camera/ECM frame: +- PSM1 gripper: index 0 +- PSM1 EEF pose (xyz + roll/pitch/yaw): indices 1-6 +- PSM2 gripper: index 7 +- PSM2 EEF pose (xyz + roll/pitch/yaw): indices 8-13 + +Orientation uses Euler RPY (roll, pitch, yaw) in radians. We assume `xyz` extrinsic. + +--- + +## Modality Mapping + +Use `open_h/embodiments/stanford_dvrk_real/modality_real_robot.json`: +- State: `psm1_pose`, `psm1_gripper`, `psm2_pose`, `psm2_gripper` +- Action: `psm1_pose`, `psm1_gripper`, `psm2_pose`, `psm2_gripper` +- Video: `endoscope_left` (mono; `endoscope_right` is defined in `modality_real_robot.json` but disabled in the current config) +- Language: `task` (from `task_index`) + +--- + +## Embodiment Configuration + +Embodiment tag: `stanford_dvrk_real` + +Action configuration: +- `REL_XYZ_ROT6D` for EEF pose (Euler RPY -> rot6d) +- `ABSOLUTE` for grippers +- `ACTION_HORIZON = 50` (30Hz -> 1.67s) + +See `open_h/embodiments/stanford_dvrk_real/stanford_dvrk_real_config.py`. + +--- + +## Split Policy (Keep Recovery, Exclude Fail) + +Needle Transfer and Peg Transfer include `recovery` and `fail` splits in `info.json`. +We **keep recovery** but **exclude fail** using runtime split filtering: + +In training configs (YAML): +```yaml +exclude_splits: ["fail"] +``` + +This applies to both training and stats generation without copying datasets. + +--- + +## Stats Generation + +Generate stats files (copies modality, generates `stats.json` and `temporal_stats.json`): + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag STANFORD_DVRK_REAL \ + --modality-json open_h/embodiments/stanford_dvrk_real/modality_real_robot.json \ + /path/to/stanford_dataset +``` + +Repeat for each task subset (Needle Transfer, Tissue Retraction, Peg Transfer). + +--- + +## Notes +- `next.done` is **not** present in Stanford real-robot datasets, so no terminal padding filter is needed. +- If joint states become useful later, extend modality config to include joint groups from `observation.state`. diff --git a/open_h/embodiments/stanford_dvrk_real/modality_real_robot.json b/open_h/embodiments/stanford_dvrk_real/modality_real_robot.json new file mode 100644 index 0000000..45d4a81 --- /dev/null +++ b/open_h/embodiments/stanford_dvrk_real/modality_real_robot.json @@ -0,0 +1,59 @@ +{ + "state": { + "psm1_pose": { + "start": 7, + "end": 13, + "original_key": "observation.state" + }, + "psm1_gripper": { + "start": 6, + "end": 7, + "original_key": "observation.state" + }, + "psm2_pose": { + "start": 20, + "end": 26, + "original_key": "observation.state" + }, + "psm2_gripper": { + "start": 19, + "end": 20, + "original_key": "observation.state" + } + }, + "action": { + "psm1_pose": { + "start": 1, + "end": 7, + "original_key": "action" + }, + "psm1_gripper": { + "start": 0, + "end": 1, + "original_key": "action" + }, + "psm2_pose": { + "start": 8, + "end": 14, + "original_key": "action" + }, + "psm2_gripper": { + "start": 7, + "end": 8, + "original_key": "action" + } + }, + "video": { + "endoscope_left": { + "original_key": "observation.images.camera_left" + }, + "endoscope_right": { + "original_key": "observation.images.camera_right" + } + }, + "annotation": { + "task": { + "original_key": "task_index" + } + } +} diff --git a/open_h/embodiments/stanford_dvrk_real/stanford_dvrk_real_config.py b/open_h/embodiments/stanford_dvrk_real/stanford_dvrk_real_config.py new file mode 100644 index 0000000..31b908c --- /dev/null +++ b/open_h/embodiments/stanford_dvrk_real/stanford_dvrk_real_config.py @@ -0,0 +1,111 @@ +""" +Stanford real-robot (dVRK) modality configuration for GR00T N1.6. + +This configuration supports the Stanford dVRK real-robot datasets: +- Needle Transfer +- Tissue Retraction +- Peg Transfer + +Data Format: +- State: 12D Cartesian EEF (2 arms Γ— [xyz + roll/pitch/yaw]) + 2 grippers +- Action: 12D Cartesian absolute (same as state poses) + 2 grippers + +Action Representation: +- EEF poses are converted to REL_XYZ_ROT6D using Euler RPY inputs (xyz + rpy) +- Grippers are ABSOLUTE jaw angles +- Euler convention: roll/pitch/yaw in radians with `xyz` extrinsic rotation + +Video: +- Stereo endoscope views: camera_left, camera_right (540x960 @ 30Hz) +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (30 FPS -> 50 steps = 1.67 seconds) +ACTION_HORIZON = 50 + +stanford_real_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "endoscope_left", + # "endoscope_right", # Mono Only + ], + ), + "state": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + mean_std_embedding_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(ACTION_HORIZON)), + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + action_configs=[ + # PSM1 pose: REL_XYZ_ROT6D with Euler RPY inputs + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm1_pose", + normalization_type="temporal_meanstd", + input_rotation_format="euler", + reference_rotation_format="euler", + ), + # PSM1 gripper: ABSOLUTE jaw angle + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + # PSM2 pose: REL_XYZ_ROT6D with Euler RPY inputs + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm2_pose", + normalization_type="temporal_meanstd", + input_rotation_format="euler", + reference_rotation_format="euler", + ), + # PSM2 gripper: ABSOLUTE jaw angle + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["task"], + ), +} + +# Register with Stanford real-robot embodiment tag +register_modality_config(stanford_real_config, embodiment_tag=EmbodimentTag.STANFORD_DVRK_REAL) diff --git a/open_h/embodiments/tud_tundra_ur5e/README.md b/open_h/embodiments/tud_tundra_ur5e/README.md new file mode 100644 index 0000000..fa23673 --- /dev/null +++ b/open_h/embodiments/tud_tundra_ur5e/README.md @@ -0,0 +1,63 @@ +# TUD TUNDRA (UR5e) - Open-H Embodiment + +TUD Dresden teleoperation dataset for surgical assistance using a Universal Robots UR5e with stereo laparoscope. + +## Embodiment Configuration + +| Property | Value | +|----------|-------| +| **Embodiment Tag** | `TUD_TUNDRA_UR5E` | +| **Config File** | `open_h/embodiments/tud_tundra_ur5e/tud_tundra_ur5e_config.py` | + +## Data Summary + +| Subset | Episodes | FPS | Task | +|--------|----------|-----|------| +| `grasping_retraction` | 146 | 30 Hz | Grasping and tissue retraction during in-vivo porcine surgery | + +## Data Format + +### State + +- `joint_position` (6) -- embedded via mean/std +- `eef_pose` (7) -- base-frame XYZ + quaternion, pass-through reference for REL_XYZ_ROT6D + +### Action + +- `eef_pose` (7) -- absolute EEF pose from `observation.state` at future timesteps (t+1..t+H) +- `gripper` (1) -- from `open_gripper` (`action[4]`), binary + +Action representation: `REL_XYZ_ROT6D` for EEF pose, `ABSOLUTE` for gripper. +Action horizon: 50 (~1.7s at 30 Hz). `delta_indices = 1..H`. + +The dataset action column contains delta commands, but REL_XYZ_ROT6D uses absolute +EEF poses from `observation.state` with a +1 offset for action timesteps. + +### Cameras + +- `laparoscope_left` (960x540 @ 30 Hz) +- Note: `laparoscope_right` is available in the dataset but only the left view is used in the config (mono only). + +### Language + +- `task` -- mapped from `task_index` via `tasks.jsonl` + +## Modality Mapping + +- `modality_grasping_retraction.json` + +## Dataset Preparation + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag TUD_TUNDRA_UR5E \ + --modality-json open_h/embodiments/tud_tundra_ur5e/modality_grasping_retraction.json \ + /path/to/tud_dataset +``` + +## Notes + +- Endoscope guidance is not configured yet; only grasping_retraction is supported in this config. +- Endoscope guidance can be included with a small additional config that drops the gripper action, adding ~50 episodes to the TUD dataset. +- The `gripper_opened` state value is constant (0) in grasping_retraction; use the `open_gripper` action (`action[4]`) as the only valid gripper signal. +- The action pipeline uses `delta_indices = 1..H`. However, the gripper is sourced from `action` while the pose is sourced from `observation.state` where `action[t] = state[t+1]`, so there is some temporal misalignment for the gripper. diff --git a/open_h/embodiments/tud_tundra_ur5e/modality_endoscope_guidance.json b/open_h/embodiments/tud_tundra_ur5e/modality_endoscope_guidance.json new file mode 100644 index 0000000..9cb4410 --- /dev/null +++ b/open_h/embodiments/tud_tundra_ur5e/modality_endoscope_guidance.json @@ -0,0 +1,64 @@ +{ + "state": { + "relative_tip_position": { + "start": 0, + "end": 3, + "original_key": "observation.state" + }, + "joint_position": { + "start": 3, + "end": 9, + "original_key": "observation.state" + }, + "joint_velocity": { + "start": 9, + "end": 15, + "original_key": "observation.state" + }, + "joint_effort": { + "start": 15, + "end": 21, + "original_key": "observation.state" + }, + "eef_pose": { + "start": 21, + "end": 28, + "original_key": "observation.state" + }, + "initial_tip_position": { + "start": 28, + "end": 31, + "original_key": "observation.state" + }, + "rcm_position": { + "start": 31, + "end": 34, + "original_key": "observation.state" + }, + "roll_angle": { + "start": 34, + "end": 35, + "original_key": "observation.state" + } + }, + "action": { + "delta_tip_pose": { + "start": 0, + "end": 4, + "original_key": "action" + } + }, + "video": { + "laparoscope_left": { + "original_key": "observation.images.laparoscope_left" + }, + "laparoscope_right": { + "original_key": "observation.images.laparoscope_right" + } + }, + "annotation": { + "task": { + "original_key": "task_index" + } + } +} diff --git a/open_h/embodiments/tud_tundra_ur5e/modality_grasping_retraction.json b/open_h/embodiments/tud_tundra_ur5e/modality_grasping_retraction.json new file mode 100644 index 0000000..c2e7bac --- /dev/null +++ b/open_h/embodiments/tud_tundra_ur5e/modality_grasping_retraction.json @@ -0,0 +1,39 @@ +{ + "state": { + "joint_position": { + "start": 8, + "end": 14, + "original_key": "observation.state" + }, + "eef_pose": { + "start": 26, + "end": 33, + "original_key": "observation.state" + } + }, + "action": { + "eef_pose": { + "start": 26, + "end": 33, + "original_key": "observation.state" + }, + "gripper": { + "start": 4, + "end": 5, + "original_key": "action" + } + }, + "video": { + "laparoscope_left": { + "original_key": "observation.images.laparoscope_left" + }, + "laparoscope_right": { + "original_key": "observation.images.laparoscope_right" + } + }, + "annotation": { + "task": { + "original_key": "task_index" + } + } +} diff --git a/open_h/embodiments/tud_tundra_ur5e/tud_tundra_ur5e_config.py b/open_h/embodiments/tud_tundra_ur5e/tud_tundra_ur5e_config.py new file mode 100644 index 0000000..22af58c --- /dev/null +++ b/open_h/embodiments/tud_tundra_ur5e/tud_tundra_ur5e_config.py @@ -0,0 +1,86 @@ +""" +TUD TUNDRA modality configuration for GR00T N1.6 (grasping_retraction only). + +This configuration supports the TUD TUNDRA UR5e grasping/retraction dataset with: +- Stereo laparoscope video (left/right) +- Joint-position state embedding (EEF pose available as pass-through reference) +- Absolute EEF pose actions derived from observation.state (t+1..t+H) +- REL_XYZ_ROT6D conversion for EEF pose actions +- Gripper command actions from open_gripper (binary) +- Task strings from tasks.jsonl via task_index + +Representation note: +The dataset action column contains delta commands (dx, dy, dz, droll). For +REL_XYZ_ROT6D we instead use the absolute base-frame EEF pose from observation.state +at future timesteps as the action source, following the CMR-style +1 offset. +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +# 30 Hz -> 50 frames ~= 1.7 seconds +ACTION_HORIZON = 50 + +tud_tundra_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "laparoscope_left", + # "laparoscope_right", # Mono Only + ], + ), + "state": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "joint_position", + "eef_pose", + ], + mean_std_embedding_keys=[ + "joint_position", + ], + pass_through_keys=[ + "eef_pose", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(1, ACTION_HORIZON + 1)), + modality_keys=[ + "eef_pose", + "gripper", + ], + action_configs=[ + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="eef_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + reference_quat_order="xyzw", + ), + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["task"], + ), +} + +# Register with TUD_TUNDRA_UR5E tag for TUNDRA UR5e surgical assistance data +register_modality_config(tud_tundra_config, embodiment_tag=EmbodimentTag.TUD_TUNDRA_UR5E) diff --git a/open_h/embodiments/tum_sonata_franka/README.md b/open_h/embodiments/tum_sonata_franka/README.md new file mode 100644 index 0000000..b98aa55 --- /dev/null +++ b/open_h/embodiments/tum_sonata_franka/README.md @@ -0,0 +1,117 @@ +# TUM SonATA Ultrasound Dataset + +## Overview + +SonATA is a robotic sonography dataset from TUM's Computer Aided Medical Procedures (CAMP) Lab. It integrates synchronized ultrasound imaging, external visual data, contact force measurements, robot motion data, and textual instructions collected from abdomen, thyroid, and arm phantoms. + +## Embodiment Configuration + +| Property | Value | +|----------|-------| +| **Embodiment Tag** | `TUM_SONATA_FRANKA` | +| **Config File** | `open_h/embodiments/tum_sonata_franka/tum_sonata_franka_config.py` | +| **Modality Mapping** | `open_h/embodiments/tum_sonata_franka/modality.json` | + +## Subsets + +| Subset | Episodes | Frames | Tasks | Status | +|--------|----------|--------|-------|--------| +| SonATA_abdomen | 1,533 | 325,634 | 287 | Ready | +| SonATA_arm | ~1,107 | - | - | Ready | +| SonATA_thyroid | ~780 | - | - | Ready | + +## Data Format + +Already in **LeRobot v2.1** format with proper metadata. + +### Raw Dataset Features + +| Feature | Shape | Type | Description | +|---------|-------|------|-------------| +| `action` | 6D | float32 | End-effector pose absolute (x, y, z, roll, pitch, yaw) | +| `observation.state` | 7D | float32 | Joint angles (joint_1 through joint_7) | +| `observation.images.tpv_camera` | 480x640x3 | video | Third-person view camera (AV1 codec) | +| `observation.images.wrist_camera` | 480x640x3 | video | Wrist-mounted camera (AV1 codec) | +| `observation.images.ultrasound` | 480x640x3 | video | Ultrasound image (AV1 codec) | +| `observation.meta.force_torque` | 6D | float32 | Force/torque (fx, fy, fz, tx, ty, tz) | +| `observation.meta.probe_type` | 1 | string | Probe model (linear or convex) | +| `observation.meta.probe_acquisition_param` | 6D | float32 | Ultrasound acquisition parameters | +| `observation.meta.probe_cali_mtx` | 7D | float32 | Probe calibration (xyz + quaternion) | +| `instruction.text` | 1 | string | Natural language instruction | +| `instruction.task` | 1 | string | Task category | +| `instruction.sub_task` | 1 | string | Sub-task description | + +### State + +- `joint_angles` (7D) -- joint angles from `observation.state[0:7]`, embedded via mean/std +- `force_torque` (6D) -- force/torque from `observation.meta.force_torque[0:6]`, embedded via mean/std +- `eef_pose` (6D) -- EEF pose from `action[0:6]`, pass-through reference for REL_XYZ_ROT6D (not embedded) + +### Action + +- `eef_pose` (6D) -- absolute EEF pose (xyz + roll/pitch/yaw Euler angles) + +Action representation: `REL_XYZ_ROT6D` (Euler angle input -> relative xyz + rot6d output = 9D per timestep). +Action horizon: 50 (~1.67s at 30 Hz). `delta_indices = 1..H`. + +Position: relative to reference EEF position (delta xyz). +Rotation: Euler angles (RPY) converted to rot6d relative to reference. + +### Cameras + +- `tpv_camera` -- third-person view (scene context) +- `wrist_camera` -- wrist-mounted (probe positioning) +- `ultrasound` -- ultrasound image + +### Language + +- `task` -- mapped from `task_index` via `tasks.jsonl` + +### Robot Configuration + +| Property | Value | +|----------|-------| +| Robot | Franka Panda | +| FPS | 30 Hz | +| Video Codec | AV1 | +| Resolution | 480x640 | + +### Data Splits (SonATA_abdomen) + +| Split | Episodes | +|-------|----------| +| Train | 0-1071 | +| Val | 1071-1225 | +| Test | 1225-1533 | + +## Task Types + +1. **Placement** - Position probe on target anatomy +2. **Scanning** - Sweep probe across anatomical region +3. **Navigation** - Move probe to specific landmarks + +## Unique Characteristics + +- **Multimodal**: Ultrasound + RGB cameras + force/torque +- **Language-conditioned**: Rich natural language instructions with multiple phrasings per task +- **Calibration data**: Camera and probe calibration matrices included +- **Ultrasound metadata**: Probe parameters (frequency, depth, FOV) per episode + +## Example Instructions + +``` +"Place the probe on the middle of the abdomen." +"Perform a transverse sweep across the aorta and IVC." +"Glide across the abdomen transversely covering the aorta and IVC." +``` + +## Dataset Preparation + +Use the shared preparation script to copy the modality JSON into each dataset's `meta/` folder and generate normalization statistics: + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag TUM_SONATA_FRANKA \ + --modality-json open_h/embodiments/tum_sonata_franka/modality.json \ + /path/to/dataset +``` diff --git a/open_h/embodiments/tum_sonata_franka/modality.json b/open_h/embodiments/tum_sonata_franka/modality.json new file mode 100644 index 0000000..694a027 --- /dev/null +++ b/open_h/embodiments/tum_sonata_franka/modality.json @@ -0,0 +1,15 @@ +{ + "video": { + "tpv_camera": {"original_key": "observation.images.tpv_camera"}, + "wrist_camera": {"original_key": "observation.images.wrist_camera"}, + "ultrasound": {"original_key": "observation.images.ultrasound"} + }, + "state": { + "joint_angles": {"start": 0, "end": 7}, + "force_torque": {"start": 0, "end": 6, "original_key": "observation.meta.force_torque"}, + "eef_pose": {"start": 0, "end": 6, "original_key": "action"} + }, + "action": { + "eef_pose": {"start": 0, "end": 6} + } +} diff --git a/open_h/embodiments/tum_sonata_franka/tum_sonata_franka_config.py b/open_h/embodiments/tum_sonata_franka/tum_sonata_franka_config.py new file mode 100644 index 0000000..d332913 --- /dev/null +++ b/open_h/embodiments/tum_sonata_franka/tum_sonata_franka_config.py @@ -0,0 +1,96 @@ +""" +TUM SonATA Ultrasound modality configuration for GR00T N1.6. + +This configuration supports the TUM SonATA robotic ultrasound sonography dataset with: +- Franka Panda robot with ultrasound probe end-effector +- REL_XYZ_ROT6D action representation with Euler angle input (RPY -> rot6d) +- 3 camera views (TPV, wrist-mounted, ultrasound) +- Joint angle state + force/torque sensor data + +Data Format: +- State: 7D joint angles + 6D force/torque = 13D embedded + + 6D EEF pose (pass-through for REL_XYZ_ROT6D reference) +- Action: 6D absolute EEF pose (xyz + roll/pitch/yaw) -> 9D REL_XYZ_ROT6D + +REL_XYZ_ROT6D Conversion: +- Position: relative to reference EEF position (delta xyz) +- Rotation: Euler angles (RPY) converted to rot6d relative to reference +- Output: 9D per timestep (xyz_rel + rot6d_rel) + +Camera Views (in model input order): +1. tpv_camera: Third-person view - scene context +2. wrist_camera: Wrist-mounted view - probe positioning +3. ultrasound: Primary imaging modality - task-relevant feedback + +Dataset: TUM SonATA (Computer Aided Medical Procedures Lab) +- SonATA_abdomen: 1,533 episodes, 325k frames +- SonATA_arm: 369 episodes +- SonATA_thyroid: 260 episodes +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +# 30 FPS -> 50 frames = 1.67 seconds of prediction +ACTION_HORIZON = 50 + +tum_sonata_config = { + "video": ModalityConfig( + delta_indices=[0], + # Camera order: external context first, then progressively more task-specific + modality_keys=["tpv_camera", "wrist_camera", "ultrasound"], + ), + "state": ModalityConfig( + delta_indices=[0], # Single reference state for REL_XYZ_ROT6D + modality_keys=[ + # Embedded state keys (sent to model) + "joint_angles", # 7D joint angles - mean_std normalized + "force_torque", # 6D force/torque sensor - mean_std normalized + # Pass-through keys (loaded but not embedded) + "eef_pose", # 6D EEF pose from action column (for REL_XYZ_ROT6D reference) + ], + # Mean-std normalization for continuous values + mean_std_embedding_keys=[ + "joint_angles", + "force_torque", + ], + # Pass-through keys: loaded for REL_XYZ_ROT6D, never embedded to model + # eef_pose is extracted from action column (via modality.json original_key) + pass_through_keys=[ + "eef_pose", + ], + ), + "action": ModalityConfig( + # Start at 1 (index 0 is state reference), go to ACTION_HORIZON + delta_indices=list(range(1, ACTION_HORIZON + 1)), + modality_keys=["eef_pose"], # 6D: xyz + roll/pitch/yaw (Euler) + action_configs=[ + # EEF pose: REL_XYZ_ROT6D action with Euler input + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, # Output: xyz_rel + rot6d_rel = 9D + state_key="eef_pose", # Reference from action at t=0 (via pass_through) + normalization_type="temporal_meanstd", + input_rotation_format="euler", # Euler angles (roll, pitch, yaw) + reference_rotation_format="euler", # Reference is also Euler + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["task"], # Maps to tasks.jsonl (287 unique task descriptions) + ), +} + +# Register with TUM_SONATA_FRANKA tag for ultrasound robotic sonography +register_modality_config(tum_sonata_config, embodiment_tag=EmbodimentTag.TUM_SONATA_FRANKA) diff --git a/open_h/embodiments/turin_mitic_ex_vivo/README.md b/open_h/embodiments/turin_mitic_ex_vivo/README.md new file mode 100644 index 0000000..ae45653 --- /dev/null +++ b/open_h/embodiments/turin_mitic_ex_vivo/README.md @@ -0,0 +1,67 @@ +# Turin MITIC Ex Vivo Embodiment + +Dual-arm dVRK (PSM1/PSM2) ex vivo surgical demonstrations covering knot tying, needle manipulation, and tissue lift tasks, recorded with stereo endoscope views. + +## Embodiment Configuration + +| Property | Value | +|----------|-------| +| **Embodiment Tag** | `TURIN_MITIC_EX_VIVO` | +| **Config File** | `open_h/embodiments/turin_mitic_ex_vivo/turin_mitic_ex_vivo_config.py` | +| **Modality Mapping** | `open_h/embodiments/turin_mitic_ex_vivo/modality.json` | + +## Data Summary + +| Property | Value | +|----------|-------| +| Source | Turin (MITIC) ex vivo surgical demonstrations | +| Robot | dVRK (dual-arm PSM1/PSM2) | +| Tasks | Knot tying, needle manipulation, tissue lift | +| Episodes | 799 | +| Frames | 388,690 | +| FPS | 30 | +| Video | 1080x1920, AV1 | + +## Data Format + +### State + +- `psm1_joints` (6D) -- PSM1 joint angles from `observation.state[0:6]`, embedded via mean/std +- `psm2_joints` (6D) -- PSM2 joint angles from `observation.state[6:12]`, embedded via mean/std +- `psm1_pose` (7D) -- PSM1 EEF pose (xyz + quat) from `action[0:7]`, pass-through reference for REL_XYZ_ROT6D (not embedded) +- `psm2_pose` (7D) -- PSM2 EEF pose (xyz + quat) from `action[7:14]`, pass-through reference for REL_XYZ_ROT6D (not embedded) + +### Action + +- `psm1_pose` (7D) -- PSM1 absolute EEF pose (xyz + quaternion) +- `psm2_pose` (7D) -- PSM2 absolute EEF pose (xyz + quaternion) + +Action representation: `REL_XYZ_ROT6D` for both arms (quaternion input -> relative xyz + rot6d output = 9D per arm per timestep). +Action horizon: 50 (~1.67s at 30 Hz). `delta_indices = 1..H`. + +### Cameras + +- `endoscope_left` +- Note: `endoscope_right` is available in the dataset but only the left view is used in the config (mono only). + +### Language + +- `annotation.instruction` (mapped from `instruction.text`) + +## Action Horizon + +- 50 steps (30 Hz, ~1.67 s of future prediction) + +## Dataset Preparation + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag TURIN_MITIC_EX_VIVO \ + --modality-json open_h/embodiments/turin_mitic_ex_vivo/modality.json \ + /path/to/turin_dataset +``` + +## Notes +- No EEF pose is present in `observation.state`. The action at `t=0` is used as the REL_XYZ_ROT6D reference pose via pass-through keys. +- Dataset synchronization uses the right endoscope timestamp; `action` comes from measured Cartesian pose topics aligned to that frame, so `action[t]` is expected to be t-aligned (not t+1) with `observation.state[t]` and images. +- Only endoscope views are available; no wrist cameras. diff --git a/open_h/embodiments/turin_mitic_ex_vivo/modality.json b/open_h/embodiments/turin_mitic_ex_vivo/modality.json new file mode 100644 index 0000000..57c4beb --- /dev/null +++ b/open_h/embodiments/turin_mitic_ex_vivo/modality.json @@ -0,0 +1,46 @@ +{ + "video": { + "endoscope_left": { + "original_key": "observation.images.endoscope.left" + }, + "endoscope_right": { + "original_key": "observation.images.endoscope.right" + } + }, + "state": { + "psm1_joints": { + "start": 0, + "end": 6 + }, + "psm2_joints": { + "start": 6, + "end": 12 + }, + "psm1_pose": { + "start": 0, + "end": 7, + "original_key": "action" + }, + "psm2_pose": { + "start": 7, + "end": 14, + "original_key": "action" + } + }, + "action": { + "psm1_pose": { + "start": 0, + "end": 7 + }, + "psm2_pose": { + "start": 7, + "end": 14 + } + }, + "annotation": { + "instruction": { + "original_key": "instruction.text", + "is_text": true + } + } +} diff --git a/open_h/embodiments/turin_mitic_ex_vivo/turin_mitic_ex_vivo_config.py b/open_h/embodiments/turin_mitic_ex_vivo/turin_mitic_ex_vivo_config.py new file mode 100644 index 0000000..c02bef8 --- /dev/null +++ b/open_h/embodiments/turin_mitic_ex_vivo/turin_mitic_ex_vivo_config.py @@ -0,0 +1,94 @@ +""" +Turin MITIC ex vivo modality configuration for GR00T N1.6. + +This configuration supports the Turin MITIC ex vivo dataset with: +- Dual-arm dVRK (PSM1/PSM2) +- REL_XYZ_ROT6D EEF action representation (xyz + quaternion input) +- Joint-angle state embeddings (12D) with pass-through EEF pose references +- Stereo endoscope video (left/right) + +Data format: +- State: 12D joint angles (6 per arm) +- Action: 14D absolute EEF pose (xyz + quat) for PSM1 + PSM2 + +REL_XYZ_ROT6D conversion: +- Reference poses are taken from action at t=0 (pass-through in state) +- Actions are predicted for t+1..t+50 at 30 Hz (~1.67s horizon) +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +ACTION_HORIZON = 50 + +turin_mitic_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "endoscope_left", + # "endoscope_right", # Mono Only + ], + ), + "state": ModalityConfig( + delta_indices=[0], + modality_keys=[ + # Embedded state keys (sent to model) + "psm1_joints", + "psm2_joints", + # Pass-through keys (loaded but not embedded) + "psm1_pose", + "psm2_pose", + ], + mean_std_embedding_keys=[ + "psm1_joints", + "psm2_joints", + ], + pass_through_keys=[ + "psm1_pose", + "psm2_pose", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(1, ACTION_HORIZON + 1)), + modality_keys=[ + "psm1_pose", + "psm2_pose", + ], + action_configs=[ + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm1_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + ), + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm2_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["annotation.instruction"], + ), +} + +# Register with the Turin MITIC ex vivo embodiment tag. +register_modality_config(turin_mitic_config, embodiment_tag=EmbodimentTag.TURIN_MITIC_EX_VIVO) diff --git a/open_h/embodiments/ucb_dvrk/README.md b/open_h/embodiments/ucb_dvrk/README.md new file mode 100644 index 0000000..6818f57 --- /dev/null +++ b/open_h/embodiments/ucb_dvrk/README.md @@ -0,0 +1,78 @@ +# UCBerkeley dVRK + +Debridement dataset collected from the da Vinci Research Kit (dVRK) with stereo endoscope views and cartesian EEF pose actions (REL_XYZ_ROT6D). + +## Dataset Statistics + +| Property | Value | +|----------|-------| +| **Episodes** | 589 | +| **Total Frames** | 221,950 | +| **Avg Frames/Episode** | ~377 | + +## Embodiment Configuration + +| Property | Value | +|----------|-------| +| **Embodiment Tag** | `ucb_dvrk` | +| **Projector Index** | 5 | +| **Config File** | `open_h/embodiments/ucb_dvrk/ucb_dvrk_config.py` | +| **Action Horizon** | 50 (delta_indices `[0..49]`) | + +## Data Format + +### State (Cartesian + Joints) +Cartesian EEF state for both PSM arms plus joint angles: + +| Arm | Indices | Dimensions | Description | +|-----|---------|------------|-------------| +| PSM1 pose | 0-6 | 7D | `[x, y, z, qx, qy, qz, qw]` | +| PSM1 gripper | 7 | 1D | `jaw` | +| PSM2 pose | 8-14 | 7D | Same format as PSM1 pose | +| PSM2 gripper | 15 | 1D | `jaw` | + +Joint-angle state (sourced from `observation.state`): +| Arm | Indices | Dimensions | Description | +|-----|---------|------------|-------------| +| PSM1 joints | 0-6 | 7D | `[outer_yaw, outer_pitch, outer_insertion, outer_roll, outer_wrist_pitch, outer_wrist_yaw, jaw]` | +| PSM2 joints | 7-13 | 7D | Same format as PSM1 | + +**Pass-through keys:** `psm1_pose`, `psm2_pose` (used as REL_XYZ_ROT6D reference frames, not tokenized as state). +**Model state keys (tokenized):** `psm1_joints`, `psm1_gripper`, `psm2_joints`, `psm2_gripper`. + +### Action (16D) +Same format as state, representing cartesian setpoints. Converted to **REL_XYZ_ROT6D** during training: +- `action[t] == state[t+1]` (setpoint for the next timestep) +- Relative conversion uses state[t] as the reference frame for the action horizon + +### Video (1 view used for training) +| Camera | Original Key | Note | +|--------|--------------|------| +| `camera_left` | `observation.images.left` | Active in config | +| `camera_right` | `observation.images.right` | Available in modality.json but commented out in config (mono only) | + +### Language +| Key | +|-----| +| `task` | + +## Dataset Preparation + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag UCB_DVRK \ + --modality-json open_h/embodiments/ucb_dvrk/modality.json \ + /path/to/ucb_dvrk_dataset +``` + +This generates `stats.json` and `temporal_stats.json` inside the dataset meta folder. + + +## File Structure + +``` +open_h/embodiments/ucb_dvrk/ +β”œβ”€β”€ README.md # This file +β”œβ”€β”€ ucb_dvrk_config.py # Built-in config module (auto-registered) +└── modality.json # Data key mappings (copied to dataset/meta/ by script) +``` diff --git a/open_h/embodiments/ucb_dvrk/modality.json b/open_h/embodiments/ucb_dvrk/modality.json new file mode 100644 index 0000000..cf24ba6 --- /dev/null +++ b/open_h/embodiments/ucb_dvrk/modality.json @@ -0,0 +1,20 @@ +{ + "video": { + "camera_left": {"original_key": "observation.images.left"}, + "camera_right": {"original_key": "observation.images.right"} + }, + "state": { + "psm1_pose": {"start": 0, "end": 7, "original_key": "observation.cartesian_state"}, + "psm1_gripper": {"start": 7, "end": 8, "original_key": "observation.cartesian_state"}, + "psm2_pose": {"start": 8, "end": 15, "original_key": "observation.cartesian_state"}, + "psm2_gripper": {"start": 15, "end": 16, "original_key": "observation.cartesian_state"}, + "psm1_joints": {"start": 0, "end": 7, "original_key": "observation.state"}, + "psm2_joints": {"start": 7, "end": 14, "original_key": "observation.state"} + }, + "action": { + "psm1_pose": {"start": 0, "end": 7, "original_key": "action.cartesian_state"}, + "psm1_gripper": {"start": 7, "end": 8, "original_key": "action.cartesian_state"}, + "psm2_pose": {"start": 8, "end": 15, "original_key": "action.cartesian_state"}, + "psm2_gripper": {"start": 15, "end": 16, "original_key": "action.cartesian_state"} + } +} diff --git a/open_h/embodiments/ucb_dvrk/ucb_dvrk_config.py b/open_h/embodiments/ucb_dvrk/ucb_dvrk_config.py new file mode 100644 index 0000000..adb802d --- /dev/null +++ b/open_h/embodiments/ucb_dvrk/ucb_dvrk_config.py @@ -0,0 +1,132 @@ +""" +dVRK UCBerkeley modality configuration for GR00T N1.6. + +This configuration supports the UCBerkeley debridement dataset with: +- Cartesian EEF pose actions (16D) using REL_XYZ_ROT6D +- Absolute gripper actions per arm +- Joint-angle state channels for normalization and dropout control +- 2 camera views (left, right stereo) instead of 4 + +Data Format: +- State: cartesian state (xyz + quat + jaw) plus joint angles, split by arm: + - psm1_pose (7D): [x, y, z, qx, qy, qz, qw] + - psm1_gripper (1D): jaw + - psm2_pose (7D): [x, y, z, qx, qy, qz, qw] + - psm2_gripper (1D): jaw +- Joint angles: + - psm1_joints (7D) + - psm2_joints (7D) +- Action: 16D cartesian setpoints (same format as state) + +REL_XYZ_ROT6D Conversion: +- Pose actions are converted to relative translation/rotation from the current state. +- action[t] corresponds to state[t+1] (setpoint for the next timestep), so relative + conversion uses state[t] as the reference frame for the action horizon. + +Dataset: UCBerkeley debridement (589 episodes, 221,950 frames) +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +# Matches existing dVRK config for consistency +ACTION_HORIZON = 50 + +dvrk_ucb_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "camera_left", + # "camera_right", # Mono Only + ], + ), + "state": ModalityConfig( + delta_indices=[0], # Single reference state at timestep t + modality_keys=[ + "psm1_joints", + "psm1_gripper", + "psm2_joints", + "psm2_gripper", + "psm1_pose", + "psm2_pose", + ], + mean_std_embedding_keys=[ + "psm1_joints", + "psm1_gripper", + "psm2_joints", + "psm2_gripper", + ], + # Pass-through cartesian pose/gripper state for REL_XYZ_ROT6D reference + pass_through_keys=[ + "psm1_pose", + "psm2_pose", + ], + ), + "action": ModalityConfig( + delta_indices=list(range(ACTION_HORIZON)), # [0, 1, 2, ..., 49] consecutive + modality_keys=[ + "psm1_pose", + "psm1_gripper", + "psm2_pose", + "psm2_gripper", + ], + action_configs=[ + # PSM1 pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm1_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="xyzw", + reference_rotation_format="quat", + reference_quat_order="xyzw", + ), + # PSM1 gripper: absolute jaw angle + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + # PSM2 pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm2_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="xyzw", + reference_rotation_format="quat", + reference_quat_order="xyzw", + ), + # PSM2 gripper: absolute jaw angle + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + state_key=None, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["task"], # Uses tasks.jsonl: "surgical_debridement" + ), +} + +# Register with UCB_DVRK tag for UCBerkeley cartesian EEF surgical robot data +register_modality_config(dvrk_ucb_config, embodiment_tag=EmbodimentTag.UCB_DVRK) diff --git a/open_h/embodiments/ucsd_dvrk/README.md b/open_h/embodiments/ucsd_dvrk/README.md new file mode 100644 index 0000000..490dcee --- /dev/null +++ b/open_h/embodiments/ucsd_dvrk/README.md @@ -0,0 +1,123 @@ +# UCSD Surgical Learning Dataset (dVRK) + +## Overview + +UCSD surgical learning data in LeRobot v2.1 format with stereo endoscope video, +dVRK kinematics, and two-arm action streams. The datasets include absolute +end-effector (EEF) pose in state and delta EEF actions for retraction and +cutting tools. + +## Subsets + +| Subset | Episodes | Frames | FPS | Tasks | Min Episode Length | +|--------|----------|--------|-----|-------|--------------------| +| surgical_learning_dataset | 912 | 288,604 | 30 | 2 | 73 | +| surgical_learning_dataset2 | 200 | 26,313 | 30 | 1 | 66 | + +All episodes are at least 50 frames, so a 50-step action horizon is valid. + +## Task List + +- `surgical_learning_dataset`: retraction, dissection +- `surgical_learning_dataset2`: Retraction + +## Cameras + +Stereo endoscope views: + +| Camera | Original Key | Resolution | FPS | Codec | Note | +|--------|--------------|------------|-----|-------|------| +| Left | `observation.images.left` | 480x640 | 30 | AV1 | Active in config (`camera_left`) | +| Right | `observation.images.right` | 480x640 | 30 | AV1 | Available in modality JSON but commented out in config (mono only) | + +## Kinematics + +### Actions (both datasets) + +16D delta EEF pose + gripper for each arm: + +- Retraction arm: `dPSM_RETRACTION_x,y,z,qw,qx,qy,qz,gripper` +- Cutter arm: `dPSM_CUTTER_x,y,z,qw,qx,qy,qz,gripper` +- Quaternion order is **wxyz** (scalar-first). + +### State + +`surgical_learning_dataset` (62D): +- Joint positions, velocities, and efforts for retraction + cutter PSMs +- Gripper position and effort for each tool +- Absolute EEF pose for both arms: `*_ee_x,y,z,qw,qx,qy,qz` +- Target points: `target_0..3_x,y` + +`surgical_learning_dataset2` (28D): +- Joint positions + gripper positions for both PSMs +- Absolute EEF pose for both arms: `*_ee_x,y,z,qw,qx,qy,qz` + +## Notes + +- Absolute EEF state is available for both arms, enabling REL_XYZ_ROT6D action + conversion using the state at timestep t as reference. +- Quaternion order is **wxyz** (qw, qx, qy, qz) for both state and action. +- `task` entries are available in `tasks.jsonl` and are used for language. + +## Embodiment Configuration + +| Property | Value | +|----------|-------| +| Embodiment Tag | `ucsd_dvrk` | +| Config File | `open_h/embodiments/ucsd_dvrk/ucsd_dvrk_config.py` | +| Action Horizon | 50 (delta_indices `[1..50]`, 30 Hz) | + +## Modality Mapping + +The modality JSON filenames intentionally mirror the dataset directory names: +`surgical_learning_dataset` uses +`modality_surgical_learning_dataset.json`, and +`surgical_learning_dataset2` uses +`modality_surgical_learning_dataset2.json`. + +This naming is important because the two UCSD datasets are not schema-identical. +They represent the same dVRK embodiment, but their recorded +`observation.state` layouts differ slightly, so each dataset needs its own +modality mapping to align the raw dataset fields with the common training keys. + +| Dataset | Modality File | +|---------|---------------| +| `surgical_learning_dataset` | `open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset.json` | +| `surgical_learning_dataset2` | `open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset2.json` | + +### State / Action Keys + +Both modality files expose the same high-level keys: + +**State keys:** +- `psm_retraction_pose` (7D: xyz + qw,qx,qy,qz) +- `psm_retraction_gripper` (1D) +- `psm_cutter_pose` (7D: xyz + qw,qx,qy,qz) +- `psm_cutter_gripper` (1D) + +**Action keys (same names, REL_XYZ_ROT6D for pose, ABSOLUTE for gripper):** +- `psm_retraction_pose` (7D) +- `psm_retraction_gripper` (1D) +- `psm_cutter_pose` (7D) +- `psm_cutter_gripper` (1D) + +**Language key:** `task` + +## Utilities + +Run `prepare_datasets.sh` separately for each sub-dataset with the matching +modality file (the script copies the modality JSON into `meta/` automatically): + +```bash +# surgical_learning_dataset +bash open_h/prepare_datasets.sh \ + --embodiment-tag UCSD_DVRK \ + --modality-json open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset.json \ + /path/to/UCSD/surgical_learning_dataset + +# surgical_learning_dataset2 +bash open_h/prepare_datasets.sh \ + --embodiment-tag UCSD_DVRK \ + --modality-json open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset2.json \ + /path/to/UCSD/surgical_learning_dataset2 +``` diff --git a/open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset.json b/open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset.json new file mode 100644 index 0000000..bd10da5 --- /dev/null +++ b/open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset.json @@ -0,0 +1,18 @@ +{ + "video": { + "camera_left": { "original_key": "observation.images.left" }, + "camera_right": { "original_key": "observation.images.right" } + }, + "state": { + "psm_retraction_pose": { "start": 40, "end": 47 }, + "psm_retraction_gripper": { "start": 18, "end": 19 }, + "psm_cutter_pose": { "start": 47, "end": 54 }, + "psm_cutter_gripper": { "start": 38, "end": 39 } + }, + "action": { + "psm_retraction_pose": { "start": 0, "end": 7 }, + "psm_retraction_gripper": { "start": 7, "end": 8 }, + "psm_cutter_pose": { "start": 8, "end": 15 }, + "psm_cutter_gripper": { "start": 15, "end": 16 } + } +} diff --git a/open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset2.json b/open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset2.json new file mode 100644 index 0000000..0641cba --- /dev/null +++ b/open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset2.json @@ -0,0 +1,18 @@ +{ + "video": { + "camera_left": { "original_key": "observation.images.left" }, + "camera_right": { "original_key": "observation.images.right" } + }, + "state": { + "psm_retraction_pose": { "start": 14, "end": 21 }, + "psm_retraction_gripper": { "start": 6, "end": 7 }, + "psm_cutter_pose": { "start": 21, "end": 28 }, + "psm_cutter_gripper": { "start": 13, "end": 14 } + }, + "action": { + "psm_retraction_pose": { "start": 0, "end": 7 }, + "psm_retraction_gripper": { "start": 7, "end": 8 }, + "psm_cutter_pose": { "start": 8, "end": 15 }, + "psm_cutter_gripper": { "start": 15, "end": 16 } + } +} diff --git a/open_h/embodiments/ucsd_dvrk/ucsd_dvrk_config.py b/open_h/embodiments/ucsd_dvrk/ucsd_dvrk_config.py new file mode 100644 index 0000000..3595d07 --- /dev/null +++ b/open_h/embodiments/ucsd_dvrk/ucsd_dvrk_config.py @@ -0,0 +1,115 @@ +""" +UCSD Surgical Learning modality configuration for GR00T N1.6. + +This config targets the UCSD surgical learning datasets: +- surgical_learning_dataset (912 episodes, 30 Hz) +- surgical_learning_dataset2 (200 episodes, 30 Hz) + +Key design choices: +- Use absolute EEF pose + gripper in state +- Use REL_XYZ_ROT6D for EEF pose actions and ABSOLUTE for gripper +- Use wxyz quaternion ordering (qw, qx, qy, qz) +- Use `task` language key from tasks.jsonl + +Note: The two UCSD datasets have different observation.state layouts. Use the +appropriate modality.json for each dataset: +- open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset.json +- open_h/embodiments/ucsd_dvrk/modality_surgical_learning_dataset2.json +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +# 30 FPS -> 50 frames = 1.67 seconds of prediction +ACTION_HORIZON = 50 + +ucsd_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "camera_left", + # "camera_right", # Mono Only + ], + ), + "state": ModalityConfig( + delta_indices=[0], # Single reference state for REL_XYZ_ROT6D + modality_keys=[ + "psm_retraction_pose", + "psm_retraction_gripper", + "psm_cutter_pose", + "psm_cutter_gripper", + ], + mean_std_embedding_keys=[ + "psm_retraction_pose", + "psm_retraction_gripper", + "psm_cutter_pose", + "psm_cutter_gripper", + ], + ), + "action": ModalityConfig( + # Start at 1 (index 0 is state reference), go to ACTION_HORIZON + delta_indices=list(range(1, ACTION_HORIZON + 1)), + modality_keys=[ + "psm_retraction_pose", + "psm_retraction_gripper", + "psm_cutter_pose", + "psm_cutter_gripper", + ], + action_configs=[ + # Retraction pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm_retraction_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="wxyz", + reference_rotation_format="quat", + reference_quat_order="wxyz", + ), + # Retraction gripper: absolute jaw value + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + # Cutter pose: REL_XYZ_ROT6D EEF action + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="psm_cutter_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + input_quat_order="wxyz", + reference_rotation_format="quat", + reference_quat_order="wxyz", + ), + # Cutter gripper: absolute jaw value + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["task"], + ), +} + +# Register with UCSD_DVRK tag for surgical robot finetuning +register_modality_config(ucsd_config, embodiment_tag=EmbodimentTag.UCSD_DVRK) diff --git a/open_h/embodiments/ustc_torin_tuodao/README.md b/open_h/embodiments/ustc_torin_tuodao/README.md new file mode 100644 index 0000000..badeb55 --- /dev/null +++ b/open_h/embodiments/ustc_torin_tuodao/README.md @@ -0,0 +1,119 @@ +# USTC Surgical Dataset (Torin) + +## Overview + +USTC provides surgical demonstrations recorded on the Torin platform with stereo endoscope video. The merged dataset contains seven task families: `exvivo_liver_sep`, `grasp_on_liver`, `invivo_liver_sep`, `knot_tying`, `Needle_handover`, `needle_pickup`, and `tissue_lifting`. All subsets are in **LeRobot v2.1** format. + +### Note + +This integration uses `action.cartesian_absolute` (xyz + quaternion + gripper per arm) +for pose targets and REL_XYZ_ROT6D for action modeling. Gripper is modeled as an +absolute target. `action.cartesian_absolute` is written alongside `absolute_action` +during conversion, and invalid rotations are repaired by carrying forward the last +valid rotation (identity if none). + +## Embodiment Configuration + +| Property | Value | +|---------|-------| +| Embodiment Tag | `ustc_torin_tuodao` | +| Config File | `open_h/embodiments/ustc_torin_tuodao/ustc_torin_tuodao_config.py` | +| Modality Mapping | `open_h/embodiments/ustc_torin_tuodao/modality.json` | +| Action Horizon | 50 (delta_indices `[0..49]`, 24 Hz) | + +## Tasks and Subsets + +### Subset Statistics (merged) + +| Subset | Episodes | Frames | FPS | Hours | +|--------|----------|--------|-----|-------| +| exvivo_liver_sep | 666 | 121,922 | 24 | 1.41 | +| grasp_on_liver | 817 | 63,538 | 24 | 0.74 | +| invivo_liver_sep | 199 | 39,899 | 24 | 0.46 | +| knot_tying | 1,098 | 182,836 | 24 | 2.12 | +| Needle_handover | 260 | 34,990 | 24 | 0.41 | +| needle_pickup | 616 | 57,172 | 24 | 0.66 | +| tissue_lifting | 110 | 11,673 | 24 | 0.14 | + +### Task Totals + +| Task | Episodes | Frames | FPS | Hours | +|------|----------|--------|-----|-------| +| **Overall** | **3,766** | **512,030** | **24** | **5.94** | + +## Data Format + +### Cameras + +| Camera Key | Modality Key | Resolution | Codec | FPS | Note | +|-----------|-------------|------------|-------|-----|------| +| `observation.images.endoscope.left` | `endoscope_left` | 1080x1920 | H.264 | 24 | Active in config | +| `observation.images.endoscope.right` | `endoscope_right` | 1080x1920 | H.264 | 24 | Available in modality.json but commented out in config (mono only) | + +### Kinematics + +**State (`observation.state`, 14D):** Joint angles for both arms. + +- Left arm: `left_joint_1` ... `left_joint_7` +- Right arm: `right_joint_8` ... `right_joint_14` + +**Action (`action`, 14D):** Cartesian **delta** commands per arm using Euler angles (plus gripper). + +- Left arm: `left_x`, `left_y`, `left_z`, `left_roll`, `left_pitch`, `left_yaw`, `left_gripper` +- Right arm: `right_x`, `right_y`, `right_z`, `right_roll`, `right_pitch`, `right_yaw`, `right_gripper` + +**Absolute Pose (`action.cartesian_absolute`, 16D):** Cartesian absolute pose per arm +(xyz + quaternion + gripper). + +- Left arm: `left_x`, `left_y`, `left_z`, `left_qx`, `left_qy`, `left_qz`, `left_qw`, `left_gripper` +- Right arm: `right_x`, `right_y`, `right_z`, `right_qx`, `right_qy`, `right_qz`, `right_qw`, `right_gripper` + +**Current Pose (`observation.current_target_psm`, 16D):** Cartesian absolute pose per arm +(xyz + quaternion + gripper) at time t. + +**Energy (`energy`, 1D):** Not present in the current USTC parquets; a future update +should append energy into the action vector once regenerated. + +### Other Metadata + +- `instruction.text` (string) +- `observation.meta.tool` (string, `left_right`) +- `timestamp`, `frame_index`, `episode_index`, `task_index` + +## Modality Keys (config.py) + +**State keys (tokenized):** `left_joints` (7D), `right_joints` (7D) +**Pass-through keys:** `left_pose` (7D from `observation.current_target_psm[0:7]`), `right_pose` (7D from `observation.current_target_psm[8:15]`) -- used as REL_XYZ_ROT6D reference frames, not tokenized as state. + +**Action keys:** +- `left_pose` (7D, REL_XYZ_ROT6D from `action.cartesian_absolute[0:7]`) +- `left_gripper` (1D, ABSOLUTE from `action.cartesian_absolute[7:8]`) +- `right_pose` (7D, REL_XYZ_ROT6D from `action.cartesian_absolute[8:15]`) +- `right_gripper` (1D, ABSOLUTE from `action.cartesian_absolute[15:16]`) + +**Video key:** `endoscope_left` +**Language key:** `annotation.instruction` (maps to `instruction.text`, raw text strings) + +## Episode Length Check (50-Step Horizon) + +All subsets are 24 Hz (50 steps β‰ˆ 2.08 seconds). Only one episode in `knot_tying/2` has length 2; every other episode across the dataset is at least 54 steps long. The dataloader automatically skips episodes that are too short for the requested action horizon, so no manual removal is required. + +## Dataset Preparation + +Use the shared preparation script to copy the modality JSON into each dataset's `meta/` folder and generate normalization statistics: + +```bash +bash open_h/prepare_datasets.sh \ + --embodiment-tag ustc_torin_tuodao \ + --modality-json open_h/embodiments/ustc_torin_tuodao/modality.json \ + /path/to/dataset +``` + +## Integration Notes + +- `observation.state` contains joint angles only; Cartesian EEF pose is sourced from `action.cartesian_absolute`. +- Pose actions use REL_XYZ_ROT6D conversion with the reference pose from `observation.current_target_psm`. +- Gripper is modeled as an ABSOLUTE action and appended after pose keys. +- Language uses per-frame `instruction.text` strings (mapped via `annotation.instruction`). +- **Warning: Data quality consideration** β€” Some subsets have all-zero rotation matrices in the raw data for one arm; the LeRobot conversion script + used replaced invalid rotations with the previous valid rotation (identity if none). diff --git a/open_h/embodiments/ustc_torin_tuodao/modality.json b/open_h/embodiments/ustc_torin_tuodao/modality.json new file mode 100644 index 0000000..c8dd8eb --- /dev/null +++ b/open_h/embodiments/ustc_torin_tuodao/modality.json @@ -0,0 +1,21 @@ +{ + "video": { + "endoscope_left": {"original_key": "observation.images.endoscope.left"}, + "endoscope_right": {"original_key": "observation.images.endoscope.right"} + }, + "state": { + "left_joints": {"start": 0, "end": 7}, + "right_joints": {"start": 7, "end": 14}, + "left_pose": {"start": 0, "end": 7, "original_key": "observation.current_target_psm"}, + "right_pose": {"start": 8, "end": 15, "original_key": "observation.current_target_psm"} + }, + "action": { + "left_pose": {"start": 0, "end": 7, "original_key": "action.cartesian_absolute"}, + "left_gripper": {"start": 7, "end": 8, "original_key": "action.cartesian_absolute"}, + "right_pose": {"start": 8, "end": 15, "original_key": "action.cartesian_absolute"}, + "right_gripper": {"start": 15, "end": 16, "original_key": "action.cartesian_absolute"} + }, + "annotation": { + "instruction": {"original_key": "instruction.text", "is_text": true} + } +} diff --git a/open_h/embodiments/ustc_torin_tuodao/ustc_torin_tuodao_config.py b/open_h/embodiments/ustc_torin_tuodao/ustc_torin_tuodao_config.py new file mode 100644 index 0000000..25ce6ad --- /dev/null +++ b/open_h/embodiments/ustc_torin_tuodao/ustc_torin_tuodao_config.py @@ -0,0 +1,129 @@ +""" +USTC Torin modality configuration for GR00T N1.6. + +This configuration supports the USTC merged surgical dataset with: +- Stereo endoscope video (left/right) +- 14D joint-angle state (7 joints per arm) +- Cartesian absolute EEF pose from `action.cartesian_absolute` (xyz + quat + gripper per arm) +- REL_XYZ_ROT6D pose actions (xyz + rot6d) with absolute gripper + +Important representation note: +The parquet files include `observation.state` (joint angles), +`observation.current_target_psm` (absolute pose per arm at t), +`action` (EEF deltas), and +`action.cartesian_absolute` (absolute pose per arm at t+1). +The modality mapping slices out pose-only values (7D per arm) and skips the +gripper slots in the cartesian absolute columns. REL_XYZ_ROT6D uses +`observation.current_target_psm` as the reference pose (t) and predicts pose +deltas toward `action.cartesian_absolute` (t+1). Gripper is modeled as an +ABSOLUTE action. + +Data update note: +The raw conversion script writes `action.cartesian_absolute` alongside +`absolute_action` and repairs invalid rotations by carrying forward the last +valid rotation (identity if none), preserving translation. + +Language: +Use per-frame `instruction.text` strings from the parquet files. This is mapped via +`annotation.instruction` in `open_h/embodiments/ustc_torin_tuodao/modality.json` with `is_text: true`, which +instructs the loader to pass through the raw text instead of using task indices. +""" + +from gr00t.configs.data.embodiment_configs import register_modality_config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + + +# Action horizon (how many future action steps to predict) +# 24 Hz -> 50 frames ~= 2.08 seconds +ACTION_HORIZON = 50 + +ustc_config = { + "video": ModalityConfig( + delta_indices=[0], + modality_keys=[ + "endoscope_left", + # "endoscope_right", # Mono Only + ], + ), + "state": ModalityConfig( + delta_indices=[0], # Single reference state for the current timestep + modality_keys=[ + "left_joints", + "right_joints", + "left_pose", + "right_pose", + ], + # Joint angles are continuous - mean/std normalization is appropriate + mean_std_embedding_keys=[ + "left_joints", + "right_joints", + ], + # Pose keys are consumed by REL_XYZ_ROT6D and should pass through + pass_through_keys=[ + "left_pose", + "right_pose", + ], + ), + "action": ModalityConfig( + # Start at +1 so index 0 is the state reference (CMR-style offset) + delta_indices=list(range(0, ACTION_HORIZON)), + modality_keys=[ + "left_pose", + "left_gripper", + "right_pose", + "right_gripper", + ], + action_configs=[ + # Left pose: REL_XYZ_ROT6D from reference pose + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="left_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + reference_quat_order="xyzw", + ), + # Left gripper: absolute + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + # Right pose: REL_XYZ_ROT6D from reference pose + ActionConfig( + rep=ActionRepresentation.REL_XYZ_ROT6D, + type=ActionType.EEF, + format=ActionFormat.XYZ_ROT6D, + state_key="right_pose", + normalization_type="temporal_meanstd", + input_rotation_format="quat", + reference_rotation_format="quat", + reference_quat_order="xyzw", + ), + # Right gripper: absolute + ActionConfig( + rep=ActionRepresentation.ABSOLUTE, + type=ActionType.NON_EEF, + format=ActionFormat.DEFAULT, + normalization_type="temporal_meanstd", + ), + ], + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["annotation.instruction"], # Uses instruction.text (raw strings) + ), +} + +# Register with USTC_TORIN_TUODAO tag for Torin surgical robot data +register_modality_config(ustc_config, embodiment_tag=EmbodimentTag.USTC_TORIN_TUODAO) diff --git a/open_h/gr00t_h_config.yaml b/open_h/gr00t_h_config.yaml new file mode 100644 index 0000000..02d2453 --- /dev/null +++ b/open_h/gr00t_h_config.yaml @@ -0,0 +1,317 @@ +# GR00T-H Pre-Training Configuration +# +# This is the training recipe used to produce the released GR00T-H checkpoint. +# It trains on 14 surgical embodiments with per-embodiment state dropout (1.0), +# meaning the model operates purely from vision + language (no proprioceptive state). +# +# Embodiment Mix Summary: +# | Embodiment | # Datasets | Frames | Mix Ratio | % of Total | +# |--------------------------|------------|-------------|-----------|------------| +# | cmr_versius | 4 | 105,945,933 | 1.6000 | 19.20% | +# | jhu_imerse_dvrk | 12 | 5,269,530 | 2.9144 | 34.98% | +# | jhu_imerse_dvrk_mono | 1 | 516,334 | 0.3102 | 3.72% | +# | jhu_lscr_dvrk_smarts | 3 | 103,025 | 0.0619 | 0.74% | +# | stanford_dvrk_real | 3 | 874,437 | 0.5253 | 6.30% | +# | obuda_dvrk | 11 | 1,156,946 | 0.6949 | 8.34% | +# | rob_surgical_bitrack | 1 | 1,003,887 | 0.6031 | 7.24% | +# | jhu_imerse_star_il | 1 | 117,247 | 0.0704 | 0.85% | +# | ustc_torin_tuodao | 7 | 512,030 | 0.3075 | 3.69% | +# | hamlyn_dvrk_30hz | 6 | 544,573 | 0.3272 | 3.93% | +# | ucsd_dvrk | 2 | 314,917 | 0.1892 | 2.27% | +# | ucb_dvrk | 1 | 221,950 | 0.1333 | 1.60% | +# | turin_mitic_ex_vivo | 4 | 997,835 | 0.5661 | 6.80% | +# | tud_tundra_ur5e | 1 | 47,753 | 0.0287 | 0.34% | +# | **Total** | **58** | | **8.3322**| **100.0%** | +# +# Usage (4 nodes x 8 GPUs): +# uv run torchrun --nnodes=4 --nproc_per_node=8 \ +# --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \ +# gr00t/experiment/launch_train.py \ +# --load-config-path open_h/gr00t_h_config.yaml +# +# IMPORTANT: replace every REPLACE_WITH_OPEN_H_DATA_PATH entry below with the +# absolute root directory of your local Open-H datasets before launching. +# These placeholder strings are documentation only; they are not expanded by YAML loading. + +data: + datasets: + # ========================================================================= + # CMR Versius (4 datasets, embodiment_tag: cmr_versius) + # Total frames: 105,945,933 Mix ratio: 1.6000 (~19.20%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/cmr-surgical-60hz/cholecystectomy_phase_gesture_prompts_v1 + - REPLACE_WITH_OPEN_H_DATA_PATH/cmr-surgical-60hz/hysterectomy_phase_prompts_v1 + - REPLACE_WITH_OPEN_H_DATA_PATH/cmr-surgical-60hz/inguinal_hernia + - REPLACE_WITH_OPEN_H_DATA_PATH/cmr-surgical-60hz/prostatectomy + mix_ratio: 1.6000 + embodiment_tag: cmr_versius + + # ========================================================================= + # JHU dVRK (12 datasets, embodiment_tag: jhu_imerse_dvrk) + # Total frames: 5,269,530 Mix ratio: 2.9144 (~34.98%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/LSCR/ARCADE/Cholecystectomy + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/LSCR/ARCADE/cautery + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/srth_porcine_chole + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/cao_cautery_combined + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/srt_needle_pickup+handover + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/srt_tissue_lift + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/Imerse/Wound_Closure/point_labeled + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/suturebot + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/Imerse/NephFat_extracted/nephfat + mix_ratio: 2.9144 + embodiment_tag: jhu_imerse_dvrk + exclude_splits: + - missing_videos + + # ========================================================================= + # JHU dVRK Monocular (1 dataset, embodiment_tag: jhu_imerse_dvrk_mono) + # Used to keep compatible with Cosmos-Surg-dVRK evals + # Total frames: 516,334 Mix ratio: 0.3102 (~3.72%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/suturebot + mix_ratio: 0.3102 + embodiment_tag: jhu_imerse_dvrk_mono + + # ========================================================================= + # JHU LSCR SMARTS (3 datasets, embodiment_tag: jhu_lscr_dvrk_smarts) + # Total frames: 103,025 Mix ratio: 0.0619 (~0.74%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/LSCR/SMARTS/offline_recorder_extracted/offline_data_part1 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/LSCR/SMARTS/offline_recorder_extracted/offline_data_part2 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/LSCR/SMARTS/offline_recorder_extracted/offline_data_part3 + mix_ratio: 0.0619 + embodiment_tag: jhu_lscr_dvrk_smarts + + # ========================================================================= + # Stanford Real Robot dVRK (3 datasets, embodiment_tag: stanford_dvrk_real) + # Total frames: 874,437 Mix ratio: 0.5253 (~6.30%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Stanford/Collaborative Haptics and Robotics in Medicine Lab/Real Robot (dVRK)/Needle Transfer + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Stanford/Collaborative Haptics and Robotics in Medicine Lab/Real Robot (dVRK)/Tissue Retraction + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Stanford/Collaborative Haptics and Robotics in Medicine Lab/Real Robot (dVRK)/Peg Transfer + mix_ratio: 0.5253 + embodiment_tag: stanford_dvrk_real + exclude_splits: + - fail + - bad_frames + + # ========================================================================= + # Obuda dVRK (11 datasets, embodiment_tag: obuda_dvrk) + # Total frames: 1,156,946 Mix ratio: 0.6949 (~8.34%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Obuda/FRS_Dome_1 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Obuda/NeedleThreading_1 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Obuda/PegTransfer_1 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Obuda/Rollercoaster_1 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Obuda/Seaspike_1 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Obuda/NeedleThreading_2 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Obuda/PegTransfer_2 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Obuda/Pork_1 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Obuda/Seaspike_2 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Obuda/Seaspike_3 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Obuda/Skinphantom_1 + mix_ratio: 0.6949 + embodiment_tag: obuda_dvrk + + # ========================================================================= + # Rob Surgical BiTrack (1 dataset, embodiment_tag: rob_surgical_bitrack) + # Total frames: 1,003,887 Mix ratio: 0.6031 (~7.24%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Rob Surgical/all_merged_data + mix_ratio: 0.6031 + embodiment_tag: rob_surgical_bitrack + + # ========================================================================= + # JHU IMERSE star_IL (1 dataset, embodiment_tag: jhu_imerse_star_il) + # Total frames: 117,247 Mix ratio: 0.0704 (~0.85%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/JHU/Imerse/star_IL_extracted/star_IL + mix_ratio: 0.0704 + embodiment_tag: jhu_imerse_star_il + exclude_splits: + - MISSING_VIDEOS + + # ========================================================================= + # USTC Torin/Tuodao (7 datasets, embodiment_tag: ustc_torin_tuodao) + # Total frames: 512,030 Mix ratio: 0.3075 (~3.69%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/USTC/exvivo_liver_sep + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/USTC/grasp_on_liver + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/USTC/invivo_liver_sep + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/USTC/ustc_knot_tying + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/USTC/Needle_handover + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/USTC/needle_pickup + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/USTC/tissue_lifting + mix_ratio: 0.3075 + embodiment_tag: ustc_torin_tuodao + + # ========================================================================= + # Hamlyn Centre dVRK (6 datasets, embodiment_tag: hamlyn_dvrk_30hz) + # Total frames: 544,573 Mix ratio: 0.3272 (~3.93%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Hamlyn/Suturing-2 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Hamlyn/peg_transfer + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Hamlyn/Suturing-1 + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Hamlyn/needle_grasp_and_handover + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Hamlyn/knot_tying + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Hamlyn/Tissue_Retraction + mix_ratio: 0.3272 + embodiment_tag: hamlyn_dvrk_30hz + exclude_splits: + - failure + + # ========================================================================= + # UCSD Surgical Learning (2 datasets, embodiment_tag: ucsd_dvrk) + # Total frames: 314,917 Mix ratio: 0.1892 (~2.27%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/UCSD/surgical_learning_dataset + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/UCSD/surgical_learning_dataset2 + mix_ratio: 0.1892 + embodiment_tag: ucsd_dvrk + + # ========================================================================= + # UCBerkeley Debridement (1 dataset, embodiment_tag: ucb_dvrk) + # Total frames: 221,950 Mix ratio: 0.1333 (~1.60%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/UCBerkeley/debridement_lerobot + mix_ratio: 0.1333 + embodiment_tag: ucb_dvrk + + # ========================================================================= + # Turin MITIC Ex Vivo (4 datasets, embodiment_tag: turin_mitic_ex_vivo) + # Total frames: 997,835 Mix ratio: 0.5661 (~6.80%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Turin/mitic_lerobot_ex_vivo + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Turin/mitic_lerobot_plastic_pad + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Turin/mitic_lerobot_plastic_pad_3DMED + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/Turin/mitic_lerobot_plastic_tube + mix_ratio: 0.5661 + embodiment_tag: turin_mitic_ex_vivo + exclude_splits: + - failure + - MISSING_VIDEOS + + # ========================================================================= + # TUD TUNDRA (1 dataset, embodiment_tag: tud_tundra_ur5e) + # Total frames: 47,753 Mix ratio: 0.0287 (~0.34%) + # ========================================================================= + - dataset_paths: + - REPLACE_WITH_OPEN_H_DATA_PATH/Surgical/TUD/260131_TUNDRA_dataset/grasping_retraction + mix_ratio: 0.0287 + embodiment_tag: tud_tundra_ur5e + + # Sharded dataset settings + shard_size: 1024 + episode_sampling_rate: 0.1 + num_shards_per_epoch: 100000 + +model: + # Tunable components + tune_llm: true + tune_visual: true + tune_projector: true + tune_diffusion_model: true + + # Image processing β€” 95% random crop + image_target_size: [236, 414] # Intermediate padded size (height, width) + image_crop_size: [224, 392] # Final crop size seen by the model + + # Augmentations + random_rotation_angle: 5 + color_jitter_params: + brightness: 0.12 + contrast: 0.15 + saturation: 0.15 + hue: 0.02 + + # State dropout: controls whether the model sees proprioceptive state inputs. + # The two modes use different mechanisms: + # + # - state_dropout_prob: global probability applied uniformly to all embodiments. + # Stochastic during training, disabled at inference. When triggered, the encoded + # state features are replaced with a learned mask_token embedding. + # + # - state_dropout_prob_per_embodiment: per-embodiment overrides. When triggered, + # the raw state vector is zeroed out BEFORE the state encoder, so the model + # learns that encoder(zeros) = "no state available". At inference, prob >= 1.0 + # applies this zeroing deterministically so the model matches training behavior. + # Embodiments not listed fall back to state_dropout_prob. + state_dropout_prob: 0.0 + state_dropout_prob_per_embodiment: + cmr_versius: 1.0 + jhu_imerse_dvrk: 1.0 + jhu_imerse_dvrk_mono: 1.0 + jhu_lscr_dvrk_smarts: 1.0 + stanford_dvrk_real: 1.0 + obuda_dvrk: 1.0 + rob_surgical_bitrack: 1.0 + ustc_torin_tuodao: 1.0 + turin_mitic_ex_vivo: 1.0 + tud_tundra_ur5e: 1.0 + hamlyn_dvrk_30hz: 1.0 + ucsd_dvrk: 1.0 + ucb_dvrk: 1.0 + jhu_imerse_star_il: 1.0 + + # Model settings + model_name: nvidia/Eagle-Block2A-2B-v2 + eagle_collator: true + use_relative_action: true + load_bf16: false + reproject_vision: false + backbone_trainable_params_fp32: true + + # Action horizon β€” padded to 50 for mixed embodiments + action_horizon: 50 + +training: + # Base model + start_from_checkpoint: nvidia/GR00T-N1.6-3B + output_dir: output/gr00t_h + + # Training hyperparameters + global_batch_size: 1024 + learning_rate: 0.00003 + weight_decay: 0.00001 + warmup_ratio: 0.00 + max_steps: 85000 + + # Optimizer + optim: adamw_torch + lr_scheduler_type: cosine + gradient_accumulation_steps: 1 + max_grad_norm: 1.0 + + # Checkpointing + save_steps: 1000 + save_total_limit: 100 + + # Distributed training + num_gpus: 32 + deepspeed_stage: 2 + + # Logging + use_wandb: true + wandb_project: finetune-gr00t-h + logging_steps: 10 + + # Data loading + dataloader_num_workers: 4 + + # Precision + bf16: true + tf32: true diff --git a/open_h/prepare_datasets.sh b/open_h/prepare_datasets.sh new file mode 100644 index 0000000..1df7b10 --- /dev/null +++ b/open_h/prepare_datasets.sh @@ -0,0 +1,107 @@ +#!/bin/bash +# Prepare one or more LeRobot datasets for GR00T by generating normalization stats. +# +# This script will: +# 1) Copy a provided modality JSON into each dataset's meta/ folder +# 2) Generate stats.json via gr00t/data/stats.py +# 3) Generate temporal_stats.json (norm stats for REL_XYZ_ROT6D actions over the action chunk) +# +# Run from repo root: +# bash open_h/prepare_datasets.sh \ +# --embodiment-tag \ +# --modality-json \ +# /path/to/dataset_a /path/to/dataset_b + +set -euo pipefail + +print_usage() { + cat <<'EOF' +Usage: + bash open_h/prepare_datasets.sh \ + --embodiment-tag \ + --modality-json \ + [ ...] + +Required arguments: + --embodiment-tag Embodiment tag to pass to stats and finetune scripts + --modality-json Path to modality JSON copied into each dataset meta/ folder + DATASET_PATH One or more dataset paths to process +EOF +} + +die() { + echo "ERROR: $*" >&2 + exit 1 +} + +if [[ ! -f "gr00t/data/stats.py" ]]; then + die "Run this script from the repo root." +fi + +EMBODIMENT_TAG="" +MODALITY_FILE="" +DATASET_PATHS=() + +while [[ $# -gt 0 ]]; do + case "$1" in + --embodiment-tag) + [[ $# -ge 2 ]] || die "Missing value for --embodiment-tag" + EMBODIMENT_TAG="$2" + shift 2 + ;; + --modality-json) + [[ $# -ge 2 ]] || die "Missing value for --modality-json" + MODALITY_FILE="$2" + shift 2 + ;; + -h|--help) + print_usage + exit 0 + ;; + --) + shift + while [[ $# -gt 0 ]]; do + DATASET_PATHS+=("$1") + shift + done + ;; + -*) + die "Unknown option: $1" + ;; + *) + DATASET_PATHS+=("$1") + shift + ;; + esac +done + +[[ -n "${EMBODIMENT_TAG}" ]] || die "Missing required --embodiment-tag" +[[ -n "${MODALITY_FILE}" ]] || die "Missing required --modality-json" +[[ ${#DATASET_PATHS[@]} -gt 0 ]] || die "Provide at least one DATASET_PATH" +[[ -f "${MODALITY_FILE}" ]] || die "Modality file does not exist: ${MODALITY_FILE}" + +for dataset_path in "${DATASET_PATHS[@]}"; do + echo "=== Processing ${dataset_path} ===" + + [[ -d "${dataset_path}" ]] || die "Dataset path does not exist: ${dataset_path}" + + if [[ ! -d "${dataset_path}/meta" ]]; then + echo "Creating meta directory: ${dataset_path}/meta" + mkdir -p "${dataset_path}/meta" + fi + + echo "Copying ${MODALITY_FILE} to ${dataset_path}/meta/modality.json" + cp "${MODALITY_FILE}" "${dataset_path}/meta/modality.json" + + echo "Generating stats..." + uv run python gr00t/data/stats.py \ + --dataset-path "${dataset_path}" \ + --embodiment-tag "${EMBODIMENT_TAG}" + + echo "Generating temporal stats..." + uv run python gr00t/experiment/launch_finetune.py \ + --base-model-path nvidia/GR00T-N1.6-3B \ + --dataset-path "${dataset_path}" \ + --embodiment-tag "${EMBODIMENT_TAG}" \ + --calculate-norm-stats +done diff --git a/pyproject.toml b/pyproject.toml index 988a83b..8335fa3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,7 +64,7 @@ gpu = [ [tool.setuptools.packages.find] where = ["."] -include = ["gr00t*"] +include = ["gr00t*", "open_h*"] [tool.uv.extra-build-dependencies] flash-attn = ["torch==2.7.1", "numpy==1.26.4", "triton==3.3.1"]