"""Pick and place in sim with MuJoCo as the world model (oracle). Milestone 0. # batch: metrics.csv + summary.txt + compute.txt (+ an mp4 per episode with --video: room view with the purple # plan and the sample spaghetti | gripper camera with the plan projected into it) uv run python scripts/pick_sim.py --episodes 12 uv run python scripts/pick_sim.py --episodes 3 --video # live, watch it think (macOS: mjpython for the viewer) uv run mjpython scripts/pick_sim.py --viewer Episodes cycle through the (object, container) pairs of configs/tasks/pick_place.yaml. """ from __future__ import annotations import argparse import sys import time from datetime import datetime from pathlib import Path import mujoco import numpy as np import yaml ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from cablepilot.compute import ComputeMeter # noqa: E402 from cablepilot.runner import Sim, run_episode, save_episode, write_summary # noqa: E402 from cablepilot.sim.scene import load_config # noqa: E402 from cablepilot.viz import draw_scene, overlay_camera # noqa: E402 def hud(f) -> list[str]: tg = f.extra["target"] lines = [f"t {f.time_s:5.1f}s phase {f.phase.upper()}", f"target conf {tg.confidence:.2f}{' SEEN' if tg.visible else ''}", f"plan {f.extra['plan_ms']:.0f} ms ESS {f.info['ess']:.0f}"] top = sorted(f.info["task_costs"].items(), key=lambda kv: -abs(kv[1]))[:3] lines.append(" ".join(f"{k} {v:.0f}" for k, v in top)) return lines class Camera: """Projection for an offscreen render of a MuJoCo camera (the gripper camera) at W x H.""" def __init__(self, model, data, cam_id, W, H): self.d, self.cam, self.W, self.H = data, cam_id, W, H self.fy = 0.5 * H / np.tan(np.radians(model.cam_fovy[cam_id]) / 2) def __call__(self, pts): R = self.d.cam_xmat[self.cam].reshape(3, 3) rel = (np.atleast_2d(pts) - self.d.cam_xpos[self.cam]) @ R z = -rel[:, 2] return (self.W / 2 + self.fy * rel[:, 0] / np.maximum(z, 1e-6), self.H / 2 - self.fy * rel[:, 1] / np.maximum(z, 1e-6), z) class VideoSink: def __init__(self, sim: Sim, path: Path, fps: int, W=640, H=360): import imageio.v2 as imageio m, d = sim.plant.model, sim.plant.data self.sim = sim self.r_room = mujoco.Renderer(m, H, W) self.r_grip = mujoco.Renderer(m, H, W) self.cam = mujoco.MjvCamera() self.cam.distance, self.cam.azimuth, self.cam.elevation = 2.6, 225.0, -28.0 self.cam.lookat[:] = [0, 0, 0.6] self.gcam = m.camera("gripper_cam").id self.proj = Camera(m, d, self.gcam, W, H) self.opt = mujoco.MjvOption() self.opt.geomgroup[2] = 0 # hide the walls in the room view self.gopt = mujoco.MjvOption() # gripper camera: what the real camera would see for o in (self.opt, self.gopt): o.sitegroup[:] = 0 # sites (the grasp point marker) are bookkeeping, not objects o.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 target = 0.5 * (f.nominal[0] + f.nominal[-1]) self.cam.lookat[:] = 0.9 * self.cam.lookat + 0.1 * target 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() self.r_grip.update_scene(d, camera=self.gcam, scene_option=self.gopt) grip = overlay_camera(self.r_grip.render(), self.proj, f.nominal, f.samples, hud(f)) self.writer.append_data(np.concatenate([room, grip], axis=1)) def close(self): self.writer.close() self.r_room.close() self.r_grip.close() class ViewerSink: def __init__(self, sim: Sim, viewer, realtime: bool): self.sim, self.v, self.realtime = sim, viewer, realtime self.t_wall0 = self.t_sim0 = None def __call__(self, f): v = self.v if not v.is_running(): raise KeyboardInterrupt with v.lock(): draw_scene(v.user_scn, f.nominal, f.samples, f.extra["target"], f.phase) v.set_texts([(mujoco.mjtFontScale.mjFONTSCALE_150, mujoco.mjtGridPos.mjGRID_TOPLEFT, "\n".join(hud(f)), None)]) v.sync() if self.realtime: now = time.perf_counter() if self.t_wall0 is None or f.time_s < self.t_sim0: self.t_wall0, self.t_sim0 = now, f.time_s ahead = (f.time_s - self.t_sim0) - (now - self.t_wall0) if ahead > 0: time.sleep(ahead) def main(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) 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/pick_place.yaml")) ap.add_argument("--object", default=None, help="only this object") ap.add_argument("--container", default=None, help="only this container") ap.add_argument("--viewer", action="store_true") ap.add_argument("--fast", action="store_true") ap.add_argument("--video", action="store_true") ap.add_argument("--out", default=None) ap.add_argument("--duration", type=float, 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() sim = Sim(cfg, planner, task_cfg) stamp = datetime.now().strftime("%Y%m%d-%H%M%S") out = Path(args.out or ROOT / "runs" / f"{stamp}_pick_place_oracle") out.mkdir(parents=True, exist_ok=True) (out / "config.yaml").write_text(yaml.safe_dump({"args": vars(args), "planner": planner, "task": task_cfg})) pairs = [p for p in task_cfg["episode"]["pairs"] if (args.object in (None, p[0])) and (args.container in (None, p[1]))] or [[args.object, args.container]] viewer = None if args.viewer: import mujoco.viewer viewer = mujoco.viewer.launch_passive(sim.plant.model, sim.plant.data, show_left_ui=False, show_right_ui=False) meter = ComputeMeter(budget_s=sim.pcfg.dt, log_path=out / "compute.jsonl") all_metrics = [] try: for ep in range(args.episodes): obj, box = pairs[ep % len(pairs)] sink = None if viewer is not None: sink = ViewerSink(sim, viewer, realtime=not args.fast) elif args.video: sink = VideoSink(sim, out / f"episode_{ep:02d}_{obj}_{box}.mp4", fps=round(3 / sim.pcfg.dt)) t0 = time.perf_counter() log, mt = run_episode(sim, seed=args.seed + 1000 * ep, obj=obj, container=box, on_frame=sink, meter=meter) if isinstance(sink, VideoSink): sink.close() save_episode(log, out / f"episode_{ep:02d}.npz") all_metrics.append(mt) events = " | ".join(f"{e['time_s']:.1f}s {e['why']}" for e in log["events"]) print(f"episode {ep:02d} {obj:6s}->{box:3s}: success {mt['success']:.0f} grasp {mt['time_to_grasp_s']:.1f}s " f"place {mt['time_to_place_s']:.1f}s closes {mt['close_attempts']} plan {mt['plan_ms_mean']:.0f} ms " f"tilt max {mt['max_tilt_deg']:.1f} T max {mt['max_tension_N']:.0f} N ({time.perf_counter() - t0:.0f} s wall)" f"\n [{events}]", flush=True) except KeyboardInterrupt: print("stopped") finally: if viewer is not None: viewer.close() if all_metrics: header = f"task: pick_place model: oracle (MuJoCo) episodes: {len(all_metrics)} seed: {args.seed}" print(write_summary(out, all_metrics, header)) table = meter.table() (out / "compute.txt").write_text(table + "\n") print(table) print(f"-> {out}") meter.close() if __name__ == "__main__": main()