From c19d95a3ce1359771d3b59d594317302eaafd625 Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Fri, 3 Jul 2026 15:04:40 -0700 Subject: [PATCH 01/13] feat(groot): add OpenCV train-time color jitter --- .../policies/groot/configuration_groot.py | 6 + src/lerobot/policies/groot/processor_groot.py | 128 ++++++++++++++++-- .../groot/test_groot_train_color_jitter.py | 122 +++++++++++++++++ 3 files changed, 247 insertions(+), 9 deletions(-) create mode 100644 tests/policies/groot/test_groot_train_color_jitter.py diff --git a/src/lerobot/policies/groot/configuration_groot.py b/src/lerobot/policies/groot/configuration_groot.py index 97e08bb763e..03fa73cb81a 100644 --- a/src/lerobot/policies/groot/configuration_groot.py +++ b/src/lerobot/policies/groot/configuration_groot.py @@ -325,6 +325,12 @@ class GrootConfig(PreTrainedConfig): # Set to True only after installing a flash-attn build matching your torch/CUDA env. use_flash_attention: bool = False + # Optional OSS-compatible train-time image color jitter. Values are non-negative magnitudes for + # brightness, contrast, saturation, and hue (hue <= 0.5), e.g. + # {"brightness": 0.3, "contrast": 0.4, "saturation": 0.5, "hue": 0.08}. + # The GR00T preprocessor samples once per sample and replays across all timesteps/camera views. + color_jitter_params: dict[str, float] | None = None + # Enable GR00T-style state-relative action chunks (action chunk expressed relative to the current # observation state). use_relative_actions: bool = False diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 20b3518a306..18020300b46 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -1245,6 +1245,8 @@ def make_groot_pre_post_processors( crop_fraction = None use_albumentations = checkpoint_assets.use_albumentations if checkpoint_assets is not None else False letter_box_transform = checkpoint_assets.letter_box_transform if checkpoint_assets is not None else False + color_jitter_params = config.color_jitter_params + use_albumentations = use_albumentations or color_jitter_params is not None input_steps: list[ProcessorStep] = [ RenameObservationsProcessorStep(rename_map={}), @@ -1256,6 +1258,7 @@ def make_groot_pre_post_processors( image_target_size=image_target_size, shortest_image_edge=shortest_image_edge, crop_fraction=crop_fraction, + color_jitter_params=color_jitter_params, use_albumentations=use_albumentations, letter_box_transform=letter_box_transform, training=dataset_meta is not None, @@ -1361,6 +1364,105 @@ def _align_video_horizon(video: np.ndarray, horizon: int | None) -> np.ndarray: return np.concatenate([pad, video], axis=1) +def _sample_n1_7_color_jitter_params(magnitudes: dict[str, float]) -> dict[str, Any]: + """Sample parameters with the same ranges and RNG order as OSS A.ColorJitter.""" + + supported = {"brightness", "contrast", "saturation", "hue"} + unknown = set(magnitudes) - supported + if unknown: + raise ValueError(f"Unsupported GR00T color jitter parameters: {sorted(unknown)}") + + parsed: dict[str, float] = {} + for name in supported: + value = magnitudes.get(name, 0.0) + if isinstance(value, bool): + raise TypeError(f"GR00T color jitter '{name}' must be a float, got bool") + try: + amount = float(value) + except (TypeError, ValueError) as exc: + raise TypeError(f"GR00T color jitter '{name}' must be a float, got {value!r}") from exc + if not np.isfinite(amount) or amount < 0: + raise ValueError(f"GR00T color jitter '{name}' must be finite and non-negative, got {amount}") + if name == "hue" and amount > 0.5: + raise ValueError(f"GR00T color jitter 'hue' must be <= 0.5, got {amount}") + parsed[name] = amount + + brightness = parsed["brightness"] + contrast = parsed["contrast"] + saturation = parsed["saturation"] + hue = parsed["hue"] + sampled = { + "brightness": random.uniform(max(0.0, 1.0 - brightness), 1.0 + brightness), + "contrast": random.uniform(max(0.0, 1.0 - contrast), 1.0 + contrast), + "saturation": random.uniform(max(0.0, 1.0 - saturation), 1.0 + saturation), + "hue": random.uniform(-hue, hue), + } + order = np.arange(4) + # Albumentations 1.4.18 creates a NumPy RandomState from one Python-random + # seed for its operation shuffle. Reproduce that sequence without importing it. + np.random.RandomState(random.randint(0, (1 << 32) - 1)).shuffle(order) + sampled["order"] = order.tolist() + return sampled + + +def _n1_7_color_jitter_lut(image: np.ndarray, factor: float, value: float = 0.0) -> np.ndarray: + lut = np.arange(256, dtype=np.float32) * factor + value + return cv2.LUT(image, np.clip(lut, 0, 255).astype(np.uint8)) + + +def _apply_n1_7_color_jitter(image: np.ndarray, params: dict[str, Any]) -> np.ndarray: + """Apply sampled OSS ColorJitter parameters to one HWC RGB uint8 image.""" + + image = np.asarray(image) + if image.dtype != np.uint8 or image.ndim != 3 or image.shape[-1] != 3: + raise ValueError( + "GR00T OpenCV color jitter expects an HWC RGB uint8 image, " + f"got shape={image.shape}, dtype={image.dtype}" + ) + if not image.flags.c_contiguous: + image = np.ascontiguousarray(image) + + def adjust_brightness(frame: np.ndarray, factor: float) -> np.ndarray: + if factor == 0: + return np.zeros_like(frame) + if factor == 1: + return frame + return _n1_7_color_jitter_lut(frame, factor) + + def adjust_contrast(frame: np.ndarray, factor: float) -> np.ndarray: + if factor == 1: + return frame + mean = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY).mean() + if factor == 0: + return np.full_like(frame, int(mean + 0.5), dtype=frame.dtype) + return _n1_7_color_jitter_lut(frame, factor, mean * (1 - factor)) + + def adjust_saturation(frame: np.ndarray, factor: float) -> np.ndarray: + if factor == 1: + return frame + grayscale = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) + grayscale = cv2.cvtColor(grayscale, cv2.COLOR_GRAY2RGB) + if factor == 0: + return grayscale + adjusted = cv2.addWeighted(frame, factor, grayscale, 1 - factor, gamma=0) + return np.clip(adjusted, 0, 255).astype(frame.dtype) + + def adjust_hue(frame: np.ndarray, factor: float) -> np.ndarray: + if factor == 0: + return frame + hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV) + lut = np.arange(256, dtype=np.int16) + lut = np.mod(lut + 180 * factor, 180).astype(np.uint8) + hsv[..., 0] = cv2.LUT(hsv[..., 0], lut) + return cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB) + + transforms = (adjust_brightness, adjust_contrast, adjust_saturation, adjust_hue) + factors = (params["brightness"], params["contrast"], params["saturation"], params["hue"]) + for index in params["order"]: + image = transforms[index](image, factors[index]) + return image + + def _build_n1_7_processor(model_name: str = GROOT_N1_7_BACKBONE_MODEL) -> ProcessorMixin: require_package("transformers", extra="groot") tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) @@ -2051,6 +2153,7 @@ class GrootN17VLMEncodeStep(ProcessorStep): image_target_size: list[int] | None = None shortest_image_edge: int | None = None crop_fraction: float | None = None + color_jitter_params: dict[str, float] | None = None use_albumentations: bool = False letter_box_transform: bool = False # Runtime-only train/eval mode: True enables Isaac's train-time random crop @@ -2091,16 +2194,22 @@ def _build_sample_images( """ if self.use_albumentations: video_np = np.asarray(video) - train_crop = self.training and torch.is_grad_enabled() + train_augmentation = self.training and torch.is_grad_enabled() sample_images: list[list[Any]] = [] for batch_idx in range(batch_size): # Isaac-GR00T samples ONE crop window per sample and replays it # across every (timestep, view) frame of that sample, keeping # cross-view geometry consistent. Eval keeps the center crop. - crop_position = (random.random(), random.random()) if train_crop else None - sample_images.append( - [ - _transform_n1_7_image_for_vlm_albumentations( + crop_position = (random.random(), random.random()) if train_augmentation else None + jitter_params = ( + _sample_n1_7_color_jitter_params(self.color_jitter_params) + if train_augmentation and self.color_jitter_params is not None + else None + ) + frames: list[np.ndarray] = [] + for timestep in range(video_np.shape[1]): + for view_idx in range(video_np.shape[2]): + frame = _transform_n1_7_image_for_vlm_albumentations( video_np[batch_idx, timestep, view_idx], image_crop_size=self.image_crop_size, image_target_size=self.image_target_size, @@ -2109,10 +2218,10 @@ def _build_sample_images( letter_box_transform=self.letter_box_transform, crop_position=crop_position, ) - for timestep in range(video_np.shape[1]) - for view_idx in range(video_np.shape[2]) - ] - ) + if jitter_params is not None: + frame = _apply_n1_7_color_jitter(frame, jitter_params) + frames.append(frame) + sample_images.append(frames) return sample_images video_t = video if torch.is_tensor(video) else torch.from_numpy(np.ascontiguousarray(video)) @@ -2205,6 +2314,7 @@ def get_config(self) -> dict[str, Any]: "image_target_size": self.image_target_size, "shortest_image_edge": self.shortest_image_edge, "crop_fraction": self.crop_fraction, + "color_jitter_params": self.color_jitter_params, "use_albumentations": self.use_albumentations, "letter_box_transform": self.letter_box_transform, "device": self.device, diff --git a/tests/policies/groot/test_groot_train_color_jitter.py b/tests/policies/groot/test_groot_train_color_jitter.py new file mode 100644 index 00000000000..8a8b0034c58 --- /dev/null +++ b/tests/policies/groot/test_groot_train_color_jitter.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# 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. + +"""Isaac-GR00T N1.7 train-time color-jitter contract.""" + +import hashlib +import random + +import numpy as np +import torch + +from lerobot.policies.groot import processor_groot +from lerobot.policies.groot.processor_groot import ( + GrootN17VLMEncodeStep, + _apply_n1_7_color_jitter, + _sample_n1_7_color_jitter_params, +) + +# Generated with Isaac-GR00T's albumentations==1.4.18 A.ColorJitter using +# random.seed(1337), np.random.seed(1337), and the deterministic input below. +OSS_COLOR_JITTER_MAGNITUDES = { + "brightness": 0.3, + "contrast": 0.4, + "saturation": 0.5, + "hue": 0.08, +} +OSS_COLOR_JITTER_PARAMS = { + "brightness": 1.0706517141708825, + "contrast": 1.0266124588840007, + "saturation": 0.8658483592493755, + "hue": 0.01372597662436345, + "order": [0, 3, 1, 2], +} +OSS_INPUT_SHA256 = "df4bf2710fd2cafea9ca517db1a16850b31ac3a7b225da50e95e81f09b81b0bb" +OSS_OUTPUT_SHA256 = "27775d8567ebb38f764821f789c57e19247b8d35d19abd3be4f65d76f493b663" + + +def _make_oss_golden_input() -> np.ndarray: + shape = (2, 3, 37, 53) + values = np.arange(np.prod(shape), dtype=np.int64) + chw = ((values * 37 + 11) % 256).astype(np.uint8).reshape(shape) + return chw.transpose(0, 2, 3, 1) + + +def _sha256_chw(images: np.ndarray) -> str: + chw = np.ascontiguousarray(images.transpose(0, 3, 1, 2)) + return hashlib.sha256(chw.tobytes()).hexdigest() + + +def test_n1_7_opencv_color_jitter_matches_oss_golden_hash(): + images = _make_oss_golden_input() + assert _sha256_chw(images) == OSS_INPUT_SHA256 + + actual = np.stack([_apply_n1_7_color_jitter(image, OSS_COLOR_JITTER_PARAMS) for image in images]) + + assert actual.shape == images.shape + assert actual.dtype == np.uint8 + assert _sha256_chw(actual) == OSS_OUTPUT_SHA256 + + +def test_n1_7_color_jitter_sampling_matches_oss_seed_sequence(): + random.seed(1337) + + assert _sample_n1_7_color_jitter_params(OSS_COLOR_JITTER_MAGNITUDES) == OSS_COLOR_JITTER_PARAMS + + +def test_training_color_jitter_runs_in_vlm_preprocessor_and_eval_is_stable(monkeypatch): + images = _make_oss_golden_input() + video = images.reshape(1, 1, 2, *images.shape[1:]) + monkeypatch.setattr( + processor_groot, + "_sample_n1_7_color_jitter_params", + lambda _: OSS_COLOR_JITTER_PARAMS, + ) + + train_step = GrootN17VLMEncodeStep( + use_albumentations=True, + color_jitter_params=OSS_COLOR_JITTER_MAGNITUDES, + training=True, + ) + train_frames = train_step._build_sample_images(video, batch_size=1, target_device=None)[0] + assert _sha256_chw(np.stack(train_frames)) == OSS_OUTPUT_SHA256 + + eval_step = GrootN17VLMEncodeStep( + use_albumentations=True, + color_jitter_params=OSS_COLOR_JITTER_MAGNITUDES, + training=False, + ) + eval_frames = eval_step._build_sample_images(video, batch_size=1, target_device=None)[0] + assert _sha256_chw(np.stack(eval_frames)) == OSS_INPUT_SHA256 + + with torch.no_grad(): + no_grad_frames = train_step._build_sample_images(video, batch_size=1, target_device=None)[0] + assert _sha256_chw(np.stack(no_grad_frames)) == OSS_INPUT_SHA256 + + +def test_color_jitter_config_round_trips_but_training_mode_does_not(): + step = GrootN17VLMEncodeStep( + color_jitter_params=OSS_COLOR_JITTER_MAGNITUDES, + training=True, + ) + + serialized = step.get_config() + assert serialized["color_jitter_params"] == OSS_COLOR_JITTER_MAGNITUDES + assert "training" not in serialized + + restored = GrootN17VLMEncodeStep(**serialized) + assert restored.color_jitter_params == OSS_COLOR_JITTER_MAGNITUDES + assert restored.training is False From bf30f630219bd472de0298ff5d23d898b36a52bc Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Fri, 3 Jul 2026 15:36:56 -0700 Subject: [PATCH 02/13] refactor(transforms): run OpenCV jitter in data workers --- .../policies/groot/configuration_groot.py | 6 - src/lerobot/policies/groot/processor_groot.py | 128 ++------------- src/lerobot/transforms/__init__.py | 2 + src/lerobot/transforms/transforms.py | 153 +++++++++++++++++- tests/datasets/test_image_transforms.py | 97 +++++++++++ .../groot/test_groot_train_color_jitter.py | 122 -------------- 6 files changed, 260 insertions(+), 248 deletions(-) delete mode 100644 tests/policies/groot/test_groot_train_color_jitter.py diff --git a/src/lerobot/policies/groot/configuration_groot.py b/src/lerobot/policies/groot/configuration_groot.py index 03fa73cb81a..97e08bb763e 100644 --- a/src/lerobot/policies/groot/configuration_groot.py +++ b/src/lerobot/policies/groot/configuration_groot.py @@ -325,12 +325,6 @@ class GrootConfig(PreTrainedConfig): # Set to True only after installing a flash-attn build matching your torch/CUDA env. use_flash_attention: bool = False - # Optional OSS-compatible train-time image color jitter. Values are non-negative magnitudes for - # brightness, contrast, saturation, and hue (hue <= 0.5), e.g. - # {"brightness": 0.3, "contrast": 0.4, "saturation": 0.5, "hue": 0.08}. - # The GR00T preprocessor samples once per sample and replays across all timesteps/camera views. - color_jitter_params: dict[str, float] | None = None - # Enable GR00T-style state-relative action chunks (action chunk expressed relative to the current # observation state). use_relative_actions: bool = False diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 18020300b46..20b3518a306 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -1245,8 +1245,6 @@ def make_groot_pre_post_processors( crop_fraction = None use_albumentations = checkpoint_assets.use_albumentations if checkpoint_assets is not None else False letter_box_transform = checkpoint_assets.letter_box_transform if checkpoint_assets is not None else False - color_jitter_params = config.color_jitter_params - use_albumentations = use_albumentations or color_jitter_params is not None input_steps: list[ProcessorStep] = [ RenameObservationsProcessorStep(rename_map={}), @@ -1258,7 +1256,6 @@ def make_groot_pre_post_processors( image_target_size=image_target_size, shortest_image_edge=shortest_image_edge, crop_fraction=crop_fraction, - color_jitter_params=color_jitter_params, use_albumentations=use_albumentations, letter_box_transform=letter_box_transform, training=dataset_meta is not None, @@ -1364,105 +1361,6 @@ def _align_video_horizon(video: np.ndarray, horizon: int | None) -> np.ndarray: return np.concatenate([pad, video], axis=1) -def _sample_n1_7_color_jitter_params(magnitudes: dict[str, float]) -> dict[str, Any]: - """Sample parameters with the same ranges and RNG order as OSS A.ColorJitter.""" - - supported = {"brightness", "contrast", "saturation", "hue"} - unknown = set(magnitudes) - supported - if unknown: - raise ValueError(f"Unsupported GR00T color jitter parameters: {sorted(unknown)}") - - parsed: dict[str, float] = {} - for name in supported: - value = magnitudes.get(name, 0.0) - if isinstance(value, bool): - raise TypeError(f"GR00T color jitter '{name}' must be a float, got bool") - try: - amount = float(value) - except (TypeError, ValueError) as exc: - raise TypeError(f"GR00T color jitter '{name}' must be a float, got {value!r}") from exc - if not np.isfinite(amount) or amount < 0: - raise ValueError(f"GR00T color jitter '{name}' must be finite and non-negative, got {amount}") - if name == "hue" and amount > 0.5: - raise ValueError(f"GR00T color jitter 'hue' must be <= 0.5, got {amount}") - parsed[name] = amount - - brightness = parsed["brightness"] - contrast = parsed["contrast"] - saturation = parsed["saturation"] - hue = parsed["hue"] - sampled = { - "brightness": random.uniform(max(0.0, 1.0 - brightness), 1.0 + brightness), - "contrast": random.uniform(max(0.0, 1.0 - contrast), 1.0 + contrast), - "saturation": random.uniform(max(0.0, 1.0 - saturation), 1.0 + saturation), - "hue": random.uniform(-hue, hue), - } - order = np.arange(4) - # Albumentations 1.4.18 creates a NumPy RandomState from one Python-random - # seed for its operation shuffle. Reproduce that sequence without importing it. - np.random.RandomState(random.randint(0, (1 << 32) - 1)).shuffle(order) - sampled["order"] = order.tolist() - return sampled - - -def _n1_7_color_jitter_lut(image: np.ndarray, factor: float, value: float = 0.0) -> np.ndarray: - lut = np.arange(256, dtype=np.float32) * factor + value - return cv2.LUT(image, np.clip(lut, 0, 255).astype(np.uint8)) - - -def _apply_n1_7_color_jitter(image: np.ndarray, params: dict[str, Any]) -> np.ndarray: - """Apply sampled OSS ColorJitter parameters to one HWC RGB uint8 image.""" - - image = np.asarray(image) - if image.dtype != np.uint8 or image.ndim != 3 or image.shape[-1] != 3: - raise ValueError( - "GR00T OpenCV color jitter expects an HWC RGB uint8 image, " - f"got shape={image.shape}, dtype={image.dtype}" - ) - if not image.flags.c_contiguous: - image = np.ascontiguousarray(image) - - def adjust_brightness(frame: np.ndarray, factor: float) -> np.ndarray: - if factor == 0: - return np.zeros_like(frame) - if factor == 1: - return frame - return _n1_7_color_jitter_lut(frame, factor) - - def adjust_contrast(frame: np.ndarray, factor: float) -> np.ndarray: - if factor == 1: - return frame - mean = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY).mean() - if factor == 0: - return np.full_like(frame, int(mean + 0.5), dtype=frame.dtype) - return _n1_7_color_jitter_lut(frame, factor, mean * (1 - factor)) - - def adjust_saturation(frame: np.ndarray, factor: float) -> np.ndarray: - if factor == 1: - return frame - grayscale = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) - grayscale = cv2.cvtColor(grayscale, cv2.COLOR_GRAY2RGB) - if factor == 0: - return grayscale - adjusted = cv2.addWeighted(frame, factor, grayscale, 1 - factor, gamma=0) - return np.clip(adjusted, 0, 255).astype(frame.dtype) - - def adjust_hue(frame: np.ndarray, factor: float) -> np.ndarray: - if factor == 0: - return frame - hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV) - lut = np.arange(256, dtype=np.int16) - lut = np.mod(lut + 180 * factor, 180).astype(np.uint8) - hsv[..., 0] = cv2.LUT(hsv[..., 0], lut) - return cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB) - - transforms = (adjust_brightness, adjust_contrast, adjust_saturation, adjust_hue) - factors = (params["brightness"], params["contrast"], params["saturation"], params["hue"]) - for index in params["order"]: - image = transforms[index](image, factors[index]) - return image - - def _build_n1_7_processor(model_name: str = GROOT_N1_7_BACKBONE_MODEL) -> ProcessorMixin: require_package("transformers", extra="groot") tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) @@ -2153,7 +2051,6 @@ class GrootN17VLMEncodeStep(ProcessorStep): image_target_size: list[int] | None = None shortest_image_edge: int | None = None crop_fraction: float | None = None - color_jitter_params: dict[str, float] | None = None use_albumentations: bool = False letter_box_transform: bool = False # Runtime-only train/eval mode: True enables Isaac's train-time random crop @@ -2194,22 +2091,16 @@ def _build_sample_images( """ if self.use_albumentations: video_np = np.asarray(video) - train_augmentation = self.training and torch.is_grad_enabled() + train_crop = self.training and torch.is_grad_enabled() sample_images: list[list[Any]] = [] for batch_idx in range(batch_size): # Isaac-GR00T samples ONE crop window per sample and replays it # across every (timestep, view) frame of that sample, keeping # cross-view geometry consistent. Eval keeps the center crop. - crop_position = (random.random(), random.random()) if train_augmentation else None - jitter_params = ( - _sample_n1_7_color_jitter_params(self.color_jitter_params) - if train_augmentation and self.color_jitter_params is not None - else None - ) - frames: list[np.ndarray] = [] - for timestep in range(video_np.shape[1]): - for view_idx in range(video_np.shape[2]): - frame = _transform_n1_7_image_for_vlm_albumentations( + crop_position = (random.random(), random.random()) if train_crop else None + sample_images.append( + [ + _transform_n1_7_image_for_vlm_albumentations( video_np[batch_idx, timestep, view_idx], image_crop_size=self.image_crop_size, image_target_size=self.image_target_size, @@ -2218,10 +2109,10 @@ def _build_sample_images( letter_box_transform=self.letter_box_transform, crop_position=crop_position, ) - if jitter_params is not None: - frame = _apply_n1_7_color_jitter(frame, jitter_params) - frames.append(frame) - sample_images.append(frames) + for timestep in range(video_np.shape[1]) + for view_idx in range(video_np.shape[2]) + ] + ) return sample_images video_t = video if torch.is_tensor(video) else torch.from_numpy(np.ascontiguousarray(video)) @@ -2314,7 +2205,6 @@ def get_config(self) -> dict[str, Any]: "image_target_size": self.image_target_size, "shortest_image_edge": self.shortest_image_edge, "crop_fraction": self.crop_fraction, - "color_jitter_params": self.color_jitter_params, "use_albumentations": self.use_albumentations, "letter_box_transform": self.letter_box_transform, "device": self.device, diff --git a/src/lerobot/transforms/__init__.py b/src/lerobot/transforms/__init__.py index 6cf9699d0b1..262bf93e1b8 100644 --- a/src/lerobot/transforms/__init__.py +++ b/src/lerobot/transforms/__init__.py @@ -16,6 +16,7 @@ ImageTransformConfig, ImageTransforms, ImageTransformsConfig, + OpenCVColorJitter, RandomSubsetApply, SharpnessJitter, make_transform_from_config, @@ -25,6 +26,7 @@ "ImageTransformConfig", "ImageTransforms", "ImageTransformsConfig", + "OpenCVColorJitter", "RandomSubsetApply", "SharpnessJitter", "make_transform_from_config", diff --git a/src/lerobot/transforms/transforms.py b/src/lerobot/transforms/transforms.py index 5240619cb6c..c4c1f195403 100644 --- a/src/lerobot/transforms/transforms.py +++ b/src/lerobot/transforms/transforms.py @@ -14,10 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. import collections +import random from collections.abc import Callable, Sequence from dataclasses import dataclass, field from typing import Any +import cv2 +import numpy as np import torch from torchvision.transforms import v2 from torchvision.transforms.v2 import ( @@ -219,13 +222,16 @@ def make_transform_from_config(cfg: ImageTransformConfig): if cfg.type == "SharpnessJitter": return SharpnessJitter(**cfg.kwargs) + if cfg.type == "OpenCVColorJitter": + return OpenCVColorJitter(**cfg.kwargs) + transform_cls = getattr(v2, cfg.type, None) if isinstance(transform_cls, type) and issubclass(transform_cls, Transform): return transform_cls(**cfg.kwargs) raise ValueError( f"Transform '{cfg.type}' is not valid. It must be a class in " - f"torchvision.transforms.v2 or 'SharpnessJitter'." + "torchvision.transforms.v2 or one of: 'OpenCVColorJitter', 'SharpnessJitter'." ) @@ -258,3 +264,148 @@ def __init__(self, cfg: ImageTransformsConfig) -> None: def forward(self, *inputs: Any) -> Any: return self.tf(*inputs) + + +class OpenCVColorJitter(Transform): + """Apply Isaac-GR00T/Albumentations-compatible color jitter with OpenCV. + + The four arguments are non-negative jitter magnitudes. Brightness, + contrast, and saturation factors are sampled uniformly from + ``[max(0, 1 - magnitude), 1 + magnitude]``; hue is sampled uniformly from + ``[-hue, hue]`` and must not exceed ``0.5``. + + A single set of factors and one random operation order are sampled per + call and applied across every frame in the input. Inputs must be CPU + ``torch.uint8`` tensors shaped ``(..., 3, H, W)``. This keeps temporal + frames photometrically consistent while retaining the byte-level behavior + of Albumentations 1.4.18 on RGB uint8 images. + """ + + def __init__( + self, + brightness: float = 0.0, + contrast: float = 0.0, + saturation: float = 0.0, + hue: float = 0.0, + ) -> None: + super().__init__() + self.brightness = self._check_magnitude("brightness", brightness) + self.contrast = self._check_magnitude("contrast", contrast) + self.saturation = self._check_magnitude("saturation", saturation) + self.hue = self._check_magnitude("hue", hue, maximum=0.5) + + @staticmethod + def _check_magnitude(name: str, value: float, maximum: float | None = None) -> float: + try: + amount = float(value) + except (TypeError, ValueError) as exc: + raise TypeError(f"OpenCVColorJitter {name} must be a float, got {value!r}") from exc + if not np.isfinite(amount) or amount < 0: + raise ValueError(f"OpenCVColorJitter {name} must be finite and non-negative, got {amount}") + if maximum is not None and amount > maximum: + raise ValueError(f"OpenCVColorJitter {name} must be <= {maximum}, got {amount}") + return amount + + def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: + del flat_inputs + params: dict[str, Any] = { + "brightness": random.uniform(max(0.0, 1.0 - self.brightness), 1.0 + self.brightness), + "contrast": random.uniform(max(0.0, 1.0 - self.contrast), 1.0 + self.contrast), + "saturation": random.uniform(max(0.0, 1.0 - self.saturation), 1.0 + self.saturation), + "hue": random.uniform(-self.hue, self.hue), + } + order = np.arange(4) + # Albumentations 1.4.18 seeds a NumPy RandomState from Python random + # before shuffling the operation order. Preserve that exact sequence. + np.random.RandomState(random.randint(0, (1 << 32) - 1)).shuffle(order) + params["order"] = order.tolist() + return params + + @staticmethod + def _lut(image: np.ndarray, factor: float, value: float = 0.0) -> np.ndarray: + lut = np.arange(256, dtype=np.float32) * factor + value + return cv2.LUT(image, np.clip(lut, 0, 255).astype(np.uint8)) + + @classmethod + def apply_rgb_image(cls, image: np.ndarray, params: dict[str, Any]) -> np.ndarray: + """Apply already-sampled parameters to one HWC RGB uint8 image.""" + + image = np.asarray(image) + if image.dtype != np.uint8 or image.ndim != 3 or image.shape[-1] != 3: + raise ValueError( + "OpenCVColorJitter expects an HWC RGB uint8 image, " + f"got shape={image.shape}, dtype={image.dtype}" + ) + if not image.flags.c_contiguous: + image = np.ascontiguousarray(image) + + def adjust_brightness(frame: np.ndarray, factor: float) -> np.ndarray: + if factor == 0: + return np.zeros_like(frame) + if factor == 1: + return frame + return cls._lut(frame, factor) + + def adjust_contrast(frame: np.ndarray, factor: float) -> np.ndarray: + if factor == 1: + return frame + mean = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY).mean() + if factor == 0: + return np.full_like(frame, int(mean + 0.5), dtype=frame.dtype) + return cls._lut(frame, factor, mean * (1 - factor)) + + def adjust_saturation(frame: np.ndarray, factor: float) -> np.ndarray: + if factor == 1: + return frame + grayscale = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) + grayscale = cv2.cvtColor(grayscale, cv2.COLOR_GRAY2RGB) + if factor == 0: + return grayscale + adjusted = cv2.addWeighted(frame, factor, grayscale, 1 - factor, gamma=0) + return np.clip(adjusted, 0, 255).astype(frame.dtype) + + def adjust_hue(frame: np.ndarray, factor: float) -> np.ndarray: + if factor == 0: + return frame + hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV) + lut = np.arange(256, dtype=np.int16) + lut = np.mod(lut + 180 * factor, 180).astype(np.uint8) + hsv[..., 0] = cv2.LUT(hsv[..., 0], lut) + return cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB) + + transforms = (adjust_brightness, adjust_contrast, adjust_saturation, adjust_hue) + factors = ( + params["brightness"], + params["contrast"], + params["saturation"], + params["hue"], + ) + for index in params["order"]: + image = transforms[index](image, factors[index]) + return image + + def transform(self, inpt: Any, params: dict[str, Any]) -> torch.Tensor: + if not torch.is_tensor(inpt): + raise TypeError(f"OpenCVColorJitter expects torch.Tensor inputs, got {type(inpt).__name__}") + if inpt.device.type != "cpu": + raise ValueError(f"OpenCVColorJitter expects CPU tensors, got device={inpt.device}") + if inpt.dtype != torch.uint8: + raise ValueError(f"OpenCVColorJitter expects uint8 tensors, got dtype={inpt.dtype}") + if inpt.ndim < 3 or inpt.shape[-3] != 3: + raise ValueError( + "OpenCVColorJitter expects CHW RGB tensors shaped (..., 3, H, W), " + f"got shape={tuple(inpt.shape)}" + ) + if inpt.numel() == 0: + return inpt.clone() + + height, width = inpt.shape[-2:] + frames = inpt.detach().contiguous().reshape(-1, 3, height, width).permute(0, 2, 3, 1).numpy() + transformed = np.stack([self.apply_rgb_image(frame, params) for frame in frames]) + return torch.from_numpy(transformed).permute(0, 3, 1, 2).reshape(inpt.shape).contiguous() + + def extra_repr(self) -> str: + return ( + f"brightness={self.brightness}, contrast={self.contrast}, " + f"saturation={self.saturation}, hue={self.hue}" + ) diff --git a/tests/datasets/test_image_transforms.py b/tests/datasets/test_image_transforms.py index 4310274e4b8..6dc36ba4263 100644 --- a/tests/datasets/test_image_transforms.py +++ b/tests/datasets/test_image_transforms.py @@ -14,6 +14,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import hashlib +import random + +import numpy as np import pytest import torch from packaging import version @@ -31,6 +35,7 @@ ImageTransformConfig, ImageTransforms, ImageTransformsConfig, + OpenCVColorJitter, RandomSubsetApply, SharpnessJitter, make_transform_from_config, @@ -392,6 +397,98 @@ def test_sharpness_jitter_invalid_range_max_smaller(): SharpnessJitter((2.0, 0.1)) +OSS_COLOR_JITTER_MAGNITUDES = { + "brightness": 0.3, + "contrast": 0.4, + "saturation": 0.5, + "hue": 0.08, +} +OSS_COLOR_JITTER_PARAMS = { + "brightness": 1.0706517141708825, + "contrast": 1.0266124588840007, + "saturation": 0.8658483592493755, + "hue": 0.01372597662436345, + "order": [0, 3, 1, 2], +} +OSS_COLOR_JITTER_OUTPUT_SHA256 = "27775d8567ebb38f764821f789c57e19247b8d35d19abd3be4f65d76f493b663" + + +def _make_oss_color_jitter_input() -> torch.Tensor: + shape = (2, 3, 37, 53) + values = np.arange(np.prod(shape), dtype=np.int64) + return torch.from_numpy(((values * 37 + 11) % 256).astype(np.uint8).reshape(shape)) + + +def test_opencv_color_jitter_sampling_matches_oss_seed_sequence(): + transform = OpenCVColorJitter(**OSS_COLOR_JITTER_MAGNITUDES) + random.seed(1337) + + assert transform.make_params([]) == OSS_COLOR_JITTER_PARAMS + + +def test_opencv_color_jitter_matches_oss_golden_hash_and_preserves_layout(): + images = _make_oss_color_jitter_input() + transform = OpenCVColorJitter(**OSS_COLOR_JITTER_MAGNITUDES) + random.seed(1337) + + actual = transform(images) + + assert actual.shape == images.shape + assert actual.dtype == torch.uint8 + assert actual.is_contiguous() + assert hashlib.sha256(actual.numpy().tobytes()).hexdigest() == OSS_COLOR_JITTER_OUTPUT_SHA256 + + +def test_opencv_color_jitter_reuses_one_sample_across_temporal_frames(): + image = _make_oss_color_jitter_input()[:1] + images = image.repeat(2, 1, 1, 1) + transform = OpenCVColorJitter(**OSS_COLOR_JITTER_MAGNITUDES) + random.seed(1337) + + first, second = transform(images) + + torch.testing.assert_close(first, second) + + +def test_opencv_color_jitter_is_available_from_image_transform_config(): + cfg = ImageTransformsConfig( + enable=True, + max_num_transforms=1, + tfs={ + "color_jitter": ImageTransformConfig( + type="OpenCVColorJitter", + kwargs=OSS_COLOR_JITTER_MAGNITUDES, + ) + }, + ) + + transform = ImageTransforms(cfg) + + assert isinstance(transform.transforms["color_jitter"], OpenCVColorJitter) + + +@pytest.mark.parametrize( + "kwargs,match", + [ + ({"brightness": -0.1}, "brightness"), + ({"contrast": float("inf")}, "contrast"), + ({"hue": 0.6}, "hue"), + ], +) +def test_opencv_color_jitter_rejects_invalid_magnitudes(kwargs, match): + with pytest.raises(ValueError, match=match): + OpenCVColorJitter(**kwargs) + + +def test_opencv_color_jitter_requires_cpu_chw_rgb_uint8(): + transform = OpenCVColorJitter() + + with pytest.raises(ValueError, match="uint8"): + transform(torch.zeros(3, 8, 8)) + with pytest.raises(ValueError, match="CHW RGB"): + transform(torch.zeros(1, 8, 8, dtype=torch.uint8)) + + def test_make_transform_from_config_with_v2_resize(img_tensor_factory): img_tensor = img_tensor_factory() tf_cfg = ImageTransformConfig(type="Resize", kwargs={"size": (32, 32)}) diff --git a/tests/policies/groot/test_groot_train_color_jitter.py b/tests/policies/groot/test_groot_train_color_jitter.py deleted file mode 100644 index 8a8b0034c58..00000000000 --- a/tests/policies/groot/test_groot_train_color_jitter.py +++ /dev/null @@ -1,122 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# 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. - -"""Isaac-GR00T N1.7 train-time color-jitter contract.""" - -import hashlib -import random - -import numpy as np -import torch - -from lerobot.policies.groot import processor_groot -from lerobot.policies.groot.processor_groot import ( - GrootN17VLMEncodeStep, - _apply_n1_7_color_jitter, - _sample_n1_7_color_jitter_params, -) - -# Generated with Isaac-GR00T's albumentations==1.4.18 A.ColorJitter using -# random.seed(1337), np.random.seed(1337), and the deterministic input below. -OSS_COLOR_JITTER_MAGNITUDES = { - "brightness": 0.3, - "contrast": 0.4, - "saturation": 0.5, - "hue": 0.08, -} -OSS_COLOR_JITTER_PARAMS = { - "brightness": 1.0706517141708825, - "contrast": 1.0266124588840007, - "saturation": 0.8658483592493755, - "hue": 0.01372597662436345, - "order": [0, 3, 1, 2], -} -OSS_INPUT_SHA256 = "df4bf2710fd2cafea9ca517db1a16850b31ac3a7b225da50e95e81f09b81b0bb" -OSS_OUTPUT_SHA256 = "27775d8567ebb38f764821f789c57e19247b8d35d19abd3be4f65d76f493b663" - - -def _make_oss_golden_input() -> np.ndarray: - shape = (2, 3, 37, 53) - values = np.arange(np.prod(shape), dtype=np.int64) - chw = ((values * 37 + 11) % 256).astype(np.uint8).reshape(shape) - return chw.transpose(0, 2, 3, 1) - - -def _sha256_chw(images: np.ndarray) -> str: - chw = np.ascontiguousarray(images.transpose(0, 3, 1, 2)) - return hashlib.sha256(chw.tobytes()).hexdigest() - - -def test_n1_7_opencv_color_jitter_matches_oss_golden_hash(): - images = _make_oss_golden_input() - assert _sha256_chw(images) == OSS_INPUT_SHA256 - - actual = np.stack([_apply_n1_7_color_jitter(image, OSS_COLOR_JITTER_PARAMS) for image in images]) - - assert actual.shape == images.shape - assert actual.dtype == np.uint8 - assert _sha256_chw(actual) == OSS_OUTPUT_SHA256 - - -def test_n1_7_color_jitter_sampling_matches_oss_seed_sequence(): - random.seed(1337) - - assert _sample_n1_7_color_jitter_params(OSS_COLOR_JITTER_MAGNITUDES) == OSS_COLOR_JITTER_PARAMS - - -def test_training_color_jitter_runs_in_vlm_preprocessor_and_eval_is_stable(monkeypatch): - images = _make_oss_golden_input() - video = images.reshape(1, 1, 2, *images.shape[1:]) - monkeypatch.setattr( - processor_groot, - "_sample_n1_7_color_jitter_params", - lambda _: OSS_COLOR_JITTER_PARAMS, - ) - - train_step = GrootN17VLMEncodeStep( - use_albumentations=True, - color_jitter_params=OSS_COLOR_JITTER_MAGNITUDES, - training=True, - ) - train_frames = train_step._build_sample_images(video, batch_size=1, target_device=None)[0] - assert _sha256_chw(np.stack(train_frames)) == OSS_OUTPUT_SHA256 - - eval_step = GrootN17VLMEncodeStep( - use_albumentations=True, - color_jitter_params=OSS_COLOR_JITTER_MAGNITUDES, - training=False, - ) - eval_frames = eval_step._build_sample_images(video, batch_size=1, target_device=None)[0] - assert _sha256_chw(np.stack(eval_frames)) == OSS_INPUT_SHA256 - - with torch.no_grad(): - no_grad_frames = train_step._build_sample_images(video, batch_size=1, target_device=None)[0] - assert _sha256_chw(np.stack(no_grad_frames)) == OSS_INPUT_SHA256 - - -def test_color_jitter_config_round_trips_but_training_mode_does_not(): - step = GrootN17VLMEncodeStep( - color_jitter_params=OSS_COLOR_JITTER_MAGNITUDES, - training=True, - ) - - serialized = step.get_config() - assert serialized["color_jitter_params"] == OSS_COLOR_JITTER_MAGNITUDES - assert "training" not in serialized - - restored = GrootN17VLMEncodeStep(**serialized) - assert restored.color_jitter_params == OSS_COLOR_JITTER_MAGNITUDES - assert restored.training is False From e19f22b11e2879228f2f330474e490fbfdf61dc8 Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Sat, 4 Jul 2026 06:02:23 -0700 Subject: [PATCH 03/13] refactor(groot): rely on dataset image transforms --- src/lerobot/policies/groot/processor_groot.py | 30 +--- src/lerobot/transforms/__init__.py | 2 - src/lerobot/transforms/transforms.py | 153 +--------------- tests/datasets/test_image_transforms.py | 97 ---------- tests/policies/groot/test_groot_n1_7.py | 14 ++ .../groot/test_groot_train_random_crop.py | 169 ------------------ 6 files changed, 18 insertions(+), 447 deletions(-) delete mode 100644 tests/policies/groot/test_groot_train_random_crop.py diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 20b3518a306..5ae4346118b 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -463,7 +463,7 @@ def _set_groot_preprocessor_training( """Set the runtime-only mode of GR00T stochastic processor steps. Any dataclass step exposing a ``training`` field participates, so processor - steps can opt into train-time-only behavior (dropout, augmentation) without + steps can opt into train-time-only behavior (for example, dropout) without this helper enumerating them. """ for step in preprocessor.steps: @@ -1258,7 +1258,6 @@ def make_groot_pre_post_processors( crop_fraction=crop_fraction, use_albumentations=use_albumentations, letter_box_transform=letter_box_transform, - training=dataset_meta is not None, device=config.device, ), DeviceProcessorStep(device=config.device), @@ -1384,7 +1383,6 @@ def _transform_n1_7_image_for_vlm_albumentations( shortest_image_edge: int | None, crop_fraction: float | None, letter_box_transform: bool = False, - crop_position: tuple[float, float] | None = None, ) -> np.ndarray: """cv2/INTER_AREA eval transform mirroring Isaac-GR00T's albumentations preprocessing. @@ -1394,12 +1392,6 @@ def _transform_n1_7_image_for_vlm_albumentations( cv2/INTER_AREA resize and floored center-crop here intentionally differ from that torch path and must stay bit-exact to the upstream reference. The hot path accepts and returns numpy arrays to avoid per-frame PIL round-trips. - - ``crop_position`` selects where the ``crop_fraction`` window sits: ``None`` - keeps the deterministic center crop (eval contract), while ``(y, x)`` - fractions in [0, 1] place the window for Isaac's train-time random crop - (0.5, 0.5 == center). Training samples one position per sample and reuses - it across camera views. """ if image_target_size is None: return image @@ -1451,13 +1443,8 @@ def resize_shortest_edge(frame: np.ndarray) -> np.ndarray: height, width = image_np.shape[:2] crop_h = max(1, int(height * crop_fraction)) crop_w = max(1, int(width * crop_fraction)) - if crop_position is None: - top = max(0, (height - crop_h) // 2) - left = max(0, (width - crop_w) // 2) - else: - pos_y, pos_x = crop_position - top = int(round((height - crop_h) * min(max(pos_y, 0.0), 1.0))) - left = int(round((width - crop_w) * min(max(pos_x, 0.0), 1.0))) + top = max(0, (height - crop_h) // 2) + left = max(0, (width - crop_w) // 2) image_np = image_np[top : top + crop_h, left : left + crop_w] return resize_shortest_edge(image_np) @@ -2053,11 +2040,6 @@ class GrootN17VLMEncodeStep(ProcessorStep): crop_fraction: float | None = None use_albumentations: bool = False letter_box_transform: bool = False - # Runtime-only train/eval mode: True enables Isaac's train-time random crop - # (one window per sample, replayed across views); False keeps the - # deterministic center crop. Never serialized - reloaded pipelines default - # to eval and are re-enabled only when processors are built with dataset_meta. - training: bool = False device: str | None = None _proc: ProcessorMixin | None = field(default=None, init=False, repr=False) @@ -2091,13 +2073,8 @@ def _build_sample_images( """ if self.use_albumentations: video_np = np.asarray(video) - train_crop = self.training and torch.is_grad_enabled() sample_images: list[list[Any]] = [] for batch_idx in range(batch_size): - # Isaac-GR00T samples ONE crop window per sample and replays it - # across every (timestep, view) frame of that sample, keeping - # cross-view geometry consistent. Eval keeps the center crop. - crop_position = (random.random(), random.random()) if train_crop else None sample_images.append( [ _transform_n1_7_image_for_vlm_albumentations( @@ -2107,7 +2084,6 @@ def _build_sample_images( shortest_image_edge=self.shortest_image_edge, crop_fraction=self.crop_fraction, letter_box_transform=self.letter_box_transform, - crop_position=crop_position, ) for timestep in range(video_np.shape[1]) for view_idx in range(video_np.shape[2]) diff --git a/src/lerobot/transforms/__init__.py b/src/lerobot/transforms/__init__.py index 262bf93e1b8..6cf9699d0b1 100644 --- a/src/lerobot/transforms/__init__.py +++ b/src/lerobot/transforms/__init__.py @@ -16,7 +16,6 @@ ImageTransformConfig, ImageTransforms, ImageTransformsConfig, - OpenCVColorJitter, RandomSubsetApply, SharpnessJitter, make_transform_from_config, @@ -26,7 +25,6 @@ "ImageTransformConfig", "ImageTransforms", "ImageTransformsConfig", - "OpenCVColorJitter", "RandomSubsetApply", "SharpnessJitter", "make_transform_from_config", diff --git a/src/lerobot/transforms/transforms.py b/src/lerobot/transforms/transforms.py index c4c1f195403..5240619cb6c 100644 --- a/src/lerobot/transforms/transforms.py +++ b/src/lerobot/transforms/transforms.py @@ -14,13 +14,10 @@ # See the License for the specific language governing permissions and # limitations under the License. import collections -import random from collections.abc import Callable, Sequence from dataclasses import dataclass, field from typing import Any -import cv2 -import numpy as np import torch from torchvision.transforms import v2 from torchvision.transforms.v2 import ( @@ -222,16 +219,13 @@ def make_transform_from_config(cfg: ImageTransformConfig): if cfg.type == "SharpnessJitter": return SharpnessJitter(**cfg.kwargs) - if cfg.type == "OpenCVColorJitter": - return OpenCVColorJitter(**cfg.kwargs) - transform_cls = getattr(v2, cfg.type, None) if isinstance(transform_cls, type) and issubclass(transform_cls, Transform): return transform_cls(**cfg.kwargs) raise ValueError( f"Transform '{cfg.type}' is not valid. It must be a class in " - "torchvision.transforms.v2 or one of: 'OpenCVColorJitter', 'SharpnessJitter'." + f"torchvision.transforms.v2 or 'SharpnessJitter'." ) @@ -264,148 +258,3 @@ def __init__(self, cfg: ImageTransformsConfig) -> None: def forward(self, *inputs: Any) -> Any: return self.tf(*inputs) - - -class OpenCVColorJitter(Transform): - """Apply Isaac-GR00T/Albumentations-compatible color jitter with OpenCV. - - The four arguments are non-negative jitter magnitudes. Brightness, - contrast, and saturation factors are sampled uniformly from - ``[max(0, 1 - magnitude), 1 + magnitude]``; hue is sampled uniformly from - ``[-hue, hue]`` and must not exceed ``0.5``. - - A single set of factors and one random operation order are sampled per - call and applied across every frame in the input. Inputs must be CPU - ``torch.uint8`` tensors shaped ``(..., 3, H, W)``. This keeps temporal - frames photometrically consistent while retaining the byte-level behavior - of Albumentations 1.4.18 on RGB uint8 images. - """ - - def __init__( - self, - brightness: float = 0.0, - contrast: float = 0.0, - saturation: float = 0.0, - hue: float = 0.0, - ) -> None: - super().__init__() - self.brightness = self._check_magnitude("brightness", brightness) - self.contrast = self._check_magnitude("contrast", contrast) - self.saturation = self._check_magnitude("saturation", saturation) - self.hue = self._check_magnitude("hue", hue, maximum=0.5) - - @staticmethod - def _check_magnitude(name: str, value: float, maximum: float | None = None) -> float: - try: - amount = float(value) - except (TypeError, ValueError) as exc: - raise TypeError(f"OpenCVColorJitter {name} must be a float, got {value!r}") from exc - if not np.isfinite(amount) or amount < 0: - raise ValueError(f"OpenCVColorJitter {name} must be finite and non-negative, got {amount}") - if maximum is not None and amount > maximum: - raise ValueError(f"OpenCVColorJitter {name} must be <= {maximum}, got {amount}") - return amount - - def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: - del flat_inputs - params: dict[str, Any] = { - "brightness": random.uniform(max(0.0, 1.0 - self.brightness), 1.0 + self.brightness), - "contrast": random.uniform(max(0.0, 1.0 - self.contrast), 1.0 + self.contrast), - "saturation": random.uniform(max(0.0, 1.0 - self.saturation), 1.0 + self.saturation), - "hue": random.uniform(-self.hue, self.hue), - } - order = np.arange(4) - # Albumentations 1.4.18 seeds a NumPy RandomState from Python random - # before shuffling the operation order. Preserve that exact sequence. - np.random.RandomState(random.randint(0, (1 << 32) - 1)).shuffle(order) - params["order"] = order.tolist() - return params - - @staticmethod - def _lut(image: np.ndarray, factor: float, value: float = 0.0) -> np.ndarray: - lut = np.arange(256, dtype=np.float32) * factor + value - return cv2.LUT(image, np.clip(lut, 0, 255).astype(np.uint8)) - - @classmethod - def apply_rgb_image(cls, image: np.ndarray, params: dict[str, Any]) -> np.ndarray: - """Apply already-sampled parameters to one HWC RGB uint8 image.""" - - image = np.asarray(image) - if image.dtype != np.uint8 or image.ndim != 3 or image.shape[-1] != 3: - raise ValueError( - "OpenCVColorJitter expects an HWC RGB uint8 image, " - f"got shape={image.shape}, dtype={image.dtype}" - ) - if not image.flags.c_contiguous: - image = np.ascontiguousarray(image) - - def adjust_brightness(frame: np.ndarray, factor: float) -> np.ndarray: - if factor == 0: - return np.zeros_like(frame) - if factor == 1: - return frame - return cls._lut(frame, factor) - - def adjust_contrast(frame: np.ndarray, factor: float) -> np.ndarray: - if factor == 1: - return frame - mean = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY).mean() - if factor == 0: - return np.full_like(frame, int(mean + 0.5), dtype=frame.dtype) - return cls._lut(frame, factor, mean * (1 - factor)) - - def adjust_saturation(frame: np.ndarray, factor: float) -> np.ndarray: - if factor == 1: - return frame - grayscale = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) - grayscale = cv2.cvtColor(grayscale, cv2.COLOR_GRAY2RGB) - if factor == 0: - return grayscale - adjusted = cv2.addWeighted(frame, factor, grayscale, 1 - factor, gamma=0) - return np.clip(adjusted, 0, 255).astype(frame.dtype) - - def adjust_hue(frame: np.ndarray, factor: float) -> np.ndarray: - if factor == 0: - return frame - hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV) - lut = np.arange(256, dtype=np.int16) - lut = np.mod(lut + 180 * factor, 180).astype(np.uint8) - hsv[..., 0] = cv2.LUT(hsv[..., 0], lut) - return cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB) - - transforms = (adjust_brightness, adjust_contrast, adjust_saturation, adjust_hue) - factors = ( - params["brightness"], - params["contrast"], - params["saturation"], - params["hue"], - ) - for index in params["order"]: - image = transforms[index](image, factors[index]) - return image - - def transform(self, inpt: Any, params: dict[str, Any]) -> torch.Tensor: - if not torch.is_tensor(inpt): - raise TypeError(f"OpenCVColorJitter expects torch.Tensor inputs, got {type(inpt).__name__}") - if inpt.device.type != "cpu": - raise ValueError(f"OpenCVColorJitter expects CPU tensors, got device={inpt.device}") - if inpt.dtype != torch.uint8: - raise ValueError(f"OpenCVColorJitter expects uint8 tensors, got dtype={inpt.dtype}") - if inpt.ndim < 3 or inpt.shape[-3] != 3: - raise ValueError( - "OpenCVColorJitter expects CHW RGB tensors shaped (..., 3, H, W), " - f"got shape={tuple(inpt.shape)}" - ) - if inpt.numel() == 0: - return inpt.clone() - - height, width = inpt.shape[-2:] - frames = inpt.detach().contiguous().reshape(-1, 3, height, width).permute(0, 2, 3, 1).numpy() - transformed = np.stack([self.apply_rgb_image(frame, params) for frame in frames]) - return torch.from_numpy(transformed).permute(0, 3, 1, 2).reshape(inpt.shape).contiguous() - - def extra_repr(self) -> str: - return ( - f"brightness={self.brightness}, contrast={self.contrast}, " - f"saturation={self.saturation}, hue={self.hue}" - ) diff --git a/tests/datasets/test_image_transforms.py b/tests/datasets/test_image_transforms.py index 6dc36ba4263..4310274e4b8 100644 --- a/tests/datasets/test_image_transforms.py +++ b/tests/datasets/test_image_transforms.py @@ -14,10 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import hashlib -import random - -import numpy as np import pytest import torch from packaging import version @@ -35,7 +31,6 @@ ImageTransformConfig, ImageTransforms, ImageTransformsConfig, - OpenCVColorJitter, RandomSubsetApply, SharpnessJitter, make_transform_from_config, @@ -397,98 +392,6 @@ def test_sharpness_jitter_invalid_range_max_smaller(): SharpnessJitter((2.0, 0.1)) -OSS_COLOR_JITTER_MAGNITUDES = { - "brightness": 0.3, - "contrast": 0.4, - "saturation": 0.5, - "hue": 0.08, -} -OSS_COLOR_JITTER_PARAMS = { - "brightness": 1.0706517141708825, - "contrast": 1.0266124588840007, - "saturation": 0.8658483592493755, - "hue": 0.01372597662436345, - "order": [0, 3, 1, 2], -} -OSS_COLOR_JITTER_OUTPUT_SHA256 = "27775d8567ebb38f764821f789c57e19247b8d35d19abd3be4f65d76f493b663" - - -def _make_oss_color_jitter_input() -> torch.Tensor: - shape = (2, 3, 37, 53) - values = np.arange(np.prod(shape), dtype=np.int64) - return torch.from_numpy(((values * 37 + 11) % 256).astype(np.uint8).reshape(shape)) - - -def test_opencv_color_jitter_sampling_matches_oss_seed_sequence(): - transform = OpenCVColorJitter(**OSS_COLOR_JITTER_MAGNITUDES) - random.seed(1337) - - assert transform.make_params([]) == OSS_COLOR_JITTER_PARAMS - - -def test_opencv_color_jitter_matches_oss_golden_hash_and_preserves_layout(): - images = _make_oss_color_jitter_input() - transform = OpenCVColorJitter(**OSS_COLOR_JITTER_MAGNITUDES) - random.seed(1337) - - actual = transform(images) - - assert actual.shape == images.shape - assert actual.dtype == torch.uint8 - assert actual.is_contiguous() - assert hashlib.sha256(actual.numpy().tobytes()).hexdigest() == OSS_COLOR_JITTER_OUTPUT_SHA256 - - -def test_opencv_color_jitter_reuses_one_sample_across_temporal_frames(): - image = _make_oss_color_jitter_input()[:1] - images = image.repeat(2, 1, 1, 1) - transform = OpenCVColorJitter(**OSS_COLOR_JITTER_MAGNITUDES) - random.seed(1337) - - first, second = transform(images) - - torch.testing.assert_close(first, second) - - -def test_opencv_color_jitter_is_available_from_image_transform_config(): - cfg = ImageTransformsConfig( - enable=True, - max_num_transforms=1, - tfs={ - "color_jitter": ImageTransformConfig( - type="OpenCVColorJitter", - kwargs=OSS_COLOR_JITTER_MAGNITUDES, - ) - }, - ) - - transform = ImageTransforms(cfg) - - assert isinstance(transform.transforms["color_jitter"], OpenCVColorJitter) - - -@pytest.mark.parametrize( - "kwargs,match", - [ - ({"brightness": -0.1}, "brightness"), - ({"contrast": float("inf")}, "contrast"), - ({"hue": 0.6}, "hue"), - ], -) -def test_opencv_color_jitter_rejects_invalid_magnitudes(kwargs, match): - with pytest.raises(ValueError, match=match): - OpenCVColorJitter(**kwargs) - - -def test_opencv_color_jitter_requires_cpu_chw_rgb_uint8(): - transform = OpenCVColorJitter() - - with pytest.raises(ValueError, match="uint8"): - transform(torch.zeros(3, 8, 8)) - with pytest.raises(ValueError, match="CHW RGB"): - transform(torch.zeros(1, 8, 8, dtype=torch.uint8)) - - def test_make_transform_from_config_with_v2_resize(img_tensor_factory): img_tensor = img_tensor_factory() tf_cfg = ImageTransformConfig(type="Resize", kwargs={"size": (32, 32)}) diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index 8b74e4664f4..0a84b0d0af7 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -619,6 +619,20 @@ def test_raw_n1_7_libero_checkpoint_processors_use_checkpoint_assets(tmp_path): assert decode_actions.action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_LIBERO +def test_groot_training_keeps_vlm_image_preprocessing_deterministic(tmp_path): + model_path = tmp_path / "libero_spatial" + _write_raw_n1_7_libero_checkpoint(model_path) + config = _raw_n1_7_libero_config(model_path) + + preprocessor, _ = make_groot_pre_post_processors(config, dataset_meta=object()) + + pack_inputs = next(step for step in preprocessor.steps if isinstance(step, GrootN17PackInputsStep)) + vlm_encode = next(step for step in preprocessor.steps if isinstance(step, GrootN17VLMEncodeStep)) + + assert pack_inputs.training is True + assert not hasattr(vlm_encode, "training") + + def test_raw_n1_7_checkpoint_requires_percentile_stats_when_config_uses_percentiles(tmp_path): model_path = tmp_path / "libero_spatial" _write_raw_n1_7_libero_checkpoint(model_path) diff --git a/tests/policies/groot/test_groot_train_random_crop.py b/tests/policies/groot/test_groot_train_random_crop.py deleted file mode 100644 index adef2958b59..00000000000 --- a/tests/policies/groot/test_groot_train_random_crop.py +++ /dev/null @@ -1,169 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# 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. - -"""Isaac-GR00T N1.7 train-time random crop contract (crop geometry only). - -Isaac-GR00T crops a random ``crop_fraction`` window during training and the -deterministic center window at eval, replaying the sampled window across all -camera views of a sample (gr00t/data/transform/video.py, n1.5-release onward: -"If mode is 'train', return a random crop transform. If mode is 'eval', return -a center crop transform."). This mirrors LeRobot's own Diffusion/VQBeT -``crop_is_random`` pattern. Color jitter is intentionally out of scope here. -""" - -import random - -import numpy as np -import torch - -from lerobot.policies.groot.processor_groot import ( - GrootN17VLMEncodeStep, - _transform_n1_7_image_for_vlm_albumentations, -) - - -def _structured_image(h=480, w=640): - yy, xx = np.mgrid[0:h, 0:w] - return np.stack([(xx * 255 / w), (yy * 255 / h), ((xx + yy) * 255 / (h + w))], axis=-1).astype(np.uint8) - - -def test_crop_position_none_is_bitexact_center_crop(): - """crop_position=None must remain byte-identical to the pre-change eval path.""" - img = _structured_image() - ref = _transform_n1_7_image_for_vlm_albumentations( - img, - image_crop_size=None, - image_target_size=[256, 256], - shortest_image_edge=256, - crop_fraction=0.95, - ) - out = _transform_n1_7_image_for_vlm_albumentations( - img, - image_crop_size=None, - image_target_size=[256, 256], - shortest_image_edge=256, - crop_fraction=0.95, - crop_position=None, - ) - np.testing.assert_array_equal(ref, out) - - -def test_crop_position_center_matches_center_crop(): - img = _structured_image() - center = _transform_n1_7_image_for_vlm_albumentations( - img, - image_crop_size=None, - image_target_size=[256, 256], - shortest_image_edge=256, - crop_fraction=0.95, - crop_position=None, - ) - explicit = _transform_n1_7_image_for_vlm_albumentations( - img, - image_crop_size=None, - image_target_size=[256, 256], - shortest_image_edge=256, - crop_fraction=0.95, - crop_position=(0.5, 0.5), - ) - # int-floor center vs rounded positional center may differ by <=1 px of grid - assert center.shape == explicit.shape - diff = np.abs(center.astype(np.int16) - explicit.astype(np.int16)) - assert diff.mean() < 3.0 - - -def test_crop_position_corners_differ_from_center(): - img = _structured_image() - - def crop_at(position): - return _transform_n1_7_image_for_vlm_albumentations( - img, - image_crop_size=None, - image_target_size=[256, 256], - shortest_image_edge=256, - crop_fraction=0.95, - crop_position=position, - ) - - center = crop_at(None) - tl = crop_at((0.0, 0.0)) - br = crop_at((1.0, 1.0)) - assert not np.array_equal(center, tl) - assert not np.array_equal(tl, br) - - -def _video(img, views=2): - return np.stack([img] * views, axis=0).reshape(1, 1, views, *img.shape) - - -def _step(training): - return GrootN17VLMEncodeStep( - image_target_size=[256, 256], - shortest_image_edge=256, - crop_fraction=0.95, - use_albumentations=True, - training=training, - ) - - -def test_training_crop_replays_one_window_across_views(): - video = _video(_structured_image()) - frames = _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0] - np.testing.assert_array_equal(np.asarray(frames[0]), np.asarray(frames[1])) - - -def test_training_crop_differs_from_eval_center_crop(): - video = _video(_structured_image()) - random.seed(3) # a draw that is not the exact center - train_frame = np.asarray( - _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0] - ) - eval_frame = np.asarray( - _step(training=False)._build_sample_images(video, batch_size=1, target_device=None)[0][0] - ) - assert not np.array_equal(train_frame, eval_frame) - - -def test_training_crop_is_disabled_under_no_grad(): - video = _video(_structured_image()) - with torch.no_grad(): - no_grad_frame = np.asarray( - _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0] - ) - eval_frame = np.asarray( - _step(training=False)._build_sample_images(video, batch_size=1, target_device=None)[0][0] - ) - np.testing.assert_array_equal(no_grad_frame, eval_frame) - - -def test_training_mode_is_not_serialized(): - step = _step(training=True) - serialized = step.get_config() - assert "training" not in serialized - restored = GrootN17VLMEncodeStep(**serialized) - assert restored.training is False - - -def test_training_crop_respects_global_seed(): - video = _video(_structured_image()) - - def draw(): - random.seed(11) - return np.asarray( - _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0] - ) - - np.testing.assert_array_equal(draw(), draw()) From ae762c6a08b121a0bd2e3895f7e6d0c26fc87e6d Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Sat, 4 Jul 2026 06:11:08 -0700 Subject: [PATCH 04/13] fix(groot): preserve train-time crop position --- src/lerobot/policies/groot/processor_groot.py | 30 +++- tests/policies/groot/test_groot_n1_7.py | 14 -- .../groot/test_groot_train_random_crop.py | 169 ++++++++++++++++++ 3 files changed, 196 insertions(+), 17 deletions(-) create mode 100644 tests/policies/groot/test_groot_train_random_crop.py diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 5ae4346118b..20b3518a306 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -463,7 +463,7 @@ def _set_groot_preprocessor_training( """Set the runtime-only mode of GR00T stochastic processor steps. Any dataclass step exposing a ``training`` field participates, so processor - steps can opt into train-time-only behavior (for example, dropout) without + steps can opt into train-time-only behavior (dropout, augmentation) without this helper enumerating them. """ for step in preprocessor.steps: @@ -1258,6 +1258,7 @@ def make_groot_pre_post_processors( crop_fraction=crop_fraction, use_albumentations=use_albumentations, letter_box_transform=letter_box_transform, + training=dataset_meta is not None, device=config.device, ), DeviceProcessorStep(device=config.device), @@ -1383,6 +1384,7 @@ def _transform_n1_7_image_for_vlm_albumentations( shortest_image_edge: int | None, crop_fraction: float | None, letter_box_transform: bool = False, + crop_position: tuple[float, float] | None = None, ) -> np.ndarray: """cv2/INTER_AREA eval transform mirroring Isaac-GR00T's albumentations preprocessing. @@ -1392,6 +1394,12 @@ def _transform_n1_7_image_for_vlm_albumentations( cv2/INTER_AREA resize and floored center-crop here intentionally differ from that torch path and must stay bit-exact to the upstream reference. The hot path accepts and returns numpy arrays to avoid per-frame PIL round-trips. + + ``crop_position`` selects where the ``crop_fraction`` window sits: ``None`` + keeps the deterministic center crop (eval contract), while ``(y, x)`` + fractions in [0, 1] place the window for Isaac's train-time random crop + (0.5, 0.5 == center). Training samples one position per sample and reuses + it across camera views. """ if image_target_size is None: return image @@ -1443,8 +1451,13 @@ def resize_shortest_edge(frame: np.ndarray) -> np.ndarray: height, width = image_np.shape[:2] crop_h = max(1, int(height * crop_fraction)) crop_w = max(1, int(width * crop_fraction)) - top = max(0, (height - crop_h) // 2) - left = max(0, (width - crop_w) // 2) + if crop_position is None: + top = max(0, (height - crop_h) // 2) + left = max(0, (width - crop_w) // 2) + else: + pos_y, pos_x = crop_position + top = int(round((height - crop_h) * min(max(pos_y, 0.0), 1.0))) + left = int(round((width - crop_w) * min(max(pos_x, 0.0), 1.0))) image_np = image_np[top : top + crop_h, left : left + crop_w] return resize_shortest_edge(image_np) @@ -2040,6 +2053,11 @@ class GrootN17VLMEncodeStep(ProcessorStep): crop_fraction: float | None = None use_albumentations: bool = False letter_box_transform: bool = False + # Runtime-only train/eval mode: True enables Isaac's train-time random crop + # (one window per sample, replayed across views); False keeps the + # deterministic center crop. Never serialized - reloaded pipelines default + # to eval and are re-enabled only when processors are built with dataset_meta. + training: bool = False device: str | None = None _proc: ProcessorMixin | None = field(default=None, init=False, repr=False) @@ -2073,8 +2091,13 @@ def _build_sample_images( """ if self.use_albumentations: video_np = np.asarray(video) + train_crop = self.training and torch.is_grad_enabled() sample_images: list[list[Any]] = [] for batch_idx in range(batch_size): + # Isaac-GR00T samples ONE crop window per sample and replays it + # across every (timestep, view) frame of that sample, keeping + # cross-view geometry consistent. Eval keeps the center crop. + crop_position = (random.random(), random.random()) if train_crop else None sample_images.append( [ _transform_n1_7_image_for_vlm_albumentations( @@ -2084,6 +2107,7 @@ def _build_sample_images( shortest_image_edge=self.shortest_image_edge, crop_fraction=self.crop_fraction, letter_box_transform=self.letter_box_transform, + crop_position=crop_position, ) for timestep in range(video_np.shape[1]) for view_idx in range(video_np.shape[2]) diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index 0a84b0d0af7..8b74e4664f4 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -619,20 +619,6 @@ def test_raw_n1_7_libero_checkpoint_processors_use_checkpoint_assets(tmp_path): assert decode_actions.action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_LIBERO -def test_groot_training_keeps_vlm_image_preprocessing_deterministic(tmp_path): - model_path = tmp_path / "libero_spatial" - _write_raw_n1_7_libero_checkpoint(model_path) - config = _raw_n1_7_libero_config(model_path) - - preprocessor, _ = make_groot_pre_post_processors(config, dataset_meta=object()) - - pack_inputs = next(step for step in preprocessor.steps if isinstance(step, GrootN17PackInputsStep)) - vlm_encode = next(step for step in preprocessor.steps if isinstance(step, GrootN17VLMEncodeStep)) - - assert pack_inputs.training is True - assert not hasattr(vlm_encode, "training") - - def test_raw_n1_7_checkpoint_requires_percentile_stats_when_config_uses_percentiles(tmp_path): model_path = tmp_path / "libero_spatial" _write_raw_n1_7_libero_checkpoint(model_path) diff --git a/tests/policies/groot/test_groot_train_random_crop.py b/tests/policies/groot/test_groot_train_random_crop.py new file mode 100644 index 00000000000..adef2958b59 --- /dev/null +++ b/tests/policies/groot/test_groot_train_random_crop.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# 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. + +"""Isaac-GR00T N1.7 train-time random crop contract (crop geometry only). + +Isaac-GR00T crops a random ``crop_fraction`` window during training and the +deterministic center window at eval, replaying the sampled window across all +camera views of a sample (gr00t/data/transform/video.py, n1.5-release onward: +"If mode is 'train', return a random crop transform. If mode is 'eval', return +a center crop transform."). This mirrors LeRobot's own Diffusion/VQBeT +``crop_is_random`` pattern. Color jitter is intentionally out of scope here. +""" + +import random + +import numpy as np +import torch + +from lerobot.policies.groot.processor_groot import ( + GrootN17VLMEncodeStep, + _transform_n1_7_image_for_vlm_albumentations, +) + + +def _structured_image(h=480, w=640): + yy, xx = np.mgrid[0:h, 0:w] + return np.stack([(xx * 255 / w), (yy * 255 / h), ((xx + yy) * 255 / (h + w))], axis=-1).astype(np.uint8) + + +def test_crop_position_none_is_bitexact_center_crop(): + """crop_position=None must remain byte-identical to the pre-change eval path.""" + img = _structured_image() + ref = _transform_n1_7_image_for_vlm_albumentations( + img, + image_crop_size=None, + image_target_size=[256, 256], + shortest_image_edge=256, + crop_fraction=0.95, + ) + out = _transform_n1_7_image_for_vlm_albumentations( + img, + image_crop_size=None, + image_target_size=[256, 256], + shortest_image_edge=256, + crop_fraction=0.95, + crop_position=None, + ) + np.testing.assert_array_equal(ref, out) + + +def test_crop_position_center_matches_center_crop(): + img = _structured_image() + center = _transform_n1_7_image_for_vlm_albumentations( + img, + image_crop_size=None, + image_target_size=[256, 256], + shortest_image_edge=256, + crop_fraction=0.95, + crop_position=None, + ) + explicit = _transform_n1_7_image_for_vlm_albumentations( + img, + image_crop_size=None, + image_target_size=[256, 256], + shortest_image_edge=256, + crop_fraction=0.95, + crop_position=(0.5, 0.5), + ) + # int-floor center vs rounded positional center may differ by <=1 px of grid + assert center.shape == explicit.shape + diff = np.abs(center.astype(np.int16) - explicit.astype(np.int16)) + assert diff.mean() < 3.0 + + +def test_crop_position_corners_differ_from_center(): + img = _structured_image() + + def crop_at(position): + return _transform_n1_7_image_for_vlm_albumentations( + img, + image_crop_size=None, + image_target_size=[256, 256], + shortest_image_edge=256, + crop_fraction=0.95, + crop_position=position, + ) + + center = crop_at(None) + tl = crop_at((0.0, 0.0)) + br = crop_at((1.0, 1.0)) + assert not np.array_equal(center, tl) + assert not np.array_equal(tl, br) + + +def _video(img, views=2): + return np.stack([img] * views, axis=0).reshape(1, 1, views, *img.shape) + + +def _step(training): + return GrootN17VLMEncodeStep( + image_target_size=[256, 256], + shortest_image_edge=256, + crop_fraction=0.95, + use_albumentations=True, + training=training, + ) + + +def test_training_crop_replays_one_window_across_views(): + video = _video(_structured_image()) + frames = _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0] + np.testing.assert_array_equal(np.asarray(frames[0]), np.asarray(frames[1])) + + +def test_training_crop_differs_from_eval_center_crop(): + video = _video(_structured_image()) + random.seed(3) # a draw that is not the exact center + train_frame = np.asarray( + _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0] + ) + eval_frame = np.asarray( + _step(training=False)._build_sample_images(video, batch_size=1, target_device=None)[0][0] + ) + assert not np.array_equal(train_frame, eval_frame) + + +def test_training_crop_is_disabled_under_no_grad(): + video = _video(_structured_image()) + with torch.no_grad(): + no_grad_frame = np.asarray( + _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0] + ) + eval_frame = np.asarray( + _step(training=False)._build_sample_images(video, batch_size=1, target_device=None)[0][0] + ) + np.testing.assert_array_equal(no_grad_frame, eval_frame) + + +def test_training_mode_is_not_serialized(): + step = _step(training=True) + serialized = step.get_config() + assert "training" not in serialized + restored = GrootN17VLMEncodeStep(**serialized) + assert restored.training is False + + +def test_training_crop_respects_global_seed(): + video = _video(_structured_image()) + + def draw(): + random.seed(11) + return np.asarray( + _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0] + ) + + np.testing.assert_array_equal(draw(), draw()) From a2e9029e94fbf1cae7aac65d2c3c6bd2b5ca26c4 Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Sat, 4 Jul 2026 07:43:08 -0700 Subject: [PATCH 05/13] perf(groot): preserve uint8 images through preprocessing --- src/lerobot/policies/groot/processor_groot.py | 151 +++++++++++++++--- src/lerobot/processor/__init__.py | 2 + src/lerobot/processor/pipeline.py | 22 +++ src/lerobot/scripts/lerobot_train.py | 24 ++- tests/policies/groot/test_groot_n1_7.py | 56 ++++++- tests/processor/test_pipeline.py | 15 ++ tests/training/test_visual_validation.py | 30 +++- 7 files changed, 268 insertions(+), 32 deletions(-) diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 20b3518a306..c882bd7818d 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -25,7 +25,6 @@ import numpy as np import torch import torchvision.transforms.v2.functional as tv_functional -from einops import rearrange from torchvision.transforms import InterpolationMode from lerobot.utils.import_utils import _datasets_available, _transformers_available, require_package @@ -54,6 +53,7 @@ AbsoluteActionsProcessorStep, AddBatchDimensionProcessorStep, DeviceProcessorStep, + ImageInputFormat, PolicyAction, PolicyProcessorPipeline, ProcessorStep, @@ -551,6 +551,10 @@ def _load_groot_processor_pipelines( to_transition=policy_action_to_transition, to_output=transition_to_policy_action, ) + # Older serialized GR00T pipelines predate the raw-image contract. GR00T's + # packer consumes worker-produced uint8 directly, so upgrade them at load + # time rather than falling back to the global float compatibility default. + preprocessor.input_image_format = ImageInputFormat.UINT8_0_255 return preprocessor, postprocessor @@ -1323,6 +1327,7 @@ def make_groot_pre_post_processors( PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( steps=input_steps, name=POLICY_PREPROCESSOR_DEFAULT_NAME, + input_image_format=ImageInputFormat.UINT8_0_255, ), PolicyProcessorPipeline[PolicyAction, PolicyAction]( steps=output_steps, @@ -1336,19 +1341,51 @@ def make_groot_pre_post_processors( # GR00T specific processor steps -def _to_uint8_np_bthwc(img_t: torch.Tensor) -> np.ndarray: - # img_t: (B, C, H, W) or (B, T, C, H, W), float in [0,1] or uint8 - if img_t.dtype.is_floating_point: - img_t = (img_t.clamp(0, 1) * 255.0).to(torch.uint8) - if img_t.dim() == 4: - return rearrange(img_t.cpu().numpy(), "b c h w -> b 1 h w c") - if img_t.dim() == 5: - return rearrange(img_t.cpu().numpy(), "b t c h w -> b t h w c") - raise ValueError(f"Expected image tensor shape (B, C, H, W) or (B, T, C, H, W), got {tuple(img_t.shape)}") +@dataclass(frozen=True) +class _GrootN17VideoBatch: + """Ordered camera tensors kept in LeRobot's native (B, T, C, H, W) layout.""" + + cameras: tuple[torch.Tensor, ...] + + def __post_init__(self) -> None: + if not self.cameras: + raise ValueError("GR00T N1.7 video batches require at least one camera.") + first_shape = self.cameras[0].shape + for camera in self.cameras: + if camera.ndim != 5: + raise ValueError( + f"GR00T N1.7 camera tensors must have shape (B, T, C, H, W), got {tuple(camera.shape)}." + ) + if camera.shape[:3] != first_shape[:3]: + raise ValueError( + "GR00T N1.7 camera tensors must share batch, horizon, and channel dimensions, " + f"got {tuple(first_shape[:3])} and {tuple(camera.shape[:3])}." + ) + @property + def batch_size(self) -> int: + return int(self.cameras[0].shape[0]) + + @property + def horizon(self) -> int: + return int(self.cameras[0].shape[1]) -def _align_video_horizon(video: np.ndarray, horizon: int | None) -> np.ndarray: - """Match the checkpoint video horizon by truncating or left-padding frames.""" + +def _as_video_tensor_btchw(image: Any) -> torch.Tensor: + """Preserve a LeRobot image tensor while making its time dimension explicit.""" + + image_t = image if isinstance(image, torch.Tensor) else torch.as_tensor(image) + if image_t.ndim == 4: + return image_t.unsqueeze(1) + if image_t.ndim == 5: + return image_t + raise ValueError( + f"Expected image tensor shape (B, C, H, W) or (B, T, C, H, W), got {tuple(image_t.shape)}." + ) + + +def _align_video_horizon_tensor(video: torch.Tensor, horizon: int | None) -> torch.Tensor: + """Match the checkpoint video horizon without changing dtype or tensor layout.""" if horizon is None or horizon <= 0: return video @@ -1357,8 +1394,21 @@ def _align_video_horizon(video: np.ndarray, horizon: int | None) -> np.ndarray: return video if current > horizon: return video[:, -horizon:] - pad = np.repeat(video[:, :1], horizon - current, axis=1) - return np.concatenate([pad, video], axis=1) + pad = video[:, :1].expand(-1, horizon - current, -1, -1, -1) + return torch.cat([pad, video], dim=1) + + +def _uint8_image_tensor(image: torch.Tensor) -> torch.Tensor: + if image.dtype.is_floating_point: + return (image.clamp(0, 1) * 255.0).to(torch.uint8) + if image.dtype != torch.uint8: + return image.to(torch.uint8) + return image + + +def _uint8_image_numpy_hwc(image: torch.Tensor) -> np.ndarray: + image = _uint8_image_tensor(image).detach().cpu() + return image.permute(1, 2, 0).contiguous().numpy() def _build_n1_7_processor(model_name: str = GROOT_N1_7_BACKBONE_MODEL) -> ProcessorMixin: @@ -1845,17 +1895,28 @@ def _cache_raw_state(state: torch.Tensor) -> None: self._last_raw_state = grouped img_keys = self._ordered_image_keys(obs) + packed_video: _GrootN17VideoBatch | None = None if img_keys: - cams = [_align_video_horizon(_to_uint8_np_bthwc(obs[k]), self.video_horizon) for k in img_keys] - video = np.stack(cams, axis=2) # (B, T, V, H, W, C) - obs["video"] = video + cameras = tuple( + _align_video_horizon_tensor(_as_video_tensor_btchw(obs[key]), self.video_horizon) + for key in img_keys + ) + # Keep the pinned worker tensors in their native channels-first + # representation. The VLM step transfers each view directly and + # never creates a CPU NumPy/HWC staging buffer. + packed_video = _GrootN17VideoBatch(cameras) + obs["video"] = packed_video image_keys_to_remove = [key for key in obs if key.startswith(OBS_IMAGES)] if OBS_IMAGE in obs: image_keys_to_remove.append(OBS_IMAGE) for k in image_keys_to_remove: obs.pop(k, None) - bsz, _device = infer_n1_7_batch_size_and_device(obs, transition.get(TransitionKey.ACTION)) + if packed_video is not None: + bsz = packed_video.batch_size + _device = packed_video.cameras[0].device + else: + bsz, _device = infer_n1_7_batch_size_and_device(obs, transition.get(TransitionKey.ACTION)) comp["language"] = prepare_n1_7_language_batch( comp.get(self.language_key), bsz, @@ -2090,6 +2151,28 @@ def _build_sample_images( ``target_device`` when set) for the torchvision-backed Qwen processor. """ if self.use_albumentations: + if isinstance(video, _GrootN17VideoBatch): + train_crop = self.training and torch.is_grad_enabled() + sample_images: list[list[Any]] = [] + for batch_idx in range(batch_size): + crop_position = (random.random(), random.random()) if train_crop else None + sample_images.append( + [ + _transform_n1_7_image_for_vlm_albumentations( + _uint8_image_numpy_hwc(video.cameras[view_idx][batch_idx, timestep]), + image_crop_size=self.image_crop_size, + image_target_size=self.image_target_size, + shortest_image_edge=self.shortest_image_edge, + crop_fraction=self.crop_fraction, + letter_box_transform=self.letter_box_transform, + crop_position=crop_position, + ) + for timestep in range(video.horizon) + for view_idx in range(len(video.cameras)) + ] + ) + return sample_images + video_np = np.asarray(video) train_crop = self.training and torch.is_grad_enabled() sample_images: list[list[Any]] = [] @@ -2115,6 +2198,36 @@ def _build_sample_images( ) return sample_images + if isinstance(video, _GrootN17VideoBatch): + cameras: list[torch.Tensor] = [] + for camera in video.cameras: + camera_t = camera + if target_device is not None and camera_t.device != target_device: + camera_t = camera_t.to( + target_device, + non_blocking=(target_device.type == "cuda"), + ) + # Float observations from direct/inference callers remain + # supported, but conversion happens after transfer. Training's + # registered uint8 contract takes this branch as a no-op. + cameras.append(_uint8_image_tensor(camera_t)) + + return [ + [ + _transform_n1_7_image_for_vlm_torch( + cameras[view_idx][batch_idx, timestep], + image_crop_size=self.image_crop_size, + image_target_size=self.image_target_size, + shortest_image_edge=self.shortest_image_edge, + crop_fraction=self.crop_fraction, + letter_box_transform=self.letter_box_transform, + ) + for timestep in range(video.horizon) + for view_idx in range(len(cameras)) + ] + for batch_idx in range(batch_size) + ] + video_t = video if torch.is_tensor(video) else torch.from_numpy(np.ascontiguousarray(video)) # (B, T, V, H, W, C) uint8 -> (B, T, V, C, H, W) video_t = video_t.permute(0, 1, 2, 5, 3, 4).contiguous() @@ -2147,7 +2260,7 @@ def __call__(self, transition: EnvTransition) -> EnvTransition: if video is None: return transition - batch_size = int(video.shape[0]) + batch_size = video.batch_size if isinstance(video, _GrootN17VideoBatch) else int(video.shape[0]) languages = prepare_n1_7_language_batch( comp.get("language"), batch_size, diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index fe35af4b4c0..360c864c7ee 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -70,6 +70,7 @@ DataProcessorPipeline, DoneProcessorStep, IdentityProcessorStep, + ImageInputFormat, InfoProcessorStep, ObservationProcessorStep, PolicyActionProcessorStep, @@ -125,6 +126,7 @@ "GymHILAdapterProcessorStep", "GripperPenaltyProcessorStep", "hotswap_stats", + "ImageInputFormat", "IdentityProcessorStep", "ImageCropResizeProcessorStep", "InfoProcessorStep", diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py index b9b9c6c4350..fc3a925be9c 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -37,6 +37,7 @@ from collections.abc import Callable, Iterable, Sequence from copy import deepcopy from dataclasses import dataclass, field +from enum import StrEnum from pathlib import Path from typing import Any, TypedDict, TypeVar, cast @@ -56,6 +57,18 @@ TOutput = TypeVar("TOutput") +class ImageInputFormat(StrEnum): + """Raw image representation expected by a policy preprocessor. + + Dataset workers use uint8 images for compact IPC. The training loop uses + this contract to preserve uint8 for preprocessors that consume it directly, + while retaining the historical float32 [0, 1] input for existing policies. + """ + + FLOAT32_0_1 = "float32_0_1" + UINT8_0_255 = "uint8_0_255" + + class ProcessorStepRegistry: """A registry for ProcessorStep classes to allow instantiation from a string name. @@ -270,6 +283,11 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): steps: Sequence[ProcessorStep] = field(default_factory=list) name: str = "DataProcessorPipeline" + # This describes the raw images accepted before ``to_transition`` and the + # first processor step. It is primarily meaningful for policy preprocessors; + # the float default keeps existing and third-party pipelines compatible. + input_image_format: ImageInputFormat = ImageInputFormat.FLOAT32_0_1 + to_transition: Callable[[TInput], EnvTransition] = field( default_factory=lambda: cast(Callable[[TInput], EnvTransition], batch_to_transition), repr=False ) @@ -437,6 +455,7 @@ def get_config(self) -> dict[str, Any]: sanitized_name = self._get_sanitized_name() pipeline_config: dict[str, Any] = { "name": self.name, + "input_image_format": self.input_image_format.value, "steps": [], } @@ -741,6 +760,7 @@ def from_pretrained( pipeline = cls( steps=steps, name=loaded_config.get("name", "DataProcessorPipeline"), + input_image_format=loaded_config.get("input_image_format", ImageInputFormat.FLOAT32_0_1), to_transition=to_transition or cast(Callable[[TInput], EnvTransition], batch_to_transition), to_output=to_output or cast(Callable[[EnvTransition], TOutput], transition_to_batch), ) @@ -777,6 +797,7 @@ def from_config( pipeline = cls( steps=steps, name=config.get("name", "DataProcessorPipeline"), + input_image_format=config.get("input_image_format", ImageInputFormat.FLOAT32_0_1), to_transition=to_transition or cast(Callable[[TInput], EnvTransition], batch_to_transition), to_output=to_output or cast(Callable[[EnvTransition], TOutput], transition_to_batch), ) @@ -1533,6 +1554,7 @@ def __repr__(self) -> str: def __post_init__(self): """Validates that all provided steps are instances of `ProcessorStep`.""" + self.input_image_format = ImageInputFormat(self.input_image_format) for i, step in enumerate(self.steps): if not isinstance(step, ProcessorStep): raise TypeError(f"Step {i} ({type(step).__name__}) must inherit from ProcessorStep") diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 6e845852306..31f7514a1ff 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -55,6 +55,7 @@ from lerobot.jobs import submit_to_hf from lerobot.optim.factory import make_optimizer_and_scheduler from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors +from lerobot.processor import ImageInputFormat, PolicyProcessorPipeline from lerobot.rewards import make_reward_pre_post_processors from lerobot.utils.collate import lerobot_collate_fn from lerobot.utils.import_utils import register_third_party_plugins @@ -71,6 +72,21 @@ from .lerobot_eval import eval_policy_all +def prepare_images_for_policy( + batch: dict[str, Any], + camera_keys: list[str], + preprocessor: PolicyProcessorPipeline, +) -> None: + """Adapt worker-produced images to the policy preprocessor's raw input contract.""" + + if preprocessor.input_image_format is ImageInputFormat.UINT8_0_255: + return + + for cam_key in camera_keys: + if cam_key in batch and batch[cam_key].dtype == torch.uint8: + batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0 + + def update_policy( train_metrics: MetricsTracker, policy: PreTrainedPolicy, @@ -566,9 +582,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): for _ in range(step, cfg.steps): start_time = time.perf_counter() batch = next(dl_iter) - for cam_key in dataset.meta.camera_keys: - if cam_key in batch and batch[cam_key].dtype == torch.uint8: - batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0 + prepare_images_for_policy(batch, dataset.meta.camera_keys, preprocessor) batch = preprocessor(batch) train_tracker.dataloading_s = time.perf_counter() - start_time @@ -621,9 +635,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): n_eval_batches = 0 with torch.no_grad(), accelerator.autocast(): for eval_batch in eval_dataloader: - for cam_key in dataset.meta.camera_keys: - if cam_key in eval_batch and eval_batch[cam_key].dtype == torch.uint8: - eval_batch[cam_key] = eval_batch[cam_key].to(dtype=torch.float32) / 255.0 + prepare_images_for_policy(eval_batch, dataset.meta.camera_keys, preprocessor) eval_batch = preprocessor(eval_batch) loss, _ = policy.forward(eval_batch) eval_loss_sum += loss.item() diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index 8b74e4664f4..25154f5f17f 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -51,6 +51,7 @@ ) from lerobot.processor import ( AbsoluteActionsProcessorStep, + ImageInputFormat, PolicyProcessorPipeline, RelativeActionsProcessorStep, ) @@ -1234,11 +1235,14 @@ def test_groot_n1_7_pack_inputs_orders_video_by_checkpoint_modality_keys(): normalize_min_max=False, video_modality_keys=["image", "wrist_image"], ) + extra = torch.full((1, 3, 2, 2), 33, dtype=torch.uint8) + wrist = torch.full((1, 3, 2, 2), 22, dtype=torch.uint8) + front = torch.full((1, 3, 2, 2), 11, dtype=torch.uint8) transition = { TransitionKey.OBSERVATION: { - f"{OBS_IMAGES}.zz_extra": torch.full((1, 3, 2, 2), 33, dtype=torch.uint8), - f"{OBS_IMAGES}.image2": torch.full((1, 3, 2, 2), 22, dtype=torch.uint8), - f"{OBS_IMAGES}.image": torch.full((1, 3, 2, 2), 11, dtype=torch.uint8), + f"{OBS_IMAGES}.zz_extra": extra, + f"{OBS_IMAGES}.image2": wrist, + f"{OBS_IMAGES}.image": front, OBS_STATE: torch.zeros(1, 8), }, TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move"]}, @@ -1247,14 +1251,53 @@ def test_groot_n1_7_pack_inputs_orders_video_by_checkpoint_modality_keys(): output = step(transition) video = output[TransitionKey.OBSERVATION]["video"] - assert video.shape == (1, 1, 2, 2, 2, 3) - assert np.unique(video[0, 0, 0]).tolist() == [11] - assert np.unique(video[0, 0, 1]).tolist() == [22] + assert len(video.cameras) == 2 + assert video.cameras[0].shape == (1, 1, 3, 2, 2) + assert video.cameras[0].data_ptr() == front.data_ptr() + assert video.cameras[1].data_ptr() == wrist.data_ptr() + assert torch.unique(video.cameras[0]).tolist() == [11] + assert torch.unique(video.cameras[1]).tolist() == [22] assert f"{OBS_IMAGES}.zz_extra" not in output[TransitionKey.OBSERVATION] assert f"{OBS_IMAGES}.image" not in output[TransitionKey.OBSERVATION] assert f"{OBS_IMAGES}.image2" not in output[TransitionKey.OBSERVATION] +def test_groot_n1_7_tensor_video_path_matches_legacy_numpy_bytes(): + camera_a = torch.arange(2 * 2 * 3 * 4 * 5, dtype=torch.uint8).reshape(2, 2, 3, 4, 5) + camera_b = (camera_a.to(torch.int16) * 3 % 251).to(torch.uint8) + pack_step = GrootN17PackInputsStep( + normalize_min_max=False, + video_modality_keys=["image", "wrist_image"], + ) + packed = pack_step( + { + TransitionKey.OBSERVATION: { + f"{OBS_IMAGES}.image": camera_a, + f"{OBS_IMAGES}.image2": camera_b, + }, + TransitionKey.COMPLEMENTARY_DATA: {"task": ["a", "b"]}, + } + ) + video = packed[TransitionKey.OBSERVATION]["video"] + legacy_video = np.stack( + [ + camera_a.permute(0, 1, 3, 4, 2).numpy(), + camera_b.permute(0, 1, 3, 4, 2).numpy(), + ], + axis=2, + ) + encode_step = GrootN17VLMEncodeStep() + + tensor_frames = encode_step._build_sample_images(video, batch_size=2, target_device=None) + legacy_frames = encode_step._build_sample_images(legacy_video, batch_size=2, target_device=None) + + assert len(tensor_frames) == len(legacy_frames) == 2 + for tensor_sample, legacy_sample in zip(tensor_frames, legacy_frames, strict=True): + assert len(tensor_sample) == len(legacy_sample) == 4 + for tensor_frame, legacy_frame in zip(tensor_sample, legacy_sample, strict=True): + torch.testing.assert_close(tensor_frame, legacy_frame, rtol=0, atol=0) + + def test_groot_n1_7_postprocessor_clips_normalized_action_before_unnormalizing(): step = GrootActionUnpackUnnormalizeStep( env_action_dim=3, @@ -1664,6 +1707,7 @@ def test_groot_n1_7_processors_are_registered_lazily_without_external_gr00t(): preprocessor, _ = make_groot_pre_post_processors(config) step_types = {type(step) for step in preprocessor.steps} + assert preprocessor.input_image_format is ImageInputFormat.UINT8_0_255 assert GrootN17PackInputsStep in step_types assert GrootN17VLMEncodeStep in step_types assert "gr00t" not in sys.modules diff --git a/tests/processor/test_pipeline.py b/tests/processor/test_pipeline.py index 0e9746a6302..6d33de828d8 100644 --- a/tests/processor/test_pipeline.py +++ b/tests/processor/test_pipeline.py @@ -33,6 +33,7 @@ from lerobot.processor import ( DataProcessorPipeline, EnvTransition, + ImageInputFormat, ProcessorStep, ProcessorStepRegistry, TransitionKey, @@ -604,6 +605,20 @@ def test_save_and_load_pretrained(): assert loaded_pipeline.steps[1].counter == 10 +def test_input_image_format_is_backward_compatible_and_serialized(): + legacy_pipeline = DataProcessorPipeline.from_config({"steps": []}) + assert legacy_pipeline.input_image_format is ImageInputFormat.FLOAT32_0_1 + + pipeline = DataProcessorPipeline( + [], + input_image_format=ImageInputFormat.UINT8_0_255, + ) + loaded_pipeline = DataProcessorPipeline.from_config(pipeline.get_config()) + + assert pipeline.get_config()["input_image_format"] == "uint8_0_255" + assert loaded_pipeline.input_image_format is ImageInputFormat.UINT8_0_255 + + def test_step_without_optional_methods(): """Test pipeline with steps that don't implement optional methods.""" step = MockStepWithoutOptionalMethods(multiplier=3.0) diff --git a/tests/training/test_visual_validation.py b/tests/training/test_visual_validation.py index 1df8006b273..38dfbc8b586 100644 --- a/tests/training/test_visual_validation.py +++ b/tests/training/test_visual_validation.py @@ -30,6 +30,7 @@ import numpy as np import pytest +import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") @@ -38,7 +39,8 @@ from lerobot.configs.train import TrainPipelineConfig from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.policies.factory import make_policy_config -from lerobot.scripts.lerobot_train import train +from lerobot.processor import ImageInputFormat, PolicyProcessorPipeline +from lerobot.scripts.lerobot_train import prepare_images_for_policy, train from lerobot.utils.device_utils import auto_select_torch_device pytest.importorskip("transformers") @@ -57,6 +59,32 @@ def temp_dir(tmp_path): DEVICE = auto_select_torch_device() +def test_prepare_images_for_policy_obeys_preprocessor_contract(): + image = torch.tensor([0, 127, 255], dtype=torch.uint8) + + float_batch = {"observation.images.front": image.clone()} + prepare_images_for_policy( + float_batch, + ["observation.images.front"], + PolicyProcessorPipeline([]), + ) + assert float_batch["observation.images.front"].dtype == torch.float32 + torch.testing.assert_close( + float_batch["observation.images.front"], + torch.tensor([0.0, 127.0 / 255.0, 1.0]), + ) + + uint8_image = image.clone() + uint8_batch = {"observation.images.front": uint8_image} + prepare_images_for_policy( + uint8_batch, + ["observation.images.front"], + PolicyProcessorPipeline([], input_image_format=ImageInputFormat.UINT8_0_255), + ) + assert uint8_batch["observation.images.front"].data_ptr() == uint8_image.data_ptr() + assert uint8_batch["observation.images.front"].dtype == torch.uint8 + + def make_dummy_dataset(camera_keys, tmp_path): """Creates a minimal dummy dataset for testing rename_mapping logic.""" features = { From c379ddeebc6d867d99a8f8a6892c534b918efc1b Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Sat, 4 Jul 2026 08:09:08 -0700 Subject: [PATCH 06/13] refactor(groot): simplify uint8 image handoff --- src/lerobot/policies/groot/processor_groot.py | 203 +++++++----------- src/lerobot/policies/groot/utils.py | 2 + tests/policies/groot/test_groot_n1_7.py | 16 +- .../groot/test_groot_train_random_crop.py | 22 +- 4 files changed, 92 insertions(+), 151 deletions(-) diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index c882bd7818d..91cbad2052c 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -1341,36 +1341,6 @@ def make_groot_pre_post_processors( # GR00T specific processor steps -@dataclass(frozen=True) -class _GrootN17VideoBatch: - """Ordered camera tensors kept in LeRobot's native (B, T, C, H, W) layout.""" - - cameras: tuple[torch.Tensor, ...] - - def __post_init__(self) -> None: - if not self.cameras: - raise ValueError("GR00T N1.7 video batches require at least one camera.") - first_shape = self.cameras[0].shape - for camera in self.cameras: - if camera.ndim != 5: - raise ValueError( - f"GR00T N1.7 camera tensors must have shape (B, T, C, H, W), got {tuple(camera.shape)}." - ) - if camera.shape[:3] != first_shape[:3]: - raise ValueError( - "GR00T N1.7 camera tensors must share batch, horizon, and channel dimensions, " - f"got {tuple(first_shape[:3])} and {tuple(camera.shape[:3])}." - ) - - @property - def batch_size(self) -> int: - return int(self.cameras[0].shape[0]) - - @property - def horizon(self) -> int: - return int(self.cameras[0].shape[1]) - - def _as_video_tensor_btchw(image: Any) -> torch.Tensor: """Preserve a LeRobot image tensor while making its time dimension explicit.""" @@ -1384,6 +1354,40 @@ def _as_video_tensor_btchw(image: Any) -> torch.Tensor: ) +def _video_cameras_btchw(video: Any) -> tuple[torch.Tensor, ...]: + """Return ordered cameras in LeRobot's native (B, T, C, H, W) layout. + + New pipelines pass a tuple of worker-produced camera tensors. The ndarray + adapter keeps serialized pipelines and direct callers using the former + (B, T, V, H, W, C) representation compatible without duplicating the + downstream image path. + """ + + if isinstance(video, (list, tuple)): + cameras = tuple(_as_video_tensor_btchw(camera) for camera in video) + else: + video_t = video if torch.is_tensor(video) else torch.from_numpy(np.ascontiguousarray(video)) + if video_t.ndim != 6: + raise ValueError( + "Expected video as ordered camera tensors or shape (B, T, V, H, W, C), " + f"got {tuple(video_t.shape)}." + ) + cameras = tuple( + video_t[:, :, view_idx].permute(0, 1, 4, 2, 3) for view_idx in range(video_t.shape[2]) + ) + + if not cameras: + raise ValueError("GR00T N1.7 video batches require at least one camera.") + expected_shape = cameras[0].shape[:3] + for camera in cameras: + if camera.ndim != 5 or camera.shape[:3] != expected_shape: + raise ValueError( + "GR00T N1.7 cameras must share (B, T, C) dimensions in (B, T, C, H, W) layout; " + f"expected {tuple(expected_shape)}, got {tuple(camera.shape)}." + ) + return cameras + + def _align_video_horizon_tensor(video: torch.Tensor, horizon: int | None) -> torch.Tensor: """Match the checkpoint video horizon without changing dtype or tensor layout.""" @@ -1895,28 +1899,26 @@ def _cache_raw_state(state: torch.Tensor) -> None: self._last_raw_state = grouped img_keys = self._ordered_image_keys(obs) - packed_video: _GrootN17VideoBatch | None = None + packed_cameras: tuple[torch.Tensor, ...] = () if img_keys: - cameras = tuple( + packed_cameras = tuple( _align_video_horizon_tensor(_as_video_tensor_btchw(obs[key]), self.video_horizon) for key in img_keys ) # Keep the pinned worker tensors in their native channels-first # representation. The VLM step transfers each view directly and # never creates a CPU NumPy/HWC staging buffer. - packed_video = _GrootN17VideoBatch(cameras) - obs["video"] = packed_video + obs["video"] = packed_cameras image_keys_to_remove = [key for key in obs if key.startswith(OBS_IMAGES)] if OBS_IMAGE in obs: image_keys_to_remove.append(OBS_IMAGE) for k in image_keys_to_remove: obs.pop(k, None) - if packed_video is not None: - bsz = packed_video.batch_size - _device = packed_video.cameras[0].device + if packed_cameras: + bsz = packed_cameras[0].shape[0] else: - bsz, _device = infer_n1_7_batch_size_and_device(obs, transition.get(TransitionKey.ACTION)) + bsz, _ = infer_n1_7_batch_size_and_device(obs, transition.get(TransitionKey.ACTION)) comp["language"] = prepare_n1_7_language_batch( comp.get(self.language_key), bsz, @@ -2096,9 +2098,9 @@ def load_state_dict(self, state: dict[str, torch.Tensor]) -> None: class GrootN17VLMEncodeStep(ProcessorStep): """Tokenize N1.7's packed video-language prompt with the Qwen3-VL processor. - The packed video has shape ``(B, T, V, H, W, C)``. Each frame/view becomes - an image item in the same chat message so the resulting image tokens match - the temporal VLM packing used by Isaac-GR00T. + The packed video is an ordered tuple of ``(B, T, C, H, W)`` camera tensors. + Each frame/view becomes an image item in the same chat message so the + resulting image tokens match the temporal VLM packing used by Isaac-GR00T. Images are handed to the torchvision-backed Qwen3-VL processor as ``(C, H, W)`` uint8 tensors (no per-frame PIL roundtrip), and, when ``device`` resolves to a @@ -2141,39 +2143,18 @@ def _target_device(self) -> torch.device | None: # the CPU path, which is bit-identical, instead of crashing. return None - def _build_sample_images( - self, video: Any, batch_size: int, target_device: torch.device | None - ) -> list[list[Any]]: + def _build_sample_images(self, video: Any, target_device: torch.device | None) -> list[list[Any]]: """Return, per batch item, its ordered ``(timestep, view)`` frames. ``use_albumentations`` keeps the legacy per-frame cv2/INTER_AREA transform; otherwise frames are ``(C, H, W)`` uint8 tensors (moved to ``target_device`` when set) for the torchvision-backed Qwen processor. """ - if self.use_albumentations: - if isinstance(video, _GrootN17VideoBatch): - train_crop = self.training and torch.is_grad_enabled() - sample_images: list[list[Any]] = [] - for batch_idx in range(batch_size): - crop_position = (random.random(), random.random()) if train_crop else None - sample_images.append( - [ - _transform_n1_7_image_for_vlm_albumentations( - _uint8_image_numpy_hwc(video.cameras[view_idx][batch_idx, timestep]), - image_crop_size=self.image_crop_size, - image_target_size=self.image_target_size, - shortest_image_edge=self.shortest_image_edge, - crop_fraction=self.crop_fraction, - letter_box_transform=self.letter_box_transform, - crop_position=crop_position, - ) - for timestep in range(video.horizon) - for view_idx in range(len(video.cameras)) - ] - ) - return sample_images + cameras = _video_cameras_btchw(video) + batch_size = cameras[0].shape[0] + horizon = cameras[0].shape[1] - video_np = np.asarray(video) + if self.use_albumentations: train_crop = self.training and torch.is_grad_enabled() sample_images: list[list[Any]] = [] for batch_idx in range(batch_size): @@ -2184,7 +2165,7 @@ def _build_sample_images( sample_images.append( [ _transform_n1_7_image_for_vlm_albumentations( - video_np[batch_idx, timestep, view_idx], + _uint8_image_numpy_hwc(cameras[view_idx][batch_idx, timestep]), image_crop_size=self.image_crop_size, image_target_size=self.image_target_size, shortest_image_edge=self.shortest_image_edge, @@ -2192,66 +2173,35 @@ def _build_sample_images( letter_box_transform=self.letter_box_transform, crop_position=crop_position, ) - for timestep in range(video_np.shape[1]) - for view_idx in range(video_np.shape[2]) + for timestep in range(horizon) + for view_idx in range(len(cameras)) ] ) return sample_images - if isinstance(video, _GrootN17VideoBatch): - cameras: list[torch.Tensor] = [] - for camera in video.cameras: - camera_t = camera - if target_device is not None and camera_t.device != target_device: - camera_t = camera_t.to( - target_device, - non_blocking=(target_device.type == "cuda"), - ) - # Float observations from direct/inference callers remain - # supported, but conversion happens after transfer. Training's - # registered uint8 contract takes this branch as a no-op. - cameras.append(_uint8_image_tensor(camera_t)) - - return [ - [ - _transform_n1_7_image_for_vlm_torch( - cameras[view_idx][batch_idx, timestep], - image_crop_size=self.image_crop_size, - image_target_size=self.image_target_size, - shortest_image_edge=self.shortest_image_edge, - crop_fraction=self.crop_fraction, - letter_box_transform=self.letter_box_transform, - ) - for timestep in range(video.horizon) - for view_idx in range(len(cameras)) - ] - for batch_idx in range(batch_size) + prepared_cameras: list[torch.Tensor] = [] + for camera in cameras: + if target_device is not None and camera.device != target_device: + camera = camera.to(target_device, non_blocking=(target_device.type == "cuda")) + # Float observations from direct/inference callers remain supported, + # but conversion happens after transfer. Training uint8 is a no-op. + prepared_cameras.append(_uint8_image_tensor(camera)) + + return [ + [ + _transform_n1_7_image_for_vlm_torch( + prepared_cameras[view_idx][batch_idx, timestep], + image_crop_size=self.image_crop_size, + image_target_size=self.image_target_size, + shortest_image_edge=self.shortest_image_edge, + crop_fraction=self.crop_fraction, + letter_box_transform=self.letter_box_transform, + ) + for timestep in range(horizon) + for view_idx in range(len(prepared_cameras)) ] - - video_t = video if torch.is_tensor(video) else torch.from_numpy(np.ascontiguousarray(video)) - # (B, T, V, H, W, C) uint8 -> (B, T, V, C, H, W) - video_t = video_t.permute(0, 1, 2, 5, 3, 4).contiguous() - if target_device is not None and video_t.device != target_device: - video_t = video_t.to(target_device, non_blocking=(target_device.type == "cuda")) - - frames_per_sample: list[list[Any]] = [] - for batch_idx in range(batch_size): - sample = video_t[batch_idx] # (T, V, C, H, W) - frames_per_sample.append( - [ - _transform_n1_7_image_for_vlm_torch( - sample[timestep, view_idx], - image_crop_size=self.image_crop_size, - image_target_size=self.image_target_size, - shortest_image_edge=self.shortest_image_edge, - crop_fraction=self.crop_fraction, - letter_box_transform=self.letter_box_transform, - ) - for timestep in range(sample.shape[0]) - for view_idx in range(sample.shape[1]) - ] - ) - return frames_per_sample + for batch_idx in range(batch_size) + ] def __call__(self, transition: EnvTransition) -> EnvTransition: obs = transition.get(TransitionKey.OBSERVATION, {}) or {} @@ -2260,16 +2210,15 @@ def __call__(self, transition: EnvTransition) -> EnvTransition: if video is None: return transition - batch_size = video.batch_size if isinstance(video, _GrootN17VideoBatch) else int(video.shape[0]) + target_device = self._target_device() + sample_images = self._build_sample_images(video, target_device) + batch_size = len(sample_images) languages = prepare_n1_7_language_batch( comp.get("language"), batch_size, formalize_language=False, ) - target_device = self._target_device() - sample_images = self._build_sample_images(video, batch_size, target_device) - texts: list[str] = [] images: list[Any] = [] for batch_idx in range(batch_size): diff --git a/src/lerobot/policies/groot/utils.py b/src/lerobot/policies/groot/utils.py index 9a65404fe85..f667e72f7d5 100644 --- a/src/lerobot/policies/groot/utils.py +++ b/src/lerobot/policies/groot/utils.py @@ -227,6 +227,8 @@ def infer_n1_7_batch_size_and_device( video = obs.get("video") if isinstance(video, np.ndarray): return video.shape[0], torch.device("cpu") + if isinstance(video, (list, tuple)) and video and isinstance(video[0], torch.Tensor): + return video[0].shape[0], video[0].device return 1, torch.device("cpu") diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index 25154f5f17f..fa899cb263e 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -1251,12 +1251,12 @@ def test_groot_n1_7_pack_inputs_orders_video_by_checkpoint_modality_keys(): output = step(transition) video = output[TransitionKey.OBSERVATION]["video"] - assert len(video.cameras) == 2 - assert video.cameras[0].shape == (1, 1, 3, 2, 2) - assert video.cameras[0].data_ptr() == front.data_ptr() - assert video.cameras[1].data_ptr() == wrist.data_ptr() - assert torch.unique(video.cameras[0]).tolist() == [11] - assert torch.unique(video.cameras[1]).tolist() == [22] + assert len(video) == 2 + assert video[0].shape == (1, 1, 3, 2, 2) + assert video[0].data_ptr() == front.data_ptr() + assert video[1].data_ptr() == wrist.data_ptr() + assert torch.unique(video[0]).tolist() == [11] + assert torch.unique(video[1]).tolist() == [22] assert f"{OBS_IMAGES}.zz_extra" not in output[TransitionKey.OBSERVATION] assert f"{OBS_IMAGES}.image" not in output[TransitionKey.OBSERVATION] assert f"{OBS_IMAGES}.image2" not in output[TransitionKey.OBSERVATION] @@ -1288,8 +1288,8 @@ def test_groot_n1_7_tensor_video_path_matches_legacy_numpy_bytes(): ) encode_step = GrootN17VLMEncodeStep() - tensor_frames = encode_step._build_sample_images(video, batch_size=2, target_device=None) - legacy_frames = encode_step._build_sample_images(legacy_video, batch_size=2, target_device=None) + tensor_frames = encode_step._build_sample_images(video, target_device=None) + legacy_frames = encode_step._build_sample_images(legacy_video, target_device=None) assert len(tensor_frames) == len(legacy_frames) == 2 for tensor_sample, legacy_sample in zip(tensor_frames, legacy_frames, strict=True): diff --git a/tests/policies/groot/test_groot_train_random_crop.py b/tests/policies/groot/test_groot_train_random_crop.py index adef2958b59..c8667ef6989 100644 --- a/tests/policies/groot/test_groot_train_random_crop.py +++ b/tests/policies/groot/test_groot_train_random_crop.py @@ -121,31 +121,23 @@ def _step(training): def test_training_crop_replays_one_window_across_views(): video = _video(_structured_image()) - frames = _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0] + frames = _step(training=True)._build_sample_images(video, target_device=None)[0] np.testing.assert_array_equal(np.asarray(frames[0]), np.asarray(frames[1])) def test_training_crop_differs_from_eval_center_crop(): video = _video(_structured_image()) random.seed(3) # a draw that is not the exact center - train_frame = np.asarray( - _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0] - ) - eval_frame = np.asarray( - _step(training=False)._build_sample_images(video, batch_size=1, target_device=None)[0][0] - ) + train_frame = np.asarray(_step(training=True)._build_sample_images(video, target_device=None)[0][0]) + eval_frame = np.asarray(_step(training=False)._build_sample_images(video, target_device=None)[0][0]) assert not np.array_equal(train_frame, eval_frame) def test_training_crop_is_disabled_under_no_grad(): video = _video(_structured_image()) with torch.no_grad(): - no_grad_frame = np.asarray( - _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0] - ) - eval_frame = np.asarray( - _step(training=False)._build_sample_images(video, batch_size=1, target_device=None)[0][0] - ) + no_grad_frame = np.asarray(_step(training=True)._build_sample_images(video, target_device=None)[0][0]) + eval_frame = np.asarray(_step(training=False)._build_sample_images(video, target_device=None)[0][0]) np.testing.assert_array_equal(no_grad_frame, eval_frame) @@ -162,8 +154,6 @@ def test_training_crop_respects_global_seed(): def draw(): random.seed(11) - return np.asarray( - _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0] - ) + return np.asarray(_step(training=True)._build_sample_images(video, target_device=None)[0][0]) np.testing.assert_array_equal(draw(), draw()) From f971773ea62a84cc886c90e298b349bcb20675d7 Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Sat, 4 Jul 2026 08:16:54 -0700 Subject: [PATCH 07/13] refactor(groot): enforce camera batch contract --- src/lerobot/policies/groot/processor_groot.py | 80 +++++++------------ src/lerobot/policies/groot/utils.py | 5 -- tests/policies/groot/test_groot_n1_7.py | 52 ++++++------ .../groot/test_groot_n1_7_oss_parity.py | 5 +- .../groot/test_groot_train_random_crop.py | 3 +- 5 files changed, 57 insertions(+), 88 deletions(-) diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 91cbad2052c..d0c2fa86e8a 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -104,6 +104,8 @@ # action chunks, so processor-side horizons are capped at this value. N1_7_NATIVE_ACTION_HORIZON = 40 +type _GrootN17CameraBatch = tuple[torch.Tensor, ...] + N1_7_EMBODIMENT_MAPPING = { "oxe_droid_relative_eef_relative_joint": 24, "xdof_relative_eef_relative_joint": 27, @@ -1341,53 +1343,18 @@ def make_groot_pre_post_processors( # GR00T specific processor steps -def _as_video_tensor_btchw(image: Any) -> torch.Tensor: +def _as_video_tensor_btchw(image: torch.Tensor) -> torch.Tensor: """Preserve a LeRobot image tensor while making its time dimension explicit.""" - image_t = image if isinstance(image, torch.Tensor) else torch.as_tensor(image) - if image_t.ndim == 4: - return image_t.unsqueeze(1) - if image_t.ndim == 5: - return image_t + if image.ndim == 4: + return image.unsqueeze(1) + if image.ndim == 5: + return image raise ValueError( - f"Expected image tensor shape (B, C, H, W) or (B, T, C, H, W), got {tuple(image_t.shape)}." + f"Expected image tensor shape (B, C, H, W) or (B, T, C, H, W), got {tuple(image.shape)}." ) -def _video_cameras_btchw(video: Any) -> tuple[torch.Tensor, ...]: - """Return ordered cameras in LeRobot's native (B, T, C, H, W) layout. - - New pipelines pass a tuple of worker-produced camera tensors. The ndarray - adapter keeps serialized pipelines and direct callers using the former - (B, T, V, H, W, C) representation compatible without duplicating the - downstream image path. - """ - - if isinstance(video, (list, tuple)): - cameras = tuple(_as_video_tensor_btchw(camera) for camera in video) - else: - video_t = video if torch.is_tensor(video) else torch.from_numpy(np.ascontiguousarray(video)) - if video_t.ndim != 6: - raise ValueError( - "Expected video as ordered camera tensors or shape (B, T, V, H, W, C), " - f"got {tuple(video_t.shape)}." - ) - cameras = tuple( - video_t[:, :, view_idx].permute(0, 1, 4, 2, 3) for view_idx in range(video_t.shape[2]) - ) - - if not cameras: - raise ValueError("GR00T N1.7 video batches require at least one camera.") - expected_shape = cameras[0].shape[:3] - for camera in cameras: - if camera.ndim != 5 or camera.shape[:3] != expected_shape: - raise ValueError( - "GR00T N1.7 cameras must share (B, T, C) dimensions in (B, T, C, H, W) layout; " - f"expected {tuple(expected_shape)}, got {tuple(camera.shape)}." - ) - return cameras - - def _align_video_horizon_tensor(video: torch.Tensor, horizon: int | None) -> torch.Tensor: """Match the checkpoint video horizon without changing dtype or tensor layout.""" @@ -1899,7 +1866,7 @@ def _cache_raw_state(state: torch.Tensor) -> None: self._last_raw_state = grouped img_keys = self._ordered_image_keys(obs) - packed_cameras: tuple[torch.Tensor, ...] = () + packed_cameras: _GrootN17CameraBatch = () if img_keys: packed_cameras = tuple( _align_video_horizon_tensor(_as_video_tensor_btchw(obs[key]), self.video_horizon) @@ -1916,12 +1883,15 @@ def _cache_raw_state(state: torch.Tensor) -> None: obs.pop(k, None) if packed_cameras: - bsz = packed_cameras[0].shape[0] + batch_size = packed_cameras[0].shape[0] + batch_device = packed_cameras[0].device else: - bsz, _ = infer_n1_7_batch_size_and_device(obs, transition.get(TransitionKey.ACTION)) + batch_size, batch_device = infer_n1_7_batch_size_and_device( + obs, transition.get(TransitionKey.ACTION) + ) comp["language"] = prepare_n1_7_language_batch( comp.get(self.language_key), - bsz, + batch_size, formalize_language=self.formalize_language, ) @@ -2030,13 +2000,14 @@ def _cache_raw_state(state: torch.Tensor) -> None: comp["action_mask"] = action_mask emb_id = self.embodiment_mapping.get(self.embodiment_tag, 0) - bsz, device = infer_n1_7_batch_size_and_device(obs, transition.get(TransitionKey.ACTION)) if "action_mask" not in comp: - action_mask = torch.zeros(bsz, self.action_horizon, dtype=torch.float32, device=device) + action_mask = torch.zeros( + batch_size, self.action_horizon, dtype=torch.float32, device=batch_device + ) valid_horizon = min(self.valid_action_horizon, self.action_horizon) action_mask[:, :valid_horizon] = 1.0 comp["action_mask"] = action_mask - comp["embodiment_id"] = torch.full((bsz,), emb_id, dtype=torch.int32, device=device) + comp["embodiment_id"] = torch.full((batch_size,), emb_id, dtype=torch.int32, device=batch_device) transition[TransitionKey.OBSERVATION] = obs transition[TransitionKey.COMPLEMENTARY_DATA] = comp @@ -2143,14 +2114,17 @@ def _target_device(self) -> torch.device | None: # the CPU path, which is bit-identical, instead of crashing. return None - def _build_sample_images(self, video: Any, target_device: torch.device | None) -> list[list[Any]]: + def _build_sample_images( + self, + cameras: _GrootN17CameraBatch, + target_device: torch.device | None, + ) -> list[list[Any]]: """Return, per batch item, its ordered ``(timestep, view)`` frames. ``use_albumentations`` keeps the legacy per-frame cv2/INTER_AREA transform; otherwise frames are ``(C, H, W)`` uint8 tensors (moved to ``target_device`` when set) for the torchvision-backed Qwen processor. """ - cameras = _video_cameras_btchw(video) batch_size = cameras[0].shape[0] horizon = cameras[0].shape[1] @@ -2206,12 +2180,12 @@ def _build_sample_images(self, video: Any, target_device: torch.device | None) - def __call__(self, transition: EnvTransition) -> EnvTransition: obs = transition.get(TransitionKey.OBSERVATION, {}) or {} comp = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) or {} - video = obs.get("video") - if video is None: + cameras: _GrootN17CameraBatch | None = obs.get("video") + if cameras is None: return transition target_device = self._target_device() - sample_images = self._build_sample_images(video, target_device) + sample_images = self._build_sample_images(cameras, target_device) batch_size = len(sample_images) languages = prepare_n1_7_language_batch( comp.get("language"), diff --git a/src/lerobot/policies/groot/utils.py b/src/lerobot/policies/groot/utils.py index f667e72f7d5..e332835bf52 100644 --- a/src/lerobot/policies/groot/utils.py +++ b/src/lerobot/policies/groot/utils.py @@ -224,11 +224,6 @@ def infer_n1_7_batch_size_and_device( for value in list(obs.values()) + [action]: if isinstance(value, torch.Tensor): return value.shape[0], value.device - video = obs.get("video") - if isinstance(video, np.ndarray): - return video.shape[0], torch.device("cpu") - if isinstance(video, (list, tuple)) and video and isinstance(video[0], torch.Tensor): - return video[0].shape[0], video[0].device return 1, torch.device("cpu") diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index fa899cb263e..99057cb0a53 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -1262,40 +1262,31 @@ def test_groot_n1_7_pack_inputs_orders_video_by_checkpoint_modality_keys(): assert f"{OBS_IMAGES}.image2" not in output[TransitionKey.OBSERVATION] -def test_groot_n1_7_tensor_video_path_matches_legacy_numpy_bytes(): - camera_a = torch.arange(2 * 2 * 3 * 4 * 5, dtype=torch.uint8).reshape(2, 2, 3, 4, 5) - camera_b = (camera_a.to(torch.int16) * 3 % 251).to(torch.uint8) +def test_groot_n1_7_single_camera_is_one_element_tuple_and_preserves_bytes(): + camera = torch.arange(2 * 2 * 3 * 4 * 5, dtype=torch.uint8).reshape(2, 2, 3, 4, 5) pack_step = GrootN17PackInputsStep( normalize_min_max=False, - video_modality_keys=["image", "wrist_image"], + video_modality_keys=["image"], ) packed = pack_step( { TransitionKey.OBSERVATION: { - f"{OBS_IMAGES}.image": camera_a, - f"{OBS_IMAGES}.image2": camera_b, + f"{OBS_IMAGES}.image": camera, }, TransitionKey.COMPLEMENTARY_DATA: {"task": ["a", "b"]}, } ) - video = packed[TransitionKey.OBSERVATION]["video"] - legacy_video = np.stack( - [ - camera_a.permute(0, 1, 3, 4, 2).numpy(), - camera_b.permute(0, 1, 3, 4, 2).numpy(), - ], - axis=2, - ) - encode_step = GrootN17VLMEncodeStep() + cameras = packed[TransitionKey.OBSERVATION]["video"] - tensor_frames = encode_step._build_sample_images(video, target_device=None) - legacy_frames = encode_step._build_sample_images(legacy_video, target_device=None) + assert isinstance(cameras, tuple) + assert len(cameras) == 1 + assert cameras[0].data_ptr() == camera.data_ptr() - assert len(tensor_frames) == len(legacy_frames) == 2 - for tensor_sample, legacy_sample in zip(tensor_frames, legacy_frames, strict=True): - assert len(tensor_sample) == len(legacy_sample) == 4 - for tensor_frame, legacy_frame in zip(tensor_sample, legacy_sample, strict=True): - torch.testing.assert_close(tensor_frame, legacy_frame, rtol=0, atol=0) + frames = GrootN17VLMEncodeStep()._build_sample_images(cameras, target_device=None) + for batch_idx, sample_frames in enumerate(frames): + assert len(sample_frames) == 2 + for timestep, frame in enumerate(sample_frames): + torch.testing.assert_close(frame, camera[batch_idx, timestep], rtol=0, atol=0) def test_groot_n1_7_postprocessor_clips_normalized_action_before_unnormalizing(): @@ -1796,7 +1787,7 @@ def __call__(self, text, images, return_tensors, padding): step._proc = fake_proc transition = { TransitionKey.OBSERVATION: { - "video": np.zeros((2, 1, 1, 2, 2, 3), dtype=np.uint8), + "video": (torch.zeros((2, 1, 3, 2, 2), dtype=torch.uint8),), }, TransitionKey.COMPLEMENTARY_DATA: { "language": ["first task", "second task"], @@ -1851,15 +1842,18 @@ def __call__(self, text, images, return_tensors, padding): fake_proc = FakeProcessor() step = GrootN17VLMEncodeStep() step._proc = fake_proc - video = np.zeros((2, 2, 2, 2, 2, 3), dtype=np.uint8) + cameras = ( + torch.zeros((2, 2, 3, 2, 2), dtype=torch.uint8), + torch.zeros((2, 2, 3, 2, 2), dtype=torch.uint8), + ) image_id = 1 for batch_idx in range(2): for timestep in range(2): for view_idx in range(2): - video[batch_idx, timestep, view_idx, :, :, :] = image_id + cameras[view_idx][batch_idx, timestep] = image_id image_id += 1 transition = { - TransitionKey.OBSERVATION: {"video": video}, + TransitionKey.OBSERVATION: {"video": cameras}, TransitionKey.COMPLEMENTARY_DATA: {"language": ["task a", "task b"]}, } @@ -1980,7 +1974,9 @@ def __call__(self, text, images, return_tensors, padding): camera_a = np.arange(3 * 5 * 3, dtype=np.uint8).reshape(3, 5, 3) camera_b = (np.arange(3 * 5 * 3, dtype=np.uint16).reshape(3, 5, 3) * 3 % 251).astype(np.uint8) - video = np.stack([camera_a, camera_b], axis=0).reshape(1, 1, 2, 3, 5, 3) + cameras = tuple( + torch.from_numpy(camera).permute(2, 0, 1).unsqueeze(0).unsqueeze(0) for camera in (camera_a, camera_b) + ) fake_proc = FakeProcessor() step = GrootN17VLMEncodeStep( image_target_size=[8, 8], @@ -1992,7 +1988,7 @@ def __call__(self, text, images, return_tensors, padding): step( { - TransitionKey.OBSERVATION: {"video": video}, + TransitionKey.OBSERVATION: {"video": cameras}, TransitionKey.COMPLEMENTARY_DATA: {"language": ["move"]}, } ) diff --git a/tests/policies/groot/test_groot_n1_7_oss_parity.py b/tests/policies/groot/test_groot_n1_7_oss_parity.py index 3fced5909a8..e39c885519f 100644 --- a/tests/policies/groot/test_groot_n1_7_oss_parity.py +++ b/tests/policies/groot/test_groot_n1_7_oss_parity.py @@ -96,7 +96,10 @@ def __call__(self, **kwargs): step._proc = processor transition = { TransitionKey.OBSERVATION: { - "video": np.zeros((1, 1, 2, 480, 640, 3), dtype=np.uint8), + "video": ( + torch.zeros((1, 1, 3, 480, 640), dtype=torch.uint8), + torch.zeros((1, 1, 3, 480, 640), dtype=torch.uint8), + ), }, TransitionKey.COMPLEMENTARY_DATA: {"language": ["pick up the vial"]}, } diff --git a/tests/policies/groot/test_groot_train_random_crop.py b/tests/policies/groot/test_groot_train_random_crop.py index c8667ef6989..50499c34161 100644 --- a/tests/policies/groot/test_groot_train_random_crop.py +++ b/tests/policies/groot/test_groot_train_random_crop.py @@ -106,7 +106,8 @@ def crop_at(position): def _video(img, views=2): - return np.stack([img] * views, axis=0).reshape(1, 1, views, *img.shape) + frame = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).unsqueeze(0) + return tuple(frame.clone() for _ in range(views)) def _step(training): From a94db76027aff30a7378466ae0766bef57e42155 Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Mon, 6 Jul 2026 08:26:24 -0700 Subject: [PATCH 08/13] refactor(policy): declare input image format on policies --- src/lerobot/policies/__init__.py | 3 ++- src/lerobot/policies/groot/modeling_groot.py | 3 ++- src/lerobot/policies/groot/processor_groot.py | 6 ----- src/lerobot/policies/pretrained.py | 11 +++++++++- src/lerobot/processor/__init__.py | 2 -- src/lerobot/processor/pipeline.py | 22 ------------------- src/lerobot/scripts/lerobot_train.py | 17 ++++++++------ tests/policies/groot/test_groot_n1_7.py | 4 ++-- tests/processor/test_pipeline.py | 16 ++++---------- tests/training/test_visual_validation.py | 8 +++---- 10 files changed, 34 insertions(+), 58 deletions(-) diff --git a/src/lerobot/policies/__init__.py b/src/lerobot/policies/__init__.py index 7f0bed2e051..dbb26d09b46 100644 --- a/src/lerobot/policies/__init__.py +++ b/src/lerobot/policies/__init__.py @@ -28,7 +28,7 @@ from .pi0.configuration_pi0 import PI0Config as PI0Config from .pi0_fast.configuration_pi0_fast import PI0FastConfig as PI0FastConfig from .pi05.configuration_pi05 import PI05Config as PI05Config -from .pretrained import PreTrainedPolicy as PreTrainedPolicy +from .pretrained import ImageInputFormat as ImageInputFormat, PreTrainedPolicy as PreTrainedPolicy from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig from .utils import make_robot_action, prepare_observation_for_inference @@ -61,6 +61,7 @@ "WallXConfig", "XVLAConfig", # Base class + "ImageInputFormat", "PreTrainedPolicy", # RTC utilities "ActionInterpolator", diff --git a/src/lerobot/policies/groot/modeling_groot.py b/src/lerobot/policies/groot/modeling_groot.py index 415af930990..8a470e3f3dd 100644 --- a/src/lerobot/policies/groot/modeling_groot.py +++ b/src/lerobot/policies/groot/modeling_groot.py @@ -39,7 +39,7 @@ from lerobot.utils.constants import ACTION, OBS_IMAGES from lerobot.utils.import_utils import _transformers_available, require_package -from ..pretrained import PreTrainedPolicy +from ..pretrained import ImageInputFormat, PreTrainedPolicy from ..utils import get_device_from_parameters from .configuration_groot import ( GROOT_N1_5, @@ -67,6 +67,7 @@ class GrootPolicy(PreTrainedPolicy): name = "groot" config_class = GrootConfig + input_image_format = ImageInputFormat.UINT8_0_255 def __init__(self, config: GrootConfig, **kwargs): """Initialize Groot policy wrapper.""" diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index d0c2fa86e8a..3a043c919ae 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -53,7 +53,6 @@ AbsoluteActionsProcessorStep, AddBatchDimensionProcessorStep, DeviceProcessorStep, - ImageInputFormat, PolicyAction, PolicyProcessorPipeline, ProcessorStep, @@ -553,10 +552,6 @@ def _load_groot_processor_pipelines( to_transition=policy_action_to_transition, to_output=transition_to_policy_action, ) - # Older serialized GR00T pipelines predate the raw-image contract. GR00T's - # packer consumes worker-produced uint8 directly, so upgrade them at load - # time rather than falling back to the global float compatibility default. - preprocessor.input_image_format = ImageInputFormat.UINT8_0_255 return preprocessor, postprocessor @@ -1329,7 +1324,6 @@ def make_groot_pre_post_processors( PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( steps=input_steps, name=POLICY_PREPROCESSOR_DEFAULT_NAME, - input_image_format=ImageInputFormat.UINT8_0_255, ), PolicyProcessorPipeline[PolicyAction, PolicyAction]( steps=output_steps, diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index 702569b8c7b..3e7caf09786 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -18,10 +18,11 @@ import dataclasses import logging import os +from enum import StrEnum from importlib.resources import files from pathlib import Path from tempfile import TemporaryDirectory -from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack +from typing import TYPE_CHECKING, ClassVar, TypedDict, TypeVar, Unpack import packaging import safetensors @@ -44,6 +45,13 @@ from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata +class ImageInputFormat(StrEnum): + """Raw image representation expected by a policy before preprocessing.""" + + FLOAT32_0_1 = "float32_0_1" + UINT8_0_255 = "uint8_0_255" + + def _build_card_context( cfg: TrainPipelineConfig | None, dataset_meta: LeRobotDatasetMetadata | None, @@ -102,6 +110,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): config_class: None name: None + input_image_format: ClassVar[ImageInputFormat] = ImageInputFormat.FLOAT32_0_1 def __init__(self, config: PreTrainedConfig, *inputs, **kwargs): super().__init__() diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 360c864c7ee..fe35af4b4c0 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -70,7 +70,6 @@ DataProcessorPipeline, DoneProcessorStep, IdentityProcessorStep, - ImageInputFormat, InfoProcessorStep, ObservationProcessorStep, PolicyActionProcessorStep, @@ -126,7 +125,6 @@ "GymHILAdapterProcessorStep", "GripperPenaltyProcessorStep", "hotswap_stats", - "ImageInputFormat", "IdentityProcessorStep", "ImageCropResizeProcessorStep", "InfoProcessorStep", diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py index fc3a925be9c..b9b9c6c4350 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -37,7 +37,6 @@ from collections.abc import Callable, Iterable, Sequence from copy import deepcopy from dataclasses import dataclass, field -from enum import StrEnum from pathlib import Path from typing import Any, TypedDict, TypeVar, cast @@ -57,18 +56,6 @@ TOutput = TypeVar("TOutput") -class ImageInputFormat(StrEnum): - """Raw image representation expected by a policy preprocessor. - - Dataset workers use uint8 images for compact IPC. The training loop uses - this contract to preserve uint8 for preprocessors that consume it directly, - while retaining the historical float32 [0, 1] input for existing policies. - """ - - FLOAT32_0_1 = "float32_0_1" - UINT8_0_255 = "uint8_0_255" - - class ProcessorStepRegistry: """A registry for ProcessorStep classes to allow instantiation from a string name. @@ -283,11 +270,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): steps: Sequence[ProcessorStep] = field(default_factory=list) name: str = "DataProcessorPipeline" - # This describes the raw images accepted before ``to_transition`` and the - # first processor step. It is primarily meaningful for policy preprocessors; - # the float default keeps existing and third-party pipelines compatible. - input_image_format: ImageInputFormat = ImageInputFormat.FLOAT32_0_1 - to_transition: Callable[[TInput], EnvTransition] = field( default_factory=lambda: cast(Callable[[TInput], EnvTransition], batch_to_transition), repr=False ) @@ -455,7 +437,6 @@ def get_config(self) -> dict[str, Any]: sanitized_name = self._get_sanitized_name() pipeline_config: dict[str, Any] = { "name": self.name, - "input_image_format": self.input_image_format.value, "steps": [], } @@ -760,7 +741,6 @@ def from_pretrained( pipeline = cls( steps=steps, name=loaded_config.get("name", "DataProcessorPipeline"), - input_image_format=loaded_config.get("input_image_format", ImageInputFormat.FLOAT32_0_1), to_transition=to_transition or cast(Callable[[TInput], EnvTransition], batch_to_transition), to_output=to_output or cast(Callable[[EnvTransition], TOutput], transition_to_batch), ) @@ -797,7 +777,6 @@ def from_config( pipeline = cls( steps=steps, name=config.get("name", "DataProcessorPipeline"), - input_image_format=config.get("input_image_format", ImageInputFormat.FLOAT32_0_1), to_transition=to_transition or cast(Callable[[TInput], EnvTransition], batch_to_transition), to_output=to_output or cast(Callable[[EnvTransition], TOutput], transition_to_batch), ) @@ -1554,7 +1533,6 @@ def __repr__(self) -> str: def __post_init__(self): """Validates that all provided steps are instances of `ProcessorStep`.""" - self.input_image_format = ImageInputFormat(self.input_image_format) for i, step in enumerate(self.steps): if not isinstance(step, ProcessorStep): raise TypeError(f"Step {i} ({type(step).__name__}) must inherit from ProcessorStep") diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 31f7514a1ff..5cbed5e4d1e 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -54,8 +54,7 @@ from lerobot.envs import close_envs, make_env, make_env_pre_post_processors from lerobot.jobs import submit_to_hf from lerobot.optim.factory import make_optimizer_and_scheduler -from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors -from lerobot.processor import ImageInputFormat, PolicyProcessorPipeline +from lerobot.policies import ImageInputFormat, PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.rewards import make_reward_pre_post_processors from lerobot.utils.collate import lerobot_collate_fn from lerobot.utils.import_utils import register_third_party_plugins @@ -75,11 +74,11 @@ def prepare_images_for_policy( batch: dict[str, Any], camera_keys: list[str], - preprocessor: PolicyProcessorPipeline, + input_image_format: ImageInputFormat, ) -> None: - """Adapt worker-produced images to the policy preprocessor's raw input contract.""" + """Adapt worker-produced images to the policy's raw input contract.""" - if preprocessor.input_image_format is ImageInputFormat.UINT8_0_255: + if input_image_format is ImageInputFormat.UINT8_0_255: return for cam_key in camera_keys: @@ -312,6 +311,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): rename_map=cfg.rename_map, ) + # Capture this before Accelerate or PEFT wraps the policy. Reward models do + # not currently declare an image format and retain the historical default. + input_image_format = getattr(policy, "input_image_format", ImageInputFormat.FLOAT32_0_1) + if cfg.peft is not None: if cfg.is_reward_model_training: raise ValueError("PEFT is only supported for policy training. ") @@ -582,7 +585,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): for _ in range(step, cfg.steps): start_time = time.perf_counter() batch = next(dl_iter) - prepare_images_for_policy(batch, dataset.meta.camera_keys, preprocessor) + prepare_images_for_policy(batch, dataset.meta.camera_keys, input_image_format) batch = preprocessor(batch) train_tracker.dataloading_s = time.perf_counter() - start_time @@ -635,7 +638,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): n_eval_batches = 0 with torch.no_grad(), accelerator.autocast(): for eval_batch in eval_dataloader: - prepare_images_for_policy(eval_batch, dataset.meta.camera_keys, preprocessor) + prepare_images_for_policy(eval_batch, dataset.meta.camera_keys, input_image_format) eval_batch = preprocessor(eval_batch) loss, _ = policy.forward(eval_batch) eval_loss_sum += loss.item() diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index 99057cb0a53..f6d81ef11ee 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -27,6 +27,7 @@ from torch import nn from lerobot.configs import FeatureType, PolicyFeature +from lerobot.policies import ImageInputFormat from lerobot.policies.factory import make_policy_config, make_pre_post_processors from lerobot.policies.groot.configuration_groot import ( GROOT_ACTION_DECODE_TRANSFORM_LIBERO, @@ -51,7 +52,6 @@ ) from lerobot.processor import ( AbsoluteActionsProcessorStep, - ImageInputFormat, PolicyProcessorPipeline, RelativeActionsProcessorStep, ) @@ -1698,7 +1698,7 @@ def test_groot_n1_7_processors_are_registered_lazily_without_external_gr00t(): preprocessor, _ = make_groot_pre_post_processors(config) step_types = {type(step) for step in preprocessor.steps} - assert preprocessor.input_image_format is ImageInputFormat.UINT8_0_255 + assert GrootPolicy.input_image_format is ImageInputFormat.UINT8_0_255 assert GrootN17PackInputsStep in step_types assert GrootN17VLMEncodeStep in step_types assert "gr00t" not in sys.modules diff --git a/tests/processor/test_pipeline.py b/tests/processor/test_pipeline.py index 6d33de828d8..56d007f9cc4 100644 --- a/tests/processor/test_pipeline.py +++ b/tests/processor/test_pipeline.py @@ -33,7 +33,6 @@ from lerobot.processor import ( DataProcessorPipeline, EnvTransition, - ImageInputFormat, ProcessorStep, ProcessorStepRegistry, TransitionKey, @@ -605,18 +604,11 @@ def test_save_and_load_pretrained(): assert loaded_pipeline.steps[1].counter == 10 -def test_input_image_format_is_backward_compatible_and_serialized(): - legacy_pipeline = DataProcessorPipeline.from_config({"steps": []}) - assert legacy_pipeline.input_image_format is ImageInputFormat.FLOAT32_0_1 +def test_policy_image_format_is_not_owned_by_pipeline(): + pipeline = DataProcessorPipeline([]) - pipeline = DataProcessorPipeline( - [], - input_image_format=ImageInputFormat.UINT8_0_255, - ) - loaded_pipeline = DataProcessorPipeline.from_config(pipeline.get_config()) - - assert pipeline.get_config()["input_image_format"] == "uint8_0_255" - assert loaded_pipeline.input_image_format is ImageInputFormat.UINT8_0_255 + assert not hasattr(pipeline, "input_image_format") + assert "input_image_format" not in pipeline.get_config() def test_step_without_optional_methods(): diff --git a/tests/training/test_visual_validation.py b/tests/training/test_visual_validation.py index 38dfbc8b586..60ee00ebe13 100644 --- a/tests/training/test_visual_validation.py +++ b/tests/training/test_visual_validation.py @@ -38,8 +38,8 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.train import TrainPipelineConfig from lerobot.datasets.lerobot_dataset import LeRobotDataset +from lerobot.policies import ImageInputFormat, PreTrainedPolicy from lerobot.policies.factory import make_policy_config -from lerobot.processor import ImageInputFormat, PolicyProcessorPipeline from lerobot.scripts.lerobot_train import prepare_images_for_policy, train from lerobot.utils.device_utils import auto_select_torch_device @@ -59,14 +59,14 @@ def temp_dir(tmp_path): DEVICE = auto_select_torch_device() -def test_prepare_images_for_policy_obeys_preprocessor_contract(): +def test_prepare_images_for_policy_obeys_policy_contract(): image = torch.tensor([0, 127, 255], dtype=torch.uint8) float_batch = {"observation.images.front": image.clone()} prepare_images_for_policy( float_batch, ["observation.images.front"], - PolicyProcessorPipeline([]), + PreTrainedPolicy.input_image_format, ) assert float_batch["observation.images.front"].dtype == torch.float32 torch.testing.assert_close( @@ -79,7 +79,7 @@ def test_prepare_images_for_policy_obeys_preprocessor_contract(): prepare_images_for_policy( uint8_batch, ["observation.images.front"], - PolicyProcessorPipeline([], input_image_format=ImageInputFormat.UINT8_0_255), + ImageInputFormat.UINT8_0_255, ) assert uint8_batch["observation.images.front"].data_ptr() == uint8_image.data_ptr() assert uint8_batch["observation.images.front"].dtype == torch.uint8 From 419ae661c3c4d65bf984ceec4fb72c971247b1c3 Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Mon, 6 Jul 2026 10:22:21 -0700 Subject: [PATCH 09/13] refactor(groot): adapt legacy image dtype once --- src/lerobot/policies/groot/processor_groot.py | 37 +++++++------------ tests/policies/groot/test_groot_n1_7.py | 21 +++++++++++ 2 files changed, 35 insertions(+), 23 deletions(-) diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 3a043c919ae..8a17765c237 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -1337,16 +1337,18 @@ def make_groot_pre_post_processors( # GR00T specific processor steps -def _as_video_tensor_btchw(image: torch.Tensor) -> torch.Tensor: - """Preserve a LeRobot image tensor while making its time dimension explicit.""" +def _as_uint8_video_tensor_btchw(image: torch.Tensor) -> torch.Tensor: + """Make the time dimension explicit and adapt legacy float inference images once.""" - if image.ndim == 4: - return image.unsqueeze(1) - if image.ndim == 5: - return image - raise ValueError( - f"Expected image tensor shape (B, C, H, W) or (B, T, C, H, W), got {tuple(image.shape)}." - ) + if image.ndim not in (4, 5): + raise ValueError( + f"Expected image tensor shape (B, C, H, W) or (B, T, C, H, W), got {tuple(image.shape)}." + ) + if image.dtype.is_floating_point: + image = (image.clamp(0, 1) * 255.0).to(torch.uint8) + elif image.dtype != torch.uint8: + image = image.to(torch.uint8) + return image.unsqueeze(1) if image.ndim == 4 else image def _align_video_horizon_tensor(video: torch.Tensor, horizon: int | None) -> torch.Tensor: @@ -1363,17 +1365,8 @@ def _align_video_horizon_tensor(video: torch.Tensor, horizon: int | None) -> tor return torch.cat([pad, video], dim=1) -def _uint8_image_tensor(image: torch.Tensor) -> torch.Tensor: - if image.dtype.is_floating_point: - return (image.clamp(0, 1) * 255.0).to(torch.uint8) - if image.dtype != torch.uint8: - return image.to(torch.uint8) - return image - - def _uint8_image_numpy_hwc(image: torch.Tensor) -> np.ndarray: - image = _uint8_image_tensor(image).detach().cpu() - return image.permute(1, 2, 0).contiguous().numpy() + return image.detach().cpu().permute(1, 2, 0).contiguous().numpy() def _build_n1_7_processor(model_name: str = GROOT_N1_7_BACKBONE_MODEL) -> ProcessorMixin: @@ -1863,7 +1856,7 @@ def _cache_raw_state(state: torch.Tensor) -> None: packed_cameras: _GrootN17CameraBatch = () if img_keys: packed_cameras = tuple( - _align_video_horizon_tensor(_as_video_tensor_btchw(obs[key]), self.video_horizon) + _align_video_horizon_tensor(_as_uint8_video_tensor_btchw(obs[key]), self.video_horizon) for key in img_keys ) # Keep the pinned worker tensors in their native channels-first @@ -2151,9 +2144,7 @@ def _build_sample_images( for camera in cameras: if target_device is not None and camera.device != target_device: camera = camera.to(target_device, non_blocking=(target_device.type == "cuda")) - # Float observations from direct/inference callers remain supported, - # but conversion happens after transfer. Training uint8 is a no-op. - prepared_cameras.append(_uint8_image_tensor(camera)) + prepared_cameras.append(camera) return [ [ diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index f6d81ef11ee..6754d424d27 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -1289,6 +1289,27 @@ def test_groot_n1_7_single_camera_is_one_element_tuple_and_preserves_bytes(): torch.testing.assert_close(frame, camera[batch_idx, timestep], rtol=0, atol=0) +def test_groot_n1_7_pack_inputs_adapts_legacy_float_images_once(): + camera = torch.tensor([0.0, 0.5, 1.0], dtype=torch.float32).reshape(1, 3, 1, 1) + pack_step = GrootN17PackInputsStep(normalize_min_max=False, video_modality_keys=["image"]) + + packed = pack_step( + { + TransitionKey.OBSERVATION: {f"{OBS_IMAGES}.image": camera}, + TransitionKey.COMPLEMENTARY_DATA: {"task": ["move"]}, + } + ) + packed_camera = packed[TransitionKey.OBSERVATION]["video"][0] + + assert packed_camera.dtype == torch.uint8 + torch.testing.assert_close( + packed_camera[:, 0, :, 0, 0], + torch.tensor([[0, 127, 255]], dtype=torch.uint8), + rtol=0, + atol=0, + ) + + def test_groot_n1_7_postprocessor_clips_normalized_action_before_unnormalizing(): step = GrootActionUnpackUnnormalizeStep( env_action_dim=3, From 8576b3d0f97b3ed152fab8aa532d1bc42e8cb481 Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Mon, 6 Jul 2026 10:39:31 -0700 Subject: [PATCH 10/13] refactor: simplify policy image input handling --- src/lerobot/policies/groot/processor_groot.py | 42 ++++---------- src/lerobot/scripts/lerobot_train.py | 25 +++------ tests/policies/groot/test_groot_n1_7.py | 55 +------------------ tests/processor/test_pipeline.py | 7 --- tests/training/test_visual_validation.py | 30 +--------- 5 files changed, 23 insertions(+), 136 deletions(-) diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 8a17765c237..365a3894001 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -1338,16 +1338,11 @@ def make_groot_pre_post_processors( def _as_uint8_video_tensor_btchw(image: torch.Tensor) -> torch.Tensor: - """Make the time dimension explicit and adapt legacy float inference images once.""" - if image.ndim not in (4, 5): raise ValueError( f"Expected image tensor shape (B, C, H, W) or (B, T, C, H, W), got {tuple(image.shape)}." ) - if image.dtype.is_floating_point: - image = (image.clamp(0, 1) * 255.0).to(torch.uint8) - elif image.dtype != torch.uint8: - image = image.to(torch.uint8) + image = tv_functional.to_dtype(image, torch.uint8, scale=True) return image.unsqueeze(1) if image.ndim == 4 else image @@ -1852,30 +1847,16 @@ def _cache_raw_state(state: torch.Tensor) -> None: if grouped: self._last_raw_state = grouped + batch_size, batch_device = infer_n1_7_batch_size_and_device(obs, transition.get(TransitionKey.ACTION)) img_keys = self._ordered_image_keys(obs) - packed_cameras: _GrootN17CameraBatch = () if img_keys: - packed_cameras = tuple( + obs["video"] = tuple( _align_video_horizon_tensor(_as_uint8_video_tensor_btchw(obs[key]), self.video_horizon) for key in img_keys ) - # Keep the pinned worker tensors in their native channels-first - # representation. The VLM step transfers each view directly and - # never creates a CPU NumPy/HWC staging buffer. - obs["video"] = packed_cameras - image_keys_to_remove = [key for key in obs if key.startswith(OBS_IMAGES)] - if OBS_IMAGE in obs: - image_keys_to_remove.append(OBS_IMAGE) - for k in image_keys_to_remove: - obs.pop(k, None) - - if packed_cameras: - batch_size = packed_cameras[0].shape[0] - batch_device = packed_cameras[0].device - else: - batch_size, batch_device = infer_n1_7_batch_size_and_device( - obs, transition.get(TransitionKey.ACTION) - ) + # Preserve channels-first tensors until VLM preprocessing. + for key in [key for key in obs if key.startswith(OBS_IMAGES) or key == OBS_IMAGE]: + obs.pop(key) comp["language"] = prepare_n1_7_language_batch( comp.get(self.language_key), batch_size, @@ -2140,11 +2121,12 @@ def _build_sample_images( ) return sample_images - prepared_cameras: list[torch.Tensor] = [] - for camera in cameras: - if target_device is not None and camera.device != target_device: - camera = camera.to(target_device, non_blocking=(target_device.type == "cuda")) - prepared_cameras.append(camera) + prepared_cameras = [ + camera.to(target_device, non_blocking=(target_device.type == "cuda")) + if target_device is not None and camera.device != target_device + else camera + for camera in cameras + ] return [ [ diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 5cbed5e4d1e..5cf3b27d095 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -71,21 +71,6 @@ from .lerobot_eval import eval_policy_all -def prepare_images_for_policy( - batch: dict[str, Any], - camera_keys: list[str], - input_image_format: ImageInputFormat, -) -> None: - """Adapt worker-produced images to the policy's raw input contract.""" - - if input_image_format is ImageInputFormat.UINT8_0_255: - return - - for cam_key in camera_keys: - if cam_key in batch and batch[cam_key].dtype == torch.uint8: - batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0 - - def update_policy( train_metrics: MetricsTracker, policy: PreTrainedPolicy, @@ -585,7 +570,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): for _ in range(step, cfg.steps): start_time = time.perf_counter() batch = next(dl_iter) - prepare_images_for_policy(batch, dataset.meta.camera_keys, input_image_format) + if input_image_format is ImageInputFormat.FLOAT32_0_1: + for cam_key in dataset.meta.camera_keys: + if cam_key in batch and batch[cam_key].dtype == torch.uint8: + batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0 batch = preprocessor(batch) train_tracker.dataloading_s = time.perf_counter() - start_time @@ -638,7 +626,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): n_eval_batches = 0 with torch.no_grad(), accelerator.autocast(): for eval_batch in eval_dataloader: - prepare_images_for_policy(eval_batch, dataset.meta.camera_keys, input_image_format) + if input_image_format is ImageInputFormat.FLOAT32_0_1: + for cam_key in dataset.meta.camera_keys: + if cam_key in eval_batch and eval_batch[cam_key].dtype == torch.uint8: + eval_batch[cam_key] = eval_batch[cam_key].to(dtype=torch.float32) / 255.0 eval_batch = preprocessor(eval_batch) loss, _ = policy.forward(eval_batch) eval_loss_sum += loss.item() diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index 6754d424d27..b639dc9331e 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -1235,12 +1235,11 @@ def test_groot_n1_7_pack_inputs_orders_video_by_checkpoint_modality_keys(): normalize_min_max=False, video_modality_keys=["image", "wrist_image"], ) - extra = torch.full((1, 3, 2, 2), 33, dtype=torch.uint8) wrist = torch.full((1, 3, 2, 2), 22, dtype=torch.uint8) front = torch.full((1, 3, 2, 2), 11, dtype=torch.uint8) transition = { TransitionKey.OBSERVATION: { - f"{OBS_IMAGES}.zz_extra": extra, + f"{OBS_IMAGES}.zz_extra": torch.full((1, 3, 2, 2), 33, dtype=torch.uint8), f"{OBS_IMAGES}.image2": wrist, f"{OBS_IMAGES}.image": front, OBS_STATE: torch.zeros(1, 8), @@ -1251,65 +1250,15 @@ def test_groot_n1_7_pack_inputs_orders_video_by_checkpoint_modality_keys(): output = step(transition) video = output[TransitionKey.OBSERVATION]["video"] - assert len(video) == 2 + assert isinstance(video, tuple) and len(video) == 2 assert video[0].shape == (1, 1, 3, 2, 2) assert video[0].data_ptr() == front.data_ptr() assert video[1].data_ptr() == wrist.data_ptr() - assert torch.unique(video[0]).tolist() == [11] - assert torch.unique(video[1]).tolist() == [22] assert f"{OBS_IMAGES}.zz_extra" not in output[TransitionKey.OBSERVATION] assert f"{OBS_IMAGES}.image" not in output[TransitionKey.OBSERVATION] assert f"{OBS_IMAGES}.image2" not in output[TransitionKey.OBSERVATION] -def test_groot_n1_7_single_camera_is_one_element_tuple_and_preserves_bytes(): - camera = torch.arange(2 * 2 * 3 * 4 * 5, dtype=torch.uint8).reshape(2, 2, 3, 4, 5) - pack_step = GrootN17PackInputsStep( - normalize_min_max=False, - video_modality_keys=["image"], - ) - packed = pack_step( - { - TransitionKey.OBSERVATION: { - f"{OBS_IMAGES}.image": camera, - }, - TransitionKey.COMPLEMENTARY_DATA: {"task": ["a", "b"]}, - } - ) - cameras = packed[TransitionKey.OBSERVATION]["video"] - - assert isinstance(cameras, tuple) - assert len(cameras) == 1 - assert cameras[0].data_ptr() == camera.data_ptr() - - frames = GrootN17VLMEncodeStep()._build_sample_images(cameras, target_device=None) - for batch_idx, sample_frames in enumerate(frames): - assert len(sample_frames) == 2 - for timestep, frame in enumerate(sample_frames): - torch.testing.assert_close(frame, camera[batch_idx, timestep], rtol=0, atol=0) - - -def test_groot_n1_7_pack_inputs_adapts_legacy_float_images_once(): - camera = torch.tensor([0.0, 0.5, 1.0], dtype=torch.float32).reshape(1, 3, 1, 1) - pack_step = GrootN17PackInputsStep(normalize_min_max=False, video_modality_keys=["image"]) - - packed = pack_step( - { - TransitionKey.OBSERVATION: {f"{OBS_IMAGES}.image": camera}, - TransitionKey.COMPLEMENTARY_DATA: {"task": ["move"]}, - } - ) - packed_camera = packed[TransitionKey.OBSERVATION]["video"][0] - - assert packed_camera.dtype == torch.uint8 - torch.testing.assert_close( - packed_camera[:, 0, :, 0, 0], - torch.tensor([[0, 127, 255]], dtype=torch.uint8), - rtol=0, - atol=0, - ) - - def test_groot_n1_7_postprocessor_clips_normalized_action_before_unnormalizing(): step = GrootActionUnpackUnnormalizeStep( env_action_dim=3, diff --git a/tests/processor/test_pipeline.py b/tests/processor/test_pipeline.py index 56d007f9cc4..0e9746a6302 100644 --- a/tests/processor/test_pipeline.py +++ b/tests/processor/test_pipeline.py @@ -604,13 +604,6 @@ def test_save_and_load_pretrained(): assert loaded_pipeline.steps[1].counter == 10 -def test_policy_image_format_is_not_owned_by_pipeline(): - pipeline = DataProcessorPipeline([]) - - assert not hasattr(pipeline, "input_image_format") - assert "input_image_format" not in pipeline.get_config() - - def test_step_without_optional_methods(): """Test pipeline with steps that don't implement optional methods.""" step = MockStepWithoutOptionalMethods(multiplier=3.0) diff --git a/tests/training/test_visual_validation.py b/tests/training/test_visual_validation.py index 60ee00ebe13..1df8006b273 100644 --- a/tests/training/test_visual_validation.py +++ b/tests/training/test_visual_validation.py @@ -30,7 +30,6 @@ import numpy as np import pytest -import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") @@ -38,9 +37,8 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.train import TrainPipelineConfig from lerobot.datasets.lerobot_dataset import LeRobotDataset -from lerobot.policies import ImageInputFormat, PreTrainedPolicy from lerobot.policies.factory import make_policy_config -from lerobot.scripts.lerobot_train import prepare_images_for_policy, train +from lerobot.scripts.lerobot_train import train from lerobot.utils.device_utils import auto_select_torch_device pytest.importorskip("transformers") @@ -59,32 +57,6 @@ def temp_dir(tmp_path): DEVICE = auto_select_torch_device() -def test_prepare_images_for_policy_obeys_policy_contract(): - image = torch.tensor([0, 127, 255], dtype=torch.uint8) - - float_batch = {"observation.images.front": image.clone()} - prepare_images_for_policy( - float_batch, - ["observation.images.front"], - PreTrainedPolicy.input_image_format, - ) - assert float_batch["observation.images.front"].dtype == torch.float32 - torch.testing.assert_close( - float_batch["observation.images.front"], - torch.tensor([0.0, 127.0 / 255.0, 1.0]), - ) - - uint8_image = image.clone() - uint8_batch = {"observation.images.front": uint8_image} - prepare_images_for_policy( - uint8_batch, - ["observation.images.front"], - ImageInputFormat.UINT8_0_255, - ) - assert uint8_batch["observation.images.front"].data_ptr() == uint8_image.data_ptr() - assert uint8_batch["observation.images.front"].dtype == torch.uint8 - - def make_dummy_dataset(camera_keys, tmp_path): """Creates a minimal dummy dataset for testing rename_mapping logic.""" features = { From cb28921bfbd6d0b37b68c325e5c401320dcf3019 Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Wed, 8 Jul 2026 14:55:21 -0700 Subject: [PATCH 11/13] refactor(train): centralize camera image preparation --- src/lerobot/scripts/lerobot_train.py | 27 ++++++++++++++------- tests/training/test_visual_validation.py | 30 +++++++++++++++++++++++- 2 files changed, 48 insertions(+), 9 deletions(-) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 5cf3b27d095..f47ffd18a05 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -71,6 +71,23 @@ from .lerobot_eval import eval_policy_all +def prep_camera_images_for_policy( + batch: dict[str, Any], + camera_keys: list[str], + input_image_format: ImageInputFormat, +) -> None: + """Adapt worker-produced camera images to the policy's raw input contract.""" + + if input_image_format is ImageInputFormat.UINT8_0_255: + return + if input_image_format is not ImageInputFormat.FLOAT32_0_1: + raise ValueError(f"Unsupported policy image input format: {input_image_format}") + + for cam_key in camera_keys: + if cam_key in batch and batch[cam_key].dtype == torch.uint8: + batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0 + + def update_policy( train_metrics: MetricsTracker, policy: PreTrainedPolicy, @@ -570,10 +587,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): for _ in range(step, cfg.steps): start_time = time.perf_counter() batch = next(dl_iter) - if input_image_format is ImageInputFormat.FLOAT32_0_1: - for cam_key in dataset.meta.camera_keys: - if cam_key in batch and batch[cam_key].dtype == torch.uint8: - batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0 + prep_camera_images_for_policy(batch, dataset.meta.camera_keys, input_image_format) batch = preprocessor(batch) train_tracker.dataloading_s = time.perf_counter() - start_time @@ -626,10 +640,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): n_eval_batches = 0 with torch.no_grad(), accelerator.autocast(): for eval_batch in eval_dataloader: - if input_image_format is ImageInputFormat.FLOAT32_0_1: - for cam_key in dataset.meta.camera_keys: - if cam_key in eval_batch and eval_batch[cam_key].dtype == torch.uint8: - eval_batch[cam_key] = eval_batch[cam_key].to(dtype=torch.float32) / 255.0 + prep_camera_images_for_policy(eval_batch, dataset.meta.camera_keys, input_image_format) eval_batch = preprocessor(eval_batch) loss, _ = policy.forward(eval_batch) eval_loss_sum += loss.item() diff --git a/tests/training/test_visual_validation.py b/tests/training/test_visual_validation.py index 1df8006b273..efc86553070 100644 --- a/tests/training/test_visual_validation.py +++ b/tests/training/test_visual_validation.py @@ -30,6 +30,7 @@ import numpy as np import pytest +import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") @@ -37,8 +38,9 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.train import TrainPipelineConfig from lerobot.datasets.lerobot_dataset import LeRobotDataset +from lerobot.policies import ImageInputFormat from lerobot.policies.factory import make_policy_config -from lerobot.scripts.lerobot_train import train +from lerobot.scripts.lerobot_train import prep_camera_images_for_policy, train from lerobot.utils.device_utils import auto_select_torch_device pytest.importorskip("transformers") @@ -57,6 +59,32 @@ def temp_dir(tmp_path): DEVICE = auto_select_torch_device() +def test_prep_camera_images_for_policy_obeys_policy_contract(): + image = torch.tensor([0, 127, 255], dtype=torch.uint8) + + float_batch = {"observation.images.front": image.clone()} + prep_camera_images_for_policy( + float_batch, + ["observation.images.front"], + ImageInputFormat.FLOAT32_0_1, + ) + assert float_batch["observation.images.front"].dtype == torch.float32 + torch.testing.assert_close( + float_batch["observation.images.front"], + torch.tensor([0.0, 127.0 / 255.0, 1.0]), + ) + + uint8_image = image.clone() + uint8_batch = {"observation.images.front": uint8_image} + prep_camera_images_for_policy( + uint8_batch, + ["observation.images.front"], + ImageInputFormat.UINT8_0_255, + ) + assert uint8_batch["observation.images.front"].data_ptr() == uint8_image.data_ptr() + assert uint8_batch["observation.images.front"].dtype == torch.uint8 + + def make_dummy_dataset(camera_keys, tmp_path): """Creates a minimal dummy dataset for testing rename_mapping logic.""" features = { From dad8709295b7e974e07c7c8e2d4566e0a9d651d7 Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Thu, 9 Jul 2026 11:11:34 -0700 Subject: [PATCH 12/13] fix(eval): honor policy image input format --- src/lerobot/envs/utils.py | 15 +++++++++++---- src/lerobot/policies/pretrained.py | 9 +-------- src/lerobot/scripts/lerobot_eval.py | 12 +++++++++--- src/lerobot/types.py | 9 ++++++++- tests/envs/test_envs.py | 21 +++++++++++++++++++++ 5 files changed, 50 insertions(+), 16 deletions(-) diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index 8b9c4f94b14..995d6adf814 100644 --- a/src/lerobot/envs/utils.py +++ b/src/lerobot/envs/utils.py @@ -28,6 +28,7 @@ from torch import Tensor from lerobot.configs import FeatureType, PolicyFeature +from lerobot.types import ImageInputFormat from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR from lerobot.utils.utils import get_channel_first_image_shape @@ -65,7 +66,10 @@ def _convert_nested_dict(d): return result -def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Tensor]: +def preprocess_observation( + observations: dict[str, np.ndarray], + image_input_format: ImageInputFormat = ImageInputFormat.FLOAT32_0_1, +) -> dict[str, Tensor]: # TODO(jadechoghari, imstevenpmwork): refactor this to use features from the environment (no hardcoding) """Convert environment observation to LeRobot format observation. Args: @@ -96,10 +100,13 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten # sanity check that images are uint8 assert img_tensor.dtype == torch.uint8, f"expect torch.uint8, but instead {img_tensor.dtype=}" - # convert to channel first of type float32 in range [0,1] + # convert to channel first, keeping the policy's expected image range img_tensor = einops.rearrange(img_tensor, "b h w c -> b c h w").contiguous() - img_tensor = img_tensor.type(torch.float32) - img_tensor /= 255 + if image_input_format is ImageInputFormat.FLOAT32_0_1: + img_tensor = img_tensor.type(torch.float32) + img_tensor /= 255 + elif image_input_format is not ImageInputFormat.UINT8_0_255: + raise ValueError(f"Unsupported policy image input format: {image_input_format}") return_observations[imgkey] = img_tensor diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index 3e7caf09786..cdd04c6fd11 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -18,7 +18,6 @@ import dataclasses import logging import os -from enum import StrEnum from importlib.resources import files from pathlib import Path from tempfile import TemporaryDirectory @@ -35,6 +34,7 @@ from lerobot.__version__ import __version__ from lerobot.configs import PreTrainedConfig from lerobot.configs.train import TrainPipelineConfig +from lerobot.types import ImageInputFormat from lerobot.utils.hub import HubMixin from .utils import log_model_loading_keys @@ -45,13 +45,6 @@ from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata -class ImageInputFormat(StrEnum): - """Raw image representation expected by a policy before preprocessing.""" - - FLOAT32_0_1 = "float32_0_1" - UINT8_0_255 = "uint8_0_255" - - def _build_card_context( cfg: TrainPipelineConfig | None, dataset_meta: LeRobotDatasetMetadata | None, diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index 722763d6e79..dc8b7575073 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -82,7 +82,7 @@ make_env_pre_post_processors, preprocess_observation, ) -from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors +from lerobot.policies import ImageInputFormat, PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.processor import PolicyProcessorPipeline from lerobot.types import PolicyAction from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD @@ -260,7 +260,10 @@ def rollout( try: while not np.all(done) and step < max_steps: # Numpy array to tensor and changing dictionary keys to LeRobot policy format. - observation = preprocess_observation(observation) + observation = preprocess_observation( + observation, + image_input_format=getattr(policy, "input_image_format", ImageInputFormat.FLOAT32_0_1), + ) if return_observations: all_observations.append(deepcopy(observation)) @@ -378,7 +381,10 @@ def rollout( # Track the final observation. if return_observations: - observation = preprocess_observation(observation) + observation = preprocess_observation( + observation, + image_input_format=getattr(policy, "input_image_format", ImageInputFormat.FLOAT32_0_1), + ) all_observations.append(deepcopy(observation)) # Stack the sequence along the first dimension so that we have (batch, sequence, *) tensors. diff --git a/src/lerobot/types.py b/src/lerobot/types.py index 9de504870b0..86ba296d726 100644 --- a/src/lerobot/types.py +++ b/src/lerobot/types.py @@ -16,7 +16,7 @@ from __future__ import annotations -from enum import Enum +from enum import Enum, StrEnum from typing import Any, TypedDict import numpy as np @@ -36,6 +36,13 @@ class TransitionKey(str, Enum): COMPLEMENTARY_DATA = "complementary_data" +class ImageInputFormat(StrEnum): + """Raw image representation expected by a policy before preprocessing.""" + + FLOAT32_0_1 = "float32_0_1" + UINT8_0_255 = "uint8_0_255" + + PolicyAction = torch.Tensor RobotAction = dict[str, Any] EnvAction = np.ndarray diff --git a/tests/envs/test_envs.py b/tests/envs/test_envs.py index c6a0b077dd8..67b4e8ac6ee 100644 --- a/tests/envs/test_envs.py +++ b/tests/envs/test_envs.py @@ -31,6 +31,7 @@ _parse_hub_url, preprocess_observation, ) +from lerobot.policies import ImageInputFormat from tests.utils import require_env OBS_TYPES = ["state", "pixels", "pixels_agent_pos"] @@ -43,6 +44,26 @@ AVAILABLE_ENVS = ["aloha", "pusht"] +def test_preprocess_observation_preserves_uint8_for_uint8_policy_contract(): + image = np.zeros((4, 5, 3), dtype=np.uint8) + image[0, 0] = [0, 127, 255] + image[1, 2] = [255, 127, 0] + + obs = preprocess_observation( + {"pixels": image}, + image_input_format=ImageInputFormat.UINT8_0_255, + ) + + assert obs["observation.image"].dtype == torch.uint8 + assert obs["observation.image"].shape == (1, 3, 4, 5) + torch.testing.assert_close( + obs["observation.image"][0, :, 0, 0], torch.tensor([0, 127, 255], dtype=torch.uint8) + ) + torch.testing.assert_close( + obs["observation.image"][0, :, 1, 2], torch.tensor([255, 127, 0], dtype=torch.uint8) + ) + + @pytest.mark.parametrize("obs_type", OBS_TYPES) @pytest.mark.parametrize("env_name, env_task", ENV_TASK_PAIRS) @require_env From d652e14b552b4bd6a137d134a28f761e2eea419d Mon Sep 17 00:00:00 2001 From: Andy Wrenn Date: Thu, 9 Jul 2026 14:04:36 -0700 Subject: [PATCH 13/13] docs(types): clarify image input format --- src/lerobot/types.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/lerobot/types.py b/src/lerobot/types.py index 86ba296d726..7b81215a818 100644 --- a/src/lerobot/types.py +++ b/src/lerobot/types.py @@ -37,7 +37,12 @@ class TransitionKey(str, Enum): class ImageInputFormat(StrEnum): - """Raw image representation expected by a policy before preprocessing.""" + """Raw image dtype/range expected by a policy before preprocessing. + + This only describes the tensor dtype and numeric range. It does not encode + channel count, channel order, or layout; those are defined by dataset/env + feature conventions and processor steps. + """ FLOAT32_0_1 = "float32_0_1" UINT8_0_255 = "uint8_0_255"