"""End-to-end demo: an agent drives the capsule to the beacon and interacts. Two agents: - ``llm`` (default): any tool-calling LLM via an OpenAI-compatible endpoint. Ships configured for a local ollama (http://localhost:11434/v1), swap ``--base-url``/``--model`` for any provider. The agent perceives through the protocol tools only; vision frames are summarized into a coarse color grid so text-only models can navigate too. The loop keeps a compact CURRENT STATE note (agent-side working memory, protocol §9) and nudges the model back on track if it drifts (no tool calls, moving away, collisions). - ``scripted``: deterministic sensor-driven fallback (same tools, no LLM). - ``auto`` (default): tries the LLM for a bounded number of steps, then hands off to the scripted agent so the demo always reaches the beacon. Every tool call and result is printed to stdout as a transcript. The final first-person frame is saved to ``--frame`` (default demo_final_frame.png). Usage: python -m testbed.demo --agent llm --model gemma4:e2b python -m testbed.demo --agent scripted """ from __future__ import annotations import argparse import asyncio import base64 import io import json import math import sys from pathlib import Path from typing import Any from aicc.client import AICCClient from aicc.errors import ToolError from aicc.protocol import SessionInit, ToolSpec from aicc.transport.websocket import WebSocketClientTransport from PIL import Image DEFAULT_URL = "ws://127.0.0.1:8765" DEFAULT_BASE_URL = "http://localhost:11434/v1" DEFAULT_MODEL = "gemma4:e2b" NUDGE = ( "You have not called a tool. Continue the mission: check the CURRENT STATE note, " "then call a tool (turn or move, or interact if within 1.6 m)." ) MISSION = ( "You are a capsule AI inside a 16x16 m room. Your mission: reach the glowing " "beacon pillar and activate it with the interact tool. " "World: yaw 0 faces +z, 90 faces +x, coordinates in meters; three wooden " "crates block parts of the room. " "Perception: proprioception gives position/heading; world_query gives the room " "layout and the beacon's exact coordinates (call it once at the start); " "vision gives the first-person frame; hear tells you how close you are to the beacon. " "Every tool result is prefixed with a CURRENT STATE note: read it — it tells you " "your position, heading, distance to the beacon, and the bearing to turn toward it. " "Standard loop, one action per turn: " "1. call proprioception (and hear every few steps); " "2. compute the bearing to the beacon and turn toward it (turn yaw_deg = that bearing, max 90 per call); " "3. move forward 1.0 m; if the result says collision, turn 60-90 degrees and go around, " "then re-aim. " "Never walk into a crate or wall twice in a row. " "Call interact ONLY when proprioception puts you within 1.6 m of the beacon. " "Keep tool calls short and decisive. Stop once the beacon is active." ) COLOR_NAMES = [ ("beacon_cyan", (120, 210, 235)), ("beacon_yellow", (255, 190, 90)), ("crate_red", (178, 64, 54)), ("crate_blue", (64, 96, 178)), ("crate_olive", (128, 128, 60)), ("floor_gray", (60, 60, 66)), ("wall_gray", (96, 96, 106)), ("dark", (30, 30, 36)), ] def _nearest_color(r: float, g: float, b: float) -> str: best, best_d = "dark", 1e9 for name, (cr, cg, cb) in COLOR_NAMES: d = (r - cr) ** 2 + (g - cg) ** 2 + (b - cb) ** 2 if d < best_d: best, best_d = name, d return best class FrameRecorder: """Visual mode: saves first-person frames and a sensor-built top-down map with the capsule's path, then stitches a GIF and a summary image. Only uses data observed through protocol sensors (world_query, proprioception, vision) — the map is reconstructed the way the agent sees the world, not read from bridge internals. """ MAP_SIZE = 320 ROOM = 16.0 def __init__(self, out_dir: str): self.out_dir = Path(out_dir) self.out_dir.mkdir(parents=True, exist_ok=True) self.layout: dict[str, Any] | None = None self.path: list[tuple[float, float]] = [] self.map_frames: list[Image.Image] = [] self.view_frames: list[Image.Image] = [] self.step = 0 self._log = lambda msg: print(f"[rec] {msg}") async def snap( self, client: AICCClient, pos: tuple[float, float], yaw: float ) -> None: """Capture one visual snapshot (one vision call + sensor layout).""" if self.layout is None: wq = (await client.call_tool("world_query", {})).output self.layout = wq self._log("world layout captured from sensors") self.path.append(pos) view = await self._view_frame(client) self.view_frames.append(view) self.step += 1 map_img = self._map_frame(pos, yaw) self.map_frames.append(map_img) view.save(self.out_dir / f"step_{self.step:03d}_view.png") map_img.save(self.out_dir / f"step_{self.step:03d}_map.png") async def _view_frame(self, client: AICCClient) -> Image.Image: from PIL import Image out = (await client.call_tool("vision", {})).output self._log(f"vision({out['width']}x{out['height']}) frame captured") return Image.open(io.BytesIO(base64.b64decode(out["png_b64"]))).convert("RGB") def _map_frame(self, pos: tuple[float, float], yaw: float) -> Image.Image: from PIL import Image, ImageDraw size = self.MAP_SIZE scale = size / self.ROOM def xy(x: float, z: float) -> tuple[float, float]: return (x * scale, size - z * scale) img = Image.new("RGB", (size, size), (52, 52, 58)) d = ImageDraw.Draw(img) for i in range(int(self.ROOM) + 1): c = 60 if i % 2 == 0 else 54 d.line([xy(i, 0), xy(i, self.ROOM)], fill=(c, c, c + 6), width=1) d.line([xy(0, i), xy(self.ROOM, i)], fill=(c, c, c + 6), width=1) if self.layout: for ob in self.layout.get("obstacles", []): x0, z0 = xy(ob["x"] - ob["width"] / 2, ob["z"] - ob["depth"] / 2) x1, z1 = xy(ob["x"] + ob["width"] / 2, ob["z"] + ob["depth"] / 2) d.rectangle( [min(x0, x1), min(z0, z1), max(x0, x1), max(z0, z1)], fill=(150, 90, 60), outline=(20, 20, 26), width=2, ) b = self.layout.get("beacon", {}) bx, bz = xy(b.get("x", 12.5), b.get("z", 12.5)) d.ellipse( [bx - 8, bz - 8, bx + 8, bz + 8], fill=(120, 210, 235), outline=(20, 20, 26), width=2, ) if len(self.path) > 1: pts = [xy(x, z) for x, z in self.path] d.line(pts, fill=(255, 170, 60), width=3) cx, cz = xy(*pos) d.ellipse( [cx - 7, cz - 7, cx + 7, cz + 7], fill=(240, 240, 250), outline=(20, 20, 26), width=2, ) rad = math.radians(yaw) tip = xy(pos[0] + math.sin(rad) * 0.7, pos[1] + math.cos(rad) * 0.7) d.line([(cx, cz), tip], fill=(30, 30, 40), width=3) d.text( (8, 8), f"step {self.step} pos ({pos[0]:.1f}, {pos[1]:.1f}) yaw {yaw:.0f}", fill=(230, 230, 230), ) return img def finish(self, interacted: bool) -> Image.Image: """Write demo.gif + demo_summary.png; returns the summary image.""" from PIL import Image, ImageDraw if self.map_frames: gif = self.out_dir / "demo.gif" self.map_frames[0].save( gif, save_all=True, append_images=self.map_frames[1:], duration=350, loop=0, ) self._log(f"animation written: {gif}") summary = Image.new("RGB", (self.MAP_SIZE + 320, 240), (20, 20, 26)) if self.map_frames: summary.paste(self.map_frames[-1], (0, 0)) if self.view_frames: summary.paste(self.view_frames[-1].resize((320, 240)), (self.MAP_SIZE, 0)) d = ImageDraw.Draw(summary) status = "BEACON ACTIVATED" if interacted else "mission not completed" d.text( (self.MAP_SIZE + 8, 244 - 16), status, fill=(255, 220, 120) if interacted else (255, 120, 120), ) out = self.out_dir / "demo_summary.png" summary.save(out) self._log(f"summary written: {out}") return summary def vision_digest(png_b64: str, cols: int = 12, rows: int = 9) -> str: """Coarse color-grid summary of a vision frame (agent-side preprocessing).""" try: from PIL import Image img = Image.open(io.BytesIO(base64.b64decode(png_b64))).convert("RGB") except Exception: # noqa: BLE001 - digest is best-effort return "(no digest: unreadable frame)" small = img.resize((cols, rows)) grid: list[list[tuple[int, int, int]]] = [] for r in range(rows): row: list[tuple[int, int, int]] = [] for c in range(cols): px = small.getpixel((c, r)) if isinstance(px, (tuple, list)) and len(px) >= 3: row.append((int(px[0]), int(px[1]), int(px[2]))) elif isinstance(px, (int, float)): v = int(px) row.append((v, v, v)) else: row.append((0, 0, 0)) grid.append(row) lines = [] for r in range(rows): lines.append(" ".join(_nearest_color(*grid[r][c]) for c in range(cols))) return f"frame digest ({cols}x{rows}, center = where you face):\n" + "\n".join( lines ) # --------------------------------------------------------------------------- # LLM agent (OpenAI-compatible) # --------------------------------------------------------------------------- def to_openai_tools(tools: list[ToolSpec]) -> list[dict[str, Any]]: out = [] for t in tools: if t.id in ( "echo", "boom", "bump", ): # conformance-only tools: not for the agent continue out.append( { "type": "function", "function": { "name": t.id, "description": t.description, "parameters": t.input_schema, }, } ) return out async def run_llm_agent( client: AICCClient, manifest: SessionInit, *, base_url: str, api_key: str, model: str, max_steps: int, recorder: FrameRecorder | None = None, ) -> dict[str, Any]: from openai import AsyncOpenAI ac = AsyncOpenAI(base_url=base_url, api_key=api_key) tools = to_openai_tools(manifest.tools) messages: list[dict[str, Any]] = [ {"role": "system", "content": MISSION}, { "role": "user", "content": "Begin your mission. Use the tools to reach and activate the beacon.", }, ] summary: dict[str, Any] = { "steps": 0, "tool_calls": 0, "result": None, "interacted": False, } # Agent-side working memory: a compact summary of state the model has # already observed via sensors, prefixed to tool results so small models # don't drift (protocol §9: agent memory is the agent's concern). beacon_pos: tuple[float, float] | None = None pos: tuple[float, float] | None = None yaw: float = 0.0 def state_hint() -> str: if pos is None: return "" if beacon_pos is not None: dist = math.hypot(beacon_pos[0] - pos[0], beacon_pos[1] - pos[1]) bearing = ( math.degrees(math.atan2(beacon_pos[0] - pos[0], beacon_pos[1] - pos[1])) - yaw ) % 360.0 return ( f"CURRENT STATE: you are at ({pos[0]:.2f}, {pos[1]:.2f}), heading {yaw:.1f} deg; " f"the beacon is {dist:.2f} m away at relative bearing {bearing:.0f} deg. " f"Action: turn yaw_deg={bearing:.0f} to face it, then move forward." ) return f"CURRENT STATE: you are at ({pos[0]:.2f}, {pos[1]:.2f}), heading {yaw:.1f} deg." def observe(output: dict[str, Any]) -> None: nonlocal pos, yaw if isinstance(output.get("position"), dict): pos = (output["position"]["x"], output["position"]["z"]) if isinstance(output.get("rotation"), dict): yaw = output["rotation"]["yaw_deg"] % 360.0 def log(role: str, msg: str) -> None: print(f"[{role}] {msg}") # Progress monitor: if the capsule keeps moving away from the beacon, # reflect that back to the model as a correction (framework-style loop). dist_history: list[float] = [] corrections = 0 last_correct_step = -99 last_collision_step = -99 def distance_to_beacon() -> float: if pos is None or beacon_pos is None: return float("nan") return math.hypot(beacon_pos[0] - pos[0], beacon_pos[1] - pos[1]) def maybe_correct(step: int) -> None: nonlocal corrections, last_correct_step if pos is None or beacon_pos is None: return dist = distance_to_beacon() if not math.isfinite(dist) or dist < 2.0: return dist_history.append(dist) recent = dist_history[-4:] moving_away = len(recent) == 4 and all( recent[i] > recent[i + 1] for i in range(3) ) bearing = ( math.degrees(math.atan2(beacon_pos[0] - pos[0], beacon_pos[1] - pos[1])) - yaw ) % 360.0 deviation = min(bearing, 360.0 - bearing) facing_wrong = dist > 4.0 and deviation > 40.0 and step - last_correct_step > 3 detouring = step - last_collision_step < 3 if (moving_away or (facing_wrong and not detouring)) and corrections < 5: corrections += 1 last_correct_step = step hint = state_hint() msg = ( "You are not making progress toward the beacon. Stop and follow the CURRENT " f"STATE note exactly: {hint} Call turn with that yaw_deg, then move forward." ) messages.append({"role": "user", "content": msg}) log("agent", "(correction: re-aim at the beacon)") dist_history.clear() for step in range(max_steps): summary["steps"] = step + 1 if recorder is not None and pos is not None: await recorder.snap(client, pos, yaw) try: resp = await ac.chat.completions.create( model=model, messages=messages, tools=tools, tool_choice="auto" ) except Exception as exc: # noqa: BLE001 - provider errors end the demo log("agent", f"LLM error: {exc}") summary["result"] = f"LLM error: {exc}" return summary choice = resp.choices[0] if choice.message.content: log("agent", choice.message.content[:400]) messages.append({"role": "assistant", "content": choice.message.content}) calls = choice.message.tool_calls if not calls: # Nudge a model that drifted into plain text back to acting. silent = (not choice.message.content) or len(choice.message.content) < 8 already_nudged = bool(messages) and messages[-1].get("content") == NUDGE if silent or already_nudged: summary["result"] = "model produced no tool call" return summary messages.append({"role": "user", "content": NUDGE}) log("agent", "(nudge: no tool call; continue the mission)") continue messages.append( { "role": "assistant", "content": choice.message.content or "", "tool_calls": [ { "id": c.id, "type": "function", "function": { "name": c.function.name, "arguments": c.function.arguments, }, } for c in calls ], } ) for tc in calls: name = tc.function.name try: args = json.loads(tc.function.arguments or "{}") except json.JSONDecodeError: args = {} log("tool", f"{name}({json.dumps(args)})") try: result = await client.call_tool(name, args) output = result.output extra = "" if name == "vision" and isinstance(output.get("png_b64"), str): extra = "\n" + vision_digest(output["png_b64"]) output = dict(output, png_b64="") log("bridge", f"ok {json.dumps(output)[:500]}{extra}") messages.append( { "role": "tool", "tool_call_id": tc.id, "content": json.dumps(result.output) + extra, } ) summary["tool_calls"] += 1 if name == "world_query": b = output.get("beacon") if isinstance(b, dict): beacon_pos = (b["x"], b["z"]) observe(output) hint = state_hint() if hint: messages[-1] = dict( messages[-1], content=hint + "\n" + messages[-1]["content"] ) maybe_correct(step) if name == "move" and output.get("collision"): last_collision_step = step obj = output.get("collision_normal") hint_msg = ( f"You collided with an obstacle (normal {obj}). Turn yaw_deg=90 and move " "forward twice to get around it, then follow the CURRENT STATE note to re-aim." ) if not (messages and messages[-1].get("content") == hint_msg): messages.append({"role": "user", "content": hint_msg}) log("agent", "(collision: go around)") if name == "interact" and output.get("success"): summary["interacted"] = True summary["result"] = output.get("message", "beacon activated") return summary except ToolError as e: log("bridge", f"error {e.code}: {e.message}") messages.append( { "role": "tool", "tool_call_id": tc.id, "content": json.dumps({"error": e.code, "message": e.message}), } ) summary["tool_calls"] += 1 summary["result"] = f"exceeded {max_steps} steps" return summary # --------------------------------------------------------------------------- # Scripted agent (sensor-driven fallback, no LLM) # --------------------------------------------------------------------------- async def run_scripted_agent( client: AICCClient, manifest: SessionInit, max_steps: int, recorder: FrameRecorder | None = None, ) -> dict[str, Any]: summary: dict[str, Any] = { "steps": 0, "tool_calls": 0, "result": None, "interacted": False, } def log(role: str, msg: str) -> None: print(f"[{role}] {msg}") wq = (await client.call_tool("world_query", {})).output beacon = wq["beacon"] bx, bz = beacon["x"], beacon["z"] log("agent", f"mission: reach beacon at ({bx}, {bz})") consecutive_collisions = 0 detour_steps = 0 # >0: escaping an obstacle before re-aiming at the beacon detour_turns = 0 for step in range(max_steps): summary["steps"] = step + 1 prop = (await client.call_tool("proprioception", {})).output summary["tool_calls"] += 1 pos = prop["position"] if recorder is not None: await recorder.snap( client, (pos["x"], pos["z"]), prop["rotation"]["yaw_deg"] ) dist = math.hypot(bx - pos["x"], bz - pos["z"]) log( "agent", f"at ({pos['x']:.2f}, {pos['z']:.2f}) heading {prop['rotation']['yaw_deg']:.1f} deg, {dist:.2f} m to beacon", ) if dist <= 1.5: la = (await client.call_tool("look_at", {"target": "beacon"})).output summary["tool_calls"] += 1 log("tool", 'look_at({"target": "beacon"})') log("bridge", f"ok {json.dumps(la)}") frame = (await client.call_tool("vision", {})).output summary["tool_calls"] += 1 log("tool", "vision({})") log("bridge", f"ok {frame['width']}x{frame['height']} frame captured") res = (await client.call_tool("interact", {})).output summary["tool_calls"] += 1 log("tool", 'interact({"target": "beacon"})') log("bridge", f"ok {json.dumps(res)}") if res["success"]: summary["interacted"] = True summary["result"] = res["message"] return summary log("agent", "not in reach yet; continue") if detour_steps > 0: # Escape the obstacle before re-aiming: keep the detour heading. mv = (await client.call_tool("move", {"forward": 0.8})).output summary["tool_calls"] += 1 log("tool", 'move({"forward": 0.8}) [detour]') log("bridge", f"ok {json.dumps(mv)}") if mv.get("moved", 0.0) < 0.2 and detour_turns < 6: detour_turns += 1 t = (await client.call_tool("turn", {"yaw_deg": 45.0})).output summary["tool_calls"] += 1 log("tool", 'turn({"yaw_deg": 45.0}) [detour]') log("bridge", f"ok {json.dumps(t)}") else: detour_steps -= 1 continue target_yaw = math.degrees(math.atan2(bx - pos["x"], bz - pos["z"])) % 360.0 cur_yaw = prop["rotation"]["yaw_deg"] % 360.0 delta = (target_yaw - cur_yaw + 540.0) % 360.0 - 180.0 # Keep off the walls so the capsule can round corners. if pos["x"] < 1.2: delta = min(delta, -60.0) elif pos["x"] > 15.0: delta = max(delta, 60.0) if pos["z"] < 1.2: delta = max(delta, 60.0) elif pos["z"] > 15.0: delta = min(delta, -60.0) if abs(delta) > 4.0: turn = (await client.call_tool("turn", {"yaw_deg": delta * 0.8})).output summary["tool_calls"] += 1 log("tool", f'turn({{"yaw_deg": {delta * 0.8:.1f}}})') log("bridge", f"ok {json.dumps(turn)}") else: mv = (await client.call_tool("move", {"forward": 0.8})).output summary["tool_calls"] += 1 log("tool", 'move({"forward": 0.8})') log("bridge", f"ok {json.dumps(mv)}") if mv.get("collision"): consecutive_collisions += 1 angle = 40.0 if consecutive_collisions < 4 else 90.0 t = (await client.call_tool("turn", {"yaw_deg": angle})).output summary["tool_calls"] += 1 log("tool", f'turn({{"yaw_deg": {angle}}})') log("bridge", f"ok {json.dumps(t)}") detour_steps = 3 detour_turns = 0 else: consecutive_collisions = 0 summary["result"] = f"exceeded {max_steps} steps" return summary # --------------------------------------------------------------------------- # Shared runner # --------------------------------------------------------------------------- async def run_demo(args: argparse.Namespace) -> dict[str, Any]: transport = WebSocketClientTransport(args.url) async with AICCClient(transport) as client: manifest = await client.handshake() print(f"[handshake] session {manifest.session_id} world {manifest.world.name}") print(f"[handshake] tools: {[t.id for t in manifest.tools]}\n") if args.frames_dir: recorder = FrameRecorder(args.frames_dir) else: recorder = None if args.agent == "llm": summary = await run_llm_agent( client, manifest, base_url=args.base_url, api_key=args.api_key, model=args.model, max_steps=args.max_steps, recorder=recorder, ) elif args.agent == "scripted": summary = await run_scripted_agent( client, manifest, args.max_steps, recorder=recorder ) else: # auto from openai import AsyncOpenAI probe = AsyncOpenAI(base_url=args.base_url, api_key=args.api_key) try: await asyncio.wait_for(probe.models.list(), timeout=5.0) llm_steps = min(args.max_steps, 15) summary = await run_llm_agent( client, manifest, base_url=args.base_url, api_key=args.api_key, model=args.model, max_steps=llm_steps, recorder=recorder, ) if summary.get("interacted"): return summary print( f"\n[auto] LLM did not finish in {llm_steps} steps " f"({summary.get('result')}); handing off to the scripted agent\n" ) summary = await run_scripted_agent( client, manifest, args.max_steps, recorder=recorder ) except Exception as exc: # noqa: BLE001 - fall back to scripted print( f"\n[auto] LLM endpoint unavailable ({type(exc).__name__}: {exc}); falling back to scripted agent" ) summary = await run_scripted_agent( client, manifest, args.max_steps, recorder=recorder ) if recorder is not None: recorder.finish(summary.get("interacted", False)) frame = (await client.call_tool("vision", {})).output if args.frame: raw = base64.b64decode(frame["png_b64"]) def _write_frame() -> None: with open(args.frame, "wb") as fh: fh.write(raw) await asyncio.to_thread(_write_frame) print(f"\n[final frame saved] {args.frame}") print( f"\n[demo done] steps={summary['steps']} tool_calls={summary['tool_calls']} " f"interacted={summary['interacted']} result={summary['result']}" ) return summary def main() -> int: parser = argparse.ArgumentParser( description="AICC capsule demo: LLM or scripted agent drives the capsule to the beacon." ) parser.add_argument( "--url", default=DEFAULT_URL, help=f"bridge WebSocket URL (default {DEFAULT_URL})", ) parser.add_argument( "--agent", choices=["llm", "scripted", "auto"], default="auto", help="agent driver (default auto: LLM if reachable, else scripted)", ) parser.add_argument( "--base-url", default=DEFAULT_BASE_URL, help="OpenAI-compatible endpoint" ) parser.add_argument( "--model", default=DEFAULT_MODEL, help="model id on the endpoint" ) parser.add_argument("--api-key", default="ollama", help="API key for the endpoint") parser.add_argument("--max-steps", type=int, default=60, help="max agent steps") parser.add_argument( "--frame", default="demo_final_frame.png", help="where to save the final vision frame", ) parser.add_argument( "--frames-dir", default=None, help="visual mode: save first-person + top-down map frames each step " "into this directory, plus demo.gif animation and demo_summary.png", ) args = parser.parse_args() summary = asyncio.run(run_demo(args)) return 0 if summary.get("interacted") else 1 if __name__ == "__main__": sys.exit(main())