754 lines
25 KiB
Python
754 lines
25 KiB
Python
#!/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()
|