File size: 1,255 Bytes
87d0bc2 | 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 | """Build the 19-d proprio array in EXACTLY the same order as the pixel frame cache."""
import argparse, json
from pathlib import Path
import numpy as np
from proprio_spec import from_mapping, PROPRIO_DIM
ap = argparse.ArgumentParser()
ap.add_argument("--root", default="data/bucket/ogbench-cube-success-1000-v1")
ap.add_argument("--budget", type=int, required=True)
ap.add_argument("--cachedir", default="cache")
a = ap.parse_args()
root = Path(a.root); tag = f"b{a.budget:04d}"
rows = [json.loads(l) for l in (root / "subsets" / f"subset-{a.budget:04d}.jsonl").read_text().splitlines()]
out = []
for row in rows: # same order as render_cache.py
with np.load(root / "data" / row["archive_file"], allow_pickle=False) as z:
T = len(z["action"]) - 1
for t in range(T): # same inner loop
out.append(from_mapping(z, t))
P = np.stack(out).astype(np.float32)
n_frames = len(np.load(f"{a.cachedir}/pix_{tag}_frames.npy", mmap_mode="r"))
assert P.shape == (n_frames, PROPRIO_DIM), f"ordering mismatch: {P.shape} vs frames {n_frames}"
np.save(f"{a.cachedir}/prop_{tag}.npy", P)
print(json.dumps({"budget": a.budget, "proprio": list(P.shape), "aligned_with_frames": True}))
|