"""Record a search mission and render a three-view GIF for the README. Panes (left -> right): top-down debug map (world state: crates, marker, path, capsule) first person (the actual frames the model received) chat (the model's messages, reasoning, and tool calls) Usage: python -m testbed.record_search --out search_mission.gif """ from __future__ import annotations import argparse import asyncio import base64 import io import sys from pathlib import Path from aicc.client import AICCClient from aicc.transport.websocket import WebSocketClientTransport, WebSocketServer from PIL import Image, ImageDraw, ImageFont from testbed.bridge import build_bridge from testbed.llm_agent import ( SEARCH_MISSION, LLMController, resolve_provider, run_llm_agent_loop, ) PANEL = 400 ROOM = 16.0 MAX_CHAT_LINES = 21 CHAT_LINE_MAX = 46 # chars per chat line ROLE_COLORS = { "agent": (220, 220, 230), "think": (150, 165, 200), "tool": (255, 184, 96), "bridge": (150, 210, 150), "mission": (255, 220, 120), "vision": (130, 210, 250), } def _font(size: int = 13) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: try: return ImageFont.load_default(size=size) except TypeError: # older Pillow return ImageFont.load_default() def _wrap(text: str, width: int = CHAT_LINE_MAX) -> list[str]: words = text.split() lines: list[str] = [] cur = "" for w in words: if len(cur) + len(w) + 1 > width: if cur: lines.append(cur) cur = w else: cur = (cur + " " + w).strip() if cur: lines.append(cur) return lines[:3] # keep the pane readable class ChatPane: """Accumulates the mission chat and renders the last lines.""" def __init__(self) -> None: self.lines: list[tuple[str, str]] = [] def add(self, role: str, msg: str) -> None: for piece in _wrap(msg): self.lines.append((role, piece)) if len(self.lines) > MAX_CHAT_LINES * 2: self.lines = self.lines[-MAX_CHAT_LINES:] def render(self) -> Image.Image: img = Image.new("RGB", (PANEL, PANEL), (12, 12, 18)) draw = ImageDraw.Draw(img) font = _font(13) draw.text((8, 6), "chat — what the model says", fill=(160, 170, 190), font=font) y = 26 for role, piece in self.lines[-MAX_CHAT_LINES:]: color = ROLE_COLORS.get(role, (200, 200, 200)) draw.text((8, y), piece[:CHAT_LINE_MAX], fill=color, font=font) y += 17 return img def _topdown_frame( bridge, path: list[tuple[float, float]], step: int, marker_found: bool ) -> Image.Image: from testbed.room.render import render_topdown img = render_topdown(bridge.world, PANEL) draw = ImageDraw.Draw(img) scale = PANEL / ROOM if len(path) > 1: pts = [(x * scale, PANEL - z * scale) for x, z in path] draw.line(pts, fill=(255, 170, 60), width=3) m = bridge.world.marker_world_pos() if m is not None: mx, mz = m[0] * scale, PANEL - m[2] * scale r = 9 draw.polygon( [(mx, mz - r), (mx - r, mz + r * 0.8), (mx + r, mz + r * 0.8)], outline=(255, 255, 255), width=2, ) status = f"step {step} {'FOUND' if marker_found else 'searching...'}" draw.text((10, PANEL - 24), status, fill=(255, 220, 120), font=_font(13)) return img def _compose(map_img, view_img, chat_img) -> Image.Image: canvas = Image.new("RGB", (PANEL * 3, PANEL), (18, 18, 26)) canvas.paste(map_img, (0, 0)) canvas.paste(view_img, (PANEL, 0)) canvas.paste(chat_img, (PANEL * 2, 0)) draw = ImageDraw.Draw(canvas) draw.text((8, 6), "top-down", fill=(160, 170, 190), font=_font(12)) draw.text((PANEL + 8, 6), "first person", fill=(160, 170, 190), font=_font(12)) draw.text((PANEL * 2 + 8, 6), "chat", fill=(160, 170, 190), font=_font(12)) for x in (PANEL, PANEL * 2): draw.line([(x, 0), (x, PANEL)], fill=(60, 60, 80)) return canvas async def run(args: argparse.Namespace) -> int: base_url, api_key, model = resolve_provider( provider=args.provider, base_url=args.base_url, api_key=args.api_key, model=args.model, ) from testbed.llm_agent import MISSION bridge = build_bridge() if args.goal == "search": # Fixed marker so the GIF tells a clear story (crate_blue, east face). bridge.world.marker = { "box": "crate_blue", "face": "x1", "u": 0.5, "v": 0.55, "size": 0.42, "color": (255, 150, 40), } m = bridge.world.marker_world_pos() print(f"[record] triangle marker on crate_blue east face at {tuple(round(v, 2) for v in m)}") frames: list[Image.Image] = [] chat = ChatPane() marker_seen = False samples: list[tuple[float, float]] = [] # capsule positions per captured frame def log(role: str, msg: str) -> None: print(f" [{role}] {msg}") if role in ("agent", "think", "tool", "mission", "vision"): chat.add(role, msg) async with WebSocketServer(bridge, port=8765): transport = WebSocketClientTransport("ws://127.0.0.1:8765") async with AICCClient(transport) as client: manifest = await client.handshake() ctl = LLMController( client, manifest, base_url=base_url, api_key=api_key, model=model, system_prompt=SEARCH_MISSION if args.goal == "search" else MISSION, log=log, ) if args.goal == "search": ctl.messages = [ ctl.messages[0], {"role": "user", "content": ( "MISSION: explore the room, find the orange triangle on the " "back of a crate, and report it. Keep exploring until found." )}, ] chat.add("mission", "MISSION: find the orange triangle") else: ctl.messages.append({ "role": "user", "content": ( "MISSION: reach the beacon and activate it. " "Keep calling tools until it is done. Move smoothly (use " "duration on move/turn). The capsule cruises while you think." ), }) chat.add("mission", "MISSION: reach and activate the beacon") if args.goal == "search" and not args.capture_hz: # capture at frame-attach events (one frame per step) orig_frame = ctl._frame_message def record_frame(png: str, note: str): nonlocal marker_seen view = ( Image.open(io.BytesIO(base64.b64decode(png))) .convert("RGB") .resize((PANEL, PANEL)) ) top = _topdown_frame(bridge, ctl.path, len(frames), marker_seen) frames.append(_compose(top, view, chat.render())) samples.append((bridge.world.capsule.x, bridge.world.capsule.z)) return orig_frame(png, note) ctl._frame_message = record_frame # type: ignore[method-assign] def is_success(turn): nonlocal marker_seen for c in turn.calls: if c.name == "report" and c.ok and c.output and c.output.get("verified"): marker_seen = True return True return False async def observer(agent_done: asyncio.Event) -> None: """Sample the world at a fixed rate so the GIF shows the actual gliding motion (mid-animation frames), not just step endpoints.""" obs_transport = WebSocketClientTransport("ws://127.0.0.1:8765") async with AICCClient(obs_transport) as obs: await obs.handshake() while not agent_done.is_set(): try: out = (await obs.call_tool("vision", {})).output except Exception: # noqa: BLE001 - session teardown break view = ( Image.open(io.BytesIO(base64.b64decode(out["png_b64"]))) .convert("RGB") .resize((PANEL, PANEL)) ) samples.append((bridge.world.capsule.x, bridge.world.capsule.z)) top = _topdown_frame(bridge, ctl.path, len(frames), marker_seen) frames.append(_compose(top, view, chat.render())) if len(frames) > 350: # keep the GIF bounded frames[:150] = [] await asyncio.sleep(1.0 / args.capture_hz) agent_done = asyncio.Event() obs_task = None if args.capture_hz: obs_task = asyncio.create_task(observer(agent_done)) try: summary = await run_llm_agent_loop( ctl, args.steps, log=log, look_every=2 if args.goal == "beacon" else 1, cruise=args.cruise, autonomous=args.goal == "search", is_success=is_success if args.goal == "search" else None, ) finally: agent_done.set() if obs_task is not None: await obs_task # Final frame: end state of the map. if marker_seen: chat.add("mission", f"DONE: {summary.get('result', 'triangle found')}") top = _topdown_frame(bridge, ctl.path, len(frames), marker_seen) last_view = ( frames[-1].crop((PANEL, 0, PANEL * 2, PANEL)) if frames else Image.new("RGB", (PANEL, PANEL), (0, 0, 0)) ) frames.append(_compose(top, last_view, chat.render())) # Smoothness check: mean movement between consecutive captured frames. if len(samples) > 2: from itertools import pairwise deltas = [ ((b[0] - a[0]) ** 2 + (b[1] - a[1]) ** 2) ** 0.5 for a, b in pairwise(samples) ] moving = [d for d in deltas if d > 0.02] if moving: print(f"[record] motion between frames: mean {sum(moving)/len(moving):.2f} m " f"(max {max(moving):.2f}) over {len(moving)} moving samples") out = Path(args.out) if frames: frames[0].save( out, save_all=True, append_images=frames[1:], duration=args.duration, loop=0, ) frames[-1].save(out.with_suffix(".png")) print(f"[record] wrote {out} ({len(frames)} frames, {args.duration} ms/frame)") print(f"[record] mission result: {summary.get('result')}") return 0 if (marker_seen or summary.get("interacted")) else 1 def main() -> int: parser = argparse.ArgumentParser( description="Record a mission into a three-view GIF (map | first person | chat)." ) parser.add_argument("--out", default="search_mission.gif") parser.add_argument( "--goal", choices=["search", "beacon"], default="search", help="which mission to record (default search)", ) parser.add_argument("--provider", default=None) parser.add_argument("--base-url", default=None) parser.add_argument("--api-key", default=None) parser.add_argument("--model", default=None) parser.add_argument("--steps", type=int, default=45) parser.add_argument("--cruise", type=float, default=0.6) parser.add_argument("--duration", type=int, default=150, help="ms per GIF frame") parser.add_argument( "--capture-hz", type=float, default=0.0, help="sample frames at this rate (Hz) with an observer client to show " "gliding motion; 0 = one frame per step (search default)", ) args = parser.parse_args() try: return asyncio.run(run(args)) except KeyboardInterrupt: print("\n[record] interrupted — the partial GIF was NOT written") return 130 if __name__ == "__main__": sys.exit(main())