src/validate_render_cache.py

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