File size: 7,214 Bytes
168e232 | 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 134 135 136 137 138 139 | """Clear the floor of a real room (imported from naavox) in sim: many cubes, one overhead camera, the bin at the named
place, MuJoCo as the world model (oracle).
uv run python scripts/clear_sim.py --scenario configs/scenarios/clear_bedroom.yaml --episodes 3
uv run python scripts/clear_sim.py --scenario configs/scenarios/clear_playroom.yaml --episodes 1 --video
Video: room view with the purple plan | the overhead camera (what the robot sees: detections boxed, the target ringed,
the plan projected) | the gripper camera (for the eye only: perception does not use it).
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from datetime import datetime
from pathlib import Path
import mujoco
import numpy as np
import yaml
from PIL import Image, ImageDraw
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from cablepilot.clear_floor import run_clear_floor # noqa: E402
from cablepilot.compute import ComputeMeter # noqa: E402
from cablepilot.runner import Sim, write_summary # noqa: E402
from cablepilot.sim.scene import load_config # noqa: E402
from cablepilot.viz import draw_scene, overlay_camera # noqa: E402
class VideoSink:
def __init__(self, sim: Sim, path: Path, fps: int, H=360):
import imageio.v2 as imageio
m, d = sim.plant.model, sim.plant.data
self.sim, self.H = sim, H
self.r_room = mujoco.Renderer(m, H, 480)
self.r_grip = mujoco.Renderer(m, H, 480)
self.cam = mujoco.MjvCamera()
self.cam.distance, self.cam.azimuth, self.cam.elevation = 3.0, 225.0, -35.0
self.gcam = m.camera("gripper_cam").id
self.opt = mujoco.MjvOption()
self.opt.geomgroup[2] = 0
for o in (self.opt,):
o.sitegroup[:] = 0
o.flags[mujoco.mjtVisFlag.mjVIS_RANGEFINDER] = 0
self.gopt = mujoco.MjvOption()
self.gopt.sitegroup[:] = 0
self.gopt.flags[mujoco.mjtVisFlag.mjVIS_RANGEFINDER] = 0
self.writer = imageio.get_writer(path, fps=fps, codec="libx264", quality=7, macro_block_size=8)
def __call__(self, f):
d = self.sim.plant.data
tgt = 0.5 * (f.nominal[0] + f.nominal[-1])
self.cam.lookat[:] = tgt if not self.cam.lookat.any() else 0.9 * self.cam.lookat + 0.1 * tgt
self.r_room.update_scene(d, camera=self.cam, scene_option=self.opt)
draw_scene(self.r_room.scene, f.nominal, f.samples, f.extra["target"], f.phase, reset=False)
room = self.r_room.render()
cam = f.extra["cam"]
ov = cam.last if cam.last is not None else cam.render()
ov = overlay_camera(ov, cam.project, f.nominal, f.samples,
[f"t {f.time_s:5.1f}s pick {f.extra['picks'] + 1} {f.phase.upper()}",
f"plan {f.extra['plan_ms']:.0f} ms blobs {len(f.extra['detections'])}"])
im = Image.fromarray(ov)
dr = ImageDraw.Draw(im)
for det in f.extra["detections"]:
u, v, _ = cam.project(det["centre"][None])
dr.rectangle([u[0] - 6, v[0] - 6, u[0] + 6, v[0] + 6], outline=(255, 255, 0) if not det["partial"] else (255, 120, 0))
tg = f.extra["target"]
if tg.centre is not None:
u, v, _ = cam.project(tg.centre[None])
dr.ellipse([u[0] - 11, v[0] - 11, u[0] + 11, v[0] + 11], outline=(0, 255, 120), width=2)
ov = np.asarray(im.resize((self.H, self.H)))
self.r_grip.update_scene(d, camera=self.gcam, scene_option=self.gopt)
grip = self.r_grip.render()
self.writer.append_data(np.concatenate([room, ov, grip], axis=1))
def close(self):
self.writer.close()
def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--scenario", default=str(ROOT / "configs/scenarios/clear_bedroom.yaml"))
ap.add_argument("--episodes", type=int, default=1)
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--planner", default=str(ROOT / "configs/planner.yaml"))
ap.add_argument("--task", default=str(ROOT / "configs/tasks/clear_floor.yaml"))
ap.add_argument("--duration", type=float, default=None)
ap.add_argument("--video", action="store_true")
ap.add_argument("--out", default=None)
args = ap.parse_args()
planner = yaml.safe_load(Path(args.planner).read_text())
task_cfg = yaml.safe_load(Path(args.task).read_text())
if args.duration:
task_cfg["episode"]["duration_s"] = args.duration
cfg = load_config(scenario=str(Path(args.scenario).resolve().relative_to(ROOT)))
sim = Sim(cfg, planner, task_cfg)
name = Path(args.scenario).stem
out = Path(args.out or ROOT / "runs" / f"{datetime.now().strftime('%Y%m%d-%H%M%S')}_{name}")
out.mkdir(parents=True, exist_ok=True)
(out / "config.yaml").write_text(yaml.safe_dump({"args": vars(args), "planner": planner, "task": task_cfg}))
meter = ComputeMeter(budget_s=sim.pcfg.dt, log_path=out / "compute.jsonl")
allm = []
for ep in range(args.episodes):
sink = VideoSink(sim, out / f"episode_{ep:02d}.mp4", fps=round(3 / sim.pcfg.dt)) if args.video else None
t0 = time.perf_counter()
res, mt = run_clear_floor(sim, seed=args.seed + 1000 * ep, on_frame=sink, meter=meter)
if sink:
sink.close()
allm.append(mt)
(out / f"episode_{ep:02d}_picks.json").write_text(json.dumps(res["picks"], indent=1, default=float))
lg = res["log"]
np.savez_compressed(out / f"episode_{ep:02d}.npz", time_s=np.asarray(lg["time_s"]), state=np.asarray(lg["state"]),
phase=np.asarray(lg["phase"], dtype="U8"), pick=np.asarray(lg["pick"]),
tensions=np.asarray(lg["tensions"]), plan_ms=np.asarray(lg["plan_ms"]),
cube_names=np.asarray(list(res["cube_start"])), cube_start=np.asarray(list(res["cube_start"].values())),
cube_end=np.asarray(list(res["cube_end"].values())))
print(f"episode {ep:02d}: binned {mt['binned']}/{mt['cubes']} left {mt['left_on_floor']} picks {mt['picks_done']}/"
f"{mt['picks_tried']} closes {mt['close_attempts']} {mt['duration_s']:.0f} s sim ({mt['s_per_binned_cube']:.0f} s/cube, "
f"pick median {mt['pick_s_median']:.0f} s) plan {mt['plan_ms_mean']:.0f} ms perception {mt['perception_ms_mean']:.0f} ms "
f"T max {mt['max_tension_N']:.0f} N furniture-contact ticks {mt['furniture_contact_ticks']} ({time.perf_counter() - t0:.0f} s wall)", flush=True)
for lf in res["left"]:
print(" LEFT", lf)
for p in res["picks"]:
print(" ", f"{p['t0']:6.1f}s", "done" if p["done"] else "GAVE UP", f"{p['duration_s']:.0f}s",
" | ".join(f"{e['time_s']:.0f}s {e['why']}" for e in p["events"][-3:]))
print(write_summary(out, allm, f"clear_floor {name} model: {sim.model_name} episodes {len(allm)} seed {args.seed}"))
(out / "compute.txt").write_text(meter.table() + "\n")
print(meter.table())
meter.close()
print("->", out)
if __name__ == "__main__":
main()
|