From 20c0f07858ddcf1acb8cc03ccf47a35d92882914 Mon Sep 17 00:00:00 2001 From: johnnynunez Date: Thu, 2 Jul 2026 00:54:20 +0200 Subject: [PATCH] feat(groot): activate checkpoint-configured N1.7 raw-state dropout during training Isaac-GR00T applies dual state regularization during fine-tuning: raw-state zeroing driven by the processor sidecar's state_dropout_prob (0.2 for the inspected N1.7 checkpoint) plus encoded-feature dropout. Baseline LeRobot kept the processor in deterministic mode, so the raw-state dropout never activated (RCA Tier-2 contributor to the LeRobot-trained SO-101 failures). - GrootN17PackInputsStep: runtime-only 'training' flag + state_dropout_prob; whole-sample state zeroing gated on torch.is_grad_enabled() so eval and no_grad validation paths are unaffected - sidecar loader reads state_dropout_prob from processor_config.json - state_dropout_prob serializes with the step; the training flag intentionally does not (reloaded pipelines default to eval, re-enabled only when processors are rebuilt with dataset_meta) - _set_groot_preprocessor_training toggles any dataclass step exposing a 'training' field on serialized-pipeline reloads Verification: tests/policies/groot/test_groot_state_dropout.py (4 passed) on RTX PRO 6000 / CUDA 13.3. --- src/lerobot/policies/groot/processor_groot.py | 36 +++++++ .../groot/test_groot_state_dropout.py | 100 ++++++++++++++++++ 2 files changed, 136 insertions(+) create mode 100644 tests/policies/groot/test_groot_state_dropout.py diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 5856f5ff1..50bf1d843 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -15,6 +15,7 @@ # limitations under the License. import logging +import random from copy import copy, deepcopy from dataclasses import dataclass, field, fields, is_dataclass from pathlib import Path @@ -136,6 +137,7 @@ class _GrootN17CheckpointProcessorAssets: video_horizon: int | None use_percentiles: bool use_relative_action: bool + state_dropout_prob: float clip_outliers: bool video_modality_keys: list[str] | None image_crop_size: list[int] | None @@ -182,6 +184,9 @@ def _load_n1_7_checkpoint_processor_assets(config: GrootConfig) -> _GrootN17Chec modality_config = {} use_relative_action = bool(processor_kwargs.get("use_relative_action", False)) + state_dropout_prob = as_optional_float(processor_kwargs.get("state_dropout_prob")) + if state_dropout_prob is None: + state_dropout_prob = 0.0 stats = _load_n1_7_checkpoint_stats( checkpoint_path, processor_kwargs, @@ -222,6 +227,7 @@ def _load_n1_7_checkpoint_processor_assets(config: GrootConfig) -> _GrootN17Chec video_horizon=video_horizon, use_percentiles=bool(processor_kwargs.get("use_percentiles", False)), use_relative_action=use_relative_action, + state_dropout_prob=state_dropout_prob, clip_outliers=clip_outliers, video_modality_keys=video_modality_keys, image_crop_size=as_int_pair(processor_kwargs.get("image_crop_size")), @@ -445,6 +451,22 @@ def _apply_groot_step_overrides( post_init() +def _set_groot_preprocessor_training( + preprocessor: PolicyProcessorPipeline, + *, + training: bool, +) -> None: + """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 + this helper enumerating them. + """ + for step in preprocessor.steps: + if is_dataclass(step) and any(f.name == "training" for f in fields(step)): + setattr(step, "training", training) + + def make_groot_pre_post_processors_from_pretrained( config: GrootConfig, pretrained_path: str, @@ -493,6 +515,7 @@ def make_groot_pre_post_processors_from_pretrained( _reconnect_groot_relative_absolute_steps(preprocessor, postprocessor) _reconnect_groot_n1_7_pack_decode_steps(preprocessor, postprocessor) _apply_groot_action_decode_transform(postprocessor, config.action_decode_transform) + _set_groot_preprocessor_training(preprocessor, training=dataset_meta is not None) return preprocessor, postprocessor @@ -1056,6 +1079,7 @@ def _build_n1_7_relative_action_processor_assets( video_horizon=base_assets.video_horizon if base_assets is not None else None, use_percentiles=use_percentiles, use_relative_action=True, + state_dropout_prob=base_assets.state_dropout_prob if base_assets is not None else 0.0, clip_outliers=base_assets.clip_outliers if base_assets is not None else True, video_modality_keys=video_modality_keys, image_crop_size=base_assets.image_crop_size if base_assets is not None else None, @@ -1160,6 +1184,8 @@ def make_groot_pre_post_processors( embodiment_tag=config.embodiment_tag, embodiment_mapping=embodiment_mapping, normalize_min_max=True, + training=dataset_meta is not None, + state_dropout_prob=(checkpoint_assets.state_dropout_prob if checkpoint_assets is not None else 0.0), stats=padded_stats, clip_outliers=clip_outliers, video_modality_keys=video_modality_keys, @@ -1472,6 +1498,8 @@ class GrootN17PackInputsStep(ProcessorStep): embodiment_tag: str = "new_embodiment" embodiment_mapping: dict[str, int] = field(default_factory=lambda: dict(N1_7_EMBODIMENT_MAPPING)) normalize_min_max: bool = True + training: bool = False + state_dropout_prob: float = 0.0 stats: dict[str, dict[str, Any]] | None = None clip_outliers: bool = True use_percentiles: bool = False @@ -1803,6 +1831,13 @@ class GrootN17PackInputsStep(ProcessorStep): if dim < self.max_state_dim: pad = torch.zeros(bsz, 1, self.max_state_dim - dim, dtype=state.dtype, device=state.device) state = torch.cat([state, pad], dim=2) + if self.training and torch.is_grad_enabled() and self.state_dropout_prob > 0: + drop_state = torch.tensor( + [random.random() < self.state_dropout_prob for _ in range(bsz)], + dtype=torch.bool, + device=state.device, + ).view(bsz, 1, 1) + state = state.masked_fill(drop_state, 0) obs["state"] = state action = transition.get(TransitionKey.ACTION) @@ -1914,6 +1949,7 @@ class GrootN17PackInputsStep(ProcessorStep): "embodiment_tag": self.embodiment_tag, "embodiment_mapping": self.embodiment_mapping, "normalize_min_max": self.normalize_min_max, + "state_dropout_prob": self.state_dropout_prob, "clip_outliers": self.clip_outliers, "use_percentiles": self.use_percentiles, "video_modality_keys": self.video_modality_keys, diff --git a/tests/policies/groot/test_groot_state_dropout.py b/tests/policies/groot/test_groot_state_dropout.py new file mode 100644 index 000000000..fcfebcb8f --- /dev/null +++ b/tests/policies/groot/test_groot_state_dropout.py @@ -0,0 +1,100 @@ +#!/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 raw-state dropout training contract. + +Isaac-GR00T zeroes the entire proprioceptive state of a sample with probability +``state_dropout_prob`` (configured in the checkpoint's processor sidecar) during +training only. Baseline LeRobot kept the processor deterministic, so this +regularization never activated. These tests pin the train/eval split. +""" + +import torch + +from lerobot.policies.groot.processor_groot import GrootN17PackInputsStep +from lerobot.types import TransitionKey +from lerobot.utils.constants import OBS_STATE + + +def _make_transition(): + return { + TransitionKey.OBSERVATION: {OBS_STATE: torch.tensor([[1.0, 2.0], [3.0, 4.0]])}, + TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move", "Move"]}, + } + + +def test_groot_n1_7_training_applies_raw_state_dropout_before_encoder(): + step = GrootN17PackInputsStep( + max_state_dim=4, + max_action_dim=4, + normalize_min_max=False, + training=True, + state_dropout_prob=1.0, + ) + + output = step(_make_transition()) + + expected = torch.zeros(2, 1, 4) + torch.testing.assert_close(output[TransitionKey.OBSERVATION]["state"], expected) + + +def test_groot_n1_7_training_state_dropout_is_disabled_under_no_grad(): + step = GrootN17PackInputsStep( + max_state_dim=4, + max_action_dim=4, + normalize_min_max=False, + training=True, + state_dropout_prob=1.0, + ) + + with torch.no_grad(): + output = step(_make_transition()) + + expected = torch.tensor([[[1.0, 2.0, 0.0, 0.0]], [[3.0, 4.0, 0.0, 0.0]]]) + torch.testing.assert_close(output[TransitionKey.OBSERVATION]["state"], expected) + + +def test_groot_n1_7_eval_mode_state_dropout_is_inactive(): + step = GrootN17PackInputsStep( + max_state_dim=4, + max_action_dim=4, + normalize_min_max=False, + training=False, + state_dropout_prob=1.0, + ) + + output = step(_make_transition()) + + expected = torch.tensor([[[1.0, 2.0, 0.0, 0.0]], [[3.0, 4.0, 0.0, 0.0]]]) + torch.testing.assert_close(output[TransitionKey.OBSERVATION]["state"], expected) + + +def test_groot_n1_7_pack_step_serializes_dropout_prob_but_not_training_mode(): + step = GrootN17PackInputsStep( + max_state_dim=4, + max_action_dim=4, + normalize_min_max=False, + training=True, + state_dropout_prob=0.2, + ) + + serialized = step.get_config() + restored = GrootN17PackInputsStep(**serialized) + + assert "training" not in serialized + assert serialized["state_dropout_prob"] == 0.2 + assert restored.training is False + assert restored.state_dropout_prob == 0.2