diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index f91ebcdda..789c86057 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -139,6 +139,13 @@ class _GrootN17CheckpointProcessorAssets: use_albumentations: bool +@dataclass(frozen=True) +class _GrootN17ActionGroup: + key: str + indices: list[int] + relative: bool + + def _load_n1_7_checkpoint_processor_assets(config: GrootConfig) -> _GrootN17CheckpointProcessorAssets | None: """Load N1.7 processor settings from checkpoint sidecar JSON files. @@ -548,16 +555,252 @@ def _reconnect_groot_n1_7_pack_decode_steps( step.pack_step = pack_step -def _resolve_action_feature_names_from_dataset_meta(dataset_meta: Any | None) -> list[str] | None: +def _resolve_feature_names_from_dataset_meta(dataset_meta: Any | None, feature_key: str) -> list[str] | None: features = getattr(dataset_meta, "features", {}) or {} - action_feature = features.get(ACTION) if isinstance(features, dict) else None - if isinstance(action_feature, dict): - names = action_feature.get("names") - else: - names = getattr(action_feature, "names", None) + feature = features.get(feature_key) if isinstance(features, dict) else None + names = feature.get("names") if isinstance(feature, dict) else getattr(feature, "names", None) return list(names) if names is not None else None +def _resolve_action_feature_names_from_dataset_meta(dataset_meta: Any | None) -> list[str] | None: + return _resolve_feature_names_from_dataset_meta(dataset_meta, ACTION) + + +def _resolve_visual_modality_keys_from_dataset_meta(dataset_meta: Any | None) -> list[str] | None: + features = getattr(dataset_meta, "features", {}) or {} + if not isinstance(features, dict): + return None + + keys: list[str] = [] + for key, value in features.items(): + dtype = value.get("dtype") if isinstance(value, dict) else getattr(value, "dtype", None) + feature_type = value.get("type") if isinstance(value, dict) else getattr(value, "type", None) + is_visual = dtype in {"image", "video"} or str(feature_type).upper().endswith("VISUAL") + if not is_visual or not isinstance(key, str) or not key.startswith(f"{OBS_IMAGES}."): + continue + keys.append(key.removeprefix(f"{OBS_IMAGES}.")) + return keys or None + + +def _slice_stats_entry(stats: dict[str, Any], indices: list[int]) -> dict[str, list[float]]: + if not indices: + return {} + + max_index = max(indices) + sliced: dict[str, list[float]] = {} + for stat_name, value in stats.items(): + tensor = torch.as_tensor(value, dtype=torch.float32).flatten() + if tensor.numel() <= max_index: + continue + sliced[stat_name] = [float(tensor[index].item()) for index in indices] + + if "min" in sliced and "max" in sliced: + if "mean" not in sliced: + sliced["mean"] = [ + (low + high) * 0.5 for low, high in zip(sliced["min"], sliced["max"], strict=True) + ] + if "std" not in sliced: + sliced["std"] = [ + abs(high - low) * 0.5 for low, high in zip(sliced["min"], sliced["max"], strict=True) + ] + return sliced + + +def _feature_group_key(name: str) -> str: + base = name.removesuffix(".pos").split(".")[-1] + return base.replace(" ", "_") or "action" + + +def _infer_n1_7_action_groups( + action_names: list[str], + *, + action_dim: int, + exclude_joints: list[str], +) -> list[_GrootN17ActionGroup]: + if not action_names or action_dim <= 0: + return [] + + names = list(action_names[:action_dim]) + exclude_tokens = [str(token).lower() for token in exclude_joints if token] + groups: list[_GrootN17ActionGroup] = [] + current_indices: list[int] = [] + + def flush_relative_group() -> None: + if not current_indices: + return + key = ( + "single_arm" + if not any(group.key == "single_arm" for group in groups) + else f"single_arm_{len(groups)}" + ) + groups.append(_GrootN17ActionGroup(key=key, indices=list(current_indices), relative=True)) + current_indices.clear() + + for index, name in enumerate(names): + lowered = str(name).lower() + is_excluded = any(token == lowered or token in lowered for token in exclude_tokens) + if is_excluded: + flush_relative_group() + groups.append( + _GrootN17ActionGroup(key=_feature_group_key(str(name)), indices=[index], relative=False) + ) + else: + current_indices.append(index) + + flush_relative_group() + return groups + + +def _group_stats_by_action_groups( + stats: dict[str, Any], groups: list[_GrootN17ActionGroup] +) -> dict[str, dict[str, list[float]]]: + return {group.key: _slice_stats_entry(stats, group.indices) for group in groups} + + +def _grouped_stats_support_percentiles( + raw_stats: dict[str, Any], + modality_config: dict[str, Any], + *, + use_relative_action: bool, +) -> bool: + state_keys = modality_config.get("state", {}).get("modality_keys", []) + for key in state_keys: + stats = raw_stats.get("state", {}).get(key, {}) + if "q01" not in stats or "q99" not in stats: + return False + + action_cfg = modality_config.get("action", {}) + action_keys = action_cfg.get("modality_keys", []) + action_configs = action_cfg.get("action_configs", []) + for idx, key in enumerate(action_keys): + cfg = action_configs[idx] if idx < len(action_configs) else {} + is_relative = ( + use_relative_action and isinstance(cfg, dict) and config_value(cfg.get("rep")) == "relative" + ) + if is_relative: + continue + stats = raw_stats.get("action", {}).get(key, {}) + if "q01" not in stats or "q99" not in stats: + return False + return True + + +def _build_n1_7_relative_action_processor_assets( + config: GrootConfig, + dataset_stats: dict[str, dict[str, Any]] | None, + dataset_meta: Any | None, + *, + base_assets: _GrootN17CheckpointProcessorAssets | None = None, +) -> _GrootN17CheckpointProcessorAssets | None: + if not config.use_relative_actions or not dataset_stats: + return None + + try: + action_dim = int(config.output_features[ACTION].shape[0]) + except Exception: + return None + + action_names = _resolve_action_feature_names_from_dataset_meta(dataset_meta) + if not action_names: + return None + + groups = _infer_n1_7_action_groups( + action_names, + action_dim=action_dim, + exclude_joints=list(config.relative_exclude_joints or []), + ) + if not groups or not any(group.relative for group in groups): + return None + + meta_stats = getattr(dataset_meta, "stats", None) or {} + state_stats = (meta_stats.get(OBS_STATE) if isinstance(meta_stats, dict) else None) or dataset_stats.get( + OBS_STATE, {} + ) + absolute_action_stats = ( + meta_stats.get(ACTION) if isinstance(meta_stats, dict) else None + ) or dataset_stats.get(ACTION, {}) + relative_action_stats = dataset_stats.get(ACTION, {}) + if not state_stats or not absolute_action_stats or not relative_action_stats: + return None + + raw_stats: dict[str, Any] = { + "state": _group_stats_by_action_groups(state_stats, groups), + "action": _group_stats_by_action_groups(absolute_action_stats, groups), + "relative_action": { + group.key: _slice_stats_entry(relative_action_stats, group.indices) + for group in groups + if group.relative + }, + } + + action_configs = [ + { + "rep": "RELATIVE" if group.relative else "ABSOLUTE", + "type": "NON_EEF", + "format": "DEFAULT", + "state_key": None, + } + for group in groups + ] + action_horizon = min(config.chunk_size, 40) + modality_config: dict[str, Any] = { + "state": {"modality_keys": [group.key for group in groups]}, + "action": { + "modality_keys": [group.key for group in groups], + "action_configs": action_configs, + "delta_indices": list(range(action_horizon)), + }, + } + video_modality_keys = ( + base_assets.video_modality_keys if base_assets is not None else None + ) or _resolve_visual_modality_keys_from_dataset_meta(dataset_meta) + if video_modality_keys: + modality_config["video"] = { + "modality_keys": list(video_modality_keys), + "delta_indices": [0], + } + + use_percentiles = _grouped_stats_support_percentiles(raw_stats, modality_config, use_relative_action=True) + flat_stats = { + OBS_STATE: flatten_n1_7_modality_stats( + embodiment_stats=raw_stats, + embodiment_config=modality_config, + modality="state", + use_percentiles=use_percentiles, + use_relative_action=True, + ), + ACTION: flatten_n1_7_modality_stats( + embodiment_stats=raw_stats, + embodiment_config=modality_config, + modality="action", + use_percentiles=use_percentiles, + use_relative_action=True, + ), + } + + return _GrootN17CheckpointProcessorAssets( + stats=flat_stats, + raw_stats=raw_stats, + modality_config=modality_config, + embodiment_mapping=base_assets.embodiment_mapping + if base_assets is not None + else dict(N1_7_EMBODIMENT_MAPPING), + formalize_language=base_assets.formalize_language if base_assets is not None else True, + valid_action_horizon=action_horizon, + max_action_horizon=action_horizon, + video_horizon=base_assets.video_horizon if base_assets is not None else None, + use_percentiles=use_percentiles, + use_relative_action=True, + 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, + image_target_size=base_assets.image_target_size if base_assets is not None else None, + shortest_image_edge=base_assets.shortest_image_edge if base_assets is not None else None, + crop_fraction=base_assets.crop_fraction if base_assets is not None else None, + use_albumentations=base_assets.use_albumentations if base_assets is not None else False, + ) + + def make_groot_pre_post_processors( config: GrootConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, @@ -592,6 +835,20 @@ def make_groot_pre_post_processors( """ checkpoint_assets = _load_n1_7_checkpoint_processor_assets(config) + checkpoint_stats = checkpoint_assets.stats if checkpoint_assets is not None else None + checkpoint_has_stats = has_modality_stats(checkpoint_stats) + if config.use_relative_actions and not checkpoint_has_stats: + relative_assets = _build_n1_7_relative_action_processor_assets( + config, + dataset_stats, + dataset_meta, + base_assets=checkpoint_assets, + ) + if relative_assets is not None: + checkpoint_assets = relative_assets + checkpoint_stats = checkpoint_assets.stats + checkpoint_has_stats = has_modality_stats(checkpoint_stats) + action_horizon = ( checkpoint_assets.max_action_horizon if checkpoint_assets is not None and checkpoint_assets.max_action_horizon is not None @@ -602,8 +859,6 @@ def make_groot_pre_post_processors( if checkpoint_assets is not None and checkpoint_assets.valid_action_horizon is not None else action_horizon ) - checkpoint_stats = checkpoint_assets.stats if checkpoint_assets is not None else None - checkpoint_has_stats = has_modality_stats(checkpoint_stats) padded_stats = checkpoint_stats if checkpoint_has_stats else (dataset_stats or {}) embodiment_mapping = ( checkpoint_assets.embodiment_mapping @@ -669,8 +924,11 @@ def make_groot_pre_post_processors( ), DeviceProcessorStep(device=config.device), ] + uses_native_relative_actions = bool( + checkpoint_assets is not None and checkpoint_assets.use_relative_action + ) relative_step: RelativeActionsProcessorStep | None = None - if config.use_relative_actions: + if config.use_relative_actions and not uses_native_relative_actions: relative_step = RelativeActionsProcessorStep( enabled=True, exclude_joints=list(config.relative_exclude_joints or []), @@ -981,9 +1239,92 @@ class GrootN17PackInputsStep(ProcessorStep): ) return ordered + def _state_groups_from_tensor(self, state: torch.Tensor) -> dict[str, torch.Tensor]: + if self.modality_config is None or self.raw_stats is None: + return {} + state_config = self.modality_config.get("state", {}) + if not isinstance(state_config, dict): + return {} + state_keys = state_config.get("modality_keys", []) + if not isinstance(state_keys, list): + return {} + + grouped: dict[str, torch.Tensor] = {} + start_idx = 0 + for key in state_keys: + if not isinstance(key, str): + continue + key_stats = self.raw_stats.get("state", {}).get(key, {}) + dim = stat_dim_from_entry(key_stats) if isinstance(key_stats, dict) else 0 + if dim <= 0: + continue + grouped[key] = state[:, start_idx : start_idx + dim] + start_idx += dim + return grouped + + def _convert_relative_action_groups_for_training( + self, action: torch.Tensor, state: torch.Tensor + ) -> torch.Tensor: + if self.modality_config is None or self.raw_stats is None: + return action + + action_config = self.modality_config.get("action", {}) + if not isinstance(action_config, dict): + return action + action_keys = action_config.get("modality_keys", []) + action_configs = action_config.get("action_configs", []) + if not isinstance(action_keys, list) or not isinstance(action_configs, list): + return action + + state_groups = self._state_groups_from_tensor(state) + if not state_groups: + return action + + converted = action + start_idx = 0 + cloned = False + for idx, key in enumerate(action_keys): + if not isinstance(key, str): + continue + key_stats = self.raw_stats.get("action", {}).get(key, {}) + dim = stat_dim_from_entry(key_stats) if isinstance(key_stats, dict) else 0 + if dim <= 0: + continue + end_idx = start_idx + dim + if end_idx > action.shape[-1]: + break + + cfg = ( + action_configs[idx] + if idx < len(action_configs) and isinstance(action_configs[idx], dict) + else {} + ) + if config_value(cfg.get("rep")) == "relative": + action_type = config_value(cfg.get("type")) + if action_type != "non_eef": + raise ValueError(f"Unsupported relative N1.7 action config for '{key}': {cfg}") + state_key = cfg.get("state_key") or key + reference = state_groups.get(state_key) + if reference is None: + raise KeyError(f"Missing raw state group '{state_key}' for relative N1.7 action '{key}'") + if reference.shape[-1] != dim: + raise ValueError( + f"Relative N1.7 action group '{key}' has dim {dim}, but state group " + f"'{state_key}' has dim {reference.shape[-1]}." + ) + if not cloned: + converted = action.clone() + cloned = True + converted[..., start_idx:end_idx] -= reference[:, None, :] + + start_idx = end_idx + + return converted + def __call__(self, transition: EnvTransition) -> EnvTransition: obs = transition.get(TransitionKey.OBSERVATION, {}) or {} comp = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) or {} + raw_state_for_action: torch.Tensor | None = None def _align_vec(vec: Any, target_dim: int, *, default: float) -> torch.Tensor: t = torch.as_tensor(vec) @@ -1068,6 +1409,7 @@ class GrootN17PackInputsStep(ProcessorStep): if dim > self.max_state_dim: raise ValueError(f"State dimension {dim} exceeds max_state_dim {self.max_state_dim}.") _cache_raw_state(state) + raw_state_for_action = state if self.normalize_min_max: state = _min_max_norm(state, OBS_STATE) state = state.unsqueeze(1) @@ -1090,6 +1432,8 @@ class GrootN17PackInputsStep(ProcessorStep): raise ValueError(f"Action horizon {horizon} exceeds action_horizon {self.action_horizon}.") if dim > self.max_action_dim: raise ValueError(f"Action dimension {dim} exceeds max_action_dim {self.max_action_dim}.") + if raw_state_for_action is not None: + action = self._convert_relative_action_groups_for_training(action, raw_state_for_action) if self.normalize_min_max: flat = _min_max_norm(action.reshape(bsz * horizon, dim), ACTION) action = flat.view(bsz, horizon, dim) diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index b1ef56495..32c972d89 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -30,7 +30,6 @@ from lerobot.configs import FeatureType, PolicyFeature from lerobot.policies.factory import make_policy_config, make_pre_post_processors from lerobot.policies.groot.configuration_groot import ( GROOT_ACTION_DECODE_TRANSFORM_LIBERO, - GROOT_N1_7, GROOT_N1_7_BASE_MODEL, GrootConfig, infer_groot_n1_7_action_execution_horizon, @@ -370,7 +369,7 @@ def test_groot_n1_7_accepts_named_action_decode_transform(): def test_groot_n1_7_rejects_legacy_libero_gripper_action_decode_transform(legacy_transform): with pytest.raises(ValueError, match="Unsupported GR00T N1.7 action decode transform"): GrootConfig( - action_decode_transform=legacy_transform, + action_decode_transform=legacy_transform, device="cpu", ) @@ -378,7 +377,7 @@ def test_groot_n1_7_rejects_legacy_libero_gripper_action_decode_transform(legacy def test_groot_config_rejects_mismatched_n1_5_path_for_n1_7(): with pytest.raises(ValueError, match="does not match base_model_path"): GrootConfig( - base_model_path="nvidia/GR00T-N1.5-3B", + base_model_path="nvidia/GR00T-N1.5-3B", device="cpu", ) @@ -504,7 +503,7 @@ def test_groot_from_pretrained_rejects_mismatched_caller_config(tmp_path): # so construction itself raises before from_pretrained is reached. with pytest.raises(ValueError, match="does not match base_model_path"): config = GrootConfig( - base_model_path="nvidia/GR00T-N1.5-3B", + base_model_path="nvidia/GR00T-N1.5-3B", input_features=input_features, output_features=output_features, device="cpu", @@ -1004,6 +1003,77 @@ def test_groot_n1_7_pack_inputs_normalizes_action_chunk_per_dimension_before_pad assert action_mask[0, :, 3:].sum().item() == 0 +def test_groot_n1_7_pack_inputs_trains_native_relative_groups_with_absolute_gripper(): + step = GrootN17PackInputsStep( + action_horizon=2, + valid_action_horizon=2, + max_state_dim=6, + max_action_dim=6, + normalize_min_max=True, + clip_outliers=False, + stats={ + OBS_STATE: { + "min": [-100.0, -100.0, -100.0, -100.0, -100.0, 0.0], + "max": [100.0, 100.0, 100.0, 100.0, 100.0, 100.0], + }, + ACTION: { + "min": [-10.0, -10.0, -10.0, -10.0, -10.0, 0.0], + "max": [10.0, 10.0, 10.0, 10.0, 10.0, 100.0], + }, + }, + raw_stats={ + "state": { + "single_arm": {"min": [-100.0] * 5, "max": [100.0] * 5}, + "gripper": {"min": [0.0], "max": [100.0]}, + }, + "action": { + "single_arm": {"min": [-100.0] * 5, "max": [100.0] * 5}, + "gripper": {"min": [0.0], "max": [100.0]}, + }, + "relative_action": { + "single_arm": {"min": [-10.0] * 5, "max": [10.0] * 5}, + }, + }, + modality_config={ + "state": {"modality_keys": ["single_arm", "gripper"]}, + "action": { + "modality_keys": ["single_arm", "gripper"], + "action_configs": [ + {"rep": "RELATIVE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None}, + {"rep": "ABSOLUTE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None}, + ], + "delta_indices": [0, 1], + }, + }, + ) + transition = { + TransitionKey.OBSERVATION: { + OBS_STATE: torch.tensor([[10.0, 20.0, 30.0, 40.0, 50.0, 25.0]]), + }, + TransitionKey.ACTION: torch.tensor( + [ + [ + [12.0, 18.0, 35.0, 30.0, 55.0, 0.0], + [9.0, 21.0, 27.0, 43.0, 50.0, 100.0], + ] + ] + ), + TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move"]}, + } + + output = step(transition) + + expected_actions = torch.tensor( + [ + [ + [0.2, -0.2, 0.5, -1.0, 0.5, -1.0], + [-0.1, 0.1, -0.3, 0.3, 0.0, 1.0], + ] + ] + ) + torch.testing.assert_close(output[TransitionKey.ACTION], expected_actions) + + def test_groot_n1_7_pack_inputs_adds_inference_action_horizon_mask(): step = GrootN17PackInputsStep( action_horizon=40, @@ -1323,7 +1393,7 @@ def test_groot_from_pretrained_rejects_caller_config_mismatch_from_local_config( # so construction itself raises before from_pretrained is reached. with pytest.raises(ValueError, match="does not match base_model_path"): config = GrootConfig( - base_model_path="nvidia/GR00T-N1.5-3B", + base_model_path="nvidia/GR00T-N1.5-3B", input_features=input_features, output_features=output_features, device="cpu", @@ -1736,9 +1806,7 @@ def test_groot_n1_7_saved_processors_reload_through_factory_preserves_saved_stat assert unpack_step.env_action_dim == 7 - - -def test_groot_n1_7_relative_action_training_processors_save_relative_action_stats(tmp_path): +def test_groot_n1_7_relative_action_training_processors_save_native_grouped_stats(tmp_path): input_features, output_features = _groot_features(state_dim=6, action_dim=6) action_names = [ "shoulder_pan.pos", @@ -1817,26 +1885,41 @@ def test_groot_n1_7_relative_action_training_processors_save_relative_action_sta postprocessor.save_pretrained(tmp_path) preprocessor_config = json.loads((tmp_path / "policy_preprocessor.json").read_text()) - assert any(step.get("registry_name") == "relative_actions_processor" for step in preprocessor_config["steps"]) + assert not any( + step.get("registry_name") == "relative_actions_processor" for step in preprocessor_config["steps"] + ) pack_entry = next( step for step in preprocessor_config["steps"] if step.get("registry_name") == "groot_n1_7_pack_inputs_v1" ) + pack_config = pack_entry["config"] + assert pack_config["modality_config"]["action"]["modality_keys"] == ["single_arm", "gripper"] + assert pack_config["modality_config"]["action"]["action_configs"] == [ + {"rep": "RELATIVE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None}, + {"rep": "ABSOLUTE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None}, + ] + assert pack_config["raw_stats"]["relative_action"]["single_arm"]["min"] == [-2.0, -3.0, -4.0, -5.0, -6.0] + assert pack_config["raw_stats"]["action"]["gripper"]["min"] == [0.0] + assert pack_config["raw_stats"]["action"]["gripper"]["max"] == [100.0] + pack_state = load_file(tmp_path / pack_entry["state_file"]) torch.testing.assert_close(pack_state[f"{ACTION}.min"], expected_relative_action_stats["min"]) torch.testing.assert_close(pack_state[f"{ACTION}.max"], expected_relative_action_stats["max"]) postprocessor_config = json.loads((tmp_path / "policy_postprocessor.json").read_text()) - assert any(step.get("registry_name") == "absolute_actions_processor" for step in postprocessor_config["steps"]) - unpack_entry = next( + assert not any( + step.get("registry_name") == "absolute_actions_processor" for step in postprocessor_config["steps"] + ) + decode_entry = next( step for step in postprocessor_config["steps"] - if step.get("registry_name", "").startswith("groot_action_unpack_unnormalize") + if step.get("registry_name") == "groot_n1_7_action_decode_v1" ) - unpack_state = load_file(tmp_path / unpack_entry["state_file"]) - torch.testing.assert_close(unpack_state[f"{ACTION}.min"], expected_relative_action_stats["min"]) - torch.testing.assert_close(unpack_state[f"{ACTION}.max"], expected_relative_action_stats["max"]) + decode_config = decode_entry["config"] + assert decode_config["use_relative_action"] is True + assert decode_config["raw_stats"]["relative_action"]["single_arm"]["max"] == [2.0, 3.0, 4.0, 5.0, 6.0] + assert decode_config["raw_stats"]["action"]["gripper"]["max"] == [100.0] def test_groot_policy_selects_n1_7_model_class(monkeypatch):