Fix GROOT relative action training stats

src/lerobot/policies/groot/processor_groot.py

@@ -581,27 +581,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
 
 
@@ -887,6 +891,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,
     )
 
@@ -1179,6 +1184,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
@@ -1327,6 +1333,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 {}
@@ -1441,8 +1514,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:
@@ -1500,6 +1577,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,

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
 

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
@@ -228,6 +229,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."""
 
@@ -271,6 +329,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."""
 
@@ -296,6 +355,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)
@@ -316,19 +376,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
 
 
@@ -472,6 +549,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 = {}