480 lines
17 KiB
Python
480 lines
17 KiB
Python
"""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.protocol import SessionInit
|
|
from aicc.transport.websocket import WebSocketClientTransport
|
|
from PIL import Image
|
|
|
|
from testbed.llm_agent import resolve_provider
|
|
|
|
DEFAULT_URL = "ws://127.0.0.1:8765"
|
|
|
|
|
|
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 testbed.room.mapview import draw_sensor_map
|
|
|
|
return draw_sensor_map(self.layout, pos, yaw, self.path, step=self.step)
|
|
|
|
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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# LLM agent (OpenAI-compatible)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def run_llm_agent(
|
|
client: AICCClient,
|
|
manifest: SessionInit,
|
|
*,
|
|
base_url: str,
|
|
api_key: str,
|
|
model: str,
|
|
max_steps: int,
|
|
recorder: FrameRecorder | None = None,
|
|
vision: bool | None = None,
|
|
digest: bool | None = None,
|
|
look_every: int = 0,
|
|
cruise: float = 0.0,
|
|
) -> dict[str, Any]:
|
|
"""Autonomous LLM run: controller + nudge/correct loop (see llm_agent)."""
|
|
from testbed.llm_agent import LLMController, run_llm_agent_loop
|
|
|
|
controller = LLMController(
|
|
client,
|
|
manifest,
|
|
base_url=base_url,
|
|
api_key=api_key,
|
|
model=model,
|
|
log=lambda role, msg: print(f"[{role}] {msg}"),
|
|
multimodal=vision,
|
|
digest=digest,
|
|
)
|
|
try:
|
|
return await run_llm_agent_loop(
|
|
controller, max_steps, log=print_log, recorder=recorder
|
|
)
|
|
except RuntimeError as exc:
|
|
return {"steps": 0, "tool_calls": 0, "result": str(exc), "interacted": False}
|
|
|
|
|
|
def print_log(role: str, msg: str) -> None:
|
|
print(f"[{role}] {msg}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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, "duration": 0.6})).output
|
|
summary["tool_calls"] += 1
|
|
log("tool", 'move({"forward": 0.8, "duration": 0.6}) [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, "duration": 0.6})).output
|
|
summary["tool_calls"] += 1
|
|
log("tool", 'move({"forward": 0.8, "duration": 0.6})')
|
|
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,
|
|
vision=args.vision,
|
|
digest=args.digest,
|
|
look_every=args.look_every,
|
|
cruise=args.cruise,
|
|
)
|
|
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,
|
|
vision=args.vision,
|
|
digest=args.digest,
|
|
look_every=args.look_every,
|
|
cruise=args.cruise,
|
|
)
|
|
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=None,
|
|
help="OpenAI-compatible endpoint (default per --provider)",
|
|
)
|
|
parser.add_argument(
|
|
"--provider",
|
|
choices=["polza", "openai", "ollama"],
|
|
default=None,
|
|
help="provider preset: endpoint + key from env (POLZA_API_KEY/OPENAI_API_KEY)",
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
default=None,
|
|
help="model id on the endpoint (default per provider: polza -> openai/gpt-5.6-luna)",
|
|
)
|
|
parser.add_argument(
|
|
"--api-key",
|
|
default=None,
|
|
help="API key (default: $POLZA_API_KEY / $OPENAI_API_KEY per provider)",
|
|
)
|
|
parser.add_argument("--max-steps", type=int, default=60, help="max agent steps")
|
|
parser.add_argument(
|
|
"--vision",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=None,
|
|
help="pass real camera frames to the model as images (auto-detected for local ollama)",
|
|
)
|
|
parser.add_argument(
|
|
"--digest",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=None,
|
|
help="always include the color-grid digest alongside images (off by default for multimodal models)",
|
|
)
|
|
parser.add_argument(
|
|
"--look-every",
|
|
type=int,
|
|
default=3,
|
|
help="attach a fresh camera frame every N agent steps (0 disables; default 3)",
|
|
)
|
|
parser.add_argument(
|
|
"--cruise",
|
|
type=float,
|
|
default=1.0,
|
|
help="proactive motion: glide forward (meters) while the LLM thinks (0 disables)",
|
|
)
|
|
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()
|
|
try:
|
|
args.base_url, args.api_key, args.model = resolve_provider(
|
|
provider=args.provider,
|
|
base_url=args.base_url,
|
|
api_key=args.api_key,
|
|
model=args.model,
|
|
)
|
|
if args.provider:
|
|
print(f"[demo] provider: {args.provider} -> {args.base_url} model: {args.model}")
|
|
except ValueError as exc:
|
|
print(f"[demo] {exc}")
|
|
return 1
|
|
try:
|
|
summary = asyncio.run(run_demo(args))
|
|
except KeyboardInterrupt:
|
|
return 2
|
|
except Exception as exc: # noqa: BLE001 - friendly failure instead of a traceback
|
|
print(f"\n[demo] failed: {type(exc).__name__}: {exc}")
|
|
print("[demo] is the bridge running? scripts/run_bridge.sh")
|
|
return 1
|
|
return 0 if summary.get("interacted") else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|