Move GROOT relative stats out of train script

src/lerobot/policies/factory.py

@@ -538,6 +538,7 @@ def make_policy(
         set_dataset_feature_metadata = getattr(cfg, "set_dataset_feature_metadata", None)
         if callable(set_dataset_feature_metadata):
             set_dataset_feature_metadata(ds_meta.features)
+        cfg._runtime_dataset_meta = ds_meta
 
     kwargs["config"] = cfg
 

src/lerobot/policies/groot/processor_groot.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 
 import logging
-from copy import copy
+from copy import copy, deepcopy
 from dataclasses import dataclass, field, fields, is_dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -55,6 +55,7 @@ from lerobot.processor import (
     RenameObservationsProcessorStep,
     batch_to_transition,
     policy_action_to_transition,
+    to_relative_actions,
     transition_to_batch,
     transition_to_policy_action,
 )
@@ -343,25 +344,22 @@ def _load_n1_7_checkpoint_video_modality_keys(
     return keys or None
 
 
-# GR00T normalizes state/action inside its own processor steps and so deliberately has no
-# NormalizerProcessorStep/UnnormalizerProcessorStep (see GrootConfig.normalization_mapping, which is
-# IDENTITY for every feature). lerobot-train nonetheless emits these standard override keys
-# unconditionally, so for a GR00T pipeline they legitimately match no step. They are dropped up front
-# by _drop_groot_absent_standard_overrides so they neither break loading nor mask genuine typos.
-_GROOT_ABSENT_STANDARD_OVERRIDE_KEYS = frozenset({"normalizer_processor", "unnormalizer_processor"})
+# GR00T normalizes and represents actions inside its own processor steps, so it deliberately has no
+# standard NormalizerProcessorStep/UnnormalizerProcessorStep or generic relative/absolute action steps.
+# ``lerobot-train`` can still emit those generic override keys; for a GR00T pipeline they legitimately
+# match no step, so drop them up front without masking unrelated typo keys.
+_GROOT_ABSENT_STANDARD_OVERRIDE_KEYS = frozenset(
+    {
+        "absolute_actions_processor",
+        "normalizer_processor",
+        "relative_actions_processor",
+        "unnormalizer_processor",
+    }
+)
 
 
 def _drop_groot_absent_standard_overrides(overrides: dict[str, Any] | None) -> dict[str, Any] | None:
-    """Strip standard normalization override keys that a GR00T pipeline has no step for.
-
-    ``lerobot-train`` emits ``normalizer_processor``/``unnormalizer_processor`` overrides
-    unconditionally, but GR00T normalizes inside its own steps and has no such step (see
-    ``GrootConfig.normalization_mapping``). Both override-application paths reject keys that match no
-    step — ``_apply_groot_step_overrides`` raises for the freshly built raw-checkpoint pipeline, and
-    ``PolicyProcessorPipeline.from_pretrained`` raises via its used-override validation for the
-    serialized pipeline — so these keys are removed before either path runs. Any other unknown key
-    (e.g. a typo) is left in place and still raises.
-    """
+    """Strip standard override keys that a GR00T pipeline has no step for."""
 
     if not overrides:
         return overrides
@@ -581,6 +579,234 @@ def _resolve_visual_modality_keys_from_dataset_meta(dataset_meta: Any | None) ->
     return keys or None
 
 
+def _as_int(value: Any) -> int:
+    if isinstance(value, torch.Tensor):
+        return int(value.item())
+    item = getattr(value, "item", None)
+    if callable(item):
+        return int(item())
+    return int(value)
+
+
+def _to_float_tensor(value: Any, *, key: str) -> torch.Tensor:
+    if value is None:
+        raise ValueError(f"Cannot compute relative action statistics: sample is missing '{key}'.")
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu().float()
+    return torch.as_tensor(value, dtype=torch.float32)
+
+
+def _state_reference_batch(state: torch.Tensor) -> torch.Tensor:
+    if state.ndim == 1:
+        return state.unsqueeze(0)
+    if state.ndim == 2:
+        return state
+    if state.ndim > 2:
+        return state.reshape(-1, state.shape[-1])[-1:].contiguous()
+    raise ValueError(f"observation.state must have at least 1 dimension, got shape {tuple(state.shape)}.")
+
+
+def _action_training_batch(action: torch.Tensor, state_batch: torch.Tensor) -> torch.Tensor:
+    if action.ndim == 1:
+        return action.unsqueeze(0)
+    if action.ndim == 2:
+        if state_batch.shape[0] == action.shape[0] and state_batch.shape[0] > 1:
+            return action
+        return action.unsqueeze(0)
+    if action.ndim == 3:
+        return action
+    raise ValueError(f"action must be (D,), (T, D), (B, D), or (B, T, D), got {tuple(action.shape)}.")
+
+
+def _relative_action_chunks_by_horizon(
+    relative_action: torch.Tensor, pad_mask: Any | None
+) -> list[list[np.ndarray]]:
+    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):
+    ensure_reader = getattr(dataset, "_ensure_reader", None)
+    if callable(ensure_reader):
+        reader = ensure_reader()
+        if reader.hf_dataset is None:
+            reader.load_and_activate()
+        delta_indices = getattr(reader, "delta_indices", None)
+        for idx in range(len(dataset)):
+            item = reader.hf_dataset[idx]
+            action = item.get(ACTION)
+            state = item.get(OBS_STATE)
+            pad_mask = None
+            if delta_indices is not None and ACTION in delta_indices:
+                ep_idx = _as_int(item["episode_index"])
+                abs_idx = _as_int(item["index"])
+                query_indices, padding = reader._get_query_indices(abs_idx, ep_idx)
+                action = reader._query_hf_dataset({ACTION: query_indices[ACTION]})[ACTION]
+                pad_mask = padding.get(f"{ACTION}_is_pad")
+            yield action, state, pad_mask
+        return
+
+    for idx in range(len(dataset)):
+        item = dataset[idx]
+        yield item.get(ACTION), item.get(OBS_STATE), item.get(f"{ACTION}_is_pad")
+
+
+def _make_relative_action_training_stats(
+    dataset: Any,
+    *,
+    exclude_joints: list[str] | None,
+    action_names: list[str] | None,
+    preserve_action_horizon: bool = True,
+) -> dict[str, dict[str, Any]]:
+    try:
+        dataset_len = len(dataset)
+    except TypeError as exc:
+        raise ValueError(
+            "Cannot compute relative action statistics for a dataset without a finite length. "
+            "Disable streaming or provide precomputed relative action statistics."
+        ) from exc
+
+    if dataset_len == 0:
+        raise ValueError("Cannot compute relative action statistics for an empty dataset.")
+
+    relative_step = RelativeActionsProcessorStep(
+        enabled=True,
+        exclude_joints=list(exclude_joints or []),
+        action_names=action_names,
+    )
+    stats = deepcopy(getattr(getattr(dataset, "meta", None), "stats", {}) or {})
+    chunks_by_horizon: list[list[np.ndarray]] | None = None
+    num_vectors = 0
+
+    for action_value, state_value, pad_mask in _iter_action_state_training_samples(dataset):
+        action = _to_float_tensor(action_value, key=ACTION)
+        state = _to_float_tensor(state_value, key=OBS_STATE)
+        state_batch = _state_reference_batch(state)
+        action_batch = _action_training_batch(action, state_batch)
+        if action_batch.shape[0] != state_batch.shape[0]:
+            if state_batch.shape[0] == 1:
+                state_batch = state_batch.expand(action_batch.shape[0], -1)
+            else:
+                raise ValueError(
+                    "Cannot compute relative action statistics: action and state batch sizes differ "
+                    f"({action_batch.shape[0]} vs {state_batch.shape[0]})."
+                )
+
+        relative_action = to_relative_actions(
+            action_batch,
+            state_batch,
+            relative_step._build_mask(action_batch.shape[-1]),
+        )
+        if not preserve_action_horizon:
+            relative_action = relative_action.reshape(-1, relative_action.shape[-1]).unsqueeze(0)
+            pad_mask = None
+        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)
+
+    if num_vectors < 2:
+        raise ValueError("Cannot compute relative action statistics from fewer than 2 unpadded action vectors.")
+
+    stats[ACTION] = _compute_horizon_relative_action_stats(chunks_by_horizon or [])
+    return stats
+
+
+def _stats_preserve_action_horizon(stats: dict[str, dict[str, Any]] | None) -> bool:
+    if not stats or ACTION not in stats:
+        return False
+    action_stats = stats.get(ACTION) or {}
+    for stat_name in ("min", "max", "mean", "std", "q01", "q99"):
+        value = action_stats.get(stat_name)
+        if value is None:
+            continue
+        return torch.as_tensor(value).ndim >= 2
+    return False
+
+
+def _make_relative_action_training_stats_from_dataset_meta(
+    config: GrootConfig, dataset_meta: Any | None
+) -> dict[str, dict[str, Any]] | None:
+    repo_id = getattr(dataset_meta, "repo_id", None)
+    root = getattr(dataset_meta, "root", None)
+    fps = getattr(dataset_meta, "fps", None)
+    if dataset_meta is None or repo_id is None or root is None or fps is None:
+        return None
+
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+    delta_timestamps = {ACTION: [index / fps for index in config.action_delta_indices]}
+    dataset = LeRobotDataset(
+        repo_id,
+        root=root,
+        delta_timestamps=delta_timestamps,
+        revision=getattr(dataset_meta, "revision", None),
+        download_videos=False,
+        return_uint8=True,
+    )
+    return _make_relative_action_training_stats(
+        dataset,
+        exclude_joints=list(config.relative_exclude_joints or []),
+        action_names=_resolve_action_feature_names_from_dataset_meta(dataset_meta),
+        preserve_action_horizon=True,
+    )
+
+
 def _slice_stats_entry(stats: dict[str, Any], indices: list[int]) -> dict[str, Any]:
     if not indices:
         return {}
