File size: 4,294 Bytes
94f697a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """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() # statement: bet.train() returns None, do not chain
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")
# verify: reload + forward with zeros
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()
|