| """Re-render dataset trajectories INSIDE the OGBench eval env for domain-matched pixels. |
| |
| The release stores full sim state (qpos/qvel); we replay it into the evaluation |
| environment and call get_pixel_observation(), so training images come from the |
| same renderer/camera as evaluation. Also stores the per-episode goal frame |
| (final state) and the 28-d state obs for the proprioception ablation. |
| """ |
| import argparse, json, time |
| from pathlib import Path |
| import numpy as np, gymnasium as gym, ogbench, mujoco |
|
|
| def main(): |
| 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("--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(parents=True, exist_ok=True) |
|
|
| env = gym.make(a.env); u = env.unwrapped |
| env.reset(seed=0, options=dict(task_id=1, render_goal=False)) |
| probe = np.asarray(u.get_pixel_observation()) |
| H, W, C = probe.shape |
| print(f"[render] env={a.env} frame={probe.shape} dtype={probe.dtype}", flush=True) |
|
|
| 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) |
| frames = np.zeros((total, H, W, C), np.uint8) |
| goals = np.zeros((total, H, W, C), np.uint8) |
| acts = np.zeros((total, 5), np.float32) |
| states = np.zeros((total, 28), 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 = z["qpos"], z["qvel"] |
| act, sobs = z["action"], z["observation"] |
| T = len(act) - 1 |
| |
| u.data.qpos[:] = qpos[-1]; u.data.qvel[:] = qvel[-1] |
| mujoco.mj_forward(u.model, u.data) |
| gframe = np.asarray(u.get_pixel_observation()).copy() |
| 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() |
| goals[i] = gframe |
| acts[i] = act[t]; states[i] = sobs[t]; tids[i] = ti |
| i += 1 |
| if (ti + 1) % 100 == 0: |
| print(f" {ti+1}/{len(rows)} episodes, {i} frames, {i/(time.time()-t0):.0f} fps", flush=True) |
| assert i == total, (i, total) |
| tag = f"b{a.budget:04d}" |
| np.save(out / f"pix_{tag}_frames.npy", frames) |
| np.save(out / f"pix_{tag}_goals.npy", goals) |
| np.savez(out / f"pix_{tag}_meta.npz", action=acts, state=states, traj_id=tids) |
| print(json.dumps({"budget": a.budget, "transitions": total, "frame": [H, W, C], |
| "seconds": round(time.time() - t0, 1), |
| "frames_gb": round(frames.nbytes / 1e9, 3)}), flush=True) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