@@ -840,20 +1066,28 @@ def make_groot_pre_post_processors(
         Tuple of (preprocessor, postprocessor) pipelines
     """
 
+    dataset_meta = dataset_meta or getattr(config, "_runtime_dataset_meta", None)
     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_dataset_stats = dataset_stats
+        if not _stats_preserve_action_horizon(relative_dataset_stats):
+            relative_dataset_stats = _make_relative_action_training_stats_from_dataset_meta(config, dataset_meta)
         relative_assets = _build_n1_7_relative_action_processor_assets(
             config,
-            dataset_stats,
+            relative_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)
+        if relative_assets is None:
+            raise ValueError(
+                "GR00T relative-action training requires horizon-preserving relative action statistics. "
+                "Pass dataset_meta with a local LeRobot dataset root, or pass precomputed relative dataset_stats."
+            )
+        checkpoint_assets = relative_assets
+        checkpoint_stats = checkpoint_assets.stats
+        checkpoint_has_stats = has_modality_stats(checkpoint_stats)
 
     action_horizon = (
         checkpoint_assets.max_action_horizon

src/lerobot/scripts/lerobot_train.py

@@ -22,14 +22,12 @@ import dataclasses
 import logging
 import time
 from contextlib import nullcontext
-from copy import deepcopy
 from pprint import pformat
 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
@@ -56,7 +54,6 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
 from lerobot.rewards import make_reward_pre_post_processors
 from lerobot.utils.collate import lerobot_collate_fn
-from lerobot.utils.constants import ACTION, OBS_STATE
 from lerobot.utils.import_utils import register_third_party_plugins
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
 from lerobot.utils.random_utils import set_seed
@@ -174,241 +171,6 @@ def update_policy(
     return train_metrics, output_dict
 
 
-def _as_int(value: Any) -> int:
-    if isinstance(value, torch.Tensor):
-        return int(value.item())
-    item = getattr(value, "item", None)
-    if callable(item):
-        return int(item())
-    return int(value)
-
-
-def _to_float_tensor(value: Any, *, key: str) -> torch.Tensor:
-    if value is None:
-        raise ValueError(f"Cannot compute relative action statistics: sample is missing '{key}'.")
-    if isinstance(value, torch.Tensor):
-        return value.detach().cpu().float()
-    return torch.as_tensor(value, dtype=torch.float32)
-
-
-def _state_reference_batch(state: torch.Tensor) -> torch.Tensor:
-    if state.ndim == 1:
-        return state.unsqueeze(0)
-    if state.ndim == 2:
-        return state
-    if state.ndim > 2:
-        return state.reshape(-1, state.shape[-1])[-1:].contiguous()
-    raise ValueError(f"observation.state must have at least 1 dimension, got shape {tuple(state.shape)}.")
-
-
-def _action_training_batch(action: torch.Tensor, state_batch: torch.Tensor) -> torch.Tensor:
-    if action.ndim == 1:
-        return action.unsqueeze(0)
-    if action.ndim == 2:
-        # A single training sample uses (T, D) action chunks with a single (1, D) state reference.
-        # Batched callers may pass (B, D); keep that shape when the state batch makes it unambiguous.
-        if state_batch.shape[0] == action.shape[0] and state_batch.shape[0] > 1:
-            return action
-        return action.unsqueeze(0)
-    if action.ndim == 3:
-        return action
-    raise ValueError(f"action must be (D,), (T, D), (B, D), or (B, T, D), got {tuple(action.shape)}.")
-
-
-def _unpadded_relative_action_vectors(relative_action: torch.Tensor, pad_mask: Any | None) -> torch.Tensor:
-    if pad_mask is None:
-        return relative_action.reshape(-1, relative_action.shape[-1])
-
-    keep = ~torch.as_tensor(pad_mask, dtype=torch.bool).cpu()
-    if relative_action.ndim == 3 and keep.ndim == 1 and relative_action.shape[0] == 1:
-        return relative_action[0, keep]
-    if relative_action.ndim == 3 and keep.ndim == 2 and tuple(keep.shape) == tuple(relative_action.shape[:2]):
-        return relative_action[keep]
-    if relative_action.ndim == 2 and keep.ndim == 1 and keep.numel() == relative_action.shape[0]:
-        return relative_action[keep]
-    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."""
-
-    ensure_reader = getattr(dataset, "_ensure_reader", None)
-    if callable(ensure_reader):
-        reader = ensure_reader()
-        if reader.hf_dataset is None:
-            reader.load_and_activate()
-        delta_indices = getattr(reader, "delta_indices", None)
-        for idx in range(len(dataset)):
-            item = reader.hf_dataset[idx]
-            action = item.get(ACTION)
-            state = item.get(OBS_STATE)
-            pad_mask = None
-            if delta_indices is not None and ACTION in delta_indices:
-                ep_idx = _as_int(item["episode_index"])
-                abs_idx = _as_int(item["index"])
-                query_indices, padding = reader._get_query_indices(abs_idx, ep_idx)
-                action = reader._query_hf_dataset({ACTION: query_indices[ACTION]})[ACTION]
-                pad_mask = padding.get(f"{ACTION}_is_pad")
-            yield action, state, pad_mask
-        return
-
-    for idx in range(len(dataset)):
-        item = dataset[idx]
-        yield item.get(ACTION), item.get(OBS_STATE), item.get(f"{ACTION}_is_pad")
-
-
-def _resolve_action_feature_names(dataset: Any) -> list[str] | None:
-    features = getattr(getattr(dataset, "meta", None), "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)
-    return list(names) if names is not None else None
-
-
-def _make_relative_action_training_stats(
-    dataset: Any,
-    *,
-    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."""
-
-    from lerobot.datasets.compute_stats import RunningQuantileStats
-    from lerobot.processor.relative_action_processor import RelativeActionsProcessorStep, to_relative_actions
-
-    try:
-        dataset_len = len(dataset)
-    except TypeError as exc:
-        raise ValueError(
-            "Cannot compute relative action statistics for a dataset without a finite length. "
-            "Disable streaming or provide precomputed relative action statistics."
-        ) from exc
-
-    if dataset_len == 0:
-        raise ValueError("Cannot compute relative action statistics for an empty dataset.")
-
-    stats = deepcopy(getattr(getattr(dataset, "meta", None), "stats", {}) or {})
-    running_stats = RunningQuantileStats()
-    relative_step = RelativeActionsProcessorStep(
-        enabled=True,
-        exclude_joints=list(exclude_joints or []),
-        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)
-        state = _to_float_tensor(state_value, key=OBS_STATE)
-        state_batch = _state_reference_batch(state)
-        action_batch = _action_training_batch(action, state_batch)
-        if action_batch.shape[0] != state_batch.shape[0]:
-            if state_batch.shape[0] == 1:
-                state_batch = state_batch.expand(action_batch.shape[0], -1)
-            else:
-                raise ValueError(
-                    "Cannot compute relative action statistics: action and state batch sizes differ "
-                    f"({action_batch.shape[0]} vs {state_batch.shape[0]})."
-                )
-
-        relative_action = to_relative_actions(
-            action_batch,
-            state_batch,
-            relative_step._build_mask(action_batch.shape[-1]),
-        )
-        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] = (
-        _compute_horizon_relative_action_stats(chunks_by_horizon or [])
-        if preserve_action_horizon
-        else running_stats.get_statistics()
-    )
-    return stats
-
-
 @parser.wrap()
 def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     """
@@ -541,29 +303,19 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     active_cfg = cfg.trainable_config
     processor_pretrained_path = active_cfg.pretrained_path
-    processor_stats = dataset.meta.stats
-    if not cfg.is_reward_model_training and getattr(active_cfg, "use_relative_actions", False):
-        if is_main_process:
-            logging.info("Computing relative-action output statistics for processor normalization")
-        processor_stats = _make_relative_action_training_stats(
-            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 = {}
     if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path:
-        processor_kwargs["dataset_stats"] = processor_stats
+        processor_kwargs["dataset_stats"] = dataset.meta.stats
 
-    if cfg.is_reward_model_training or getattr(active_cfg, "use_relative_actions", False):
+    if cfg.is_reward_model_training:
         processor_kwargs["dataset_meta"] = dataset.meta
 
     if not cfg.is_reward_model_training and processor_pretrained_path is not None:
         preprocessor_overrides = {
             "device_processor": {"device": device.type},
             "normalizer_processor": {
-                "stats": processor_stats,
+                "stats": dataset.meta.stats,
                 "features": {**policy.config.input_features, **policy.config.output_features},
                 "norm_map": policy.config.normalization_mapping,
             },
@@ -571,7 +323,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         }
         postprocessor_overrides = {
             "unnormalizer_processor": {
-                "stats": processor_stats,
+                "stats": dataset.meta.stats,
                 "features": policy.config.output_features,
                 "norm_map": policy.config.normalization_mapping,
             },
@@ -580,7 +332,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             preprocessor_overrides["relative_actions_processor"] = {
                 "enabled": True,
                 "exclude_joints": getattr(active_cfg, "relative_exclude_joints", []),
-                "action_names": _resolve_action_feature_names(dataset),
+                "action_names": getattr(active_cfg, "action_feature_names", None),
             }
             postprocessor_overrides["absolute_actions_processor"] = {"enabled": True}
         processor_kwargs["preprocessor_overrides"] = preprocessor_overrides

tests/policies/groot/test_groot_n1_7.py

@@ -41,6 +41,7 @@ from lerobot.policies.groot.processor_groot import (
     GrootN17ActionDecodeStep,
     GrootN17PackInputsStep,
     GrootN17VLMEncodeStep,
+    _make_relative_action_training_stats,
     _transform_n1_7_image_for_vlm_albumentations,
     make_groot_pre_post_processors,
 )
@@ -49,7 +50,6 @@ from lerobot.processor import (
     PolicyProcessorPipeline,
     RelativeActionsProcessorStep,
 )
-from lerobot.scripts.lerobot_train import _make_relative_action_training_stats
 from lerobot.types import TransitionKey
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
@@ -1990,6 +1990,101 @@ def test_groot_n1_7_relative_action_training_processors_save_native_grouped_stat
     assert decode_config["raw_stats"]["action"]["gripper"]["max"] == [100.0]
 
 
+def test_groot_n1_7_relative_action_processors_compute_stats_from_runtime_dataset_meta(
+    monkeypatch, tmp_path
+):
+    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=2,
+        n_action_steps=2,
+        use_relative_actions=True,
+        relative_exclude_joints=["gripper"],
+    )
+    absolute_dataset_stats = {
+        OBS_STATE: {
+            "min": torch.tensor([-50.0, -60.0, -70.0, -80.0, -90.0, 0.0]),
+            "max": torch.tensor([50.0, 60.0, 70.0, 80.0, 90.0, 100.0]),
+        },
+        ACTION: {
+            "min": torch.tensor([-100.0, -110.0, -120.0, -130.0, -140.0, 0.0]),
+            "max": torch.tensor([100.0, 110.0, 120.0, 130.0, 140.0, 100.0]),
+        },
+    }
+    samples = [
+        {
+            OBS_STATE: torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 0.0]),
+            ACTION: torch.tensor(
+                [
+                    [8.0, 17.0, 26.0, 35.0, 44.0, 0.0],
+                    [12.0, 23.0, 34.0, 45.0, 56.0, 100.0],
+                ]
+            ),
+        },
+        {
+            OBS_STATE: torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 50.0]),
+            ACTION: torch.tensor(
+                [
+                    [-1.0, -2.0, -3.0, -4.0, -5.0, 25.0],
+                    [1.0, 2.0, 3.0, 4.0, 5.0, 75.0],
+                ]
+            ),
+        },
+    ]
+    runtime_meta = SimpleNamespace(
+        repo_id="local/relative",
+        root=tmp_path,
+        revision="main",
+        fps=30,
+        stats=absolute_dataset_stats,
+        features={ACTION: {"names": action_names}},
+    )
+
+    class _RelativeStatsDataset:
+        meta = runtime_meta
+
+        def __len__(self):
+            return len(samples)
+
+        def __getitem__(self, idx):
+            return samples[idx]
+
+    def _fake_lerobot_dataset(repo_id, **kwargs):
+        assert repo_id == runtime_meta.repo_id
+        assert kwargs["root"] == runtime_meta.root
+        assert kwargs["revision"] == runtime_meta.revision
+        assert kwargs["download_videos"] is False
+        assert kwargs["delta_timestamps"][ACTION] == [0.0, 1 / runtime_meta.fps]
+        return _RelativeStatsDataset()
+
+    monkeypatch.setattr("lerobot.datasets.lerobot_dataset.LeRobotDataset", _fake_lerobot_dataset)
+    config._runtime_dataset_meta = runtime_meta
+
+    preprocessor, postprocessor = make_groot_pre_post_processors(config, dataset_stats=absolute_dataset_stats)
+
+    assert not any(isinstance(step, RelativeActionsProcessorStep) for step in preprocessor.steps)
+    assert isinstance(postprocessor.steps[0], GrootN17ActionDecodeStep)
+    pack_step = next(step for step in preprocessor.steps if isinstance(step, GrootN17PackInputsStep))
+    assert pack_step.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_step.raw_stats["relative_action"]["single_arm"]["count"] == [2, 2]
+    assert pack_step.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 = [