519 lines
17 KiB
Python
519 lines
17 KiB
Python
"""The testbed bridge: connects the AICC protocol to the room world.
|
|
|
|
Registers the world tools (proprioception, vision, depth, hear, world_query,
|
|
move, turn, look_at, interact) plus the conformance tools the scenario suite
|
|
expects (echo, boom, bump). Subclasses ``aicc.Bridge`` only to learn which
|
|
session is handling the current tool call, so tools can push async events
|
|
(collision, audio) via ``emit_event``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import copy
|
|
import io
|
|
import math
|
|
from typing import Any
|
|
|
|
from aicc.bridge import Bridge
|
|
from aicc.protocol import TickMode, ToolClass
|
|
from pydantic import BaseModel
|
|
|
|
from testbed.room.render import Raycaster
|
|
from testbed.room.world import Room
|
|
|
|
VISION_WIDTH = 160
|
|
VISION_HEIGHT = 120
|
|
DEPTH_WIDTH = VISION_WIDTH // 4
|
|
DEPTH_HEIGHT = VISION_HEIGHT // 4
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Typed tool outputs (advertised as output_schema in the manifest)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class Vec3(BaseModel):
|
|
x: float
|
|
y: float
|
|
z: float
|
|
|
|
|
|
class Rotation(BaseModel):
|
|
yaw_deg: float
|
|
pitch_deg: float
|
|
|
|
|
|
class ProprioceptionOutput(BaseModel):
|
|
position: Vec3
|
|
rotation: Rotation
|
|
velocity: Vec3
|
|
health: float
|
|
tick: int
|
|
|
|
|
|
class VisionOutput(BaseModel):
|
|
png_b64: str
|
|
width: int
|
|
height: int
|
|
tick: int
|
|
|
|
|
|
class DepthOutput(BaseModel):
|
|
width: int
|
|
height: int
|
|
max_depth: float
|
|
depth: list[list[float]]
|
|
tick: int
|
|
|
|
|
|
class Sound(BaseModel):
|
|
kind: str
|
|
direction_deg: float
|
|
intensity: float
|
|
|
|
|
|
class HearOutput(BaseModel):
|
|
sounds: list[Sound]
|
|
tick: int
|
|
|
|
|
|
class MoveOutput(BaseModel):
|
|
moved: float
|
|
collision: bool
|
|
collision_normal: Vec3 | None
|
|
position: Vec3
|
|
tick: int
|
|
duration: float = 0.0
|
|
|
|
|
|
class TurnOutput(BaseModel):
|
|
rotation: Rotation
|
|
tick: int
|
|
|
|
|
|
class LookAtOutput(BaseModel):
|
|
target: str
|
|
rotation: Rotation
|
|
tick: int
|
|
|
|
|
|
class InteractOutput(BaseModel):
|
|
target: str
|
|
success: bool
|
|
message: str
|
|
tick: int
|
|
|
|
|
|
class WorldQueryOutput(BaseModel):
|
|
room: dict[str, Any]
|
|
obstacles: list[dict[str, Any]]
|
|
beacon: dict[str, Any]
|
|
tick: int
|
|
|
|
|
|
class ReportOutput(BaseModel):
|
|
discovery: str
|
|
success: bool
|
|
verified: bool
|
|
distance: float
|
|
message: str
|
|
tick: int
|
|
|
|
|
|
class EchoOutput(BaseModel):
|
|
value: str
|
|
|
|
|
|
class BumpOutput(BaseModel):
|
|
bumped: bool
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bridge
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class RoomBridge(Bridge):
|
|
"""AICC bridge over one Room. Event-driven ticks: the world advances
|
|
only on tool calls (protocol tick_mode = "event")."""
|
|
|
|
def __init__(
|
|
self,
|
|
world: Room | None = None,
|
|
name: str = "testbed_room_01",
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.world = world or Room()
|
|
self.renderer = Raycaster(VISION_WIDTH, VISION_HEIGHT)
|
|
self._session_id: str | None = None
|
|
kwargs.setdefault("tick_mode", TickMode.EVENT)
|
|
kwargs.setdefault("tick_rate_hz", 10.0)
|
|
super().__init__(name=name, kind="3d", **kwargs)
|
|
self._register_tools()
|
|
|
|
# -- session plumbing for event emission from tool handlers ------------
|
|
|
|
async def _handle_tool_call(self, session_id: str, call) -> dict[str, Any]:
|
|
self._session_id = session_id
|
|
return await super()._handle_tool_call(session_id, call)
|
|
|
|
def emit(
|
|
self,
|
|
topic: str,
|
|
payload: dict[str, Any] | None = None,
|
|
*,
|
|
tick: int | None = None,
|
|
) -> None:
|
|
"""Push an async event to the session currently handling a tool call."""
|
|
if self._session_id is not None:
|
|
self.emit_event(self._session_id, topic, payload, tick=tick)
|
|
|
|
# -- world helpers ------------------------------------------------------
|
|
|
|
def _tick(self) -> int:
|
|
return self.world.advance_tick()
|
|
|
|
async def _animate_turn(
|
|
self, yaw_deg: float, pitch_deg: float, duration: float
|
|
) -> None:
|
|
"""Rotate smoothly over ``duration`` seconds (observable mid-turn)."""
|
|
n = max(1, min(int(duration / 0.1), 100))
|
|
dt = duration / n
|
|
for i in range(n):
|
|
self.world.turn(yaw_deg / n, pitch_deg / n)
|
|
if i < n - 1:
|
|
await asyncio.sleep(dt)
|
|
|
|
async def _animate_move(
|
|
self, forward: float, duration: float
|
|
) -> tuple[float, Any | None]:
|
|
"""Smoothly displace the capsule over ``duration`` seconds.
|
|
|
|
The position advances in small timed steps, so concurrent clients
|
|
(e.g. the live viewer) observe genuine gliding motion instead of a
|
|
teleport. Collision resolution is unchanged.
|
|
"""
|
|
world = self.world
|
|
n = max(1, min(int(duration / 0.1), 200))
|
|
step = forward / n
|
|
dt = duration / n
|
|
total = 0.0
|
|
hit: Any | None = None
|
|
for i in range(n):
|
|
moved, h = world.move_forward(step)
|
|
total += moved
|
|
if h is not None or moved < step * 0.01:
|
|
hit = h
|
|
break
|
|
if i < n - 1:
|
|
await asyncio.sleep(dt)
|
|
return total, hit
|
|
|
|
def _vec3(self, x: float, y: float, z: float) -> Vec3:
|
|
return Vec3(x=round(x, 3), y=round(y, 3), z=round(z, 3))
|
|
|
|
def _rotation(self) -> Rotation:
|
|
c = self.world.capsule
|
|
return Rotation(
|
|
yaw_deg=round(c.yaw_deg % 360.0, 2), pitch_deg=round(c.pitch_deg, 2)
|
|
)
|
|
|
|
def _bearing_to(self, dx: float, dz: float) -> float:
|
|
"""Compass bearing of (dx, dz) relative to the capsule's heading."""
|
|
abs_bearing = math.degrees(math.atan2(dx, dz))
|
|
return (abs_bearing - self.world.capsule.yaw_deg) % 360.0
|
|
|
|
# -- tool registration --------------------------------------------------
|
|
|
|
def _register_tools(self) -> None:
|
|
world = self.world
|
|
renderer = self.renderer
|
|
|
|
@self.tool(
|
|
cls=ToolClass.SENSOR,
|
|
description="Return the capsule's own state: position, rotation, velocity, health.",
|
|
)
|
|
async def proprioception() -> ProprioceptionOutput:
|
|
tick = self._tick()
|
|
c = world.capsule
|
|
return ProprioceptionOutput(
|
|
position=self._vec3(c.x, 0.0, c.z),
|
|
rotation=self._rotation(),
|
|
velocity=self._vec3(0.0, 0.0, c.speed),
|
|
health=round(c.health, 1),
|
|
tick=tick,
|
|
).model_dump()
|
|
|
|
@self.tool(
|
|
cls=ToolClass.SENSOR,
|
|
description="Return the current first-person RGB frame as base64 PNG (160x120).",
|
|
limits={"calls_per_minute": 60},
|
|
)
|
|
async def vision() -> VisionOutput:
|
|
tick = self._tick()
|
|
# Render off the event loop so frames never stall the bridge under
|
|
# load; deepcopy keeps the snapshot consistent with this tool call.
|
|
snapshot = copy.deepcopy(world)
|
|
img, _ = await asyncio.to_thread(renderer.render, snapshot)
|
|
|
|
def _encode() -> str:
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG")
|
|
return base64.b64encode(buf.getvalue()).decode("ascii")
|
|
|
|
return VisionOutput(
|
|
png_b64=await asyncio.to_thread(_encode),
|
|
width=VISION_WIDTH,
|
|
height=VISION_HEIGHT,
|
|
tick=tick,
|
|
).model_dump()
|
|
|
|
@self.tool(
|
|
cls=ToolClass.SENSOR,
|
|
description="Return a depth map (meters, 40x30) aligned to the vision frame.",
|
|
)
|
|
async def depth(max_depth: float = 10.0) -> DepthOutput:
|
|
tick = self._tick()
|
|
snapshot = copy.deepcopy(world)
|
|
_, grid = await asyncio.to_thread(renderer.render, snapshot, max_depth)
|
|
return DepthOutput(
|
|
width=DEPTH_WIDTH,
|
|
height=DEPTH_HEIGHT,
|
|
max_depth=round(max_depth, 2),
|
|
depth=grid,
|
|
tick=tick,
|
|
).model_dump()
|
|
|
|
@self.tool(
|
|
cls=ToolClass.SENSOR,
|
|
description="Return audio events since the last call: beacon hum when near, collision thuds.",
|
|
)
|
|
async def hear() -> HearOutput:
|
|
tick = self._tick()
|
|
sounds = world.hear_now()
|
|
return HearOutput(
|
|
sounds=[Sound(**s) for s in sounds],
|
|
tick=tick,
|
|
).model_dump()
|
|
|
|
@self.tool(
|
|
cls=ToolClass.SENSOR,
|
|
description="Return static room layout: bounds, obstacle positions, beacon location.",
|
|
)
|
|
async def world_query() -> WorldQueryOutput:
|
|
tick = self._tick()
|
|
d = world.describe()
|
|
return WorldQueryOutput(**d, tick=tick).model_dump()
|
|
|
|
@self.tool(
|
|
cls=ToolClass.ACTUATOR,
|
|
description=(
|
|
"Move the capsule forward along its heading by the given distance in meters. "
|
|
"duration (seconds) animates the motion smoothly instead of teleporting."
|
|
),
|
|
)
|
|
async def move(forward: float = 1.0, duration: float = 0.0) -> MoveOutput:
|
|
tick = self._tick()
|
|
if forward < 0.0 or forward > 5.0:
|
|
raise ValueError(f"forward must be in [0, 5] m, got {forward}")
|
|
if duration < 0.0 or duration > 10.0:
|
|
raise ValueError(f"duration must be in [0, 10] s, got {duration}")
|
|
if duration > 0.0:
|
|
moved, hit = await self._animate_move(forward, duration)
|
|
else:
|
|
moved, hit = world.move_forward(forward)
|
|
c = world.capsule
|
|
normal = None
|
|
if hit is not None:
|
|
normal = self._vec3(hit.normal_x, 0.0, hit.normal_z)
|
|
self.emit(
|
|
"collision",
|
|
{
|
|
"other": hit.other,
|
|
"normal": normal.model_dump(),
|
|
"impulse": round(hit.impulse, 3),
|
|
},
|
|
tick=tick,
|
|
)
|
|
world.queue_audio(
|
|
"thud",
|
|
self._bearing_to(-hit.normal_x, -hit.normal_z),
|
|
min(1.0, 0.4 + hit.impulse * 0.15),
|
|
)
|
|
return MoveOutput(
|
|
moved=round(moved, 3),
|
|
collision=hit is not None,
|
|
collision_normal=normal,
|
|
position=self._vec3(c.x, 0.0, c.z),
|
|
tick=tick,
|
|
duration=round(duration, 3),
|
|
).model_dump()
|
|
|
|
@self.tool(
|
|
cls=ToolClass.ACTUATOR,
|
|
description=(
|
|
"Rotate the camera: yaw_deg turns left/right, pitch_deg tilts up/down. "
|
|
"duration (seconds) animates the rotation smoothly."
|
|
),
|
|
)
|
|
async def turn(
|
|
yaw_deg: float = 0.0, pitch_deg: float = 0.0, duration: float = 0.0
|
|
) -> TurnOutput:
|
|
tick = self._tick()
|
|
if duration < 0.0 or duration > 10.0:
|
|
raise ValueError(f"duration must be in [0, 10] s, got {duration}")
|
|
if duration > 0.0:
|
|
await self._animate_turn(yaw_deg, pitch_deg, duration)
|
|
else:
|
|
world.turn(yaw_deg, pitch_deg)
|
|
return TurnOutput(rotation=self._rotation(), tick=tick).model_dump()
|
|
|
|
@self.tool(
|
|
cls=ToolClass.ACTUATOR,
|
|
description=(
|
|
"Orient the camera toward a named target: 'beacon' or a crate "
|
|
"(crate_red, crate_blue, crate_olive)."
|
|
),
|
|
)
|
|
async def look_at(target: str) -> LookAtOutput:
|
|
tick = self._tick()
|
|
box = next((b for b in world.boxes if b.id == target), None)
|
|
if target == "beacon":
|
|
yaw, pitch = world.look_at_beacon()
|
|
elif box is not None:
|
|
dx = box.cx - world.capsule.x
|
|
dz = box.cz - world.capsule.z
|
|
dist = math.hypot(dx, dz) or 1.0
|
|
yaw = math.degrees(math.atan2(dx, dz))
|
|
pitch = math.degrees(
|
|
math.atan2(box.height / 2 - world.capsule.eye_height, dist)
|
|
)
|
|
world.face(yaw, pitch)
|
|
else:
|
|
raise ValueError(
|
|
f"unknown target {target!r}; available: beacon, crate_red, crate_blue, crate_olive"
|
|
)
|
|
return LookAtOutput(
|
|
target=target,
|
|
rotation=Rotation(
|
|
yaw_deg=round(yaw % 360.0, 2), pitch_deg=round(pitch, 2)
|
|
),
|
|
tick=tick,
|
|
).model_dump()
|
|
|
|
@self.tool(
|
|
cls=ToolClass.ACTUATOR,
|
|
description="Use the named interactable object (e.g. 'beacon'). Must be within reach.",
|
|
)
|
|
async def interact(target: str = "beacon") -> InteractOutput:
|
|
tick = self._tick()
|
|
if target != "beacon":
|
|
raise ValueError(
|
|
f"unknown target {target!r}; available targets: beacon"
|
|
)
|
|
success, message = world.interact_beacon()
|
|
if success and world.beacon.active:
|
|
world.queue_audio(
|
|
"beacon_chime",
|
|
self._bearing_to(
|
|
world.beacon.x - world.capsule.x,
|
|
world.beacon.z - world.capsule.z,
|
|
),
|
|
1.0,
|
|
)
|
|
return InteractOutput(
|
|
target=target, success=success, message=message, tick=tick
|
|
).model_dump()
|
|
|
|
@self.tool(
|
|
cls=ToolClass.ACTUATOR,
|
|
description=(
|
|
"Report a discovery to the operator (used in search missions): "
|
|
"say what you found. The report is verified against your position."
|
|
),
|
|
)
|
|
async def report(discovery: str = "triangle") -> ReportOutput:
|
|
tick = self._tick()
|
|
dist = world.distance_to_marker()
|
|
verified = world.marker is not None and dist <= 6.0
|
|
if verified:
|
|
message = f"report verified: {discovery} sighted {dist:.1f} m away"
|
|
else:
|
|
message = (
|
|
f"not verified: nothing matching {discovery!r} within 6 m "
|
|
f"(closest check {dist:.1f} m) — keep exploring"
|
|
)
|
|
return ReportOutput(
|
|
discovery=discovery,
|
|
success=True,
|
|
verified=verified,
|
|
distance=round(dist, 2),
|
|
message=message,
|
|
tick=tick,
|
|
).model_dump()
|
|
|
|
# -- conformance tools (design doc: register alongside world tools) --
|
|
|
|
@self.tool(
|
|
cls=ToolClass.SENSOR,
|
|
description="Echo a value back to the caller.",
|
|
)
|
|
async def echo(value: str = "") -> EchoOutput:
|
|
return EchoOutput(value=value).model_dump()
|
|
|
|
@self.tool(
|
|
cls=ToolClass.ACTUATOR,
|
|
description="Always raises an exception.",
|
|
)
|
|
async def boom() -> dict[str, Any]:
|
|
raise RuntimeError("kaboom")
|
|
|
|
@self.tool(
|
|
cls=ToolClass.ACTUATOR,
|
|
description="Returns ok and emits a collision event.",
|
|
)
|
|
async def bump() -> BumpOutput:
|
|
tick = self._tick()
|
|
self.emit(
|
|
"collision",
|
|
{
|
|
"other": "wall",
|
|
"normal": {"x": 0.0, "y": 0.0, "z": 1.0},
|
|
"impulse": 1.0,
|
|
},
|
|
tick=tick,
|
|
)
|
|
return BumpOutput(bumped=True).model_dump()
|
|
|
|
|
|
def build_bridge(world: Room | None = None, **kwargs: Any) -> RoomBridge:
|
|
"""Construct a fully-wired testbed bridge."""
|
|
return RoomBridge(world=world, **kwargs)
|
|
|
|
|
|
async def serve(host: str = "127.0.0.1", port: int = 8765) -> None:
|
|
"""Run the bridge over WebSocket until interrupted."""
|
|
from aicc.transport.websocket import WebSocketServer
|
|
|
|
bridge = build_bridge()
|
|
async with WebSocketServer(bridge, host=host, port=port):
|
|
print(f"AICC capsule testbed listening on ws://{host}:{port}")
|
|
await bridge.serve_forever()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description="Run the AICC capsule testbed bridge over WebSocket."
|
|
)
|
|
parser.add_argument("--host", default="127.0.0.1")
|
|
parser.add_argument("--port", type=int, default=8765)
|
|
args = parser.parse_args()
|
|
try:
|
|
asyncio.run(serve(args.host, args.port))
|
|
except KeyboardInterrupt:
|
|
pass
|