diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml new file mode 100644 index 000000000..079a30421 --- /dev/null +++ b/.github/workflows/docker-build.yml @@ -0,0 +1,36 @@ +name: Build and Push Gr00t Images + +on: + push: + branches: + - main-positronic + +jobs: + build-and-push: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Free disk space + run: | + sudo rm -rf /usr/share/dotnet + sudo rm -rf /usr/local/lib/android + sudo rm -rf /opt/ghc + sudo rm -rf "${AGENT_TOOLSDIRECTORY:-/opt/hostedtoolcache}/CodeQL" + sudo rm -rf "${AGENT_TOOLSDIRECTORY:-/opt/hostedtoolcache}/node" + sudo rm -rf "${AGENT_TOOLSDIRECTORY:-/opt/hostedtoolcache}/go" + sudo rm -rf /usr/lib/jvm + sudo docker system prune -af || true + df -h + + - name: Log in to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKER_USERNAME }} + password: ${{ secrets.DOCKER_PASSWORD }} + + - name: Build and Push + working-directory: docker + run: | + make push diff --git a/README.md b/README.md index 6d86ea8af..eb935ad8e 100644 --- a/README.md +++ b/README.md @@ -255,7 +255,7 @@ See the [Orin setup guide](scripts/deployment/README.md#jetson-orin-setup) for D > per-platform Docker and bare-metal setup. -For a containerized setup that avoids system-level dependency conflicts, see our [Docker Setup Guide](docker/README.md). The recommended container workflow is to start the image first, then clone or pull the repo inside the running container so your checkout uses the image's prebuilt dependency environment. +For a containerized setup that avoids system-level dependency conflicts, see our [Docker Setup Guide](docker/README.md). The image includes this fork at `/gr00t` with its prebuilt dependency environment. --- diff --git a/docker/Dockerfile b/docker/Dockerfile index 34fe6d6f1..1d7f7a47a 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -225,4 +225,10 @@ ENV MUJOCO_GL="egl" \ # bootstrap_wheels.sh before sync (then committed back) if missing. flash-attn # comes from official release URLs — no source build needed. +# Install the fork alongside its dependency environment for Positronic subprocesses. +WORKDIR /gr00t +COPY . /gr00t +RUN uv pip install --python /opt/gr00t-venv/bin/python --no-deps -e /gr00t +ENV PYTHONPATH=/gr00t + CMD ["/bin/bash"] diff --git a/docker/Makefile b/docker/Makefile new file mode 100644 index 000000000..a5035665e --- /dev/null +++ b/docker/Makefile @@ -0,0 +1,79 @@ +.PHONY: all build tag push clean prune help + +# Image configuration +IMAGE_NAME_GROOT := positro/gr00t-base + +# Extract version from pyproject.toml (first literal version entry) +VERSION := $(shell sed -n 's/^version = "\([^"]*\)"/\1/p' ../pyproject.toml | head -n 1) + +GIT_SHA := $(shell git rev-parse --short HEAD 2>/dev/null || echo "unknown") + +REGISTRY_URL ?= docker.io + +TAG_GROOT_LATEST := $(IMAGE_NAME_GROOT):latest +TAG_GROOT_VERSION := $(IMAGE_NAME_GROOT):v$(VERSION) +TAG_GROOT_SHA := $(IMAGE_NAME_GROOT):$(GIT_SHA) +LOCAL_TAG_GROOT := $(IMAGE_NAME_GROOT):local + +help: + @echo "gr00t Groot Docker Build System" + @echo "" + @echo "Configuration:" + @echo " Groot Image: $(IMAGE_NAME_GROOT)" + @echo " Version: $(VERSION)" + @echo " Git SHA: $(GIT_SHA)" + @echo " Registry URL: $(REGISTRY_URL)" + @echo "" + @echo "Targets:" + @echo " make build Build the groot image" + @echo " make tag Tag the groot image" + @echo " make push Push groot tags to Docker Hub" + @echo " make clean Remove local groot images" + @echo " make prune Remove dangling/unused Docker images" + @echo " make help Show this help message" + @echo "" + +build: + @echo "Building $(IMAGE_NAME_GROOT) image..." + @if [ -z "$(VERSION)" ]; then \ + echo "Error: Could not extract version from pyproject.toml"; \ + exit 1; \ + fi + docker build --platform linux/amd64 \ + -f Dockerfile \ + -t $(LOCAL_TAG_GROOT) .. + +tag: build + @echo "Tagging groot image with multiple tags..." + docker tag $(LOCAL_TAG_GROOT) $(TAG_GROOT_LATEST) + docker tag $(LOCAL_TAG_GROOT) $(TAG_GROOT_VERSION) + docker tag $(LOCAL_TAG_GROOT) $(TAG_GROOT_SHA) + @echo "Tagged with:" + @echo " - $(TAG_GROOT_LATEST)" + @echo " - $(TAG_GROOT_VERSION)" + @echo " - $(TAG_GROOT_SHA)" + +push: tag + @echo "Pushing groot images to Docker Hub..." + docker push $(TAG_GROOT_LATEST) + docker push $(TAG_GROOT_VERSION) + docker push $(TAG_GROOT_SHA) + @echo "" + @echo "Successfully pushed groot images to Docker Hub!" + @echo "To use on cloud instances, set:" + @echo " export IMAGE_REGISTRY=$(REGISTRY_URL)/" + +all: push + +clean: + @echo "Removing local groot images..." + -docker rmi $(LOCAL_TAG_GROOT) + -docker rmi $(TAG_GROOT_LATEST) + -docker rmi $(TAG_GROOT_VERSION) + -docker rmi $(TAG_GROOT_SHA) + @echo "Cleanup complete." + +prune: + @echo "Pruning dangling and unused Docker images..." + docker image prune -f + @echo "Prune complete." diff --git a/docker/README.md b/docker/README.md index e4c62d7c7..c8cf40d51 100644 --- a/docker/README.md +++ b/docker/README.md @@ -18,11 +18,38 @@ From the repository root: bash docker/build.sh ``` -This builds from `nvidia/cuda:12.8.0-devel-ubuntu24.04` and installs all dependencies into `/opt/gr00t-venv`. The image does not include a working source checkout; for normal use, start the image and then clone or pull the repo you want to run inside the container. +This builds from `nvidia/cuda:12.8.0-devel-ubuntu24.04` and installs all dependencies into `/opt/gr00t-venv`. The image includes this fork at `/gr00t`, installed into `/opt/gr00t-venv`. + +## Positronic base image + +```bash +make -C docker build +make -C docker push +``` + +The Makefile publishes `positro/gr00t-base` with `latest`, version, and commit tags. +It targets `linux/amd64`, including when built on an Apple Silicon Mac. +For a native ARM build, use `bash docker/build.sh` on the target host. +Positronic's `GROOT_BASE_IMAGE` selects an existing base image for its adapter build. + +Fine-tuning defaults to the base checkpoint's saved model and modality configuration. +The Python launcher accepts `--video-keys` to select the fine-tuning camera layout; omitted, it retains the checkpoint's views. +When using `examples/finetune.sh`, put this option after the script's `--` passthrough delimiter: + +```bash +bash examples/finetune.sh \ + --base-model-path nvidia/GR00T-N1.7-DROID \ + --dataset-path /data/droid \ + --embodiment-tag oxe_droid_relative_eef_relative_joint \ + --output-dir /data/checkpoints \ + -- --video-keys exterior_image_1_left exterior_image_2_left wrist_image_left +``` + +The policy server accepts `--model-path hf://nvidia/GR00T-N1.7-DROID` and downloads that snapshot. ## Running the Container -**Recommended workflow: run the image, then clone or update the repo inside it.** +**Run the included fork from `/gr00t`.** Start an interactive shell: @@ -35,9 +62,7 @@ docker run -it --rm --gpus all \ Then, inside the container: ```bash -git clone --recurse-submodules https://github.com/NVIDIA/Isaac-GR00T /workspace/Isaac-GR00T -cd /workspace/Isaac-GR00T -export PYTHONPATH="$PWD${PYTHONPATH:+:$PYTHONPATH}" +cd /gr00t python -c "import gr00t; print('GR00T ready')" ``` diff --git a/gr00t/configs/finetune_config.py b/gr00t/configs/finetune_config.py index 15918e206..6b1d3ce66 100644 --- a/gr00t/configs/finetune_config.py +++ b/gr00t/configs/finetune_config.py @@ -41,10 +41,13 @@ class FinetuneConfig: modality_config_path: str | None = None """ - Path to a Python file defining the modality configuration for the given embodiment. - If None, use the pre-registered modality config in `gr00t/configs/data/embodiment_configs.py`. + Path to a Python file defining the modality configuration for the given embodiment. + If None, retain the modality configuration saved in the base checkpoint. """ + video_keys: list[str] | None = None + """Explicit camera keys for the fine-tuning dataset; omission retains the checkpoint's cameras.""" + # --- Model Tuning Flags --- tune_llm: bool = False """If True, fine-tune the language model (LLM) backbone during training.""" diff --git a/gr00t/data/types.py b/gr00t/data/types.py index 7bd6d5fe5..379241a3c 100644 --- a/gr00t/data/types.py +++ b/gr00t/data/types.py @@ -22,6 +22,10 @@ from gr00t.data.embodiment_tags import EmbodimentTag +VIDEO = "video" +LANGUAGE = "language" + + class MessageType(Enum): START_OF_EPISODE = "start_of_episode" END_OF_EPISODE = "end_of_episode" diff --git a/gr00t/experiment/launch_finetune.py b/gr00t/experiment/launch_finetune.py index 00cb3ef1f..5901b1d67 100644 --- a/gr00t/experiment/launch_finetune.py +++ b/gr00t/experiment/launch_finetune.py @@ -16,15 +16,19 @@ # Launch finetuning for N1.7 on "single node". # This script tries to provide a similar user experience as current OSS. +import copy import json import os from pathlib import Path +from transformers.utils import cached_file import tyro from gr00t.configs.base_config import get_default_config from gr00t.configs.finetune_config import FinetuneConfig -from gr00t.experiment.experiment import run +from gr00t.configs.model.gr00t_n1d7 import Gr00tN1d7Config +from gr00t.data.embodiment_tags import EmbodimentTag +from gr00t.data.types import ModalityConfig # Make sure the user provided modality config is registered. @@ -41,16 +45,9 @@ def load_modality_config(modality_config_path: str): raise FileNotFoundError(f"Modality config path does not exist: {modality_config_path}") -if __name__ == "__main__": - # 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 - ft_config = tyro.cli(FinetuneConfig, description=__doc__) - from gr00t.data.embodiment_tags import EmbodimentTag - - ft_config.embodiment_tag = EmbodimentTag.resolve(ft_config.embodiment_tag) - embodiment_tag = ft_config.embodiment_tag.value +def build_config(ft_config: FinetuneConfig): + """Inherit model and modality contracts from the checkpoint before applying training overrides.""" + embodiment_tag = EmbodimentTag.resolve(ft_config.embodiment_tag).value # all rank workers should register for the modality config if ft_config.modality_config_path is not None: @@ -73,6 +70,24 @@ def load_modality_config(modality_config_path: str): } ) config.load_config_path = None + config.model = Gr00tN1d7Config.from_pretrained(ft_config.base_model_path) + checkpoint = Path(ft_config.base_model_path) + processor_root = checkpoint / "processor" if (checkpoint / "processor").is_dir() else checkpoint + processor_file = cached_file(str(processor_root), "processor_config.json") + with open(processor_file) as f: + processor_kwargs = json.load(f)["processor_kwargs"] + config.model.use_relative_action = processor_kwargs["use_relative_action"] + if ft_config.modality_config_path is None: + modalities = processor_kwargs["modality_configs"][embodiment_tag] + config.data.modality_configs = { + embodiment_tag: {name: ModalityConfig(**value) for name, value in modalities.items()} + } + else: + config.data.modality_configs = copy.deepcopy(config.data.modality_configs) + if ft_config.video_keys is not None: + if not ft_config.video_keys or len(set(ft_config.video_keys)) != len(ft_config.video_keys): + raise ValueError("video_keys must be nonempty and unique") + config.data.modality_configs[embodiment_tag]["video"].modality_keys = ft_config.video_keys # overwrite with finetune config supplied by the user config.model.tune_llm = ft_config.tune_llm @@ -80,14 +95,17 @@ 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.random_rotation_angle = ft_config.random_rotation_angle - config.model.color_jitter_params = ft_config.color_jitter_params + if ft_config.random_rotation_angle is not None: + config.model.random_rotation_angle = ft_config.random_rotation_angle + if ft_config.color_jitter_params is not None: + config.model.color_jitter_params = ft_config.color_jitter_params config.model.use_percentiles = ft_config.use_percentiles if (ft_config.shortest_image_edge is None) != (ft_config.crop_fraction is None): raise ValueError("shortest_image_edge and crop_fraction must be set together") if ft_config.shortest_image_edge is not None: config.model.shortest_image_edge = ft_config.shortest_image_edge config.model.crop_fraction = ft_config.crop_fraction + if config.model.shortest_image_edge is not None and config.model.crop_fraction is not None: config.model.image_crop_size = None config.model.image_target_size = None if ft_config.extra_augmentation_config: @@ -99,7 +117,6 @@ def load_modality_config(modality_config_path: str): config.model.reproject_vision = False config.model.model_name = "nvidia/Cosmos-Reason2-2B" config.model.backbone_trainable_params_fp32 = True - config.model.use_relative_action = True config.training.experiment_name = ft_config.experiment_name config.training.start_from_checkpoint = ft_config.base_model_path @@ -127,4 +144,11 @@ def load_modality_config(modality_config_path: str): config.training.resume_from_checkpoint = ft_config.resume_from_checkpoint config.training.skip_weight_loading = ft_config.skip_weight_loading - run(config) + return config + + +if __name__ == "__main__": + from gr00t.experiment.experiment import run + + os.environ.setdefault("LOGURU_LEVEL", "INFO") + run(build_config(tyro.cli(FinetuneConfig, description=__doc__))) diff --git a/gr00t/model/gr00t_n1d7/processing_gr00t_n1d7.py b/gr00t/model/gr00t_n1d7/processing_gr00t_n1d7.py index ec0d59f77..4be4715ec 100644 --- a/gr00t/model/gr00t_n1d7/processing_gr00t_n1d7.py +++ b/gr00t/model/gr00t_n1d7/processing_gr00t_n1d7.py @@ -872,6 +872,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | Path, **kwargs): "exclude_state", "state_dropout_prob", "use_mean_std", + "use_percentiles", + "extra_augmentation_config", "model_name", "model_type", "max_action_horizon", diff --git a/gr00t/policy/gr00t_policy.py b/gr00t/policy/gr00t_policy.py index 6f5a46b10..173d56e25 100644 --- a/gr00t/policy/gr00t_policy.py +++ b/gr00t/policy/gr00t_policy.py @@ -23,13 +23,14 @@ from pathlib import Path from typing import Any +from huggingface_hub import snapshot_download import numpy as np import torch from transformers import AutoModel, AutoProcessor from gr00t.data.embodiment_tags import FINETUNE_ONLY_TAGS, POSTTRAIN_TAGS, EmbodimentTag from gr00t.data.interfaces import BaseProcessor -from gr00t.data.types import MessageType, ModalityConfig, VLAStepData +from gr00t.data.types import LANGUAGE, VIDEO, MessageType, ModalityConfig, VLAStepData from .policy import BasePolicy, PolicyWrapper @@ -103,7 +104,10 @@ def __init__( super().__init__(strict=strict) if isinstance(embodiment_tag, str): embodiment_tag = EmbodimentTag.resolve(embodiment_tag) - model_dir = Path(model_path) + if str(model_path).startswith("hf://"): + model_dir = Path(snapshot_download(str(model_path).removeprefix("hf://"))) + else: + model_dir = Path(model_path) # Load the pretrained model and move to target device with bfloat16 precision model = AutoModel.from_pretrained(model_dir) @@ -251,6 +255,12 @@ def check_observation(self, observation: dict[str, Any]) -> None: bs = -1 # ===== VIDEO VALIDATION ===== + expected_cameras = set(self.modality_configs[VIDEO].modality_keys) + if set(observation[VIDEO]) - expected_cameras: + raise ValueError( + f"Checkpoint cameras {sorted(expected_cameras)} do not match " + f"observation cameras {sorted(observation[VIDEO])}" + ) # Validate each video stream defined in the modality config for video_key in self.modality_configs["video"].modality_keys: assert video_key in observation["video"], ( @@ -332,8 +342,8 @@ def check_observation(self, observation: dict[str, Any]) -> None: ) # ===== LANGUAGE VALIDATION ===== - # Validate each language stream defined in the modality config - for language_key in self.modality_configs["language"].modality_keys: + # Inference uses the first instruction key; additional keys are training paraphrases. + for language_key in self.modality_configs[LANGUAGE].modality_keys[:1]: # Check that the expected language key exists in the observation # (must happen before indexing — see video validation above) assert language_key in observation["language"], ( @@ -607,8 +617,8 @@ def check_observation(self, observation: dict[str, Any]) -> None: ) # ===== LANGUAGE VALIDATION ===== - # Check language modalities (special handling for DC environment compatibility) - for language_key in modality_configs["language"].modality_keys: + # Validate the instruction selected for inference, including its DC environment alias. + for language_key in modality_configs[LANGUAGE].modality_keys[:1]: # PATCH: Legacy compatibility for DC environments # DC envs use 'annotation.human.coarse_action' instead of 'task' if language_key == "task" and "annotation.human.coarse_action" in observation: @@ -663,6 +673,8 @@ def _get_action( for modality in ["video", "state", "language"]: new_obs[modality] = {} for key in self.policy.modality_configs[modality].modality_keys: + if modality == LANGUAGE and key != self.policy.language_key: + continue if modality == "language": # PATCH: Legacy compatibility for DC environments if key == "task" and "annotation.human.coarse_action" in observation: diff --git a/tests/fixtures/droid/config.json b/tests/fixtures/droid/config.json new file mode 100644 index 000000000..11a846e5a --- /dev/null +++ b/tests/fixtures/droid/config.json @@ -0,0 +1,87 @@ +{ + "action_horizon": 40, + "add_pos_embed": true, + "apply_sincos_state_encoding": false, + "architectures": [ + "Gr00tN1d7" + ], + "attn_dropout": 0.2, + "attn_implementation": null, + "backbone_embedding_dim": 2048, + "color_jitter_params": { + "brightness": 0.3, + "contrast": 0.4, + "hue": 0.08, + "saturation": 0.5 + }, + "crop_fraction": 0.95, + "diffusion_model_cfg": { + "attention_head_dim": 48, + "dropout": 0.2, + "final_dropout": true, + "interleave_self_attention": true, + "norm_type": "ada_norm", + "num_attention_heads": 32, + "num_layers": 32, + "output_dim": 1024, + "positional_embeddings": null + }, + "dtype": "bfloat16", + "exclude_state": false, + "formalize_language": true, + "hidden_size": 1024, + "image_crop_size": [ + 230, + 230 + ], + "image_target_size": [ + 256, + 256 + ], + "letter_box_transform": false, + "load_bf16": true, + "max_action_dim": 132, + "max_num_embodiments": 32, + "max_seq_len": 1024, + "max_state_dim": 132, + "model_dtype": "bfloat16", + "model_type": "Gr00tN1d7", + "noise_beta_alpha": 1.5, + "noise_beta_beta": 1.0, + "noise_s": 0.999, + "num_inference_timesteps": 4, + "num_timestep_buckets": 1000, + "random_history_crop": true, + "random_rotation_angle": 0, + "reproject_vision": false, + "rtc_ramp_rate": 6.0, + "select_layer": 16, + "shortest_image_edge": 256, + "state_dropout_prob": 0.2, + "state_gaussian_noise_std": 0.0, + "transformers_version": "4.57.1", + "tune_diffusion_model": true, + "tune_linear": true, + "tune_llm": true, + "tune_projector": true, + "tune_top_llm_layers": 0, + "tune_visual": true, + "tune_vlln": true, + "use_albumentations": true, + "use_alternate_vl_dit": true, + "use_flash_attention": true, + "use_future_tokens": false, + "use_mean_std": false, + "use_percentiles": true, + "use_vl_self_attention": true, + "use_vlln": true, + "vl_self_attention_cfg": { + "attention_head_dim": 64, + "dropout": 0.2, + "final_dropout": true, + "num_attention_heads": 32, + "num_layers": 4, + "positional_embeddings": null + }, + "model_name": "nvidia/Cosmos-Reason2-2B" +} diff --git a/tests/fixtures/droid/processor_config.json b/tests/fixtures/droid/processor_config.json new file mode 100644 index 000000000..55b4d74b3 --- /dev/null +++ b/tests/fixtures/droid/processor_config.json @@ -0,0 +1,134 @@ +{ + "processor_class": "Gr00tN1d7Processor", + "processor_kwargs": { + "modality_configs": { + "oxe_droid_relative_eef_relative_joint": { + "video": { + "delta_indices": [ + 0 + ], + "modality_keys": [ + "exterior_image_1_left", + "wrist_image_left" + ] + }, + "state": { + "delta_indices": [ + 0 + ], + "modality_keys": [ + "eef_9d", + "gripper_position", + "joint_position" + ] + }, + "action": { + "delta_indices": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + 28, + 29, + 30, + 31, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39 + ], + "modality_keys": [ + "eef_9d", + "gripper_position", + "joint_position" + ], + "action_configs": [ + { + "rep": "RELATIVE", + "type": "EEF", + "format": "XYZ_ROT6D", + "state_key": "eef_9d" + }, + { + "rep": "ABSOLUTE", + "type": "NON_EEF", + "format": "DEFAULT", + "state_key": "gripper_position" + }, + { + "rep": "RELATIVE", + "type": "NON_EEF", + "format": "DEFAULT", + "state_key": "joint_position" + } + ] + }, + "language": { + "delta_indices": [ + 0 + ], + "modality_keys": [ + "annotation.language.language_instruction" + ] + } + } + }, + "use_percentiles": true, + "use_mean_std": false, + "image_crop_size": [ + 230, + 230 + ], + "image_target_size": [ + 256, + 256 + ], + "formalize_language": true, + "max_state_dim": 132, + "max_action_dim": 132, + "apply_sincos_state_encoding": false, + "color_jitter_params": { + "brightness": 0.3, + "contrast": 0.4, + "saturation": 0.5, + "hue": 0.08 + }, + "random_rotation_angle": 0, + "letter_box_transform": false, + "exclude_state": false, + "state_dropout_prob": 0.2, + "use_albumentations": true, + "shortest_image_edge": 256, + "crop_fraction": 0.95, + "max_action_horizon": 40, + "use_relative_action": true + } +} diff --git a/tests/gr00t/experiment/test_finetune_checkpoint.py b/tests/gr00t/experiment/test_finetune_checkpoint.py new file mode 100644 index 000000000..f71420750 --- /dev/null +++ b/tests/gr00t/experiment/test_finetune_checkpoint.py @@ -0,0 +1,72 @@ +"""DROID checkpoint contracts survive fine-tuning configuration. + +Fixtures are from nvidia/GR00T-N1.7-DROID revision 05e7cc9 on Hugging Face. +""" + +import json +from pathlib import Path + +from gr00t.configs.finetune_config import FinetuneConfig +from gr00t.data.types import ActionFormat, ActionRepresentation, ActionType +from gr00t.experiment.experiment import warn_configs +from gr00t.experiment.launch_finetune import build_config +import pytest + + +CHECKPOINT = Path(__file__).parents[2] / "fixtures" / "droid" +EMBODIMENT = "oxe_droid_relative_eef_relative_joint" + + +@pytest.mark.parametrize( + "cameras", [None, ["exterior_image_1_left", "exterior_image_2_left", "wrist_image_left"]] +) +def test_finetuning_retains_droid_checkpoint_contract(cameras): + config = build_config( + FinetuneConfig( + base_model_path=str(CHECKPOINT), + dataset_path="/dataset", + embodiment_tag=EMBODIMENT, + video_keys=cameras, + resume_from_checkpoint=True, + ) + ) + saved = json.loads((CHECKPOINT / "processor_config.json").read_text())["processor_kwargs"] + modalities = config.data.modality_configs[EMBODIMENT] + assert modalities["video"].delta_indices == [0] + assert modalities["video"].modality_keys == ( + cameras or saved["modality_configs"][EMBODIMENT]["video"]["modality_keys"] + ) + assert modalities["state"].delta_indices == [0] + assert modalities["action"].delta_indices == list(range(40)) + eef, grip, joints = modalities["action"].action_configs + assert (eef.rep, eef.type, eef.format) == ( + ActionRepresentation.RELATIVE, + ActionType.EEF, + ActionFormat.XYZ_ROT6D, + ) + assert grip.rep is ActionRepresentation.ABSOLUTE + assert joints.rep is ActionRepresentation.RELATIVE + assert config.model.use_relative_action + assert config.model.diffusion_model_cfg["num_layers"] == 32 + assert config.model.select_layer == 16 + assert config.model.action_horizon == 40 + assert config.model.shortest_image_edge == saved["shortest_image_edge"] == 256 + assert config.model.crop_fraction == saved["crop_fraction"] == 0.95 + assert config.model.color_jitter_params == saved["color_jitter_params"] + assert not config.model.letter_box_transform + assert not config.model.tune_llm + assert not config.model.tune_visual + assert config.training.resume_from_checkpoint + warn_configs(config) + + +def test_camera_override_does_not_change_a_subsequent_run(): + arguments = dict( + base_model_path=str(CHECKPOINT), dataset_path="/dataset", embodiment_tag=EMBODIMENT + ) + build_config(FinetuneConfig(**arguments, video_keys=["custom_wrist", "custom_external"])) + config = build_config(FinetuneConfig(**arguments)) + assert config.data.modality_configs[EMBODIMENT]["video"].modality_keys == [ + "exterior_image_1_left", + "wrist_image_left", + ] diff --git a/tests/gr00t/model/test_gr00t_processor.py b/tests/gr00t/model/test_gr00t_processor.py index 590b0ea46..f7517775b 100644 --- a/tests/gr00t/model/test_gr00t_processor.py +++ b/tests/gr00t/model/test_gr00t_processor.py @@ -22,6 +22,7 @@ import json from pathlib import Path +import shutil import tempfile from unittest.mock import MagicMock, patch @@ -105,6 +106,27 @@ def fake_cached_file(path_or_repo_id, filename, **kwargs): } +@pytest.mark.parametrize("use_percentiles", [True, False]) +def test_checkpoint_overrides_reach_normalization_and_mask_transforms(tmp_path, use_percentiles): + from gr00t.model.gr00t_n1d7 import processing_gr00t_n1d7 as processor_module + from gr00t.model.gr00t_n1d7.image_augmentations import BackgroundNoiseTransform + + shutil.copytree(FIXTURE_DIR, tmp_path, dirs_exist_ok=True) + config_path = tmp_path / "processor_config.json" + config = json.loads(config_path.read_text()) + config["processor_kwargs"]["use_percentiles"] = not use_percentiles + config_path.write_text(json.dumps(config)) + augmentation = {"background_noise_transforms": [{"target_mask_values": [0], "p": 1.0}]} + with patch.object(processor_module, "build_processor", return_value=MagicMock()): + processor = processor_module.Gr00tN1d7Processor.from_pretrained( + tmp_path, use_percentiles=use_percentiles, extra_augmentation_config=augmentation + ) + assert processor.state_action_processor.use_percentiles is use_percentiles + assert processor.extra_augmentation_config == augmentation + assert len(processor.train_image_transform.mask_transforms) == 1 + assert isinstance(processor.train_image_transform.mask_transforms[0], BackgroundNoiseTransform) + + def _make_step_data(proc_config) -> VLAStepData: """Create synthetic VLAStepData matching the fixture config.""" import json as _json diff --git a/tests/gr00t/policy/test_gr00t_policy.py b/tests/gr00t/policy/test_gr00t_policy.py index 54de946f8..71d53eafa 100644 --- a/tests/gr00t/policy/test_gr00t_policy.py +++ b/tests/gr00t/policy/test_gr00t_policy.py @@ -22,7 +22,7 @@ from pathlib import Path from unittest.mock import MagicMock, patch -from gr00t.data.types import ModalityConfig +from gr00t.data.types import LANGUAGE, ModalityConfig import numpy as np import pytest import torch @@ -44,13 +44,12 @@ def _build_modality_configs(): "video": ModalityConfig(delta_indices=[0], modality_keys=VIDEO_KEYS), "state": ModalityConfig(delta_indices=[0], modality_keys=STATE_KEYS), "action": ModalityConfig(delta_indices=list(range(16)), modality_keys=ACTION_KEYS), - "language": ModalityConfig(delta_indices=[0], modality_keys=[LANGUAGE_KEY]), + LANGUAGE: ModalityConfig(delta_indices=[0], modality_keys=[LANGUAGE_KEY]), } } -@pytest.fixture -def policy(): +def _make_policy(model_path): mock_model = MagicMock() mock_model.eval = MagicMock() mock_model.to = MagicMock(return_value=mock_model) @@ -96,6 +95,7 @@ def fake_decode_action(action, embodiment_tag, state=None): with ( patch("gr00t.policy.gr00t_policy.AutoModel") as MockAutoModel, patch("gr00t.policy.gr00t_policy.AutoProcessor") as MockAutoProcessor, + patch("gr00t.policy.gr00t_policy.snapshot_download", return_value="/fake/path") as download, patch("pathlib.Path.is_dir", return_value=False), patch("pathlib.Path.exists", return_value=True), ): @@ -106,12 +106,27 @@ def fake_decode_action(action, embodiment_tag, state=None): p = Gr00tPolicy( embodiment_tag=EMBODIMENT, - model_path="/fake/path", + model_path=model_path, device="cpu", ) + if model_path.startswith("hf://"): + download.assert_called_once_with("owner/model") + else: + download.assert_not_called() + MockAutoModel.from_pretrained.assert_called_once_with(Path("/fake/path")) return p +@pytest.fixture +def policy(): + return _make_policy("/fake/path") + + +@pytest.mark.parametrize("model_path", ["/fake/path", "hf://owner/model"]) +def test_policy_resolves_local_and_hub_checkpoints(model_path): + _make_policy(model_path) + + def _make_observation(batch_size=1): return { "video": { @@ -123,7 +138,7 @@ def _make_observation(batch_size=1): for k in STATE_KEYS[:-1] # all except gripper } | {"gripper": np.random.randn(batch_size, 1, 2).astype(np.float32)}, - "language": { + LANGUAGE: { LANGUAGE_KEY: [["pick up the apple"]] * batch_size, }, } @@ -139,6 +154,20 @@ def test_policy_embodiment_tag(self, policy): class TestGr00tPolicyCheckObservation: + def test_inference_requires_only_the_selected_language_key(self, policy): + policy.modality_configs[LANGUAGE].modality_keys.append("training_paraphrase") + observation = _make_observation() + policy.get_action(observation) + del observation[LANGUAGE][LANGUAGE_KEY] + with pytest.raises(AssertionError, match="Language key"): + policy.check_observation(observation) + + def test_extra_camera_cannot_be_silently_ignored(self, policy): + obs = _make_observation() + obs["video"]["second_external"] = obs["video"][VIDEO_KEYS[0]] + with pytest.raises(ValueError, match="Checkpoint cameras"): + policy.check_observation(obs) + def test_valid_observation_passes(self, policy): obs = _make_observation() policy.check_observation(obs) @@ -179,11 +208,12 @@ def __init__(self): modality_keys=["state"], ), "action": ModalityConfig(delta_indices=[0], modality_keys=["action"]), - "language": ModalityConfig( + LANGUAGE: ModalityConfig( delta_indices=[0], modality_keys=["annotation.human.action.task_description"], ), } + self.language_key = self.modality_configs[LANGUAGE].modality_keys[0] self.last_observation = None def get_modality_config(self): @@ -197,10 +227,12 @@ def reset(self, options=None): return {} -def test_sim_policy_wrapper_accepts_numpy_language_batches(): +@pytest.mark.parametrize("extra_language_keys", [[], ["training_paraphrase"]]) +def test_sim_policy_wrapper_accepts_numpy_language_batches(extra_language_keys): from gr00t.policy.gr00t_policy import Gr00tSimPolicyWrapper policy = _NumpyLanguageSimPolicy() + policy.modality_configs[LANGUAGE].modality_keys.extend(extra_language_keys) wrapper = Gr00tSimPolicyWrapper(policy) observation = { "video.camera": np.zeros((1, 1, 256, 256, 3), dtype=np.uint8), @@ -210,6 +242,23 @@ def test_sim_policy_wrapper_accepts_numpy_language_batches(): action, info = wrapper.get_action(observation) - assert policy.last_observation["language"][LANGUAGE_KEY] == [["follow the instruction"]] + assert policy.last_observation[LANGUAGE][LANGUAGE_KEY] == [["follow the instruction"]] + assert list(policy.last_observation[LANGUAGE]) == [policy.language_key] assert "action.action" in action assert info == {} + + +def test_sim_policy_wrapper_requires_selected_language_key(): + from gr00t.policy.gr00t_policy import Gr00tSimPolicyWrapper + + policy = _NumpyLanguageSimPolicy() + policy.modality_configs[LANGUAGE].modality_keys.append("training_paraphrase") + wrapper = Gr00tSimPolicyWrapper(policy) + observation = { + "video.camera": np.zeros((1, 1, 256, 256, 3), dtype=np.uint8), + "state.state": np.zeros((1, 1, 3), dtype=np.float32), + "training_paraphrase": ["follow the instruction"], + } + with pytest.raises(AssertionError, match="Language key .* must be in observation"): + wrapper.get_action(observation) + assert policy.last_observation is None