diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 789c86057..c0208dbe5 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -582,27 +582,31 @@ def _resolve_visual_modality_keys_from_dataset_meta(dataset_meta: Any | None) -> return keys or None -def _slice_stats_entry(stats: dict[str, Any], indices: list[int]) -> dict[str, list[float]]: +def _slice_stats_entry(stats: dict[str, Any], indices: list[int]) -> dict[str, Any]: if not indices: return {} max_index = max(indices) - sliced: dict[str, list[float]] = {} + sliced: dict[str, Any] = {} 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] + tensor = torch.as_tensor(value, dtype=torch.float32) + if tensor.ndim >= 2: + if tensor.shape[-1] <= max_index: + continue + sliced[stat_name] = tensor[..., indices].tolist() + else: + tensor = tensor.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: + min_arr = np.asarray(sliced["min"], dtype=np.float32) + max_arr = np.asarray(sliced["max"], dtype=np.float32) if "mean" not in sliced: - sliced["mean"] = [ - (low + high) * 0.5 for low, high in zip(sliced["min"], sliced["max"], strict=True) - ] + sliced["mean"] = ((min_arr + max_arr) * 0.5).tolist() if "std" not in sliced: - sliced["std"] = [ - abs(high - low) * 0.5 for low, high in zip(sliced["min"], sliced["max"], strict=True) - ] + sliced["std"] = (np.abs(max_arr - min_arr) * 0.5).tolist() return sliced @@ -888,6 +892,7 @@ def make_groot_pre_post_processors( clip_outliers=clip_outliers, video_modality_keys=video_modality_keys, raw_stats=checkpoint_assets.raw_stats if checkpoint_assets is not None else None, + use_percentiles=checkpoint_assets.use_percentiles if checkpoint_assets is not None else False, modality_config=checkpoint_assets.modality_config if checkpoint_assets is not None else None, ) @@ -1173,6 +1178,7 @@ class GrootN17PackInputsStep(ProcessorStep): normalize_min_max: bool = True stats: dict[str, dict[str, Any]] | None = None clip_outliers: bool = True + use_percentiles: bool = False video_modality_keys: list[str] | None = None raw_stats: dict[str, Any] | None = None modality_config: dict[str, Any] | None = None @@ -1321,6 +1327,73 @@ class GrootN17PackInputsStep(ProcessorStep): return converted + def _normalize_action_groups_for_training(self, action: torch.Tensor) -> torch.Tensor | None: + if self.modality_config is None or self.raw_stats is None: + return None + + action_config = self.modality_config.get("action", {}) + if not isinstance(action_config, dict): + return None + 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 None + + normalized_groups: list[torch.Tensor] = [] + start_idx = 0 + for idx, key in enumerate(action_keys): + if not isinstance(key, str): + continue + cfg = ( + action_configs[idx] + if idx < len(action_configs) and isinstance(action_configs[idx], dict) + else {} + ) + is_relative = config_value(cfg.get("rep")) == "relative" + stats_modality = "relative_action" if is_relative else "action" + key_stats = self.raw_stats.get(stats_modality, {}).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]: + return None + + min_v, max_v = _n1_7_decode_stats_for_action( + self.raw_stats, + key, + cfg, + use_relative_action=True, + use_percentiles=self.use_percentiles, + ) + group = action[..., start_idx:end_idx] + min_t = torch.as_tensor(min_v, dtype=group.dtype, device=group.device) + max_t = torch.as_tensor(max_v, dtype=group.dtype, device=group.device) + if min_t.ndim == 1: + min_t = min_t.view(1, 1, -1) + max_t = max_t.view(1, 1, -1) + elif min_t.ndim == 2: + if group.shape[1] > min_t.shape[0]: + return None + min_t = min_t[: group.shape[1]].unsqueeze(0) + max_t = max_t[: group.shape[1]].unsqueeze(0) + else: + return None + + denom = max_t - min_t + mask = denom != 0 + safe_denom = torch.where(mask, denom, torch.ones_like(denom)) + normalized = torch.where(mask, 2 * (group - min_t) / safe_denom - 1, torch.zeros_like(group)) + if self.clip_outliers: + normalized = normalized.clamp(-1.0, 1.0) + normalized_groups.append(normalized) + start_idx = end_idx + + if not normalized_groups or start_idx != action.shape[-1]: + return None + return torch.cat(normalized_groups, dim=-1) + + def __call__(self, transition: EnvTransition) -> EnvTransition: obs = transition.get(TransitionKey.OBSERVATION, {}) or {} comp = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) or {} @@ -1435,8 +1508,12 @@ class GrootN17PackInputsStep(ProcessorStep): 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) + normalized_action = self._normalize_action_groups_for_training(action) + if normalized_action is not None: + action = normalized_action + else: + flat = _min_max_norm(action.reshape(bsz * horizon, dim), ACTION) + action = flat.view(bsz, horizon, dim) valid_dim = min(dim, self.max_action_dim) valid_horizon = min(horizon, self.valid_action_horizon, self.action_horizon) if dim < self.max_action_dim: @@ -1494,6 +1571,7 @@ class GrootN17PackInputsStep(ProcessorStep): "embodiment_mapping": self.embodiment_mapping, "normalize_min_max": self.normalize_min_max, "clip_outliers": self.clip_outliers, + "use_percentiles": self.use_percentiles, "video_modality_keys": self.video_modality_keys, "raw_stats": self.raw_stats, "modality_config": self.modality_config, diff --git a/src/lerobot/policies/groot/utils.py b/src/lerobot/policies/groot/utils.py index 7c1d03bcd..76111134d 100644 --- a/src/lerobot/policies/groot/utils.py +++ b/src/lerobot/policies/groot/utils.py @@ -111,7 +111,14 @@ def has_modality_stats(stats: dict[str, dict[str, Any]] | None) -> bool: def stat_dim_from_entry(entry: dict[str, Any]) -> int: for stat_name in ("mean", "q01", "min", "max", "std"): value = entry.get(stat_name) + if isinstance(value, torch.Tensor): + return int(value.shape[-1]) if value.ndim > 0 else 1 + if isinstance(value, np.ndarray): + return int(value.shape[-1]) if value.ndim > 0 else 1 if isinstance(value, list) and len(value) > 0: + first = value[0] + if isinstance(first, (list, tuple)) and len(first) > 0: + return len(first) return len(value) return 0 diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index eed3f1178..7d120e3b7 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -29,6 +29,7 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from accelerate import Accelerator +import numpy as np import torch from termcolor import colored from torch.optim import Optimizer @@ -218,6 +219,63 @@ def _unpadded_relative_action_vectors(relative_action: torch.Tensor, pad_mask: A return relative_action.reshape(-1, relative_action.shape[-1]) +def _relative_action_chunks_by_horizon( + relative_action: torch.Tensor, pad_mask: Any | None +) -> list[list[np.ndarray]]: + """Return per-horizon lists of valid relative action vectors.""" + + if relative_action.ndim == 2: + relative_action = relative_action.unsqueeze(0) + if relative_action.ndim != 3: + raise ValueError( + "Cannot compute horizon-preserving relative action statistics from " + f"shape {tuple(relative_action.shape)}." + ) + + batch_size, horizon, _action_dim = relative_action.shape + keep = torch.ones(batch_size, horizon, dtype=torch.bool) + if pad_mask is not None: + mask = torch.as_tensor(pad_mask, dtype=torch.bool).cpu() + if mask.ndim == 1 and batch_size == 1 and mask.numel() == horizon: + keep[0] = ~mask + elif mask.ndim == 2 and tuple(mask.shape) == (batch_size, horizon): + keep = ~mask + + chunks: list[list[np.ndarray]] = [[] for _ in range(horizon)] + relative_np = relative_action.detach().cpu().numpy() + for batch_idx in range(batch_size): + for horizon_idx in range(horizon): + if keep[batch_idx, horizon_idx]: + chunks[horizon_idx].append(relative_np[batch_idx, horizon_idx]) + return chunks + + +def _compute_horizon_relative_action_stats(chunks_by_horizon: list[list[np.ndarray]]) -> dict[str, np.ndarray]: + if not chunks_by_horizon or not any(chunks_by_horizon): + raise ValueError("Cannot compute relative action statistics without unpadded action vectors.") + + stats: dict[str, list[np.ndarray]] = {key: [] for key in ("min", "max", "mean", "std", "q01", "q99")} + counts: list[int] = [] + for horizon_idx, vectors in enumerate(chunks_by_horizon): + if len(vectors) < 2: + raise ValueError( + "Cannot compute horizon-preserving relative action statistics from fewer than 2 " + f"unpadded vectors at action timestep {horizon_idx}." + ) + values = np.stack(vectors, axis=0).astype(np.float32) + stats["min"].append(np.min(values, axis=0)) + stats["max"].append(np.max(values, axis=0)) + stats["mean"].append(np.mean(values, axis=0)) + stats["std"].append(np.std(values, axis=0)) + stats["q01"].append(np.quantile(values, 0.01, axis=0).astype(np.float32)) + stats["q99"].append(np.quantile(values, 0.99, axis=0).astype(np.float32)) + counts.append(len(vectors)) + + computed = {key: np.stack(values, axis=0) for key, values in stats.items()} + computed["count"] = np.asarray(counts, dtype=np.int64) + return computed + + def _iter_action_state_training_samples(dataset: Any): """Yield action chunks, reference states, and action padding masks without decoding videos when possible.""" @@ -261,6 +319,7 @@ def _make_relative_action_training_stats( *, exclude_joints: list[str] | None, action_names: list[str] | None, + preserve_action_horizon: bool = False, ) -> dict[str, dict[str, Any]]: """Return dataset stats whose action entry describes the relative action tensor used for training.""" @@ -286,6 +345,7 @@ def _make_relative_action_training_stats( action_names=action_names, ) num_vectors = 0 + chunks_by_horizon: list[list[np.ndarray]] | None = None for action_value, state_value, pad_mask in _iter_action_state_training_samples(dataset): action = _to_float_tensor(action_value, key=ACTION) @@ -306,19 +366,36 @@ def _make_relative_action_training_stats( state_batch, relative_step._build_mask(action_batch.shape[-1]), ) - vectors = _unpadded_relative_action_vectors(relative_action, pad_mask) - if vectors.numel() == 0: - continue - vector_count = int(vectors.reshape(-1, vectors.shape[-1]).shape[0]) - running_stats.update(vectors.numpy()) - num_vectors += vector_count + if preserve_action_horizon: + sample_chunks = _relative_action_chunks_by_horizon(relative_action, pad_mask) + if chunks_by_horizon is None: + chunks_by_horizon = [[] for _ in range(len(sample_chunks))] + if len(sample_chunks) != len(chunks_by_horizon): + raise ValueError( + "Cannot compute horizon-preserving relative action statistics from samples with " + f"different action horizons ({len(sample_chunks)} vs {len(chunks_by_horizon)})." + ) + for horizon_idx, vectors in enumerate(sample_chunks): + chunks_by_horizon[horizon_idx].extend(vectors) + num_vectors += len(vectors) + else: + vectors = _unpadded_relative_action_vectors(relative_action, pad_mask) + if vectors.numel() == 0: + continue + vector_count = int(vectors.reshape(-1, vectors.shape[-1]).shape[0]) + running_stats.update(vectors.numpy()) + num_vectors += vector_count if num_vectors < 2: raise ValueError( "Cannot compute relative action statistics from fewer than 2 unpadded action vectors." ) - stats[ACTION] = running_stats.get_statistics() + stats[ACTION] = ( + _compute_horizon_relative_action_stats(chunks_by_horizon or []) + if preserve_action_horizon + else running_stats.get_statistics() + ) return stats @@ -461,6 +538,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): dataset, exclude_joints=getattr(active_cfg, "relative_exclude_joints", []), action_names=_resolve_action_feature_names(dataset), + preserve_action_horizon=getattr(active_cfg, "type", None) == "groot", ) processor_kwargs = {} diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index 32c972d89..d37b18f90 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -1872,10 +1872,11 @@ def test_groot_n1_7_relative_action_training_processors_save_native_grouped_stat _RelativeStatsDataset(), exclude_joints=["gripper"], action_names=action_names, + preserve_action_horizon=True, ) expected_relative_action_stats = { - "min": torch.tensor([-2.0, -3.0, -4.0, -5.0, -6.0, 0.0]), - "max": torch.tensor([2.0, 3.0, 4.0, 5.0, 6.0, 100.0]), + "min": torch.tensor([-2.0, -3.0, -4.0, -5.0, -6.0, 1.0, 2.0, 3.0, 4.0, 5.0, 0.0]), + "max": torch.tensor([-1.0, -2.0, -3.0, -4.0, -5.0, 2.0, 3.0, 4.0, 5.0, 6.0, 100.0]), } preprocessor, postprocessor = make_groot_pre_post_processors( @@ -1899,7 +1900,10 @@ def test_groot_n1_7_relative_action_training_processors_save_native_grouped_stat {"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"]["relative_action"]["single_arm"]["min"] == [ + [-2.0, -3.0, -4.0, -5.0, -6.0], + [1.0, 2.0, 3.0, 4.0, 5.0], + ] assert pack_config["raw_stats"]["action"]["gripper"]["min"] == [0.0] assert pack_config["raw_stats"]["action"]["gripper"]["max"] == [100.0] @@ -1918,10 +1922,185 @@ def test_groot_n1_7_relative_action_training_processors_save_native_grouped_stat ) 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"]["relative_action"]["single_arm"]["max"] == [ + [-1.0, -2.0, -3.0, -4.0, -5.0], + [2.0, 3.0, 4.0, 5.0, 6.0], + ] assert decode_config["raw_stats"]["action"]["gripper"]["max"] == [100.0] +def test_groot_n1_7_generated_relative_stats_match_oss_gr00t_reference_numbers(): + input_features, output_features = _groot_features(state_dim=6, action_dim=6) + action_names = [ + "shoulder_pan.pos", + "shoulder_lift.pos", + "elbow_flex.pos", + "wrist_flex.pos", + "wrist_roll.pos", + "gripper.pos", + ] + config = GrootConfig( + input_features=input_features, + output_features=output_features, + device="cpu", + use_bf16=False, + action_decode_transform=None, + chunk_size=3, + n_action_steps=3, + use_relative_actions=True, + relative_exclude_joints=["gripper"], + ) + absolute_dataset_stats = { + OBS_STATE: { + "min": torch.tensor([-20.0, -30.0, -40.0, -50.0, -60.0, 0.0]), + "max": torch.tensor([80.0, 70.0, 60.0, 50.0, 40.0, 100.0]), + "mean": torch.tensor([30.0, 20.0, 10.0, 0.0, -10.0, 50.0]), + "std": torch.tensor([10.0, 10.0, 10.0, 10.0, 10.0, 10.0]), + "q01": torch.tensor([-10.0, -20.0, -30.0, -40.0, -50.0, 10.0]), + "q99": torch.tensor([70.0, 60.0, 50.0, 40.0, 30.0, 90.0]), + }, + ACTION: { + "min": torch.tensor([-5.0, -20.0, 0.0, -25.0, 10.0, 20.0]), + "max": torch.tensor([20.0, 30.0, 45.0, 60.0, 70.0, 90.0]), + "mean": torch.tensor([5.0, 5.0, 20.0, 20.0, 40.0, 55.0]), + "std": torch.tensor([5.0, 10.0, 10.0, 20.0, 20.0, 25.0]), + "q01": torch.tensor([-4.0, -19.0, 1.0, -24.0, 11.0, 20.0]), + "q99": torch.tensor([19.0, 29.0, 44.0, 59.0, 69.0, 90.0]), + }, + } + state_a = torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 25.0]) + state_b = torch.tensor([0.0, -10.0, 10.0, -20.0, 20.0, 75.0]) + action_a = torch.tensor( + [ + [11.0, 22.0, 33.0, 44.0, 55.0, 20.0], + [12.0, 24.0, 36.0, 48.0, 60.0, 80.0], + [13.0, 26.0, 39.0, 52.0, 65.0, 90.0], + ] + ) + action_b = torch.tensor( + [ + [-1.0, -8.0, 13.0, -16.0, 25.0, 30.0], + [-2.0, -6.0, 16.0, -12.0, 30.0, 40.0], + [-3.0, -4.0, 19.0, -8.0, 35.0, 50.0], + ] + ) + samples = [ + {OBS_STATE: state_a, ACTION: action_a}, + {OBS_STATE: state_b, ACTION: action_b}, + ] + + class _Dataset: + meta = SimpleNamespace( + stats=absolute_dataset_stats, + features={ACTION: {"names": action_names}}, + ) + + def __len__(self): + return len(samples) + + def __getitem__(self, idx): + return samples[idx] + + relative_dataset_stats = _make_relative_action_training_stats( + _Dataset(), + exclude_joints=["gripper"], + action_names=action_names, + preserve_action_horizon=True, + ) + + # Static reference values from OSS GR00T's JointActionChunk.relative_chunking + + # calculate_stats_for_key path: stats are computed per chunk timestep, not + # flattened over all timesteps. + oss_arm_min = torch.tensor( + [ + [-1.0, 2.0, 3.0, 4.0, 5.0], + [-2.0, 4.0, 6.0, 8.0, 10.0], + [-3.0, 6.0, 9.0, 12.0, 15.0], + ] + ) + oss_arm_max = torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0, 5.0], + [2.0, 4.0, 6.0, 8.0, 10.0], + [3.0, 6.0, 9.0, 12.0, 15.0], + ] + ) + oss_arm_mean = torch.tensor( + [ + [0.0, 2.0, 3.0, 4.0, 5.0], + [0.0, 4.0, 6.0, 8.0, 10.0], + [0.0, 6.0, 9.0, 12.0, 15.0], + ] + ) + oss_arm_std = torch.tensor( + [ + [1.0, 0.0, 0.0, 0.0, 0.0], + [2.0, 0.0, 0.0, 0.0, 0.0], + [3.0, 0.0, 0.0, 0.0, 0.0], + ] + ) + oss_arm_q01 = torch.tensor( + [ + [-0.98, 2.0, 3.0, 4.0, 5.0], + [-1.96, 4.0, 6.0, 8.0, 10.0], + [-2.94, 6.0, 9.0, 12.0, 15.0], + ] + ) + oss_arm_q99 = torch.tensor( + [ + [0.98, 2.0, 3.0, 4.0, 5.0], + [1.96, 4.0, 6.0, 8.0, 10.0], + [2.94, 6.0, 9.0, 12.0, 15.0], + ] + ) + + torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["min"][:, :5]), oss_arm_min) + torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["max"][:, :5]), oss_arm_max) + torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["mean"][:, :5]), oss_arm_mean) + torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["std"][:, :5]), oss_arm_std) + torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["q01"][:, :5]), oss_arm_q01) + torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["q99"][:, :5]), oss_arm_q99) + + preprocessor, postprocessor = make_groot_pre_post_processors( + config, + dataset_stats=relative_dataset_stats, + dataset_meta=_Dataset.meta, + ) + pack_step = next(step for step in preprocessor.steps if isinstance(step, GrootN17PackInputsStep)) + decode_step = next(step for step in postprocessor.steps if isinstance(step, GrootN17ActionDecodeStep)) + + assert pack_step.use_percentiles is True + torch.testing.assert_close( + torch.as_tensor(pack_step.raw_stats["relative_action"]["single_arm"]["min"]), + oss_arm_min, + ) + torch.testing.assert_close( + torch.as_tensor(pack_step.raw_stats["relative_action"]["single_arm"]["q99"]), + oss_arm_q99, + ) + assert pack_step.stats[ACTION]["min"] == pytest.approx([*oss_arm_min.flatten().tolist(), 20.0]) + assert pack_step.stats[ACTION]["max"] == pytest.approx([*oss_arm_max.flatten().tolist(), 90.0]) + + packed = pack_step( + { + TransitionKey.OBSERVATION: {OBS_STATE: state_a.unsqueeze(0)}, + TransitionKey.ACTION: action_a.unsqueeze(0), + TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move the vial"]}, + } + ) + expected_normalized = torch.tensor( + [ + [1.0, 0.0, 0.0, 0.0, 0.0, -1.0], + [1.0, 0.0, 0.0, 0.0, 0.0, 5.0 / 7.0], + [1.0, 0.0, 0.0, 0.0, 0.0, 1.0], + ] + ) + torch.testing.assert_close(packed[TransitionKey.ACTION][0, :3, :6], expected_normalized) + + decoded = decode_step({TransitionKey.ACTION: packed[TransitionKey.ACTION]}) + torch.testing.assert_close(decoded[TransitionKey.ACTION], action_a.unsqueeze(0), atol=1e-5, rtol=1e-5) + + def test_groot_policy_selects_n1_7_model_class(monkeypatch): from lerobot.policies.groot.groot_n1_7 import GR00TN17