| """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 |
| from cablepilot.runner import Sim, run_episode, save_episode, write_summary |
| from cablepilot.sim.scene import load_config |
| from cablepilot.viz import draw_scene, overlay_camera |
|
|
|
|
| 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 |
| self.gopt = mujoco.MjvOption() |
| for o in (self.opt, self.gopt): |
| o.sitegroup[:] = 0 |
| 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() |
|
|
