| """Assemble Thor-deployable ckpt dir from a finished VQ-BeT yor run. |
| |
| Output <run>/ckpt/: |
| policy.pt whole-module BehaviorTransformer pickle, CPU, eval mode |
| (VQ-BeT: no normalizer object; actions raw, act_scale=1.0) |
| encoder.pt whole-module DinoV2Encoder pickle, CPU, eval mode |
| (B,C,H,W)/255 RGB -> (B,P,E) |
| card.yaml views, img_size, obs_horizon, chunk, exec_n, action_hz, |
| act_fields (58, exact converter order), proprio_fields |
| |
| action_hz = measured training rate = 1 / median(STRIDE * tick_dt), |
| computed from the raw episodes' t0_ns column. |
| """ |
| import argparse |
| import csv |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| sys.path.insert(0, "/home_shared/grail_charles/patch_policy") |
|
|
| 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)] |
| ACTION_COLS = EE_L + EE_R + HAND_L + HAND_R + ["lift"] + \ |
| ["base_pose_target_0", "base_pose_target_1", "base_pose_target_2"] |
| assert len(ACTION_COLS) == 58 |
|
|
|
|
| def measured_tick_dt(raw_root: Path) -> float: |
| dts = [] |
| for ep in sorted(raw_root.glob("ep_*")): |
| with open(ep / "robot_data.csv", newline="") as f: |
| t0 = np.array([int(r["t0_ns"]) for r in csv.DictReader(f)], np.int64) |
| d = np.diff(t0) / 1e9 |
| dts.append(d[(d > 0) & (d < 1.0)]) |
| return float(np.median(np.concatenate(dts))) |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--run", required=True, help="run dir with model_final.pt") |
| ap.add_argument("--model", default="model_final.pt", |
| help="ckpt file name in the run dir (e.g. model_stopped.pt)") |
| ap.add_argument("--raw-root", |
| default=str(Path.home() / "data/yor_episodes/5c829ca0")) |
| ap.add_argument("--stride", type=int, default=3) |
| ap.add_argument("--chunk", type=int, default=5) |
| ap.add_argument("--exec-n", type=int, default=4) |
| ap.add_argument("--feature-key", default="x_norm_patchtokens", |
| help="x_norm_patchtokens (P=256) or x_norm_clstoken (P=1)") |
| args = ap.parse_args() |
|
|
| import torch |
| run = Path(args.run) |
| out = run / "ckpt" |
| out.mkdir(parents=True, exist_ok=True) |
|
|
| tick_dt = measured_tick_dt(Path(args.raw_root)) |
| action_hz = 1.0 / (args.stride * tick_dt) |
| print(f"tick_dt={tick_dt*1000:.1f}ms stride={args.stride} " |
| f"-> action_hz={action_hz:.3f}") |
|
|
| print("loading policy ...", flush=True) |
| policy = torch.load(run / args.model, map_location="cpu", |
| weights_only=False) |
| policy.to("cpu") |
| policy.eval() |
| torch.save(policy, out / "policy.pt") |
|
|
| from models.encoder.dino import DinoV2Encoder |
| print("building encoder ...", flush=True) |
| n_patches = 1 if args.feature_key == "x_norm_clstoken" else 256 |
| enc = DinoV2Encoder(name="dinov2_vits14", |
| feature_key=args.feature_key, |
| output_dim=384, n_patches=n_patches) |
| enc.to("cpu") |
| enc.eval() |
| for p in enc.parameters(): |
| p.requires_grad = False |
| torch.save(enc, out / "encoder.pt") |
|
|
| card = [ |
| "views: [odin, wl, wr]", |
| "img_size: 224", |
| "obs_horizon: 2", |
| f"chunk: {args.chunk}", |
| f"exec_n: {args.exec_n}", |
| f"action_hz: {action_hz:.3f}", |
| "act_fields: [" + ", ".join(ACTION_COLS) + "]", |
| "proprio_fields: []", |
| ] |
| (out / "card.yaml").write_text("\n".join(card) + "\n") |
|
|
| |
| print("verifying ...", flush=True) |
| p2 = torch.load(out / "policy.pt", map_location="cpu", weights_only=False) |
| p2.eval() |
| e2 = torch.load(out / "encoder.pt", map_location="cpu", weights_only=False) |
| e2.eval() |
| with torch.no_grad(): |
| assert e2(torch.rand(2, 3, 224, 224)).shape == (2, n_patches, 384) |
| pr, _, _ = p2(torch.zeros(1, 2, 768, 384), torch.ones(1, 2, 1, 1), None) |
| assert pr.shape == (1, 2, args.chunk, 58), pr.shape |
| print(f"OK -> {out}: policy{p2.__class__.__name__} " |
| f"pred {tuple(pr.shape)}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
