"""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()