diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 62ad34094..a695c7d8e 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -469,6 +469,7 @@ def make_groot_pre_post_processors_from_pretrained( # lerobot-eval or the policy server) to the freshly built steps. _apply_groot_step_overrides(preprocessor, preprocessor_overrides) _apply_groot_step_overrides(postprocessor, postprocessor_overrides) + _apply_groot_action_decode_transform(postprocessor, config.action_decode_transform) return preprocessor, postprocessor preprocessor, postprocessor = _load_groot_processor_pipelines( @@ -480,6 +481,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) return preprocessor, postprocessor @@ -552,6 +554,20 @@ def _reconnect_groot_n1_7_pack_decode_steps( step.pack_step = pack_step +def _apply_groot_action_decode_transform( + postprocessor: PolicyProcessorPipeline, + action_decode_transform: str | None, +) -> None: + use_libero_transform = action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_LIBERO + + for step in postprocessor.steps: + if isinstance(step, GrootN17ActionDecodeStep): + step.action_decode_transform = action_decode_transform + elif isinstance(step, GrootActionUnpackUnnormalizeStep): + step.libero_gripper_action = use_libero_transform + if use_libero_transform: + step.libero_gripper_binarize = True + def _resolve_feature_names_from_dataset_meta(dataset_meta: Any | None, feature_key: str) -> list[str] | None: features = getattr(dataset_meta, "features", {}) or {} @@ -1198,6 +1214,7 @@ def make_groot_pre_post_processors( stats=padded_stats, normalize_min_max=True, clip_normalized_action=True, + libero_gripper_action=config.action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_LIBERO, ) else: action_decode_step = GrootN17ActionDecodeStep( diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index 6f31b5695..f915f985c 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -1391,6 +1391,80 @@ def test_groot_n1_7_action_decode_requires_gripper_key_for_libero_transform(): step({TransitionKey.ACTION: torch.zeros(1, 1, 1)}) +def test_groot_n1_7_fallback_processors_wire_libero_transform_to_postprocessor(): + config = _groot_config() + dataset_stats = { + OBS_STATE: { + "min": torch.zeros(8), + "max": torch.ones(8), + }, + ACTION: { + "min": torch.zeros(7), + "max": torch.ones(7), + }, + } + + _, postprocessor = make_groot_pre_post_processors(config, dataset_stats=dataset_stats) + + action_decode_step = next( + step for step in postprocessor.steps if isinstance(step, GrootActionUnpackUnnormalizeStep) + ) + assert action_decode_step.libero_gripper_action is True + + +def test_groot_n1_7_loaded_fallback_postprocessor_honors_config_action_decode_transform(tmp_path): + input_features, output_features = _groot_features(state_dim=8, action_dim=7) + dataset_stats = { + OBS_STATE: { + "min": torch.zeros(8), + "max": torch.ones(8), + }, + ACTION: { + "min": torch.zeros(7), + "max": torch.ones(7), + }, + } + disabled_config = GrootConfig( + input_features=input_features, + output_features=output_features, + device="cpu", + use_bf16=False, + action_decode_transform=None, + ) + preprocessor, postprocessor = make_groot_pre_post_processors( + disabled_config, + dataset_stats=dataset_stats, + ) + save_dir = tmp_path / "saved_fallback_processors" + disabled_config.save_pretrained(save_dir) + preprocessor.save_pretrained(save_dir) + postprocessor.save_pretrained(save_dir) + + saved_postprocessor = json.loads((save_dir / "policy_postprocessor.json").read_text()) + saved_decode_config = next( + step["config"] + for step in saved_postprocessor["steps"] + if step["registry_name"] == "groot_action_unpack_unnormalize_v2" + ) + assert saved_decode_config["libero_gripper_action"] is False + + enabled_config = GrootConfig( + input_features=input_features, + output_features=output_features, + device="cpu", + use_bf16=False, + action_decode_transform=GROOT_ACTION_DECODE_TRANSFORM_LIBERO, + ) + _, loaded_postprocessor = make_pre_post_processors(enabled_config, pretrained_path=str(save_dir)) + action_decode_step = next( + step for step in loaded_postprocessor.steps if isinstance(step, GrootActionUnpackUnnormalizeStep) + ) + + assert action_decode_step.libero_gripper_action is True + output = action_decode_step({TransitionKey.ACTION: torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -1.0]])}) + torch.testing.assert_close(output[TransitionKey.ACTION][0, -1], torch.tensor(1.0)) + + def test_groot_n1_7_postprocessor_converts_libero_gripper_convention(): step = GrootActionUnpackUnnormalizeStep( env_action_dim=7,