| """Re-render the release at 224x224 inside the eval env (OpenVLA's native size). |
| |
| Raises the model's offscreen framebuffer (hard-set to 64) before swapping in a |
| 224 renderer. Stores frames as a memmap plus per-transition action and goal xy. |
| """ |
| import argparse, json, time |
| from pathlib import Path |
| import numpy as np, gymnasium as gym, ogbench, mujoco |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--root", default="data/bucket/ogbench-cube-success-1000-v1") |
| ap.add_argument("--budget", type=int, default=1000) |
| ap.add_argument("--env", default="visual-cube-single-v0") |
| ap.add_argument("--outdir", default="cache") |
| a = ap.parse_args() |
| root = Path(a.root); out = Path(a.outdir); out.mkdir(exist_ok=True) |
|
|
| env = gym.make(a.env); u = env.unwrapped |
| env.reset(seed=0, options=dict(task_id=1, render_goal=False)) |
| u.model.vis.global_.offwidth = 224 |
| u.model.vis.global_.offheight = 224 |
| u._renderer = mujoco.Renderer(u.model, height=224, width=224) |
| u._render_height = 224; u._render_width = 224 |
| assert np.asarray(u.get_pixel_observation()).shape[:2] == (224, 224) |
|
|
| rows = [json.loads(l) for l in (root / "subsets" / f"subset-{a.budget:04d}.jsonl").read_text().splitlines()] |
| total = sum(r["transition_count"] for r in rows) |
| tag = f"b{a.budget:04d}" |
| frames = np.lib.format.open_memmap(out / f"vla_{tag}_frames224.npy", mode="w+", |
| dtype=np.uint8, shape=(total, 224, 224, 3)) |
| acts = np.zeros((total, 5), np.float32) |
| goals = np.zeros((total, 2), np.float32) |
| tids = np.zeros(total, np.int32) |
|
|
| i, t0 = 0, time.time() |
| for ti, row in enumerate(rows): |
| with np.load(root / "data" / row["archive_file"], allow_pickle=False) as z: |
| qpos, qvel, act = z["qpos"], z["qvel"], z["action"] |
| gxy = np.asarray(z["variation.cube.goal_position"])[0, 0, :] |
| T = len(act) - 1 |
| for t in range(T): |
| u.data.qpos[:] = qpos[t]; u.data.qvel[:] = qvel[t] |
| mujoco.mj_forward(u.model, u.data) |
| frames[i] = u.get_pixel_observation() |
| acts[i] = act[t]; goals[i] = gxy; tids[i] = ti |
| i += 1 |
| if (ti + 1) % 200 == 0: |
| print(f" {ti+1}/{len(rows)} eps, {i} frames, {i/(time.time()-t0):.0f} fps", flush=True) |
| assert i == total |
| frames.flush() |
| np.savez(out / f"vla_{tag}_meta.npz", action=acts, goal_xy=goals, traj_id=tids) |
| print(json.dumps({"budget": a.budget, "transitions": total, "frame": [224, 224, 3], |
| "gb": round(frames.nbytes / 1e9, 2), "seconds": round(time.time() - t0, 1)})) |
|
|
