| """Audit every episode of a T-Rex-format LeRobot v2.1 dataset. |
| |
| Usage: python scripts/data/audit_episode_quality.py --dataset-root data/trex_full |
| Contact = fingertip |(fx,fy,fz)| minus the per-finger whole-episode median baseline (if >0.5 N) exceeds 0.5 N, |
| the same definition as data/sharpa/tools/compare_difficulty.py. Writes <root>/meta/episode_quality.json. |
| """ |
| import json, glob, os, sys, numpy as np, pyarrow.parquet as pq, time |
| import argparse |
| ap=argparse.ArgumentParser(description="Per-episode quality audit for T-Rex-format LeRobot datasets (duration, bimanual motion, contact ratio, force baseline).") |
| ap.add_argument("--dataset-root",required=True); ap.add_argument("--out",default=None,help="json output (default <root>/meta/episode_quality.json)") |
| args=ap.parse_args(); root=args.dataset_root |
| eps={e["episode_index"]:e for e in map(json.loads, open(f"{root}/meta/episodes.jsonl"))} |
| F_THR=0.5 |
| rows=[]; t0=time.time() |
| files=sorted(glob.glob(f"{root}/data/chunk-*/episode_*.parquet")) |
| for i,f in enumerate(files): |
| t=pq.read_table(f, columns=["episode_index","timestamp","observation.state","observation.tactile_force"]) |
| ei=int(t["episode_index"][0].as_py()) |
| ts=np.array(t["timestamp"].to_pylist(),dtype=np.float64) |
| st=np.stack(t["observation.state"].to_numpy(zero_copy_only=False)).astype(np.float32) |
| tac=np.stack(t["observation.tactile_force"].to_numpy(zero_copy_only=False)).astype(np.float32).reshape(-1,10,6) |
| fn=np.linalg.norm(tac[:,:,:3],axis=-1) |
| med=np.median(fn,0); base=np.where(med>0.5, med, 0.0) |
| fn=fn-base |
| contact=fn>F_THR |
| l=contact[:,:5].any(1); r=contact[:,5:].any(1) |
| T=len(ts); dur=float(ts[-1]-ts[0]) if T>1 else 0.0 |
| def path(a): return float(np.abs(np.diff(a,axis=0)).sum()) |
| def events(m): |
| on=0; gap=99; prev=False |
| for v in m: |
| if v and not prev and gap>9: on+=1 |
| gap=0 if v else gap+1; prev=v |
| return on |
| dq=np.abs(np.diff(st,axis=0)); idle=float((dq.max(1)<1e-3).mean()) if T>1 else 1.0 |
| rows.append(dict(ep=ei, task=eps[ei]["tasks"][0], frames=T, dur=dur, |
| contact=float((l|r).mean()), both=float((l&r).mean()), left=float(l.mean()), right=float(r.mean()), |
| larm=path(st[:,0:7]), rarm=path(st[:,29:36]), lhand=path(st[:,7:29]), rhand=path(st[:,36:58]), |
| fp99=float(np.percentile(fn,99)), fmax=float(fn.max()), fingers=int((contact.mean(0)>0.05).sum()), |
| events=events(l)+events(r), idle=idle, baseline_max=float(base.max()), |
| ts_mono=bool(np.all(np.diff(ts)>0)), dt_med=float(np.median(np.diff(ts))) if T>1 else 0)) |
| if i%500==0: print(f"{i}/{len(files)} {time.time()-t0:.0f}s", flush=True) |
| out=args.out or f"{root}/meta/episode_quality.json"; json.dump(rows, open(out,"w")); print("wrote",out) |
| d=np.array([r["dur"] for r in rows]); print("durations: min %.1f median %.1f p90 %.1f max %.1f; n>=30s: %d; 30-120s: %d"%(d.min(),np.median(d),np.percentile(d,90),d.max(),(d>=30).sum(),((d>=30)&(d<=120)).sum())) |
| cand=[r for r in rows if 30<=r["dur"]<=120 and r["larm"]>1.0 and r["rarm"]>1.0 and r["left"]>0.1 and r["right"]>0.1] |
| print("candidates (30-120 s, both arms move >1 rad, each hand in contact >10%):", len(cand)) |
| cand.sort(key=lambda r: -(r["contact"]*0.6+r["both"]*0.4)) |
| for r in cand[:15]: |
| print(f" ep {r['ep']:5d} {r['dur']:6.1f}s contact {r['contact']:.2f} both {r['both']:.2f} L {r['left']:.2f} R {r['right']:.2f} larm {r['larm']:.1f} rarm {r['rarm']:.1f} fingers {r['fingers']} fp99 {r['fp99']:.1f}N | {r['task'][:80]}") |
| import statistics |
| print("dataset medians: contact %.2f both %.2f events %d baseline_max %.1fN"%(statistics.median(r['contact'] for r in rows),statistics.median(r['both'] for r in rows),statistics.median(r['events'] for r in rows),statistics.median(r['baseline_max'] for r in rows))) |
| print("episodes with baseline offset >0.5N on some finger:", sum(r['baseline_max']>0.5 for r in rows)) |
| |
| cand2=[r for r in cand if 40<=r['dur']<=90 and r['fingers']>=6 and r['events']>=6] |
| cand2.sort(key=lambda r: -(r['contact']+r['both']+0.02*r['fingers'])) |
| print("shortlist (40-90 s, >=6 fingers, >=6 events):", len(cand2)) |
| for r in cand2[:12]: |
| print(f" ep {r['ep']:5d} {r['dur']:6.1f}s contact {r['contact']:.2f} both {r['both']:.2f} L {r['left']:.2f} R {r['right']:.2f} ev {r['events']:3d} idle {r['idle']:.2f} fingers {r['fingers']} fp99 {r['fp99']:.1f}N base {r['baseline_max']:.1f} | {r['task']}") |
| print("DONE") |
|
|
