src/proprio_cache.py

"""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}))