File size: 3,039 Bytes
87d0bc2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """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
# goal frame = final (terminal) state of this successful demo
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()
|