File size: 3,952 Bytes
5a2e445 | 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 | #!/usr/bin/env python
"""Gate 1: per-dataset SO101 FK calibration audit.
Same recorded degrees can map to different physical poses across lerobot
zero-conventions → FK garbage that would poison B/C canonical labels while
leaving A untouched (a confound AGAINST the hypothesis). Drop datasets whose
FK produces non-physical EE trajectories.
Checks per dataset (sampled frames across episodes):
- reach in [0.02, 0.40] m (SO101 max reach ~0.35)
- z above a floor (> -0.20 m; base frame origin at arm mount)
- trajectory smoothness (median consecutive EE step < 0.05 m at 30fps)
- FK(state) vs FK(action-target) diffs correlate (both go through FK, deltas
should track since action = commanded target of the same arm)
Writes ~/tinyvla_data/so101_fk_audit.json with per-dataset verdict.
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
DATA_ROOT = Path.home() / "tinyvla_data" / "so101_v3"
OUT = Path.home() / "tinyvla_data" / "so101_fk_audit.json"
def audit_one(name, root, fk, n_frames=200):
from lerobot.datasets.lerobot_dataset import LeRobotDataset
ds = LeRobotDataset(name, root=root)
anames = ds.meta.features["action"]["shape"][0]
if anames != 6:
return {"verdict": "SKIP", "reason": f"action dim {anames} != 6"}
# smoothness needs CONSECUTIVE frames → sample a few contiguous windows;
# reach/corr can use sparse knots across the whole dataset
n = len(ds)
eps = ds.meta.episodes
sparse = np.linspace(0, n - 1, min(n_frames, n)).astype(int)
pos_state, pos_act = [], []
for i in sparse:
item = ds[int(i)]
pos_state.append(fk.ee_pose(item["observation.state"].numpy())[:3, 3])
pos_act.append(fk.ee_pose(item["action"].numpy())[:3, 3])
ps = np.array(pos_state)
pa = np.array(pos_act)
reach = np.linalg.norm(ps, axis=1)
# consecutive-frame EE steps within the first episode (real per-frame motion)
e0, e1 = int(eps["dataset_from_index"][0]), int(eps["dataset_to_index"][0])
consec = []
prev = None
for i in range(e0, min(e1, e0 + 150)):
p = fk.ee_pose(ds[i]["observation.state"].numpy())[:3, 3]
if prev is not None:
consec.append(np.linalg.norm(p - prev))
prev = p
steps = np.array(consec) if consec else np.array([0.0])
# correlation of state-vs-target displacement over sampled knots
d_state = np.diff(ps, axis=0).flatten()
d_act = np.diff(pa, axis=0).flatten()
corr = float(np.corrcoef(d_state, d_act)[0, 1]) if d_state.std() > 1e-9 else 0.0
ok_reach = bool(0.02 < reach.mean() < 0.40 and reach.max() < 0.50)
ok_z = bool(ps[:, 2].min() > -0.20)
ok_smooth = bool(np.median(steps) < 0.06)
ok_corr = bool(corr > 0.5)
verdict = "KEEP" if (ok_reach and ok_z and ok_smooth and ok_corr) else "DROP"
return {
"verdict": verdict,
"reach_mean": round(float(reach.mean()), 3),
"reach_max": round(float(reach.max()), 3),
"z_min": round(float(ps[:, 2].min()), 3),
"step_median": round(float(np.median(steps)), 4),
"corr_state_target": round(corr, 3),
"flags": {"reach": ok_reach, "z": ok_z, "smooth": ok_smooth, "corr": ok_corr},
}
def main():
from tinyvla.data.kinematics_so101 import SO101FK
fk = SO101FK()
results = {}
roots = sorted(DATA_ROOT.iterdir())
for r in roots:
if not (r / "meta" / "info.json").exists():
continue
try:
res = audit_one(r.name, r, fk)
except Exception as e:
res = {"verdict": "ERROR", "reason": f"{type(e).__name__}: {str(e)[:120]}"}
results[r.name] = res
print(f"{res['verdict']:6} {r.name}: {res}")
OUT.write_text(json.dumps(results, indent=1))
keep = sum(1 for v in results.values() if v["verdict"] == "KEEP")
print(f"\nKEEP {keep}/{len(results)} -> {OUT}")
if __name__ == "__main__":
main()
|