Update Wall-X to 1.1.0 (#104)

scripts/compute_norm_stats.py

@@ -1,182 +1,752 @@
 #!/usr/bin/env python3
+"""Compute LeRobot normalization stats (mean, std, q01, q99) for training.
 
+Writes JSON in the format expected by wall-x training configs::
+
+    {"norm_stats": {
+        "observation.state": {"mean": [...], "std": [...], "q01": [...], "q99": [...]},
+        "action": {"mean": [...], "std": [...], "q01": [...], "q99": [...]}
+    }}
+
+When ``--train_config`` is provided, the script reads ``data.lerobot_config.repo_id``,
+``norm_stats_path``, ``task.dof_config``, ``task.agent_pos_config``, and
+``task.action_horizon`` from the YAML. Per-DOF slices are aggregated separately;
+keys ending with ``_relative`` use the same relative-pose logic as the LeRobot loader.
+
+Usage
+-----
+Recommended: pass a finetune YAML (paths in the config can be placeholders; override
+with CLI flags if needed)::
+
+    python scripts/compute_norm_stats.py \\
+        --train_config /path/to/train_config.yml
+
+Multi-task example::
+
+    python scripts/compute_norm_stats.py \\
+        --train_config /path/to/multitask_config.yml
+
+Override dataset or output path from the command line::
+
+    python scripts/compute_norm_stats.py \\
+        --train_config /path/to/train_config.yml \\
+        --data_root /path/to/repo_id \\
+        --output_path /path/to/norm_stats_path
+
+Without a train config (global stats only, no per-DOF relative slices)::
+
+    python scripts/compute_norm_stats.py \\
+        --data_root /path/to/lerobot_dataset \\
+        --output_path /path/to/norm_stats.json
+
+Requirements
+------------
+- Local LeRobot v3 dataset at ``--data_root`` (or ``data.lerobot_config.repo_id``)
+- ``lerobot>=0.3``, ``datasets``, ``pyarrow``, ``numpy``, ``pyyaml``, ``tqdm``
+
+After running, set ``norm_stats_path`` in your training YAML to the generated JSON.
+"""
+
+import argparse
 import json
 import logging
-from collections import defaultdict
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List
-from tqdm import tqdm
+from typing import Any
 
 import numpy as np
+import yaml
+from numba import jit, prange
+from tqdm import tqdm
 
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
+SKIP_DOF_KEYS = frozenset(
+    {"velocity_decomposed", "height", "head_actions", "action_padding"}
+)
+
+_GEOMETRY_FUNCS_LOADED = False
 
 
-def write_json(path: Path, data: Dict) -> None:
+def _ensure_geometry_funcs() -> None:
+    global _GEOMETRY_FUNCS_LOADED
+    global canonicalize_euler_zyx_batch_nb
+    global euler_to_matrix_zyx_batch_nb
+    global matrix_to_euler_zyx_batch_nb
+    global so3_to_matrix_batch_nb
+
+    if _GEOMETRY_FUNCS_LOADED:
+        return
+
+    try:
+        from wall_x._vendor.x2robot_utils.geometry import (
+            canonicalize_euler_zyx_batch_nb,
+            euler_to_matrix_zyx_batch_nb,
+            matrix_to_euler_zyx_batch_nb,
+            so3_to_matrix_batch_nb,
+        )
+    except ImportError as exc:
+        raise RuntimeError(
+            "compute_norm_stats.py requires the vendored x2robot_utils geometry "
+            "helpers. Install Wall-X first so wall_x._vendor is available."
+        ) from exc
+
+    _GEOMETRY_FUNCS_LOADED = True
+
+
+@dataclass
+class TrainConfigContext:
+    data_root: Path
+    output_path: Path
+    state_key: str
+    action_key: str
+    propri_ranges: dict[str, list[int]]
+    action_ranges: dict[str, list[int]]
+    action_chunk: int
+    dof_config: dict[str, int]
+    agent_pos_config: dict[str, int]
+
+
+def write_norm_stats(path: Path, norm_stats: dict[str, dict]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
     path.write_text(
-        json.dumps(data, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
+        json.dumps({"norm_stats": norm_stats}, ensure_ascii=False, indent=2) + "\n",
         encoding="utf-8",
     )
 
 
-def compute_action_statistics(
-    action_data_by_robot: Dict[str, Dict[str, List]]
-) -> Dict[str, Dict[str, Dict]]:
+def compute_vector_stats(values: np.ndarray) -> dict[str, list[float]]:
+    if values.ndim == 1:
+        values = values.reshape(-1, 1)
+    return {
+        "mean": np.mean(values, axis=0).tolist(),
+        "std": np.std(values, axis=0).tolist(),
+        "q01": np.quantile(values, 0.01, axis=0).tolist(),
+        "q99": np.quantile(values, 0.99, axis=0).tolist(),
+    }
+
+
+def _apply_slice_stats(
+    full_stats: dict[str, list[float]],
+    index_range: list[int],
+    slice_stats: dict[str, list[float]],
+) -> None:
+    start, end = index_range
+    for field in ("mean", "std", "q01", "q99"):
+        full_stats[field][start:end] = slice_stats[field]
+
+
+def config_to_index_ranges(config: dict[str, int]) -> dict[str, list[int]]:
+    ranges: dict[str, list[int]] = {}
+    cur = 0
+    for key, dim in config.items():
+        if key in SKIP_DOF_KEYS:
+            continue
+        ranges[key] = [cur, cur + int(dim)]
+        cur += int(dim)
+    return ranges
+
+
+def load_train_config(path: Path) -> dict[str, Any]:
+    with path.open(encoding="utf-8") as f:
+        config = yaml.load(f, Loader=yaml.SafeLoader)
+    if not isinstance(config, dict):
+        raise ValueError(f"train config must be a YAML mapping, got {type(config)}")
+    return config
+
+
+def parse_train_config(config: dict[str, Any]) -> TrainConfigContext:
+    task = config.get("task")
+    if not isinstance(task, dict):
+        raise ValueError("train config must contain a 'task' section")
+
+    dof_config = task.get("dof_config")
+    agent_pos_config = task.get("agent_pos_config")
+    if not isinstance(dof_config, dict) or not dof_config:
+        raise ValueError("task.dof_config is required in train config")
+    if not isinstance(agent_pos_config, dict) or not agent_pos_config:
+        raise ValueError("task.agent_pos_config is required in train config")
+
+    data_cfg = config.get("data")
+    if not isinstance(data_cfg, dict):
+        raise ValueError("train config must contain a 'data' section")
+
+    lerobot_config = data_cfg.get("lerobot_config")
+    if not isinstance(lerobot_config, dict):
+        raise ValueError("data.lerobot_config is required in train config")
+
+    repo_id = lerobot_config.get("repo_id")
+    if not repo_id:
+        raise ValueError("data.lerobot_config.repo_id is required in train config")
+
+    norm_stats_path = config.get("norm_stats_path") or data_cfg.get("norm_stats_path")
+    if not norm_stats_path:
+        raise ValueError(
+            "norm_stats_path is required in train config " "(top-level or under data)"
+        )
+
+    key_mappings = data_cfg.get("key_mappings")
+    if not isinstance(key_mappings, dict):
+        raise ValueError("data.key_mappings is required in train config")
+
+    state_key = key_mappings.get("state", "observation.state")
+    action_key = key_mappings.get("action", "action")
+
+    action_chunk = int(
+        task.get("action_horizon")
+        or task.get("action_horizon_flow")
+        or data_cfg.get("action_horizon")
+        or 32
+    )
+
+    return TrainConfigContext(
+        data_root=Path(repo_id),
+        output_path=Path(norm_stats_path),
+        state_key=state_key,
+        action_key=action_key,
+        propri_ranges=config_to_index_ranges(agent_pos_config),
+        action_ranges=config_to_index_ranges(dof_config),
+        action_chunk=action_chunk,
+        dof_config=dof_config,
+        agent_pos_config=agent_pos_config,
+    )
+
+
+def layout_vector_dim(layout_config: dict[str, int]) -> int:
+    """Config vector width excluding virtual padding keys."""
+    return sum(
+        int(dim) for key, dim in layout_config.items() if key not in SKIP_DOF_KEYS
+    )
+
+
+def _prepare_arrays_for_layout(
+    states: np.ndarray,
+    actions: np.ndarray,
+    agent_pos_config: dict[str, int],
+    dof_config: dict[str, int],
+) -> tuple[np.ndarray, np.ndarray]:
     """
-    Compute statistics (min, q01, q99, max) for each action type and dimension.
-
-    Args:
-        action_data_by_robot: Dict[robot_id][action_type] -> list of arrays/lists
-
-    Returns:
-        Dict[robot_id][action_type] -> {
-            "min": [min for each dim],
-            "q01": [quantile 1% for each dim],
-            "q99": [quantile 99% for each dim],
-            "max": [max for each dim],
-            "delta": [max - min for each dim]
-            "delta_q99_q01": [q99 - q01 for each dim]
-        }
+    Match LeRobot training loader: convert 3D Euler slices to 6D when config
+    uses rotation_6D keys but the dataset stores 14-dim Euler vectors.
     """
-    stats = {}
+    from wall_x.data.backends.lerobot.rotation_layout import (
+        euler_layout_dim,
+        layout_uses_6d_rotation,
+        maybe_convert_euler_to_6d,
+    )
 
-    for robot_id, action_data in action_data_by_robot.items():
-        stats[robot_id] = {}
+    convert_state = layout_uses_6d_rotation(agent_pos_config)
+    convert_action = layout_uses_6d_rotation(dof_config)
 
-        for action_type, values_list in action_data.items():
-            if not values_list:
-                continue
+    if convert_state:
+        raw_dim = euler_layout_dim(agent_pos_config)
+        target_dim = layout_vector_dim(agent_pos_config)
+        if states.shape[-1] == raw_dim:
+            states = maybe_convert_euler_to_6d(states, agent_pos_config, True)
+            logging.info(
+                "Converted state Euler->6D (%d -> %d dims)", raw_dim, states.shape[-1]
+            )
+        elif states.shape[-1] != target_dim:
+            raise ValueError(
+                f"State dim {states.shape[-1]} does not match Euler raw dim "
+                f"{raw_dim} or 6D layout dim {target_dim} from agent_pos_config"
+            )
 
-            # Convert to numpy array: shape (num_samples, num_dims)
-            try:
-                values_array = np.array(values_list)
-                if values_array.size == 0:
-                    continue
+    if convert_action:
+        raw_dim = euler_layout_dim(dof_config)
+        target_dim = layout_vector_dim(dof_config)
+        if actions.shape[-1] == raw_dim:
+            actions = maybe_convert_euler_to_6d(actions, dof_config, True)
+            logging.info(
+                "Converted action Euler->6D (%d -> %d dims)", raw_dim, actions.shape[-1]
+            )
+        elif actions.shape[-1] != target_dim:
+            raise ValueError(
+                f"Action dim {actions.shape[-1]} does not match Euler raw dim "
+                f"{raw_dim} or 6D layout dim {target_dim} from dof_config"
+            )
 
-                # Handle both 1D and 2D cases
-                if values_array.ndim == 1:
-                    values_array = values_array.reshape(-1, 1)
-                elif values_array.ndim == 2:
-                    pass
-                else:
-                    logging.warning(
-                        f"Unexpected shape for {robot_id}/{action_type}: {values_array.shape}"
-                    )
-                    continue
+    expected_state_dim = layout_vector_dim(agent_pos_config)
+    expected_action_dim = layout_vector_dim(dof_config)
+    if states.shape[-1] != expected_state_dim:
+        raise ValueError(
+            f"State dim {states.shape[-1]} != expected layout dim {expected_state_dim}"
+        )
+    if actions.shape[-1] != expected_action_dim:
+        raise ValueError(
+            f"Action dim {actions.shape[-1]} != expected layout dim {expected_action_dim}"
+        )
+    return states, actions
 
-                # Compute statistics for each dimension
-                min_vals = np.min(values_array, axis=0).tolist()
-                max_vals = np.max(values_array, axis=0).tolist()
-                q01_vals = np.quantile(values_array, 0.01, axis=0).tolist()
-                q99_vals = np.quantile(values_array, 0.99, axis=0).tolist()
-                delta_vals = (np.array(max_vals) - np.array(min_vals)).tolist()
-                delta_q99_q01_vals = (np.array(q99_vals) - np.array(q01_vals)).tolist()
 
-                stats[robot_id][action_type] = {
-                    "min": min_vals,
-                    "q01": q01_vals,
-                    "q99": q99_vals,
-                    "max": max_vals,
-                    "delta": delta_vals,
-                    "delta_q99_q01": delta_q99_q01_vals,
-                }
+def resolve_lerobot_dataset_paths(dataset_root: Path) -> tuple[str, Path]:
+    """Return ``(repo_id, root)`` for a local LeRobot dataset directory."""
+    root = dataset_root.expanduser().resolve()
+    if not root.is_dir():
+        raise FileNotFoundError(f"LeRobot dataset root not found: {root}")
+    return root.name, root
 
-            except Exception as e:
-                logging.warning(
-                    f"Error computing statistics for {robot_id}/{action_type}: {e}"
+
+def _load_parquet_state_action_only(
+    data_root: Path,
+    state_key: str,
+    action_key: str,
+):
+    """
+    Read only state/action columns from LeRobot v3 parquet.
+
+    Does not load image/video columns or decode mp4 files.
+    """
+    import pyarrow.dataset as pa_ds
+    from datasets import Dataset
+
+    root = data_root.expanduser().resolve()
+    paths = sorted((root / "data").glob("*/*.parquet"))
+    if not paths:
+        raise FileNotFoundError(f"No parquet files under {root / 'data'}")
+
+    logging.info(
+        "Reading parquet columns %r, %r only (no video/images)",
+        state_key,
+        action_key,
+    )
+    arrow_dataset = pa_ds.dataset([str(path) for path in paths], format="parquet")
+    table = arrow_dataset.to_table(columns=[state_key, action_key])
+    return Dataset(table)
+
+
+def _load_state_action_table(
+    data_root: Path,
+    state_key: str,
+    action_key: str,
+):
+    root = data_root.expanduser().resolve()
+    if root.is_dir() and (root / "meta" / "info.json").is_file():
+        table = _load_parquet_state_action_only(root, state_key, action_key)
+        return table, state_key, action_key
+
+    try:
+        from lerobot.datasets.lerobot_dataset import LeRobotDataset
+    except ImportError as exc:
+        raise RuntimeError(
+            "compute_norm_stats.py requires LeRobot. Install it first, for example "
+            "`pip install lerobot==0.4.4` or follow the repository README."
+        ) from exc
+
+    dataset = LeRobotDataset(str(data_root), root=None, video_backend="pyav")
+    non_image_columns = [
+        col for col in dataset.features if "image" not in col and col not in {"task"}
+    ]
+    if state_key not in non_image_columns or action_key not in non_image_columns:
+        raise ValueError(
+            f"Expected keys {state_key!r} and {action_key!r} in dataset columns, "
+            f"got {non_image_columns}"
+        )
+    table = dataset.hf_dataset.select_columns([state_key, action_key])
+    return table, state_key, action_key
+
+
+def _table_to_arrays(
+    table,
+    state_key: str,
+    action_key: str,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Load full columns once; avoids O(N*chunk) random row access."""
+    logging.info("Loading state/action columns into memory...")
+    try:
+        states = np.asarray(table[state_key], dtype=np.float32)
+        actions = np.asarray(table[action_key], dtype=np.float32)
+    except (KeyError, TypeError, ValueError) as exc:
+        logging.warning(
+            "Column-wise load failed (%s); falling back to per-row stack.", exc
+        )
+        states = np.stack(
+            [
+                np.asarray(table[i][state_key], dtype=np.float32)
+                for i in range(len(table))
+            ]
+        )
+        actions = np.stack(
+            [
+                np.asarray(table[i][action_key], dtype=np.float32)
+                for i in range(len(table))
+            ]
+        )
+    if states.ndim == 1:
+        states = states.reshape(-1, 1)
+    if actions.ndim == 1:
+        actions = actions.reshape(-1, 1)
+    logging.info(
+        "  frames=%d state_dim=%d action_dim=%d",
+        len(states),
+        states.shape[1],
+        actions.shape[1],
+    )
+    return states, actions
+
+
+def _collect_relative_cartesian(
+    actions: np.ndarray,
+    states: np.ndarray,
+    index_range: list[int],
+    action_chunk: int,
+) -> np.ndarray:
+    start, end = index_range
+    max_start = max(0, len(actions) - action_chunk)
+    chunks = []
+    anchor_states = states[:max_start, start:end]
+    for offset in range(action_chunk):
+        chunks.append(actions[offset : offset + max_start, start:end] - anchor_states)
+    return np.concatenate(chunks, axis=0)
+
+
+def _compute_delta_from_state_and_abs_rot(
+    rotations: np.ndarray, state: np.ndarray
+) -> np.ndarray:
+    """Relative rotation: R_rel = R_abs @ R_state^T (same convention as the loader)."""
+    _ensure_geometry_funcs()
+
+    if rotations.shape[-1] == 3:
+        rotations_matrix = euler_to_matrix_zyx_batch_nb(rotations)
+        out_is_euler = True
+    elif rotations.shape[-1] == 6:
+        rotations_matrix = so3_to_matrix_batch_nb(rotations)
+        out_is_euler = False
+    else:
+        raise ValueError(
+            f"Only 3D euler or 6D rotation supported, got {rotations.shape[-1]}D"
+        )
+
+    if state.shape[-1] == 3:
+        state_matrix = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0]
+    elif state.shape[-1] == 6:
+        state_matrix = so3_to_matrix_batch_nb(state[np.newaxis, :])[0]
+    else:
+        raise ValueError(
+            f"Only 3D euler or 6D rotation supported, got {state.shape[-1]}D"
+        )
+
+    return _abs_rot_to_delta(rotations_matrix, state_matrix, out_is_euler)
+
+
+@jit(nopython=True, parallel=True)
+def _abs_rot_to_delta(
+    rotations_matrix: np.ndarray,
+    state_matrix: np.ndarray,
+    out_is_euler: bool,
+) -> np.ndarray:
+    st = np.empty((3, 3), dtype=np.float64)
+    st[0, 0] = state_matrix[0, 0]
+    st[0, 1] = state_matrix[1, 0]
+    st[0, 2] = state_matrix[2, 0]
+    st[1, 0] = state_matrix[0, 1]
+    st[1, 1] = state_matrix[1, 1]
+    st[1, 2] = state_matrix[2, 1]
+    st[2, 0] = state_matrix[0, 2]
+    st[2, 1] = state_matrix[1, 2]
+    st[2, 2] = state_matrix[2, 2]
+
+    n = rotations_matrix.shape[0]
+    r_rel = np.empty((n, 3, 3), dtype=np.float64)
+    for i in prange(n):
+        a00 = rotations_matrix[i, 0, 0]
+        a01 = rotations_matrix[i, 0, 1]
+        a02 = rotations_matrix[i, 0, 2]
+        a10 = rotations_matrix[i, 1, 0]
+        a11 = rotations_matrix[i, 1, 1]
+        a12 = rotations_matrix[i, 1, 2]
+        a20 = rotations_matrix[i, 2, 0]
+        a21 = rotations_matrix[i, 2, 1]
+        a22 = rotations_matrix[i, 2, 2]
+
+        r_rel[i, 0, 0] = a00 * st[0, 0] + a01 * st[1, 0] + a02 * st[2, 0]
+        r_rel[i, 0, 1] = a00 * st[0, 1] + a01 * st[1, 1] + a02 * st[2, 1]
+        r_rel[i, 0, 2] = a00 * st[0, 2] + a01 * st[1, 2] + a02 * st[2, 2]
+        r_rel[i, 1, 0] = a10 * st[0, 0] + a11 * st[1, 0] + a12 * st[2, 0]
+        r_rel[i, 1, 1] = a10 * st[0, 1] + a11 * st[1, 1] + a12 * st[2, 1]
+        r_rel[i, 1, 2] = a10 * st[0, 2] + a11 * st[1, 2] + a12 * st[2, 2]
+        r_rel[i, 2, 0] = a20 * st[0, 0] + a21 * st[1, 0] + a22 * st[2, 0]
+        r_rel[i, 2, 1] = a20 * st[0, 1] + a21 * st[1, 1] + a22 * st[2, 1]
+        r_rel[i, 2, 2] = a20 * st[0, 2] + a21 * st[1, 2] + a22 * st[2, 2]
+
+    if out_is_euler:
+        d_euler = matrix_to_euler_zyx_batch_nb(r_rel)
+        return canonicalize_euler_zyx_batch_nb(d_euler)
+
+    out6 = np.empty((n, 6), dtype=np.float64)
+    for i in prange(n):
+        out6[i, 0] = r_rel[i, 0, 0]
+        out6[i, 1] = r_rel[i, 0, 1]
+        out6[i, 2] = r_rel[i, 0, 2]
+        out6[i, 3] = r_rel[i, 1, 0]
+        out6[i, 4] = r_rel[i, 1, 1]
+        out6[i, 5] = r_rel[i, 1, 2]
+    return out6
+
+
+def _collect_relative_rotation(
+    actions: np.ndarray,
+    states: np.ndarray,
+    index_range: list[int],
+    action_chunk: int,
+) -> np.ndarray:
+    """
+    Per-anchor action chunk relative to anchor state (matches lerobot loader).
+
+    Unlike cartesian relative, each anchor processes a full [chunk, dim] action
+    clip against a single proprio rotation at the anchor frame.
+    """
+    start, end = index_range
+    max_start = max(0, len(actions) - action_chunk)
+    if max_start == 0:
+        return np.empty((0, end - start), dtype=np.float32)
+
+    chunks = []
+    for anchor_idx in tqdm(
+        range(max_start),
+        desc="  relative rotation anchors",
+        leave=False,
+    ):
+        action_clip = actions[anchor_idx : anchor_idx + action_chunk, start:end]
+        proprio_clip = states[anchor_idx, start:end]
+        rel = _compute_delta_from_state_and_abs_rot(
+            action_clip.astype(np.float64), proprio_clip.astype(np.float64)
+        ).astype(np.float32)
+        chunks.append(rel)
+    return np.concatenate(chunks, axis=0)
+
+
+def collect_dof_vectors_from_arrays(
+    states: np.ndarray,
+    actions: np.ndarray,
+    propri_ranges: dict[str, list[int]],
+    action_ranges: dict[str, list[int]],
+    action_chunk: int = 32,
+) -> dict[str, np.ndarray]:
+    vectors: dict[str, np.ndarray] = {}
+
+    absolute_action_keys = {
+        key: index_range
+        for key, index_range in action_ranges.items()
+        if not key.endswith("_relative")
+    }
+    relative_action_keys = {
+        key: index_range
+        for key, index_range in action_ranges.items()
+        if key.endswith("_relative")
+    }
+
+    for sub_key, index_range in propri_ranges.items():
+        start, end = index_range
+        vectors[sub_key] = states[:, start:end]
+
+    for sub_key, index_range in absolute_action_keys.items():
+        start, end = index_range
+        vectors[sub_key] = actions[:, start:end]
+
+    if relative_action_keys:
+        logging.info(
+            "Computing relative action slices (chunk=%d, anchors=%d)...",
+            action_chunk,
+            max(0, len(actions) - action_chunk),
+        )
+        for sub_key, index_range in tqdm(
+            relative_action_keys.items(), desc="Relative action keys"
+        ):
+            if "rotation" in sub_key:
+                vectors[sub_key] = _collect_relative_rotation(
+                    actions, states, index_range, action_chunk
+                )
+            else:
+                vectors[sub_key] = _collect_relative_cartesian(
+                    actions, states, index_range, action_chunk
                 )
-                continue
 
-    return stats
+    return vectors
 
 
-def load_lerobot_dataset(
-    repo_id: str,
-    trajectory_keys: Dict,
-    base_dir: Path,
-) -> None:
+def compute_norm_stats_with_dof_config(
+    data_root: Path,
+    output_path: Path,
+    propri_ranges: dict[str, list[int]],
+    action_ranges: dict[str, list[int]],
+    state_key: str = "observation.state",
+    action_key: str = "action",
+    action_chunk: int = 32,
+    dof_config: dict[str, int] | None = None,
+    agent_pos_config: dict[str, int] | None = None,
+) -> dict[str, dict]:
+    table, state_key, action_key = _load_state_action_table(
+        data_root, state_key, action_key
+    )
+    states, actions = _table_to_arrays(table, state_key, action_key)
+    if dof_config is not None and agent_pos_config is not None:
+        states, actions = _prepare_arrays_for_layout(
+            states, actions, agent_pos_config, dof_config
+        )
+    norm_stats = {
+        state_key: compute_vector_stats(states),
+        action_key: compute_vector_stats(actions),
+    }
 
-    # Load local or remote dataset
-    dataset = LeRobotDataset(base_dir)
+    vectors = collect_dof_vectors_from_arrays(
+        states=states,
+        actions=actions,
+        propri_ranges=propri_ranges,
+        action_ranges=action_ranges,
+        action_chunk=action_chunk,
+    )
 
-    # Iterate through all data
-    frames: Dict[str, Dict[str, List]] = defaultdict(lambda: defaultdict(list))
+    for sub_key, index_range in propri_ranges.items():
+        if sub_key not in vectors:
+            logging.warning("No samples collected for propri key %s, skipping", sub_key)
+            continue
+        slice_stats = compute_vector_stats(vectors[sub_key])
+        _apply_slice_stats(norm_stats[state_key], index_range, slice_stats)
+        logging.info("  %s (agent_pos): dim=%d", sub_key, len(slice_stats["mean"]))
 
-    all_features = dataset.features
-    non_image_columns = [col for col in all_features if "image" not in col]
+    for sub_key, index_range in action_ranges.items():
+        if sub_key not in vectors:
+            logging.warning("No samples collected for action key %s, skipping", sub_key)
+            continue
+        slice_stats = compute_vector_stats(vectors[sub_key])
+        _apply_slice_stats(norm_stats[action_key], index_range, slice_stats)
+        mode = "relative" if sub_key.endswith("_relative") else "absolute"
+        logging.info("  %s (dof, %s): dim=%d", sub_key, mode, len(slice_stats["mean"]))
 
-    print(f"Reading the following fields:{non_image_columns}")
-    fast_dataset = dataset.hf_dataset.select_columns(non_image_columns)
-
-    for i in tqdm(range(len(fast_dataset))):
-        sample = fast_dataset[i]
-        action = sample["action"]  # torch.Tensor
-        propri = sample["observation.state"]
-
-        for key, action_keys in trajectory_keys.items():
-            for action_key, action_range in action_keys.items():
-                if key == "action":
-                    frames[repo_id][action_key].append(
-                        action[action_range[0] : action_range[1]].numpy().tolist()
-                    )
-                else:
-                    frames[repo_id][action_key].append(
-                        propri[action_range[0] : action_range[1]].numpy().tolist()
-                    )
-
-    return frames
+    write_norm_stats(output_path, norm_stats)
+    return norm_stats
 
 
-def compute_action_normalizer(
-    repo_id: str, trajectory_keys: Dict, base_dir: Path, output_dir: Path
-) -> None:
-    """
-    Compute action normalizer statistics for all robot_ids.
-    """
-    logging.info("Starting action normalizer computation...")
+def load_vectors(
+    data_root: Path,
+    state_key: str,
+    action_key: str,
+) -> tuple[np.ndarray, np.ndarray]:
+    table, state_key, action_key = _load_state_action_table(
+        data_root, state_key, action_key
+    )
+    return _table_to_arrays(table, state_key, action_key)
 
-    frames = load_lerobot_dataset(repo_id, trajectory_keys, base_dir)
 
-    # Compute statistics
-    stats = compute_action_statistics(frames)
+def compute_norm_stats(
+    data_root: Path,
+    output_path: Path,
+    state_key: str = "observation.state",
+    action_key: str = "action",
+    train_ctx: TrainConfigContext | None = None,
+) -> dict[str, dict]:
+    if train_ctx is None:
+        states, actions = load_vectors(data_root, state_key, action_key)
+        norm_stats = {
+            state_key: compute_vector_stats(states),
+            action_key: compute_vector_stats(actions),
+        }
+        write_norm_stats(output_path, norm_stats)
+        return norm_stats
 
-    # Save statistics for each robot_id
-    output_dir = Path(output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
+    return compute_norm_stats_with_dof_config(
+        data_root=data_root,
+        output_path=output_path,
+        propri_ranges=train_ctx.propri_ranges,
+        action_ranges=train_ctx.action_ranges,
+        state_key=state_key,
+        action_key=action_key,
+        action_chunk=train_ctx.action_chunk,
+        dof_config=train_ctx.dof_config,
+        agent_pos_config=train_ctx.agent_pos_config,
+    )
 
-    # for robot_id, robot_stats in stats.items():
-    #     output_file = output_dir / f"{robot_id}_action_stats.json"
-    #     write_json(output_file, robot_stats)
-    #     logging.info(f"Saved action statistics for {robot_id} to {output_file}")
 
-    # Also save a combined file
-    combined_output = output_dir / "all_robots_action_stats.json"
-    write_json(combined_output, stats)
-    logging.info(f"Saved combined action statistics to {combined_output}")
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Compute norm stats for a local LeRobot v3 dataset.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""\
+examples:
+  %(prog)s --train_config /path/to/train_config.yml
+  %(prog)s --train_config /path/to/multitask_config.yml \\
+      --data_root /path/to/repo_id --output_path /path/to/norm_stats_path
+  %(prog)s --data_root /path/to/lerobot_dataset --output_path /path/to/norm_stats.json
+""",
+    )
+    parser.add_argument(
+        "--train_config",
+        type=str,
+        default=None,
+        help=(
+            "Training YAML config (e.g. cvpr_example.yml). When set, reads "
+            "data.lerobot_config.repo_id, data.norm_stats_path, task.dof_config, "
+            "task.agent_pos_config and task.action_horizon. Action keys ending "
+            "with '_relative' use the same relative-pose logic as lerobot loader."
+        ),
+    )
+    parser.add_argument(
+        "--data_root",
+        type=str,
+        default=None,
+        help="Local LeRobot dataset directory (overrides train config)",
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        default=None,
+        help="Output json path (overrides train config)",
+    )
+    parser.add_argument(
+        "--state_key",
+        type=str,
+        default=None,
+        help="Dataset column for proprioception (overrides train config)",
+    )
+    parser.add_argument(
+        "--action_key",
+        type=str,
+        default=None,
+        help="Dataset column for action (overrides train config)",
+    )
+    return parser.parse_args()
 
 
 def main() -> None:
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
+    args = parse_args()
 
-    repo_id = "xxx"  # your dataset name
-    data_root_path = "/path/to/lerobot/dataset"
-    output_stats_dir = "/path/to/save/action_stats"
-    trajectory_keys = {  # your dataset keys
-        "propri": {
-            "follow_right_ee_cartesian_pos": [0, 3],
-            "follow_right_ee_rotation": [3, 6],
-            "follow_right_gripper": [6, 7],
-        },
-        "action": {
-            "master_right_ee_cartesian_pos": [0, 3],
-            "master_right_ee_rotation": [3, 6],
-            "master_right_gripper": [6, 7],
-        },
-    }
+    train_ctx: TrainConfigContext | None = None
+    if args.train_config:
+        config_path = Path(args.train_config)
+        if not config_path.exists():
+            raise FileNotFoundError(f"train config not found: {config_path}")
+        train_ctx = parse_train_config(load_train_config(config_path))
+        logging.info("train_config: %s", config_path)
+        logging.info("  dof_config keys: %s", list(train_ctx.action_ranges))
+        logging.info("  agent_pos_config keys: %s", list(train_ctx.propri_ranges))
+        logging.info("  action_chunk: %d", train_ctx.action_chunk)
 
-    compute_action_normalizer(
-        repo_id, trajectory_keys, data_root_path, output_stats_dir
+    data_root = Path(args.data_root or (train_ctx.data_root if train_ctx else ""))
+
+    output_path = Path(args.output_path or (train_ctx.output_path if train_ctx else ""))
+    state_key = args.state_key or (
+        train_ctx.state_key if train_ctx else "observation.state"
     )
-    logging.info("Action normalizer computation completed.")
+    action_key = args.action_key or (train_ctx.action_key if train_ctx else "action")
+
+    if not data_root.exists():
+        raise FileNotFoundError(f"Dataset not found: {data_root}")
+
+    logging.info("dataset: %s", data_root)
+    logging.info("output:  %s", output_path)
+
+    norm_stats = compute_norm_stats(
+        data_root=data_root,
+        output_path=output_path,
+        state_key=state_key,
+        action_key=action_key,
+        train_ctx=train_ctx,
+    )
+
+    for key, stats in norm_stats.items():
+        logging.info("  %s: dim=%d", key, len(stats["mean"]))
+
+    logging.info("Saved norm stats to %s", output_path)
 
 
 if __name__ == "__main__":