| """Generic hold-tail extender for ANY 3-PoinTr dataset (single- or multi-object). |
| |
| Appends H timesteps that REPEAT the final cum_flow (pos is constant = x0, visibility |
| copied from the last frame), so the flow model learns to HOLD at the final pose instead |
| of relaxing the points back toward frame-0 near the end. Task-agnostic: operates on the |
| existing `no_actions/<task>/<NNNN>/timestep_k/{k_pos,k_cum_flow,k_visibility}.pth` tree, |
| copies labels.pth + metadata.json, and bumps T → T+H. Per-episode T is read from the |
| data (variable-length demos handled). |
| |
| python extend_3pointr_holdtail.py \ |
| --src /data0/.../3pointr_data_open_drawer_2d_wb/no_actions/open_drawer \ |
| --dst /data0/.../3pointr_data_open_drawer_2d_wb_hold30/no_actions/open_drawer --hold_tail 30 |
| """ |
| import os, json, glob, shutil, argparse |
| from pathlib import Path |
| import torch |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--src", required=True, help="no_actions/<task> dir of the existing dataset") |
| ap.add_argument("--dst", required=True, help="no_actions/<task> dir to write the extended dataset") |
| ap.add_argument("--hold_tail", type=int, default=30) |
| a = ap.parse_args() |
| H = a.hold_tail |
| os.makedirs(a.dst, exist_ok=True) |
| eps = sorted(glob.glob(os.path.join(a.src, "[0-9][0-9][0-9][0-9]"))) |
| print(f"[extend] {len(eps)} episodes hold_tail={H} {a.src} -> {a.dst}", flush=True) |
| minN = 10**9 |
| for ep in eps: |
| name = Path(ep).name |
| dst_ep = Path(a.dst) / name |
| if dst_ep.exists(): |
| shutil.rmtree(dst_ep) |
| shutil.copytree(ep, dst_ep) |
| T = len(glob.glob(os.path.join(dst_ep, "timestep_*"))) |
| x0 = torch.load(dst_ep / "timestep_0" / "0_pos.pth", weights_only=False) |
| cf = torch.load(dst_ep / f"timestep_{T-1}" / f"{T-1}_cum_flow.pth", weights_only=False) |
| vis = torch.load(dst_ep / f"timestep_{T-1}" / f"{T-1}_visibility.pth", weights_only=False) |
| for k in range(T, T + H): |
| ts = dst_ep / f"timestep_{k}" |
| ts.mkdir(parents=True, exist_ok=True) |
| torch.save(x0, ts / f"{k}_pos.pth") |
| torch.save(cf, ts / f"{k}_cum_flow.pth") |
| torch.save(vis, ts / f"{k}_visibility.pth") |
| mp = dst_ep / "metadata.json" |
| meta = json.loads(mp.read_text()) if mp.exists() else {} |
| meta["T_tracked"] = int(T); meta["hold_tail"] = int(H); meta["T"] = int(T + H) |
| mp.write_text(json.dumps(meta, indent=2)) |
| minN = min(minN, int(x0.shape[0])) |
| print(f"[extend] done: {len(eps)} eps, each +{H} held frames. min_N = {minN} " |
| f"(set --subsample <= {minN})", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
