File size: 5,871 Bytes
8101696 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | #!/usr/bin/env python3
"""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())
|