File size: 5,476 Bytes
5d08972 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | #!/usr/bin/env python3
"""Render episodes listed in a render plan, one mp4 per camera.
The plan fixes every folder name up front, in one global pass over all runs, so a name is
unique across the whole output rather than only within the run that produced it. Assigning
names per run is what silently dropped episodes: a later run's success collided with an
earlier one's folder and was skipped as already done.
Plan entries are `{run_dir, demo, target, name}`. Shard with `--shard i --num-shards n` to
render in parallel; each episode writes a `.done` marker, so re-running skips finished work.
MUJOCO_GL=egl python render_from_plan.py --plan render_plan.json \
--source src150_fix.hdf5 --filter-key valid --out output_aug_256 --shard 0 --num-shards 6
"""
from __future__ import annotations
import argparse
import json
import os
import time
os.environ.setdefault("MUJOCO_GL", "egl")
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
import h5py
import imageio
import mujoco
import numpy as np
import robocasa # noqa: F401
import robosuite
from robosuite.controllers import load_composite_controller_config
CAMERAS = ["robot0_agentview_left", "robot0_agentview_right", "robot0_eye_in_hand"]
def source_index(path, filter_key):
out = {}
with h5py.File(path, "r") as f:
keys = ([k.decode() for k in f[f"mask/{filter_key}"][:]] if filter_key
else sorted(f["data"].keys(), key=lambda x: int(x.split("_")[1])))
for i, k in enumerate(keys):
meta = json.loads(f[f"data/{k}"].attrs["ep_meta"])
cat = next((c["info"]["cat"] for c in meta.get("object_cfgs", [])
if c.get("name") == "obj"), "unknown")
out[i] = (k, cat)
return out
def build_env(env_args, res):
kw = dict(env_args["env_kwargs"])
kw.pop("env_name", None)
kw.update(has_renderer=False, has_offscreen_renderer=True, use_camera_obs=False,
use_object_obs=True, camera_depths=False, camera_names=CAMERAS,
camera_widths=res, camera_heights=res, ignore_done=True)
kw.setdefault("controller_configs",
load_composite_controller_config(controller=None, robot="PandaOmron"))
return robosuite.make(env_name=env_args["env_name"], **kw)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--plan", required=True)
ap.add_argument("--source", required=True)
ap.add_argument("--filter-key", default=None)
ap.add_argument("--out", required=True)
ap.add_argument("--shard", type=int, default=0)
ap.add_argument("--num-shards", type=int, default=1)
ap.add_argument("--res", type=int, default=256)
ap.add_argument("--fps", type=int, default=20)
a = ap.parse_args()
plan = [e for i, e in enumerate(json.load(open(a.plan))) if i % a.num_shards == a.shard]
src = source_index(a.source, a.filter_key)
os.makedirs(a.out, exist_ok=True)
env, opened, done = None, None, 0
for e in plan:
ep_dir = os.path.join(a.out, e["name"])
if os.path.exists(os.path.join(ep_dir, ".done")):
done += 1
continue
t0 = time.time()
with h5py.File(os.path.join(e["run_dir"], "demo.hdf5"), "r") as f:
g = f[f"data/{e['demo']}"]
states = g["states"][:]
meta = json.loads(g.attrs["ep_meta"])
model_xml = g.attrs["model_file"]
si = int(g["src_demo_inds"][0]) if "src_demo_inds" in g else -1
env_args = json.loads(f["data"].attrs["env_args"])
sk, scat = src.get(si, (f"idx{si}", "unknown"))
if env is None:
env = build_env(env_args, a.res)
env.set_ep_meta(meta)
env.reset()
env.reset_from_xml_string(model_xml)
env.sim.reset()
model, data = env.sim.model._model, env.sim.data._data
opt = mujoco.MjvOption()
opt.geomgroup[:] = 0
opt.geomgroup[1] = 1 # visual geoms, as the dataset's own videos are drawn
opt.sitegroup[:] = 0
r = mujoco.Renderer(model, height=a.res, width=a.res)
cam_ids = {c: mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, c) for c in CAMERAS}
frames = {c: [] for c in CAMERAS}
for t in range(len(states)):
env.sim.set_state_from_flattened(states[t])
env.sim.forward()
for c in CAMERAS:
r.update_scene(data, camera=cam_ids[c], scene_option=opt)
frames[c].append(r.render().copy())
r.close()
for c in CAMERAS:
d = os.path.join(ep_dir, c)
os.makedirs(d, exist_ok=True)
imageio.mimwrite(os.path.join(d, f"{c}.mp4"), frames[c], fps=a.fps,
codec="libx264", macro_block_size=1,
ffmpeg_params=["-pix_fmt", "yuv420p"])
with open(os.path.join(ep_dir, ".done"), "w") as fh:
json.dump(dict(generated_demo=e["demo"], run_dir=e["run_dir"],
source_demo=sk, source_object=scat, target_object=e["target"],
frames=len(states), layout=meta.get("layout_id"),
style=meta.get("style_id"), lang=meta.get("lang")), fh, indent=1)
done += 1
print(f" [{a.shard}] {e['name']}: {len(states)}f {time.time()-t0:.0f}s", flush=True)
print(f"shard {a.shard}: {done}/{len(plan)}", flush=True)
if __name__ == "__main__":
raise SystemExit(main())
|