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 = {}

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