| |
| """Validate a converted xtrainer2 LeRobot dataset. |
| |
| Two independent checks: |
| |
| 1. Decoded rotations must reproduce the official ``norm_stats.json`` in the |
| 20-D model space, which is what proves the quaternion convention is right. |
| 2. StarVLA's dataloader must load a sample end to end. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import glob |
| import importlib.util |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
|
|
| from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotSingleDataset |
| from starVLA.dataloader.gr00t_lerobot.transform.state_action import RotationTransform |
|
|
| REPO_ROOT = Path(__file__).resolve().parents[1] |
| OFFICIAL_STATS = "/apdcephfs_nj7/share_305204761/arianliu/data/benchmark/norm_stats.json" |
| ARM_BLOCKS = ((0, 3, 7), (8, 11, 15)) |
|
|
|
|
| def load_data_config(): |
| """The examples tree is not a package, so load the registry by path.""" |
| path = REPO_ROOT / "examples/Xtrainer2/train_files/data_registry/data_config.py" |
| spec = importlib.util.spec_from_file_location("_xtrainer2_data_config", path) |
| assert spec is not None and spec.loader is not None |
| module = importlib.util.module_from_spec(spec) |
| sys.modules[spec.name] = module |
| spec.loader.exec_module(module) |
| return module.Xtrainer2BimanualEefDataConfig() |
|
|
|
|
| def to_model_space(raw: np.ndarray, scramble_back_to_xyzw: bool = False) -> np.ndarray: |
| """Expand stored 16-D poses into the 20-D layout used by the official stats.""" |
| transform = RotationTransform("quaternion", "rotation_6d") |
| columns = [] |
| for position, rotation, gripper in ARM_BLOCKS: |
| columns.append(raw[:, position : position + 3]) |
| block = raw[:, rotation : rotation + 4] |
| if scramble_back_to_xyzw: |
| block = block[:, [1, 2, 3, 0]] |
| six = transform.forward(torch.from_numpy(np.ascontiguousarray(block))) |
| columns.append(six.numpy()) |
| columns.append(raw[:, gripper : gripper + 1]) |
| return np.concatenate(columns, axis=1) |
|
|
|
|
| def _stats_error(model_space: np.ndarray, official: dict) -> tuple[float, np.ndarray]: |
| ref_mean = np.array(official["state"]["mean"])[0] |
| ref_lo = np.array(official["state"]["p_low"])[0] |
| ref_hi = np.array(official["state"]["p_high"])[0] |
| mean = model_space.mean(0) |
| lo = np.quantile(model_space, 0.01, axis=0) |
| hi = np.quantile(model_space, 0.99, axis=0) |
| error = float( |
| (np.abs(mean - ref_mean).mean() + np.abs(lo - ref_lo).mean() + np.abs(hi - ref_hi).mean()) |
| / 3 |
| ) |
| return error, mean |
|
|
|
|
| def read_states(dataset_root: Path) -> np.ndarray: |
| frames = [ |
| pd.read_parquet(path)["observation.state"] |
| for path in sorted(glob.glob(str(dataset_root / "data" / "chunk-*" / "*.parquet"))) |
| ] |
| return np.stack([np.asarray(row, dtype=np.float32) for f in frames for row in f]) |
|
|
|
|
| def check_convention(dataset_roots: list[Path]) -> tuple[float, float]: |
| """Compare the stored convention against the wrong one, pooled over all tasks. |
| |
| The official stats cover the whole benchmark, so any single task is biased |
| against them; only the pooled corpus is comparable, and only the gap between |
| the two conventions is meaningful. |
| """ |
| official = json.load(open(OFFICIAL_STATS)) |
| ref_mean = np.array(official["state"]["mean"])[0] |
| raw = np.concatenate([read_states(root) for root in dataset_roots], axis=0) |
|
|
| stored_error, mean = _stats_error(to_model_space(raw), official) |
| wrong_error, _ = _stats_error(to_model_space(raw, scramble_back_to_xyzw=True), official) |
|
|
| print(f" pooled frames={len(raw)} over {len(dataset_roots)} datasets") |
| print(f" ours mean {np.round(mean, 3)}") |
| print(f" ref mean {np.round(ref_mean, 3)}") |
| print(f" err as stored (wxyz) : {stored_error:.4f}") |
| print(f" err if misread (xyzw) : {wrong_error:.4f}") |
| return stored_error, wrong_error |
|
|
|
|
| def check_dataloader(dataset_root: Path) -> None: |
| config = load_data_config() |
| dataset = LeRobotSingleDataset( |
| dataset_path=dataset_root, |
| modality_configs=config.modality_config(), |
| embodiment_tag=config.embodiment_tag, |
| transforms=config.transform(), |
| data_cfg={"lerobot_version": "v2.0", "action_mode": "abs", "include_state": True}, |
| ) |
| sample = dataset[0] |
| print(f" episodes={len(dataset.trajectory_lengths)} samples={len(dataset)}") |
| for key, value in sample.items(): |
| if isinstance(value, (np.ndarray, torch.Tensor)): |
| print(f" {key}: {tuple(value.shape)} {value.dtype}") |
| else: |
| print(f" {key}: {value!r}") |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument( |
| "--root", |
| type=Path, |
| default=Path("/apdcephfs_nj7/share_305204761/arianliu/data/lerobot"), |
| ) |
| parser.add_argument( |
| "--min-margin", |
| type=float, |
| default=3.0, |
| help="How many times better the stored convention must fit the official stats", |
| ) |
| args = parser.parse_args() |
|
|
| dataset_roots = sorted(p for p in args.root.iterdir() if p.is_dir()) |
| failures = [] |
|
|
| print("=== rotation convention vs official norm_stats ===") |
| stored_error, wrong_error = check_convention(dataset_roots) |
| margin = wrong_error / max(stored_error, 1e-9) |
| print(f" margin: {margin:.1f}x") |
| if margin < args.min_margin: |
| failures.append(f"convention margin {margin:.1f}x < {args.min_margin}x") |
|
|
| for dataset_root in dataset_roots: |
| print(f"\n=== {dataset_root.name}: StarVLA dataloader ===") |
| check_dataloader(dataset_root) |
|
|
| if failures: |
| print("\nFAILED:") |
| for failure in failures: |
| print(f" {failure}") |
| return 1 |
| print("\nAll checks passed.") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
