VLA/scripts/compute_norm_stats.py

#!/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 dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import yaml
from numba import jit, prange
from tqdm import tqdm

SKIP_DOF_KEYS = frozenset(
    {"velocity_decomposed", "height", "head_actions", "action_padding"}
)

_GEOMETRY_FUNCS_LOADED = False


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({"norm_stats": norm_stats}, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )


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]:
    """
    Match LeRobot training loader: convert 3D Euler slices to 6D when config
    uses rotation_6D keys but the dataset stores 14-dim Euler vectors.
    """
    from wall_x.data.backends.lerobot.rotation_layout import (
        euler_layout_dim,
        layout_uses_6d_rotation,
        maybe_convert_euler_to_6d,
    )

    convert_state = layout_uses_6d_rotation(agent_pos_config)
    convert_action = layout_uses_6d_rotation(dof_config)

    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"
            )

    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"
            )

    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


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


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
                )

    return vectors


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),
    }

    vectors = collect_dof_vectors_from_arrays(
        states=states,
        actions=actions,
        propri_ranges=propri_ranges,
        action_ranges=action_ranges,
        action_chunk=action_chunk,
    )

    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"]))

    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"]))

    write_norm_stats(output_path, norm_stats)
    return norm_stats


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)


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

    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,
    )


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()

    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)

    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"
    )
    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__":
    main()