| """Build Psi0 dataset views of the LeRobot v2.1 recordings (raw data is never modified). |
| |
| Usage: <psi0 python> prepare_data.py [--cams <cam0>,<cam1>] [--force] <task> [<task> ...] (b6 basket a3 a7 a18) |
| |
| View = $G/psi0-data/<task>/ (LeRobot v2.1 layout read by psi/data/lerobot/lerobot_ext.py): |
| data/chunk-*/episode_*.parquet NEW small parquet: timestamp, frame_index, episode_index, index, task_index, |
| <state_key> (float32[S]), <action_key> (float32[A]) |
| videos/... symlinks to the raw mp4s of the selected cameras only |
| (LeRobot decodes every video key listed in info.json) |
| meta/info.json rewritten features / totals |
| meta/episodes.jsonl, tasks.jsonl copied (B6: first 100 episodes only) |
| meta/episodes_stats.jsonl per-episode stats of the two new columns (lerobot aggregate_stats input) |
| meta/stats_psi0.json min/max/mean/std/q01/q99 used by SonicActionStateTransform (bounds) |
| hb-view-receipt.json robot layout (keys, dims, body_dims, cameras) read by hb_args.py + diagnostics |
| |
| Per-robot layouts (same row t; labels are the values of row t): |
| g1 (b6 basket a3 a7) observation.hb44 = observation.state[0:29] (29 body joints, checked by name) |
| + inspire_left_hand_state 6 + inspire_right_hand_state 6 + projected_gravity 3 |
| action.hb76 = action.motion_token 64 + inspire_left 6 + inspire_right 6 (measured hands) |
| cameras ego_view (cam0) + chest; joint-perturbed dims 29 |
| h2 (a18) observation.hb78 = observation.state[0:75] (legs 12, waist 3, arms 14, head 2, hands 22+22) |
| + projected_gravity 3 |
| action.hb75 = action.body_target 31 (names == state[0:31]) + action.hands 44 (names == state[31:75]) |
| cameras head (cam0) + chest; joint-perturbed dims 31 (29 body + head 2) |
| Why derived columns: SonicRepackTransform.delta_timestamps keys windows by column name, so one column cannot be |
| both a 1-frame state and a 30-frame action chunk (G1 hands), and Psi0 takes one state / one action key here. |
| The CLIP pooled text cache is NOT precomputed here: FinetuneTrainer._warm_pooled_cache_from_tasks |
| builds it at start from meta/tasks.jsonl and saves <run>/clip_pooled_cache.pt (shipped with the run). |
| """ |
| import argparse |
| import hashlib |
| import json |
| import os |
| import shutil |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import pyarrow as pa |
| import pyarrow.parquet as pq |
|
|
| H = Path(__file__).resolve().parent |
| G = Path(os.environ.get("G", H.parent)) |
| P = G.parent |
| OUT = Path(os.environ.get("HB_DATA_ROOT", G / "psi0-data")) |
| TASKS = { |
| "b6": (P / "b6-v5t7-ready-start-60ep-20261109/data/deliveries/g1/b6-v5t7-ready-start-20261109", 100, "g1"), |
| "basket": (P / "basket-current-dynamics-100-v21-b0ff522/data", None, "g1"), |
| "trolley": (P / "trolley-push-expert-v1-20261110/data/trolley_push", 100, "g1"), |
| "b7": (P / "b7-black-clothes-hamper-v1-1bb002a/data", 100, "g1"), |
| "a3": (P / "a3-a7-20261109/a3-fridge-open-expert-v1/780fcf1e4bd7db947a8c2d8f7b92e61f24f71b45", None, "g1"), |
| "a7": (P / "a3-a7-20261109/a7-cabinet-close-expert-v1/1935a98e8f55d4d1d2d1615ddb97c2679713595c", None, "g1"), |
| "a18": (P / "a18-h2-dataset-audit-8e956a10/dataset", None, "h2"), |
| "a17": (P / "a17-unscrew-bottle-v1-407e6e2/data", None, "h2"), |
| } |
| BODY = ["left_hip_pitch_joint", "left_hip_roll_joint", "left_hip_yaw_joint", "left_knee_joint", "left_ankle_pitch_joint", |
| "left_ankle_roll_joint", "right_hip_pitch_joint", "right_hip_roll_joint", "right_hip_yaw_joint", |
| "right_knee_joint", "right_ankle_pitch_joint", "right_ankle_roll_joint", "waist_yaw_joint", "waist_roll_joint", |
| "waist_pitch_joint", "left_shoulder_pitch_joint", "left_shoulder_roll_joint", "left_shoulder_yaw_joint", |
| "left_elbow_joint", "left_wrist_roll_joint", "left_wrist_pitch_joint", "left_wrist_yaw_joint", |
| "right_shoulder_pitch_joint", "right_shoulder_roll_joint", "right_shoulder_yaw_joint", "right_elbow_joint", |
| "right_wrist_roll_joint", "right_wrist_pitch_joint", "right_wrist_yaw_joint"] |
| STATE, TOKEN, LH, RH, GRAV = ("observation.state", "action.motion_token", "observation.inspire_left_hand_state", |
| "observation.inspire_right_hand_state", "observation.projected_gravity") |
| BT, HANDS = "action.body_target", "action.hands" |
| |
| LAYOUTS = { |
| "g1": dict(state_key="observation.hb44", action_key="action.hb76", body_dims=29, |
| state_parts=[(STATE, 0, 29), (LH, 0, 6), (RH, 0, 6), (GRAV, 0, 3)], |
| action_parts=[(TOKEN, 0, 64), (LH, 0, 6), (RH, 0, 6)], |
| exact_dims={TOKEN: 64, LH: 6, RH: 6, GRAV: 3}, |
| cams={"ego": "observation.images.ego_view", "chest": "observation.images.chest"}, default_cams="ego,chest"), |
| "h2": dict(state_key="observation.hb78", action_key="action.hb75", body_dims=31, |
| state_parts=[(STATE, 0, 75), (GRAV, 0, 3)], |
| action_parts=[(BT, 0, 31), (HANDS, 0, 44)], |
| exact_dims={STATE: 75, GRAV: 3, BT: 31, HANDS: 44}, |
| cams={"head": "observation.images.head", "chest": "observation.images.chest", |
| "wrist": "observation.images.wrist", "third_person": "observation.images.third_person"}, |
| default_cams="head,chest"), |
| } |
| SCALARS = ["timestamp", "frame_index", "episode_index", "index", "task_index"] |
| TS_TOL = 1e-4 |
|
|
|
|
| def jl(path): |
| return [json.loads(x) for x in Path(path).read_text().splitlines() if x.strip()] |
|
|
|
|
| def wl(path, rows): |
| Path(path).write_text("".join(json.dumps(r) + "\n" for r in rows)) |
|
|
|
|
| def names_of(ft, dim, prefix): |
| n = ft.get("names") |
| if isinstance(n, dict): |
| n = next(iter(n.values()), None) |
| return list(n) if isinstance(n, list) and len(n) == dim else [f"{prefix}_{i}" for i in range(dim)] |
|
|
|
|
| def raw_names(ft): |
| n = ft.get("names") |
| if isinstance(n, dict): |
| n = next(iter(n.values()), None) |
| return list(n) if isinstance(n, list) else None |
|
|
|
|
| def col2d(t, key, dim): |
| c = t.column(key).combine_chunks() |
| v = np.asarray(c.flatten().to_numpy(zero_copy_only=False), dtype=np.float64) |
| if v.size != len(t) * dim: |
| raise ValueError(f"{key}: {v.size} values for {len(t)} rows x {dim}") |
| return v.reshape(len(t), dim) |
|
|
|
|
| def fsl(x): |
| x = np.ascontiguousarray(x, dtype=np.float32) |
| return pa.FixedSizeListArray.from_arrays(pa.array(x.reshape(-1)), x.shape[1]) |
|
|
|
|
| def ep_stats(x): |
| return {"min": x.min(0).tolist(), "max": x.max(0).tolist(), "mean": x.mean(0).tolist(), |
| "std": x.std(0).tolist(), "count": [int(len(x))]} |
|
|
|
|
| def full_stats(x, names): |
| return {"min": x.min(0).tolist(), "max": x.max(0).tolist(), "mean": x.mean(0).tolist(), "std": x.std(0).tolist(), |
| "q01": np.quantile(x, 0.01, axis=0).tolist(), "q99": np.quantile(x, 0.99, axis=0).tolist(), |
| "count": [int(len(x))], "names": names} |
|
|
|
|
| def video_frames(path): |
| try: |
| import av |
| with av.open(str(path)) as c: |
| s = c.streams.video[0] |
| return int(s.frames), float(s.average_rate or 0) |
| except Exception as e: |
| return None, str(e)[:200] |
|
|
|
|
| def part_names(feats, parts): |
| out = [] |
| for key, a, b in parts: |
| nm = names_of(feats[key], feats[key]["shape"][0], key.split(".")[-1]) |
| out += nm[a:b] if key == STATE else [f"{key.split('.')[-1]}.{x}" for x in nm[a:b]] |
| return out |
|
|
|
|
| def check_layout(task, robot, L, feats, snames): |
| for key, dim in L["exact_dims"].items(): |
| assert key in feats and list(feats[key]["shape"]) == [dim], (task, key, feats.get(key)) |
| for key, a, b in L["state_parts"] + L["action_parts"]: |
| assert key in feats and feats[key]["shape"][0] >= b, (task, key, feats.get(key)) |
| if robot == "g1": |
| body_idx = [snames.index(n) for n in BODY] |
| assert body_idx == list(range(29)), f"{task}: body joints are not observation.state[0:29]: {body_idx}" |
| else: |
| for key, a, b in ((BT, 0, 31), (HANDS, 31, 75)): |
| n = raw_names(feats[key]) |
| if n is None: |
| print(f"WARN {task}: {key} has no names; cannot check it mirrors observation.state[{a}:{b}]") |
| else: |
| assert n == snames[a:b], f"{task}: {key} names != observation.state[{a}:{b}] names" |
| print(f"INFO {task}: perturbed joints (state[0:31]) = {snames[:31]}") |
|
|
|
|
| def build(task, cams, force, keep_ts): |
| src, limit, robot = TASKS[task] |
| L = LAYOUTS[robot] |
| SK, AK, NB = L["state_key"], L["action_key"], L["body_dims"] |
| info = json.loads((src / "meta/info.json").read_text()) |
| feats = info["features"] |
| if info.get("codebase_version") != "v2.1": |
| print(f"WARN {task}: codebase_version={info.get('codebase_version')} (expected v2.1)") |
| fps = int(info["fps"]) |
| assert fps == 30, (task, fps) |
| snames = names_of(feats[STATE], feats[STATE]["shape"][0], "state") |
| check_layout(task, robot, L, feats, snames) |
| vkeys = [L["cams"][c] for c in cams] |
| for k in vkeys: |
| assert feats[k]["dtype"] == "video", (task, k) |
| dims = {k: feats[k]["shape"][0] for k, _, _ in L["state_parts"] + L["action_parts"]} |
|
|
| episodes = jl(src / "meta/episodes.jsonl") |
| total_src = len(episodes) |
| n = total_src if limit is None else limit |
| assert total_src >= n, (task, total_src, n) |
| episodes = episodes[:n] |
| assert [e["episode_index"] for e in episodes] == list(range(n)), f"{task}: episodes not contiguous 0..{n-1}" |
| tasks = jl(src / "meta/tasks.jsonl") |
|
|
| dst = OUT / task |
| tmp = OUT / f".{task}.partial" |
| if dst.exists(): |
| if not force: |
| raise SystemExit(f"{dst} exists (use --force to rebuild the view; raw data is untouched)") |
| shutil.rmtree(dst) |
| if tmp.exists(): |
| shutil.rmtree(tmp) |
| (tmp / "meta").mkdir(parents=True) |
| s_names, a_names = part_names(feats, L["state_parts"]), part_names(feats, L["action_parts"]) |
|
|
| S, A, ep_rows = [], [], [] |
| diag = {"ts_max_dev_s": 0.0, "ts_rewritten_episodes": [], "basket_hand_diff": None, |
| "h2_target_minus_state_absmax": None} |
| chunks = int(info["chunks_size"]) |
| for e in episodes: |
| i = e["episode_index"] |
| chunk = i // chunks |
| rel = info["data_path"].format(episode_chunk=chunk, episode_index=i) |
| t = pq.read_table(src / rel) |
| T = len(t) |
| assert T == int(e["length"]), (task, i, T, e["length"]) |
| c = {k: col2d(t, k, d) for k, d in dims.items()} |
| st = np.concatenate([c[k][:, a:b] for k, a, b in L["state_parts"]], 1).astype(np.float32) |
| ac = np.concatenate([c[k][:, a:b] for k, a, b in L["action_parts"]], 1).astype(np.float32) |
| assert np.isfinite(st).all() and np.isfinite(ac).all(), (task, i) |
| if task == "basket" and c[STATE].shape[1] >= 41: |
| d = float(np.abs(c[STATE][:, 29:41] - np.concatenate([c[LH], c[RH]], 1)).max()) |
| diag["basket_hand_diff"] = max(diag["basket_hand_diff"] or 0.0, d) |
| if robot == "h2": |
| d = np.abs(np.concatenate([c[BT], c[HANDS]], 1) - c[STATE][:, :75]).max(0) |
| prev = diag["h2_target_minus_state_absmax"] |
| diag["h2_target_minus_state_absmax"] = d if prev is None else np.maximum(prev, d) |
| cols = {k: t.column(k).combine_chunks() for k in SCALARS} |
| fi = np.asarray(cols["frame_index"].to_numpy(zero_copy_only=False)) |
| assert (fi == np.arange(T)).all(), (task, i, "frame_index") |
| assert (np.asarray(cols["episode_index"].to_numpy(zero_copy_only=False)) == i).all(), (task, i) |
| ts = np.asarray(cols["timestamp"].to_numpy(zero_copy_only=False), dtype=np.float64) |
| dev = float(np.abs(ts - fi / fps).max()) |
| diag["ts_max_dev_s"] = max(diag["ts_max_dev_s"], dev) |
| if dev > TS_TOL and not keep_ts: |
| cols["timestamp"] = pa.array((fi / fps).astype(np.float32)) |
| diag["ts_rewritten_episodes"].append(i) |
| out = pa.table({**cols, SK: fsl(st), AK: fsl(ac)}) |
| (tmp / rel).parent.mkdir(parents=True, exist_ok=True) |
| pq.write_table(out, tmp / rel) |
| for k in vkeys: |
| relv = info["video_path"].format(episode_chunk=chunk, video_key=k, episode_index=i) |
| (tmp / relv).parent.mkdir(parents=True, exist_ok=True) |
| os.symlink(src / relv, tmp / relv) |
| ep_rows.append({"episode_index": i, "stats": {SK: ep_stats(st), AK: ep_stats(ac)}}) |
| S.append(st) |
| A.append(ac) |
| S, A = np.concatenate(S), np.concatenate(A) |
| frames = int(len(S)) |
|
|
| new_feats = {k: feats[k] for k in SCALARS} |
| new_feats.update({k: feats[k] for k in vkeys}) |
| new_feats[SK] = {"dtype": "float32", "shape": [S.shape[1]], "names": s_names} |
| new_feats[AK] = {"dtype": "float32", "shape": [A.shape[1]], "names": a_names} |
| info2 = dict(info, features=new_feats, total_episodes=n, total_frames=frames, total_videos=n * len(vkeys), |
| total_chunks=(n - 1) // chunks + 1, splits={"train": f"0:{n}"}) |
| (tmp / "meta/info.json").write_text(json.dumps(info2, indent=2)) |
| wl(tmp / "meta/episodes.jsonl", episodes) |
| wl(tmp / "meta/episodes_stats.jsonl", ep_rows) |
| wl(tmp / "meta/tasks.jsonl", tasks) |
| stats = {AK: full_stats(A, a_names), SK: full_stats(S, s_names)} |
| (tmp / "meta/stats_psi0.json").write_text(json.dumps(stats, indent=1)) |
|
|
| vf = {k: [{"episode": j, "expected": int(episodes[j]["length"]), |
| "frames_rate": video_frames(src / info["video_path"].format( |
| episode_chunk=j // chunks, video_key=k, episode_index=j))} |
| for j in (0, n - 1)] for k in vkeys} |
| if diag["h2_target_minus_state_absmax"] is not None: |
| d = diag["h2_target_minus_state_absmax"] |
| diag["h2_target_minus_state_absmax"] = {"body31_max": float(d[:31].max()), "hands44_max": float(d[31:].max())} |
| rec = dict( |
| task=task, robot=robot, source=str(src), source_episodes=total_src, episodes=n, frames=frames, fps=fps, |
| image_keys=vkeys, state_key=SK, action_key=AK, state_dim=int(S.shape[1]), action_dim=int(A.shape[1]), |
| body_dims=NB, perturbed_joint_names=s_names[:NB], |
| state_layout=" + ".join(f"{k}[{a}:{b}]" for k, a, b in L["state_parts"]), |
| action_layout=" + ".join(f"{k}[{a}:{b}]" for k, a, b in L["action_parts"]), |
| raw_state_dim=len(snames), tasks=[x["task"] for x in tasks], |
| constant_dims={"action": [a_names[j] for j in np.where(A.max(0) - A.min(0) < 1e-6)[0]], |
| "state": [s_names[j] for j in np.where(S.max(0) - S.min(0) < 1e-6)[0]]}, |
| timestamps=diag, video_frames_first_last_episode=vf, |
| stats_sha256=hashlib.sha256((tmp / "meta/stats_psi0.json").read_bytes()).hexdigest(), |
| ) |
| if robot == "g1": |
| tok = A[:, :64] |
| grid = np.abs(tok * 16 - np.round(tok * 16)) < 1e-3 |
| rec["token_diag"] = {"min": float(tok.min()), "max": float(tok.max()), |
| "per_dim_min_range": [float(tok.min(0).min()), float(tok.min(0).max())], |
| "per_dim_max_range": [float(tok.max(0).min()), float(tok.max(0).max())], |
| "frac_on_1_16_grid": float(grid.mean()), "frac_abs_le_1": float((np.abs(tok) <= 1).mean())} |
| rec["hand_range"] = {"left": [float(A[:, 64:70].min()), float(A[:, 64:70].max())], |
| "right": [float(A[:, 70:76].min()), float(A[:, 70:76].max())]} |
| else: |
| rec["action_range"] = {"body_target31": [float(A[:, :31].min()), float(A[:, :31].max())], |
| "hands44": [float(A[:, 31:].min()), float(A[:, 31:].max())]} |
| (tmp / "hb-view-receipt.json").write_text(json.dumps(rec, indent=2) + "\n") |
| os.replace(tmp, dst) |
| keys = ("robot", "episodes", "frames", "state_dim", "action_dim", "image_keys", "constant_dims", |
| "token_diag", "hand_range", "action_range") |
| print("VIEW_OK", task, json.dumps({k: rec[k] for k in keys if k in rec}), flush=True) |
| if diag["ts_rewritten_episodes"]: |
| print(f"NOTE {task}: timestamps of {len(diag['ts_rewritten_episodes'])} episodes deviated up to " |
| f"{diag['ts_max_dev_s']:.2e}s from frame_index/fps and were rewritten in the view") |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("tasks", nargs="+") |
| ap.add_argument("--cams", default=os.environ.get("HB_CAMS"), |
| help="comma list, cam0 first (g1: ego[,chest]; h2: head[,chest,...]); default per robot") |
| ap.add_argument("--force", action="store_true") |
| ap.add_argument("--keep-timestamps", action="store_true") |
| a = ap.parse_args() |
| OUT.mkdir(parents=True, exist_ok=True) |
| for t in a.tasks: |
| if t not in TASKS: |
| raise SystemExit(f"unknown task {t}; known {list(TASKS)}") |
| L = LAYOUTS[TASKS[t][2]] |
| cams = [c for c in (a.cams or L["default_cams"]).split(",") if c] |
| cam0 = L["default_cams"].split(",")[0] |
| assert cams and cams[0] == cam0 and set(cams) <= set(L["cams"]), f"{t}: cameras {cams}, cam0 must be {cam0}" |
| build(t, cams, a.force, a.keep_timestamps) |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
