"""Validation report for the re-rendered frame caches used by every pixel policy. Checks, per budget: cache files/shapes/sha256; transition counts vs the release's subset manifests; nested-subset property; ordering agreement between the 64x64 cache, the proprio cache and the 224x224 cache (actions/states must be identical row-for-row); exact recomputation of the 19-d proprio vectors; reproduction of the release's 28-d `observation` after replaying qpos/qvel into the eval env; re-render determinism (cached frame == fresh render of the same state); goal frame == render of the terminal state; and the raw-release vs eval-renderer domain gap. Writes validation/render_cache_validation.json plus example PNGs. """ import argparse, hashlib, json, os, time from pathlib import Path import numpy as np, gymnasium as gym, ogbench, mujoco from proprio_spec import from_mapping, PROPRIO_DIM def sha256(p, chunk=1 << 24): h = hashlib.sha256() with open(p, "rb") as f: for b in iter(lambda: f.read(chunk), b""): h.update(b) return h.hexdigest() def obs_fn(u): for name in ["compute_observation", "_compute_observation", "get_observation", "_get_obs"]: if hasattr(u, name): return name, getattr(u, name) raise RuntimeError("no observation method on env") def make_env(size): env = gym.make("visual-cube-single-v0"); u = env.unwrapped env.reset(seed=0, options=dict(task_id=1, render_goal=False)) if size != 64: u.model.vis.global_.offwidth = size; u.model.vis.global_.offheight = size u._renderer = mujoco.Renderer(u.model, height=size, width=size) u._render_height = size; u._render_width = size return env, u def main(): ap = argparse.ArgumentParser() ap.add_argument("--root", default="data/bucket/ogbench-cube-success-1000-v1") ap.add_argument("--cachedir", default="cache") ap.add_argument("--budgets", nargs="+", type=int, default=[10, 50, 100, 1000]) ap.add_argument("--sample", type=int, default=400, help="frames re-rendered per budget for determinism") ap.add_argument("--out", default="validation") a = ap.parse_args() root, cache, out = Path(a.root), Path(a.cachedir), Path(a.out); out.mkdir(exist_ok=True) rng = np.random.default_rng(0) report = dict(generated_utc=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), root=str(root), dataset_tar_sha256=None, budgets={}, nested_subsets=None, cache224={}) tar = root.parent / "ogbench-cube-success-1000-v1.tar" if tar.exists(): report["dataset_tar_sha256"] = sha256(tar) manifests = {b: [json.loads(l) for l in (root / "subsets" / f"subset-{b:04d}.jsonl").read_text().splitlines()] for b in a.budgets} # Nested-subset property: each smaller manifest is a prefix of the next larger one. bs = sorted(a.budgets); nested = [] for s, l in zip(bs[:-1], bs[1:]): ids_s = [r["archive_file"] for r in manifests[s]]; ids_l = [r["archive_file"] for r in manifests[l]] nested.append(dict(small=s, large=l, is_prefix=ids_l[:len(ids_s)] == ids_s)) report["nested_subsets"] = nested env64, u64 = make_env(64) # 28-d state observations come from the state env (the visual env's compute_observation returns pixels). senv = gym.make("cube-single-v0"); su = senv.unwrapped senv.reset(seed=0, options=dict(task_id=1, render_goal=False)) oname, get_obs = obs_fn(su) assert (su.model.nq, su.model.nv) == (u64.model.nq, u64.model.nv) report["obs_method"] = f"cube-single-v0.{oname}" report["env"] = dict(nq=int(u64.model.nq), nv=int(u64.model.nv), frame=list(np.asarray(u64.get_pixel_observation()).shape)) for b in a.budgets: t0 = time.time(); tag = f"b{b:04d}"; rows = manifests[b] files = {k: cache / f for k, f in dict(frames=f"pix_{tag}_frames.npy", goals=f"pix_{tag}_goals.npy", meta=f"pix_{tag}_meta.npz", prop=f"prop_{tag}.npy").items()} R = dict(files={k: dict(path=str(p), exists=p.exists(), bytes=p.stat().st_size if p.exists() else None, sha256=sha256(p) if p.exists() else None) for k, p in files.items()}) frames = np.load(files["frames"], mmap_mode="r"); goals = np.load(files["goals"], mmap_mode="r") meta = np.load(files["meta"]); prop = np.load(files["prop"]) N = len(frames); expected = sum(r["transition_count"] for r in rows) R.update(episodes=len(rows), transitions_cache=int(N), transitions_manifest=int(expected), count_match=bool(N == expected), frame_shape=list(frames.shape[1:]), goals_shape_match=bool(goals.shape == frames.shape), proprio_shape=list(prop.shape), proprio_aligned=bool(prop.shape == (N, PROPRIO_DIM))) # Row-for-row recomputation from the release: actions, states, proprio; obs reproduction; goal frames. i = 0; act_ok = st_ok = prop_ok = True; obs_err = []; goal_frame_exact = 0; goal_frame_maxdiff = 0 traj_ok = True for ti, row in enumerate(rows): with np.load(root / "data" / row["archive_file"], allow_pickle=False) as z: qpos, qvel, act, sobs = z["qpos"], z["qvel"], z["action"], z["observation"] T = len(act) - 1 act_ok &= np.array_equal(meta["action"][i:i + T], act[:T].astype(np.float32)) st_ok &= np.array_equal(meta["state"][i:i + T], sobs[:T].astype(np.float32)) traj_ok &= bool((meta["traj_id"][i:i + T] == ti).all()) P = np.stack([from_mapping(z, t) for t in range(T)]) prop_ok &= np.array_equal(prop[i:i + T], P) # goal frame == fresh render of the terminal state u64.data.qpos[:] = qpos[-1]; u64.data.qvel[:] = qvel[-1]; mujoco.mj_forward(u64.model, u64.data) g = np.asarray(u64.get_pixel_observation()) d = np.abs(g.astype(int) - goals[i].astype(int)).max() goal_frame_exact += int(d == 0); goal_frame_maxdiff = max(goal_frame_maxdiff, int(d)) # 28-d observation reproduction on a few rows per episode for t in sorted(set([0, T // 2, T - 1])): su.data.qpos[:] = qpos[t]; su.data.qvel[:] = qvel[t]; mujoco.mj_forward(su.model, su.data) obs_err.append(np.abs(np.asarray(get_obs()) - sobs[t])) i += T obs_err = np.stack(obs_err) R.update(actions_match_release=bool(act_ok), states_match_release=bool(st_ok), traj_ids_match=bool(traj_ok), proprio_recomputed_exact=bool(prop_ok), goal_frames=dict(exact=goal_frame_exact, of=len(rows), max_abs_pixel_diff=goal_frame_maxdiff), obs_reproduction=dict(rows_checked=int(len(obs_err)), max_abs=float(obs_err.max()), mean_abs=float(obs_err.mean()), per_dim_max=[float(x) for x in obs_err.max(0)], worst_dims=[int(x) for x in np.argsort(-obs_err.max(0))[:5]], # dim 18 is the binary gripper_contact flag: a static mj_forward on a replayed # state does not always reproduce a contact registered during live stepping. max_abs_continuous_dims=float(np.delete(obs_err, 18, axis=1).max()), contact_flag_mismatch_fraction=float((obs_err[:, 18] > 0.5).mean()), note="continuous dims reproduce the release; only the binary contact flag " "(dim 18) can differ. Proprio is never taken from replayed states: " "training uses the release's stored live values, eval uses env info.")) # Re-render determinism on a random sample of transitions. idx = np.sort(rng.choice(N, size=min(a.sample, N), replace=False)) tid = meta["traj_id"]; exact = 0; maxd = 0; meand = [] # map traj -> (row, offset) to find qpos/qvel for each sampled transition starts = np.searchsorted(tid, np.arange(len(rows))) cur_row, z = -1, None for j in idx: r = int(tid[j]); t = int(j - starts[r]) if r != cur_row: if z is not None: z.close() z = np.load(root / "data" / rows[r]["archive_file"], allow_pickle=False); cur_row = r u64.data.qpos[:] = z["qpos"][t]; u64.data.qvel[:] = z["qvel"][t]; mujoco.mj_forward(u64.model, u64.data) f = np.asarray(u64.get_pixel_observation()).astype(int); c = frames[j].astype(int) d = np.abs(f - c); exact += int(d.max() == 0); maxd = max(maxd, int(d.max())); meand.append(float(d.mean())) if z is not None: z.close() R["rerender_determinism"] = dict(sampled=int(len(idx)), pixel_exact=exact, max_abs_pixel_diff=maxd, mean_abs_pixel_diff=float(np.mean(meand))) R["seconds"] = round(time.time() - t0, 1) report["budgets"][b] = R print(f"[validate] b{b}: N={N} match={R['count_match']} act={act_ok} state={st_ok} prop={prop_ok} " f"obs_max={R['obs_reproduction']['max_abs']:.4g} rerender_exact={exact}/{len(idx)} " f"goal_exact={goal_frame_exact}/{len(rows)} ({R['seconds']}s)", flush=True) # Domain gap: raw release pixels.front (224, collection renderer) vs eval renderer at 224 for the same states. env224, u224 = make_env(224) row = manifests[bs[0]][0] with np.load(root / "data" / row["archive_file"], allow_pickle=False) as z: qpos, qvel, raw = z["qpos"], z["qvel"], z["pixels.front"] gaps = [] for t in [0, len(qpos) // 2, len(qpos) - 1]: u224.data.qpos[:] = qpos[t]; u224.data.qvel[:] = qvel[t]; mujoco.mj_forward(u224.model, u224.data) rr = np.asarray(u224.get_pixel_observation()) gaps.append(float(np.abs(rr.astype(int) - raw[t].astype(int)).mean())) if t == 0: try: import imageio imageio.imwrite(out / "domain_gap_raw_release_224.png", raw[t]) imageio.imwrite(out / "domain_gap_eval_renderer_224.png", rr) u64.data.qpos[:] = qpos[t]; u64.data.qvel[:] = qvel[t]; mujoco.mj_forward(u64.model, u64.data) imageio.imwrite(out / "domain_gap_eval_renderer_64.png", np.asarray(u64.get_pixel_observation())) except Exception as e: print("png export skipped:", e) report["domain_gap"] = dict(episode=row["archive_file"], mean_abs_pixel_diff_raw_vs_eval_224=gaps, note="raw release frames come from the collection renderer (swm/OGBCube-v0); " "policies are trained on eval-renderer frames instead") # 224x224 cache (OpenVLA): ordering vs the 64x64 cache, goal xy, and re-render determinism. f224 = cache / "vla_b1000_frames224.npy"; m224 = cache / "vla_b1000_meta.npz" if f224.exists(): t0 = time.time() fr = np.load(f224, mmap_mode="r"); mm = np.load(m224); m64 = np.load(cache / "pix_b1000_meta.npz") rows = manifests[1000]; N = len(fr) R = dict(files={"frames224": dict(path=str(f224), bytes=f224.stat().st_size, sha256=sha256(f224)), "meta": dict(path=str(m224), bytes=m224.stat().st_size, sha256=sha256(m224))}, shape=list(fr.shape), transitions_match_64=bool(N == len(m64["action"])), actions_identical_to_64_cache=bool(np.array_equal(mm["action"], m64["action"])), traj_ids_identical_to_64_cache=bool(np.array_equal(mm["traj_id"], m64["traj_id"]))) # goal xy per trajectory and re-render determinism on a sample idx = np.sort(rng.choice(N, size=min(a.sample, N), replace=False)); tid = mm["traj_id"] starts = np.searchsorted(tid, np.arange(len(rows))); exact = 0; maxd = 0; goal_ok = True; cur_row, z = -1, None for j in idx: r = int(tid[j]); t = int(j - starts[r]) if r != cur_row: if z is not None: z.close() z = np.load(root / "data" / rows[r]["archive_file"], allow_pickle=False); cur_row = r goal_ok &= bool(np.allclose(mm["goal_xy"][j], np.asarray(z["variation.cube.goal_position"])[0, 0, :])) u224.data.qpos[:] = z["qpos"][t]; u224.data.qvel[:] = z["qvel"][t]; mujoco.mj_forward(u224.model, u224.data) d = np.abs(np.asarray(u224.get_pixel_observation()).astype(int) - fr[j].astype(int)) exact += int(d.max() == 0); maxd = max(maxd, int(d.max())) if z is not None: z.close() R["goal_xy_matches_release"] = goal_ok R["rerender_determinism"] = dict(sampled=int(len(idx)), pixel_exact=exact, max_abs_pixel_diff=maxd) R["seconds"] = round(time.time() - t0, 1) report["cache224"] = R print(f"[validate] 224: order_ok={R['actions_identical_to_64_cache']} goal_ok={goal_ok} " f"rerender_exact={exact}/{len(idx)} ({R['seconds']}s)", flush=True) # Overall verdict. ok = all(v["count_match"] and v["actions_match_release"] and v["states_match_release"] and v["proprio_recomputed_exact"] and v["traj_ids_match"] for v in report["budgets"].values()) ok &= all(n["is_prefix"] for n in nested) if report["cache224"]: ok &= report["cache224"]["actions_identical_to_64_cache"] and report["cache224"]["goal_xy_matches_release"] report["verdict"] = "PASS" if ok else "FAIL" json.dump(report, open(out / "render_cache_validation.json", "w"), indent=2) print("[validate] verdict:", report["verdict"], "->", out / "render_cache_validation.json", flush=True) if __name__ == "__main__": main()