89 lines
2.7 KiB
Python
89 lines
2.7 KiB
Python
"""Top-down map drawing from sensor data.
|
|
|
|
Shared by the demo's frame recorder and the live viewer. Only uses what an
|
|
agent can observe: the world_query layout (room, obstacles, beacon) and
|
|
proprioception (position, heading).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from PIL import Image, ImageDraw
|
|
|
|
ROOM_SIZE = 16.0
|
|
MAP_SIZE = 320
|
|
SCALE = MAP_SIZE / ROOM_SIZE
|
|
|
|
|
|
def draw_sensor_map(
|
|
layout: dict[str, Any] | None,
|
|
pos: tuple[float, float],
|
|
yaw: float,
|
|
path: list[tuple[float, float]] | None = None,
|
|
*,
|
|
step: int | None = None,
|
|
beacon_active: bool = False,
|
|
) -> Image.Image:
|
|
"""Draw the room from the agent's sensor view: floor grid, crates, the
|
|
beacon, the capsule's path and current position/heading."""
|
|
path = path or []
|
|
|
|
def xy(x: float, z: float) -> tuple[float, float]:
|
|
return (x * SCALE, MAP_SIZE - z * SCALE)
|
|
|
|
img = Image.new("RGB", (MAP_SIZE, MAP_SIZE), (52, 52, 58))
|
|
d = ImageDraw.Draw(img)
|
|
for i in range(int(ROOM_SIZE) + 1):
|
|
c = 60 if i % 2 == 0 else 54
|
|
d.line([xy(i, 0), xy(i, ROOM_SIZE)], fill=(c, c, c + 6), width=1)
|
|
d.line([xy(0, i), xy(ROOM_SIZE, i)], fill=(c, c, c + 6), width=1)
|
|
|
|
if layout:
|
|
for ob in 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 = layout.get("beacon", {})
|
|
bx, bz = xy(b.get("x", 12.5), b.get("z", 12.5))
|
|
core = (255, 200, 90) if beacon_active else (120, 210, 235)
|
|
glow = (255, 240, 180) if beacon_active else (160, 235, 250)
|
|
d.ellipse([bx - 14, bz - 14, bx + 14, bz + 14], fill=glow)
|
|
d.ellipse(
|
|
[bx - 7, bz - 7, bx + 7, bz + 7], fill=core, outline=(20, 20, 26), width=2
|
|
)
|
|
|
|
if len(path) > 1:
|
|
pts = [xy(x, z) for x, z in 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,
|
|
)
|
|
fx, fz = _forward(yaw)
|
|
tip = xy(pos[0] + fx * 0.7, pos[1] + fz * 0.7)
|
|
d.line([(cx, cz), tip], fill=(30, 30, 40), width=3)
|
|
if step is not None:
|
|
d.text(
|
|
(8, 8),
|
|
f"step {step} pos ({pos[0]:.1f}, {pos[1]:.1f}) yaw {yaw:.0f}",
|
|
fill=(230, 230, 230),
|
|
)
|
|
return img
|
|
|
|
|
|
def _forward(yaw_deg: float) -> tuple[float, float]:
|
|
import math
|
|
|
|
rad = math.radians(yaw_deg)
|
|
return math.sin(rad), math.cos(rad)
|