| """Convert raw YOR episodes (thor2 recorder format) to patch-policy training format. |
| |
| Raw (per episode): |
| robot_data.csv wide tick rows, incl. frame_idx (odin) + t1_ns clock |
| odin.mp4 head cam, one frame per unique image |
| wrist_left/right.mp4 + *_index.csv (own clock domain, same machine) |
| |
| Output (patch-policy style, like cube/pusht): |
| <out-dir>/ |
| seq_lengths.pkl list[int], subsampled lengths |
| actions.pth FloatTensor [N, Tmax, 58] zero-padded |
| obses/episode_00000.pth ... uint8 Tensor [T, V=3, H, W, C] |
| meta.json action column order, stride, img size, episode names |
| |
| Action (58D, all commanded where available): |
| [left_ee_target(7), right_ee_target(7), |
| left_hand_qpos(20, wuji-commanded), right_hand_qpos(20), |
| lift(1), base_pose_target(3)] |
| |
| Images: center-crop square -> resize 224 (DINO/14 native -> 256 patches). |
| Views: V=3 [odin, wrist_left, wrist_right], timestamp-matched per tick. |
| """ |
| import argparse |
| import csv |
| import json |
| import pickle |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
| import torch |
|
|
| EE_L = [f"left_ee_target_wxyz_xyz_{i}" for i in range(7)] |
| EE_R = [f"right_ee_target_wxyz_xyz_{i}" for i in range(7)] |
| HAND_L = [f"left_hand_qpos_{i}" for i in range(20)] |
| HAND_R = [f"right_hand_qpos_{i}" for i in range(20)] |
| LIFT = ["lift"] |
| BASE = ["base_pose_target_0", "base_pose_target_1", "base_pose_target_2"] |
| ACTION_COLS = EE_L + EE_R + HAND_L + HAND_R + LIFT + BASE |
| assert len(ACTION_COLS) == 58 |
|
|
|
|
| def read_csv_rows(path): |
| with open(path, newline="") as f: |
| r = csv.DictReader(f) |
| return r.fieldnames, list(r) |
|
|
|
|
| def decode_resize(video_path, size=224): |
| """Decode mp4 -> list of uint8 HxWxC resized frames (center-crop square first).""" |
| cap = cv2.VideoCapture(str(video_path)) |
| if not cap.isOpened(): |
| raise RuntimeError(f"cannot open {video_path}") |
| frames = [] |
| while True: |
| ok, img = cap.read() |
| if not ok: |
| break |
| h, w = img.shape[:2] |
| side = min(h, w) |
| x0 = (w - side) // 2 |
| y0 = (h - side) // 2 |
| crop = img[y0:y0 + side, x0:x0 + side] |
| small = cv2.resize(crop, (size, size), interpolation=cv2.INTER_AREA) |
| small = cv2.cvtColor(small, cv2.COLOR_BGR2RGB) |
| frames.append(small) |
| cap.release() |
| return np.stack(frames) |
|
|
|
|
| def nearest_idx(sorted_ts, queries): |
| """For each query, index of closest entry in sorted_ts (both int64 ns).""" |
| pos = np.searchsorted(sorted_ts, queries) |
| pos = np.clip(pos, 0, len(sorted_ts) - 1) |
| prev = np.clip(pos - 1, 0, len(sorted_ts) - 1) |
| pick_next = np.abs(sorted_ts[pos] - queries) < np.abs(sorted_ts[prev] - queries) |
| return np.where(pick_next, pos, prev) |
|
|
|
|
| def convert_episode(epdir, size=224, stride=3): |
| _, ticks = read_csv_rows(epdir / "robot_data.csv") |
| _, wleft_idx = read_csv_rows(epdir / "wrist_left_index.csv") |
| _, wright_idx = read_csv_rows(epdir / "wrist_right_index.csv") |
|
|
| print(f" decoding {epdir.name} ...", flush=True) |
| odin = decode_resize(epdir / "odin.mp4", size) |
| wl = decode_resize(epdir / "wrist_left.mp4", size) |
| wr = decode_resize(epdir / "wrist_right.mp4", size) |
| print(f" odin={odin.shape} wl={wl.shape} wr={wr.shape}", flush=True) |
|
|
| wleft_ts = np.array([int(r["img_recv_ts_ns"]) for r in wleft_idx], dtype=np.int64) |
| wright_ts = np.array([int(r["img_recv_ts_ns"]) for r in wright_idx], dtype=np.int64) |
|
|
| keep = list(range(0, len(ticks), stride)) |
| T = len(keep) |
| V = 3 |
| obs = np.empty((T, V, size, size, 3), dtype=np.uint8) |
| act = np.zeros((T, len(ACTION_COLS)), dtype=np.float32) |
|
|
| tick_t1 = np.array([int(ticks[k]["t1_ns"]) for k in keep], dtype=np.int64) |
| wl_match = nearest_idx(wleft_ts, tick_t1) |
| wr_match = nearest_idx(wright_ts, tick_t1) |
|
|
| prev = np.zeros(len(ACTION_COLS), dtype=np.float32) |
| for j, k in enumerate(keep): |
| row = ticks[k] |
| fi = int(float(row["frame_idx"])) |
| fi = max(0, min(fi, len(odin) - 1)) |
| obs[j, 0] = odin[fi] |
| obs[j, 1] = wl[wl_match[j]] |
| obs[j, 2] = wr[wr_match[j]] |
| vals = [] |
| for c in ACTION_COLS: |
| v = row.get(c, "") |
| try: |
| vals.append(float(v) if v != "" else np.nan) |
| except ValueError: |
| vals.append(np.nan) |
| a = np.array(vals, dtype=np.float32) |
| nan = np.isnan(a) |
| a[nan] = prev[nan] |
| prev = a |
| act[j] = a |
| return obs, act |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--raw-root", type=str, |
| default=str(Path.home() / "data/yor_episodes/5c829ca0")) |
| ap.add_argument("--out-dir", type=str, |
| default=str(Path.home() / "data/yor_dataset/yor_dataset")) |
| ap.add_argument("--stride", type=int, default=3) |
| ap.add_argument("--img-size", type=int, default=224) |
| args = ap.parse_args() |
|
|
| raw_root = Path(args.raw_root) |
| out = Path(args.out_dir) |
| (out / "obses").mkdir(parents=True, exist_ok=True) |
|
|
| epdirs = sorted(raw_root.glob("ep_*")) |
| assert epdirs, f"no episodes in {raw_root}" |
| print(f"found {len(epdirs)} episodes") |
|
|
| all_act, lengths, names = [], [], [] |
| for i, epdir in enumerate(epdirs): |
| obs, act = convert_episode(epdir, args.img_size, args.stride) |
| torch.save(torch.from_numpy(obs), |
| out / "obses" / f"episode_{i:05d}.pth") |
| all_act.append(torch.from_numpy(act)) |
| lengths.append(act.shape[0]) |
| names.append(epdir.name) |
| print(f" [{i}] {epdir.name}: T={act.shape[0]} act_mean={act.mean():.4f}") |
|
|
| Tmax = max(lengths) |
| A = len(ACTION_COLS) |
| actions = torch.zeros(len(all_act), Tmax, A, dtype=torch.float32) |
| for i, a in enumerate(all_act): |
| actions[i, :a.shape[0]] = a |
| torch.save(actions, out / "actions.pth") |
| with open(out / "seq_lengths.pkl", "wb") as f: |
| pickle.dump(lengths, f) |
| with open(out / "meta.json", "w") as f: |
| json.dump({ |
| "action_columns": ACTION_COLS, |
| "action_dim": A, |
| "views": 3, |
| "view_names": ["odin", "wrist_left", "wrist_right"], |
| "img_size": args.img_size, |
| "stride": args.stride, |
| "episodes": names, |
| "seq_lengths": lengths, |
| }, f, indent=2) |
| print(f"wrote {out}: N={len(all_act)} Tmax={Tmax} A={A}") |
| print(f" actions range [{actions.min():.3f}, {actions.max():.3f}]") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
