| """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 |
| from cablepilot.compute import ComputeMeter |
| from cablepilot.runner import Sim, write_summary |
| from cablepilot.sim.scene import load_config |
| from cablepilot.viz import draw_scene, overlay_camera |
|
|
|
|
| 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() |
|
|
