Add auto-running RL demo with PPO mini-updates

README.md

@@ -1,12 +1,15 @@
 ---
-title: Openenv Rl Demo
-emoji: 📚
-colorFrom: gray
+title: OpenENV RL Demo
+emoji: 🚗
+colorFrom: blue
 colorTo: green
 sdk: gradio
-sdk_version: 6.9.0
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# OpenENV RL — Live Policy Training
+
+Auto-runs 20 steps per episode using the openenv policy system (FlatMLPPolicy).
+PPO mini-update after each episode — rewards increase over time.

app.py

@@ -0,0 +1,303 @@
+"""
+OpenENV RL Demo — auto-runs 20 steps per episode using the openenv policy system.
+
+Policies: FlatMLPPolicy / TicketAttentionPolicy (from openenv)
+Training:  PPO mini-update after each episode — rewards increase over time
+Display:   Live step-by-step feed + episode reward history
+"""
+
+import math, time, threading
+import numpy as np
+import torch
+import torch.optim as optim
+import gradio as gr
+
+from overflow_env.server.overflow_environment import OverflowEnvironment
+from overflow_env.models import OverflowAction
+from policies.flat_mlp_policy import FlatMLPPolicy
+from policies.ticket_attention_policy import TicketAttentionPolicy
+from policies.policy_spec import build_obs, build_ticket_vector, OBS_DIM
+
+
+STEPS_PER_EPISODE = 20
+
+
+# ── Observation adapter ───────────────────────────────────────────────────────
+
+def obs_to_vec(overflow_obs) -> np.ndarray:
+    cars = overflow_obs.cars
+    if not cars:
+        return np.zeros(OBS_DIM, dtype=np.float32)
+    ego = next((c for c in cars if c.carId == 0), cars[0])
+    ego_spd = ego.speed / 4.5
+    ego_x   = ego.position.x
+    ego_y   = (ego.lane - 2) * 3.7
+    tickets = []
+    for car in cars:
+        if car.carId == 0:
+            continue
+        rx = car.position.x - ego.position.x
+        ry = (car.lane - ego.lane) * 3.7
+        cs = car.speed / 4.5
+        d  = math.sqrt(rx**2 + ry**2)
+        if d > 80:
+            continue
+        cl = max(ego_spd - cs * math.copysign(1, max(rx, 0.01)), 0.1)
+        tickets.append(build_ticket_vector(
+            severity_weight=1.0 if d < 8 else 0.75 if d < 15 else 0.5,
+            ttl=5.0, pos_x=rx, pos_y=ry, pos_z=0.0,
+            vel_x=cs, vel_y=0.0, vel_z=0.0, heading=0.0,
+            size_length=4.0, size_width=2.0, size_height=1.5,
+            distance=d, time_to_collision=min(d / cl, 30.0),
+            bearing=math.atan2(ry, max(rx, 0.01)),
+            ticket_type="collision_risk", entity_type="vehicle", confidence=1.0,
+        ))
+    tv = np.array(tickets, dtype=np.float32) if tickets else None
+    return build_obs(ego_x=ego_x, ego_y=ego_y, ego_z=0.0,
+                     ego_vx=ego_spd, ego_vy=0.0,
+                     heading=0.0, speed=ego_spd,
+                     steer=0.0, throttle=0.5, brake=0.0,
+                     ticket_vectors=tv)
+
+
+def action_to_decision(a: np.ndarray) -> str:
+    s, t, b = float(a[0]), float(a[1]), float(a[2])
+    if abs(s) > 0.35: return "lane_change_left" if s < 0 else "lane_change_right"
+    if b > 0.25:      return "brake"
+    if t > 0.20:      return "accelerate"
+    return "maintain"
+
+
+# ── Global training state ─────────────────────────────────────────────────────
+
+policy   = FlatMLPPolicy(obs_dim=OBS_DIM)
+optimizer = optim.Adam(policy.parameters(), lr=3e-4, eps=1e-5)
+
+# Rollout buffer (lightweight — one episode at a time)
+_buf_obs   = []
+_buf_acts  = []
+_buf_rews  = []
+_buf_logps = []
+_buf_vals  = []
+_buf_dones = []
+
+episode_history = []   # [{ep, steps, reward, outcome}]
+step_log        = []   # [{ep, step, decision, reward, scene}]
+_running        = False
+_lock           = threading.Lock()
+
+
+def _ppo_mini_update():
+    """Single PPO gradient step on the just-completed episode."""
+    if len(_buf_obs) < 2:
+        return
+    obs_t  = torch.tensor(np.array(_buf_obs),   dtype=torch.float32)
+    acts_t = torch.tensor(np.array(_buf_acts),  dtype=torch.float32)
+    rews_t = torch.tensor(_buf_rews,             dtype=torch.float32)
+    logp_t = torch.tensor(_buf_logps,            dtype=torch.float32)
+    vals_t = torch.tensor(_buf_vals,             dtype=torch.float32)
+    done_t = torch.tensor(_buf_dones,            dtype=torch.float32)
+
+    # GAE returns
+    gamma, lam = 0.99, 0.95
+    adv = torch.zeros_like(rews_t)
+    gae = 0.0
+    for t in reversed(range(len(rews_t))):
+        nv  = 0.0 if t == len(rews_t) - 1 else float(vals_t[t + 1])
+        d   = rews_t[t] + gamma * nv * (1 - done_t[t]) - vals_t[t]
+        gae = d + gamma * lam * (1 - done_t[t]) * gae
+        adv[t] = gae
+    ret = adv + vals_t
+    adv = (adv - adv.mean()) / (adv.std() + 1e-8)
+
+    policy.train()
+    act_mean, val = policy(obs_t)
+    val = val.squeeze(-1)
+    dist    = torch.distributions.Normal(act_mean, torch.ones_like(act_mean) * 0.3)
+    logp    = dist.log_prob(acts_t).sum(dim=-1)
+    entropy = dist.entropy().sum(dim=-1).mean()
+    ratio   = torch.exp(logp - logp_t)
+    pg      = torch.max(-adv * ratio, -adv * ratio.clamp(0.8, 1.2)).mean()
+    vf      = 0.5 * ((val - ret) ** 2).mean()
+    loss    = pg + 0.5 * vf - 0.02 * entropy
+    optimizer.zero_grad()
+    loss.backward()
+    torch.nn.utils.clip_grad_norm_(policy.parameters(), 0.5)
+    optimizer.step()
+
+
+def run_episodes_loop():
+    """Background thread — runs episodes continuously, updates policy after each."""
+    global _running
+    ep_num = 0
+    env    = OverflowEnvironment()
+
+    while _running:
+        ep_num += 1
+        obs    = env.reset()
+        ep_rew = 0.0
+        outcome = "timeout"
+
+        _buf_obs.clear(); _buf_acts.clear(); _buf_rews.clear()
+        _buf_logps.clear(); _buf_vals.clear(); _buf_dones.clear()
+
+        for step in range(1, STEPS_PER_EPISODE + 1):
+            if not _running:
+                break
+
+            obs_vec = obs_to_vec(obs)
+            policy.eval()
+            with torch.no_grad():
+                obs_t = torch.tensor(obs_vec, dtype=torch.float32).unsqueeze(0)
+                act_mean, val = policy(obs_t)
+                dist   = torch.distributions.Normal(act_mean.squeeze(0),
+                                                    torch.ones(3) * 0.3)
+                action = dist.sample().clamp(-1, 1)
+                logp   = dist.log_prob(action).sum()
+
+            decision = action_to_decision(action.numpy())
+            obs      = env.step(OverflowAction(decision=decision, reasoning=""))
+            reward   = float(obs.reward or 0.0)
+            done     = obs.done
+
+            _buf_obs.append(obs_vec)
+            _buf_acts.append(action.numpy())
+            _buf_rews.append(reward)
+            _buf_logps.append(float(logp))
+            _buf_vals.append(float(val.squeeze()))
+            _buf_dones.append(float(done))
+
+            ep_rew += reward
+
+            with _lock:
+                step_log.append({
+                    "ep": ep_num, "step": step,
+                    "decision": decision,
+                    "reward": round(reward, 2),
+                    "ep_reward": round(ep_rew, 2),
+                    "scene": obs.scene_description.split("\n")[0],
+                    "incident": obs.incident_report or "",
+                })
+
+            if done:
+                outcome = "CRASH" if "CRASH" in (obs.incident_report or "") else "GOAL"
+                break
+
+            time.sleep(0.6)   # pace so UI can show each step
+
+        _ppo_mini_update()
+
+        with _lock:
+            episode_history.append({
+                "ep": ep_num,
+                "steps": step,
+                "reward": round(ep_rew, 2),
+                "outcome": outcome,
+            })
+
+
+# ── Gradio UI ─────────────────────────────────────────────────────────────────
+
+def start_training():
+    global _running
+    if not _running:
+        _running = True
+        step_log.clear()
+        episode_history.clear()
+        t = threading.Thread(target=run_episodes_loop, daemon=True)
+        t.start()
+    return gr.update(value="Running...", interactive=False), gr.update(interactive=True)
+
+
+def stop_training():
+    global _running
+    _running = False
+    return gr.update(value="Start", interactive=True), gr.update(interactive=False)
+
+
+def get_updates():
+    """Called by gr.Timer every second — returns latest display content."""
+    with _lock:
+        logs = list(step_log[-20:])
+        eps  = list(episode_history[-30:])
+
+    # Step feed
+    lines = []
+    for e in reversed(logs):
+        flag = ""
+        if "CRASH" in e["incident"]: flag = " 💥"
+        elif "GOAL" in e["incident"]: flag = " ✓"
+        elif "NEAR MISS" in e["incident"]: flag = " ⚠"
+        lines.append(
+            f"ep {e['ep']:>3d} | step {e['step']:>2d} | "
+            f"{e['decision']:<20} | r={e['reward']:>+6.2f} | "
+            f"ep_total={e['ep_reward']:>7.2f}{flag}"
+        )
+    step_text = "\n".join(lines) if lines else "Waiting for first episode..."
+
+    # Episode summary
+    ep_lines = ["Episode | Steps | Total Reward | Outcome", "-" * 44]
+    for e in reversed(eps):
+        ep_lines.append(
+            f"  {e['ep']:>4d}  |  {e['steps']:>3d}  | "
+            f"  {e['reward']:>+8.2f}   | {e['outcome']}"
+        )
+    ep_text = "\n".join(ep_lines) if eps else "No episodes completed yet."
+
+    # Mean reward trend
+    if len(eps) >= 2:
+        rewards = [e["reward"] for e in eps]
+        n = len(rewards)
+        half = max(n // 2, 1)
+        early = sum(rewards[:half]) / half
+        late  = sum(rewards[half:]) / max(n - half, 1)
+        trend = f"Mean reward (early {half} eps): {early:+.2f}  →  (last {n-half} eps): {late:+.2f}"
+        arrow = "↑ improving" if late > early else "↓ declining"
+        trend_text = f"{trend}   {arrow}"
+    else:
+        trend_text = "Collecting data..."
+
+    status = "● RUNNING" if _running else "■ STOPPED"
+
+    return step_text, ep_text, trend_text, status
+
+
+with gr.Blocks(title="OpenENV RL Demo") as demo:
+    gr.Markdown("# OpenENV RL — Live Policy Training\n"
+                "FlatMLPPolicy runs 20 steps per episode on OverflowEnvironment. "
+                "PPO mini-update after each episode — watch rewards improve over time.")
+
+    with gr.Row():
+        start_btn  = gr.Button("Start", variant="primary")
+        stop_btn   = gr.Button("Stop",  variant="stop", interactive=False)
+        status_box = gr.Textbox(value="■ STOPPED", label="Status",
+                                interactive=False, scale=0, min_width=120)
+
+    gr.Markdown("### Live Step Feed (most recent 20 steps)")
+    step_display = gr.Textbox(
+        value="Press Start to begin...",
+        lines=22, max_lines=22, interactive=False,
+        elem_id="step_feed",
+    )
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Episode History")
+            ep_display = gr.Textbox(lines=12, interactive=False)
+        with gr.Column():
+            gr.Markdown("### Reward Trend")
+            trend_display = gr.Textbox(lines=3, interactive=False)
+
+    # Auto-refresh every 1 second
+    timer = gr.Timer(value=1.0)
+    timer.tick(
+        fn=get_updates,
+        outputs=[step_display, ep_display, trend_display, status_box],
+    )
+
+    start_btn.click(fn=start_training, outputs=[start_btn, stop_btn])
+    stop_btn.click(fn=stop_training,   outputs=[start_btn, stop_btn])
+
+
+if __name__ == "__main__":
+    demo.launch()

policies/__init__.py

@@ -0,0 +1,2 @@
+from .flat_mlp_policy import FlatMLPPolicy
+from .ticket_attention_policy import TicketAttentionPolicy

policies/base_policy.py

@@ -0,0 +1,66 @@
+"""
+BasePolicy — abstract interface all policies implement.
+
+All policies expose the same predict() and train_step() API so the
+curriculum trainer can swap them out transparently.
+"""
+
+from __future__ import annotations
+
+import abc
+from typing import Any, Dict, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class BasePolicy(nn.Module, abc.ABC):
+    """
+    Abstract base for all driving policies.
+
+    Subclasses implement:
+        forward(obs_tensor)  → action_tensor, value_tensor
+        encode_obs(obs_np)   → torch.Tensor
+    """
+
+    def __init__(self, obs_dim: int, action_dim: int = 3):
+        super().__init__()
+        self.obs_dim    = obs_dim
+        self.action_dim = action_dim
+
+    @abc.abstractmethod
+    def forward(
+        self, obs: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Returns:
+            action_mean  — shape (B, action_dim)
+            value        — shape (B, 1)
+        """
+        ...
+
+    def predict(
+        self,
+        obs: np.ndarray,
+        deterministic: bool = False,
+    ) -> np.ndarray:
+        """Numpy in, numpy out. Used by the env during rollout."""
+        self.eval()
+        with torch.no_grad():
+            t = torch.as_tensor(obs, dtype=torch.float32).unsqueeze(0)
+            mean, _ = self.forward(t)
+            if deterministic:
+                action = mean
+            else:
+                action = mean + torch.randn_like(mean) * 0.1
+        return action.squeeze(0).numpy()
+
+    @staticmethod
+    def _mlp(dims: list[int], activation=nn.Tanh) -> nn.Sequential:
+        layers = []
+        for i in range(len(dims) - 1):
+            layers.append(nn.Linear(dims[i], dims[i + 1]))
+            if i < len(dims) - 2:
+                layers.append(activation())
+        return nn.Sequential(*layers)

policies/flat_mlp_policy.py

@@ -0,0 +1,50 @@
+"""
+FlatMLPPolicy — sanity-check baseline.
+
+Concatenates the full observation (ego + all tickets flattened) and passes
+it through a standard MLP. No attention, no structure.
+
+Use this to:
+  1. Verify the reward signal and environment are working
+  2. Establish a performance floor
+  3. Confirm that TicketAttentionPolicy actually improves over this
+
+If FlatMLPPolicy can't learn Stage 1 survival, the reward or env is broken.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from .base_policy import BasePolicy
+
+
+class FlatMLPPolicy(BasePolicy):
+    """Standard 3-layer MLP over the full flat observation."""
+
+    def __init__(self, obs_dim: int, hidden: int = 256):
+        super().__init__(obs_dim)
+
+        self.actor = nn.Sequential(
+            nn.Linear(obs_dim, hidden), nn.LayerNorm(hidden), nn.Tanh(),
+            nn.Linear(hidden, hidden),                          nn.Tanh(),
+            nn.Linear(hidden, hidden // 2),                     nn.Tanh(),
+            nn.Linear(hidden // 2, 3),                          nn.Tanh(),
+        )
+        self.critic = nn.Sequential(
+            nn.Linear(obs_dim, hidden),    nn.Tanh(),
+            nn.Linear(hidden, hidden // 2), nn.Tanh(),
+            nn.Linear(hidden // 2, 1),
+        )
+        self._init_weights()
+
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.orthogonal_(m.weight, gain=1.0)
+                nn.init.zeros_(m.bias)
+        nn.init.orthogonal_(self.actor[-2].weight, gain=0.01)
+
+    def forward(self, obs: torch.Tensor):
+        return self.actor(obs), self.critic(obs)

policies/policy_spec.py

@@ -0,0 +1,409 @@
+"""
+Policy data input specifications — formal contracts for observation, action, and ticket data.
+
+This module defines the exact data shapes, normalization ranges, and semantic meaning
+of every field consumed by OpenENV policies. Use this as the reference when:
+
+  1. Building a new environment that targets these policies
+  2. Writing a bridge/adapter from a different simulator
+  3. Implementing a new policy that must interoperate with the existing set
+
+All policies share the same raw observation layout (EGO + ticket matrix).
+Specialized policies (ThreatAvoidance, SystemFailure) select subsets internally.
+
+Example usage:
+    from openenv.policies.policy_spec import ObsSpec, ActionSpec, validate_obs
+
+    spec = ObsSpec()
+    obs  = my_env.get_observation()
+    validate_obs(obs, spec)  # raises ValueError on shape/range mismatch
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+
+# ── Ego state specification ──────────────────────────────────────────────────
+
+EGO_STATE_DIM = 11
+
+@dataclass(frozen=True)
+class EgoField:
+    """Description of a single ego state field."""
+    index:       int
+    name:        str
+    unit:        str
+    raw_range:   Tuple[float, float]   # physical range before normalization
+    norm_divisor: float                # obs_value = raw_value / norm_divisor
+    description: str
+
+EGO_FIELDS: List[EgoField] = [
+    EgoField(0,  "x",             "m",     (-5000, 5000),  1000.0,  "Forward displacement from episode start"),
+    EgoField(1,  "y",             "m",     (-6.0, 6.0),    3.7,     "Lateral displacement (0 = lane center, + = left)"),
+    EgoField(2,  "z",             "m",     (-10, 10),      10.0,    "Vertical position (flat road = 0)"),
+    EgoField(3,  "vx",            "m/s",   (-20, 20),      20.0,    "Forward velocity in world frame"),
+    EgoField(4,  "vy",            "m/s",   (-20, 20),      20.0,    "Lateral velocity in world frame"),
+    EgoField(5,  "vz",            "m/s",   (0, 0),         1.0,     "Vertical velocity (always 0 on flat road)"),
+    EgoField(6,  "heading_sin",   "rad",   (-1, 1),        1.0,     "sin(heading angle), 0 = forward"),
+    EgoField(7,  "heading_cos",   "rad",   (-1, 1),        1.0,     "cos(heading angle), 1 = forward"),
+    EgoField(8,  "speed",         "m/s",   (0, 20),        20.0,    "Scalar speed = sqrt(vx^2 + vy^2)"),
+    EgoField(9,  "steer",         "norm",  (-1, 1),        1.0,     "Current steering command [-1=full left, 1=full right]"),
+    EgoField(10, "net_drive",     "norm",  (-1, 1),        1.0,     "throttle - brake [-1=full brake, 1=full throttle]"),
+]
+
+
+# ── Ticket vector specification ──────────────────────────────────────────────
+
+TICKET_VECTOR_DIM = 37   # 18 fixed + 14 type one-hot + 5 entity one-hot
+MAX_TICKETS = 16
+
+# Ticket types (14 total) — one-hot encoded starting at index 18
+TICKET_TYPES = [
+    "collision_risk", "sudden_brake", "side_impact", "head_on",
+    "merge_cut", "rear_end_risk",
+    "pedestrian_crossing", "cyclist_lane",
+    "tire_blowout", "brake_fade", "steering_loss", "sensor_occlusion",
+    "road_hazard", "weather_visibility",
+]
+
+# Entity types (5 total) — one-hot encoded after ticket types
+ENTITY_TYPES = ["vehicle", "pedestrian", "cyclist", "obstacle", "system"]
+
+# Verify dimension
+assert 18 + len(TICKET_TYPES) + len(ENTITY_TYPES) == TICKET_VECTOR_DIM, (
+    f"Ticket vector dim mismatch: 18 + {len(TICKET_TYPES)} + {len(ENTITY_TYPES)} "
+    f"!= {TICKET_VECTOR_DIM}"
+)
+
+@dataclass(frozen=True)
+class TicketField:
+    """Description of a single ticket vector field."""
+    offset:       int            # index within the TICKET_VECTOR_DIM vector
+    length:       int            # number of floats
+    name:         str
+    unit:         str
+    raw_range:    Tuple[float, float]
+    norm_divisor: float
+    description:  str
+
+TICKET_FIELDS: List[TicketField] = [
+    TicketField(0,  1, "severity_weight",   "norm",   (0, 1),       1.0,   "Severity: 0.25=LOW, 0.5=MED, 0.75=HIGH, 1.0=CRITICAL"),
+    TicketField(1,  1, "ttl_norm",          "s",      (0, 10),     10.0,   "Time-to-live remaining, clamped to [0,1]"),
+    TicketField(2,  1, "pos_x",             "m",      (-100, 100), 100.0,  "Ego-relative X (forward positive)"),
+    TicketField(3,  1, "pos_y",             "m",      (-50, 50),    50.0,  "Ego-relative Y (left positive)"),
+    TicketField(4,  1, "pos_z",             "m",      (-10, 10),    10.0,  "Ego-relative Z (up positive)"),
+    TicketField(5,  1, "vel_x",             "m/s",    (-30, 30),    30.0,  "Entity velocity X in world frame"),
+    TicketField(6,  1, "vel_y",             "m/s",    (-30, 30),    30.0,  "Entity velocity Y in world frame"),
+    TicketField(7,  1, "vel_z",             "m/s",    (-10, 10),    10.0,  "Entity velocity Z in world frame"),
+    TicketField(8,  1, "heading_sin",       "rad",    (-1, 1),      1.0,   "sin(entity heading relative to ego)"),
+    TicketField(9,  1, "heading_cos",       "rad",    (-1, 1),      1.0,   "cos(entity heading relative to ego)"),
+    TicketField(10, 1, "size_length",       "m",      (0, 10),     10.0,   "Entity bounding box length"),
+    TicketField(11, 1, "size_width",        "m",      (0, 5),       5.0,   "Entity bounding box width"),
+    TicketField(12, 1, "size_height",       "m",      (0, 4),       4.0,   "Entity bounding box height"),
+    TicketField(13, 1, "distance_norm",     "m",      (0, 100),   100.0,   "Euclidean distance to ego, clamped to [0,1]"),
+    TicketField(14, 1, "ttc_norm",          "s",      (0, 30),     30.0,   "Time-to-collision, clamped to [0,1]. 1.0 = no collision"),
+    TicketField(15, 1, "bearing_sin",       "rad",    (-1, 1),      1.0,   "sin(bearing angle from ego forward axis)"),
+    TicketField(16, 1, "bearing_cos",       "rad",    (-1, 1),      1.0,   "cos(bearing angle from ego forward axis)"),
+    TicketField(17, 1, "confidence",        "norm",   (0, 1),       1.0,   "Perception confidence [0=unreliable, 1=certain]"),
+    TicketField(18, len(TICKET_TYPES), "type_onehot",   "bool", (0, 1), 1.0, "One-hot ticket type"),
+    TicketField(18 + len(TICKET_TYPES), len(ENTITY_TYPES), "entity_onehot", "bool", (0, 1), 1.0, "One-hot entity type"),
+]
+
+
+# ── Full observation specification ───────────────────────────────────────────
+
+OBS_DIM = EGO_STATE_DIM + MAX_TICKETS * TICKET_VECTOR_DIM   # 11 + 16*37 = 603
+
+@dataclass(frozen=True)
+class ObsSpec:
+    """Complete observation space specification."""
+    ego_dim:         int = EGO_STATE_DIM
+    ticket_dim:      int = TICKET_VECTOR_DIM
+    max_tickets:     int = MAX_TICKETS
+    total_dim:       int = OBS_DIM
+    dtype:           str = "float32"
+    value_range:     Tuple[float, float] = (-1.0, 1.0)
+
+    # Layout: obs[0:ego_dim] = ego state
+    #         obs[ego_dim:] reshaped to (max_tickets, ticket_dim)
+    # Tickets are sorted by severity desc, distance asc. Zero-padded rows = empty slots.
+
+
+# ── Action specification ─────────────────────────────────────────────────────
+
+@dataclass(frozen=True)
+class ActionField:
+    index:       int
+    name:        str
+    raw_range:   Tuple[float, float]
+    description: str
+
+ACTION_DIM = 3
+
+ACTION_FIELDS: List[ActionField] = [
+    ActionField(0, "steer",    (-1.0, 1.0), "Steering command. -1=full left, +1=full right. Scaled by MAX_STEER=0.6 rad"),
+    ActionField(1, "throttle", (-1.0, 1.0), "Throttle command. Only positive values used (clipped to [0,1]). Scaled by MAX_ACCEL=4.0 m/s^2"),
+    ActionField(2, "brake",    (-1.0, 1.0), "Brake command. Only positive values used (clipped to [0,1]). Scaled by MAX_BRAKE=8.0 m/s^2"),
+]
+
+@dataclass(frozen=True)
+class ActionSpec:
+    """Action space specification."""
+    dim:          int = ACTION_DIM
+    dtype:        str = "float32"
+    value_range:  Tuple[float, float] = (-1.0, 1.0)
+
+
+# ── Policy input requirements ────────────────────────────────────────────────
+
+@dataclass(frozen=True)
+class PolicyInputSpec:
+    """Describes what a specific policy reads from the observation."""
+    name:              str
+    reads_ego:         bool
+    ego_indices:       Tuple[int, ...]           # which ego fields are used
+    reads_tickets:     bool
+    ticket_filter:     Optional[str]             # None = all, or "kinematic" / "failure"
+    max_tickets_used:  int                       # how many ticket slots the policy actually reads
+    requires_history:  bool                      # whether GRU/recurrent hidden state is needed
+    description:       str
+
+POLICY_SPECS: Dict[str, PolicyInputSpec] = {
+    "SurvivalPolicy": PolicyInputSpec(
+        name="SurvivalPolicy",
+        reads_ego=True,
+        ego_indices=tuple(range(EGO_STATE_DIM)),
+        reads_tickets=False,
+        ticket_filter=None,
+        max_tickets_used=0,
+        requires_history=False,
+        description="Stage 1 baseline. Reads only ego state (first 11 dims). "
+                    "Ticket portion of obs is ignored entirely.",
+    ),
+    "FlatMLPPolicy": PolicyInputSpec(
+        name="FlatMLPPolicy",
+        reads_ego=True,
+        ego_indices=tuple(range(EGO_STATE_DIM)),
+        reads_tickets=True,
+        ticket_filter=None,
+        max_tickets_used=MAX_TICKETS,
+        requires_history=False,
+        description="Sanity-check baseline. Reads full flat observation (ego + all tickets "
+                    "concatenated). No attention or structure.",
+    ),
+    "TicketAttentionPolicy": PolicyInputSpec(
+        name="TicketAttentionPolicy",
+        reads_ego=True,
+        ego_indices=tuple(range(EGO_STATE_DIM)),
+        reads_tickets=True,
+        ticket_filter=None,
+        max_tickets_used=MAX_TICKETS,
+        requires_history=False,
+        description="Main policy (Stage 2+). Cross-attention: ego queries ticket set. "
+                    "Order-invariant over tickets. Padding mask on zero-rows.",
+    ),
+    "ThreatAvoidancePolicy": PolicyInputSpec(
+        name="ThreatAvoidancePolicy",
+        reads_ego=True,
+        ego_indices=tuple(range(EGO_STATE_DIM)),
+        reads_tickets=True,
+        ticket_filter="kinematic",
+        max_tickets_used=1,
+        requires_history=False,
+        description="Specialist for kinematic threats (collision_risk, sudden_brake, "
+                    "side_impact, head_on, merge_cut, rear_end_risk). Extracts the "
+                    "highest-severity kinematic ticket and gates between brake/evade branches.",
+    ),
+    "SystemFailurePolicy": PolicyInputSpec(
+        name="SystemFailurePolicy",
+        reads_ego=True,
+        ego_indices=tuple(range(EGO_STATE_DIM)),
+        reads_tickets=True,
+        ticket_filter="failure",
+        max_tickets_used=1,
+        requires_history=False,
+        description="Specialist for onboard failures (tire_blowout, brake_fade, steering_loss). "
+                    "Mixture-of-experts with one expert per failure type. Initialized with "
+                    "domain-correct response priors.",
+    ),
+    "RecurrentPolicy": PolicyInputSpec(
+        name="RecurrentPolicy",
+        reads_ego=True,
+        ego_indices=tuple(range(EGO_STATE_DIM)),
+        reads_tickets=True,
+        ticket_filter=None,
+        max_tickets_used=MAX_TICKETS,
+        requires_history=True,
+        description="GRU-based policy for partial observability (Stage 4+). Carries hidden "
+                    "state across timesteps. Requires h_prev to be tracked by caller.",
+    ),
+}
+
+
+# ── Validation helpers ───────────────────────────────────────────────────────
+
+def validate_obs(obs: np.ndarray, spec: Optional[ObsSpec] = None) -> None:
+    """
+    Validate an observation array against the spec.
+    Raises ValueError with a descriptive message on any mismatch.
+    """
+    spec = spec or ObsSpec()
+    if obs.ndim != 1:
+        raise ValueError(f"Observation must be 1D, got shape {obs.shape}")
+    if obs.shape[0] != spec.total_dim:
+        raise ValueError(
+            f"Observation dim mismatch: expected {spec.total_dim}, got {obs.shape[0]}. "
+            f"Check ego_dim ({spec.ego_dim}) + max_tickets ({spec.max_tickets}) "
+            f"* ticket_dim ({spec.ticket_dim})"
+        )
+    if obs.dtype != np.float32:
+        raise ValueError(f"Observation dtype must be float32, got {obs.dtype}")
+
+
+def validate_action(action: np.ndarray) -> None:
+    """Validate an action array."""
+    if action.shape != (ACTION_DIM,):
+        raise ValueError(f"Action shape mismatch: expected ({ACTION_DIM},), got {action.shape}")
+    if np.any(action < -1.0) or np.any(action > 1.0):
+        raise ValueError(f"Action values must be in [-1, 1], got min={action.min()}, max={action.max()}")
+
+
+def build_obs(
+    ego_x: float, ego_y: float, ego_z: float,
+    ego_vx: float, ego_vy: float,
+    heading: float, speed: float,
+    steer: float, throttle: float, brake: float,
+    ticket_vectors: Optional[np.ndarray] = None,
+    max_tickets: int = MAX_TICKETS,
+) -> np.ndarray:
+    """
+    Build a valid observation vector from raw values.
+
+    This is the primary entry point for external environments that want to
+    produce observations compatible with OpenENV policies.
+
+    Parameters
+    ----------
+    ego_x       : forward displacement from episode start (metres)
+    ego_y       : lateral displacement from lane center (metres, + = left)
+    ego_z       : vertical position (metres)
+    ego_vx      : forward velocity (m/s)
+    ego_vy      : lateral velocity (m/s)
+    heading     : heading angle (radians, 0 = forward)
+    speed       : scalar speed (m/s)
+    steer       : current steering command [-1, 1]
+    throttle    : current throttle command [0, 1]
+    brake       : current brake command [0, 1]
+    ticket_vectors : (N, TICKET_VECTOR_DIM) array of ticket vectors, or None.
+                     Use EventTicket.to_vector() or build_ticket_vector() to create these.
+    max_tickets : number of ticket slots (must match policy expectation, default 16)
+
+    Returns
+    -------
+    obs : np.ndarray of shape (EGO_STATE_DIM + max_tickets * TICKET_VECTOR_DIM,)
+    """
+    import math
+
+    ego = np.array([
+        ego_x / 1000.0,
+        ego_y / 3.7,        # ROAD_HALF_WIDTH
+        ego_z / 10.0,
+        ego_vx / 20.0,      # MAX_SPEED
+        ego_vy / 20.0,
+        0.0,                 # vz (flat road)
+        math.sin(heading),
+        math.cos(heading),
+        speed / 20.0,
+        steer,
+        throttle - brake,    # net drive signal
+    ], dtype=np.float32)
+
+    ticket_matrix = np.zeros((max_tickets, TICKET_VECTOR_DIM), dtype=np.float32)
+    if ticket_vectors is not None:
+        n = min(len(ticket_vectors), max_tickets)
+        ticket_matrix[:n] = ticket_vectors[:n]
+
+    return np.concatenate([ego, ticket_matrix.flatten()])
+
+
+def build_ticket_vector(
+    severity_weight: float,
+    ttl: float,
+    pos_x: float, pos_y: float, pos_z: float,
+    vel_x: float, vel_y: float, vel_z: float,
+    heading: float,
+    size_length: float, size_width: float, size_height: float,
+    distance: float,
+    time_to_collision: Optional[float],
+    bearing: float,
+    ticket_type: str,
+    entity_type: str,
+    confidence: float = 1.0,
+) -> np.ndarray:
+    """
+    Build a single ticket vector from raw values without needing the full
+    EventTicket class. Use this when adapting a different simulator.
+
+    Parameters
+    ----------
+    severity_weight    : 0.25 (LOW), 0.5 (MEDIUM), 0.75 (HIGH), 1.0 (CRITICAL)
+    ttl                : seconds remaining until ticket expires
+    pos_x/y/z          : ego-relative position (metres)
+    vel_x/y/z          : entity velocity in world frame (m/s)
+    heading            : entity heading relative to ego (radians)
+    size_length/width/height : entity bounding box (metres)
+    distance           : euclidean distance to ego (metres)
+    time_to_collision  : seconds until collision, or None if no collision course
+    bearing            : angle from ego forward axis (radians)
+    ticket_type        : one of TICKET_TYPES (e.g., "collision_risk")
+    entity_type        : one of ENTITY_TYPES (e.g., "vehicle")
+    confidence         : perception confidence [0, 1]
+
+    Returns
+    -------
+    vec : np.ndarray of shape (TICKET_VECTOR_DIM,) = (37,)
+    """
+    import math
+
+    ttc_norm = min((time_to_collision if time_to_collision is not None else 30.0) / 30.0, 1.0)
+
+    type_oh = [0.0] * len(TICKET_TYPES)
+    entity_oh = [0.0] * len(ENTITY_TYPES)
+
+    if ticket_type in TICKET_TYPES:
+        type_oh[TICKET_TYPES.index(ticket_type)] = 1.0
+    else:
+        raise ValueError(f"Unknown ticket_type '{ticket_type}'. Must be one of {TICKET_TYPES}")
+
+    if entity_type in ENTITY_TYPES:
+        entity_oh[ENTITY_TYPES.index(entity_type)] = 1.0
+    else:
+        raise ValueError(f"Unknown entity_type '{entity_type}'. Must be one of {ENTITY_TYPES}")
+
+    vec = [
+        severity_weight,
+        min(ttl / 10.0, 1.0),
+        pos_x / 100.0,
+        pos_y / 50.0,
+        pos_z / 10.0,
+        vel_x / 30.0,
+        vel_y / 30.0,
+        vel_z / 10.0,
+        math.sin(heading),
+        math.cos(heading),
+        size_length / 10.0,
+        size_width / 5.0,
+        size_height / 4.0,
+        min(distance / 100.0, 1.0),
+        ttc_norm,
+        math.sin(bearing),
+        math.cos(bearing),
+        confidence,
+        *type_oh,
+        *entity_oh,
+    ]
+    return np.array(vec, dtype=np.float32)

policies/ticket_attention_policy.py

@@ -0,0 +1,227 @@
+"""
+TicketAttentionPolicy — the main policy (Stage 2+).
+
+Architecture: two-pass "reflective" cross-attention.
+
+    Pass 1: ego queries tickets → raw threat context
+    Pass 2: (ego + raw context) queries tickets again → refined context
+    This forces the policy to "think twice" — first perceive, then plan.
+
+    [ego | refined_context] → steer head  → steer action
+                            → drive head  → throttle, brake
+                            → critic head → value
+
+Why two-pass:
+  The first pass gathers what threats exist. The second pass re-examines
+  tickets knowing what the overall threat picture looks like. This prevents
+  the impulsive single-shot responses that cause wild oscillation.
+
+Why separate heads:
+  Steering requires smooth, conservative output (off-road = death).
+  Throttle/brake can be more aggressive. Separate heads + separate
+  noise levels let each dimension learn at its own pace.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .base_policy import BasePolicy
+EGO_STATE_DIM = 11
+MAX_TICKETS = 16
+TICKET_VECTOR_DIM = 37
+
+
+class TicketAttentionPolicy(BasePolicy):
+    """
+    Two-pass reflective attention policy.
+
+    Pass 1: perceive — what threats exist?
+    Pass 2: plan    — given what I see, which threats matter most?
+    Output: separate steer head (conservative) + drive head (throttle/brake)
+    """
+
+    def __init__(
+        self,
+        obs_dim:      int,
+        ego_embed:    int = 64,
+        ticket_embed: int = 64,
+        n_heads:      int = 4,
+        hidden:       int = 256,
+    ):
+        super().__init__(obs_dim)
+        assert ego_embed % n_heads == 0
+        assert ticket_embed == ego_embed
+
+        self.ego_embed    = ego_embed
+        self.max_tickets  = MAX_TICKETS
+        self.ticket_dim   = TICKET_VECTOR_DIM
+
+        # ── Encoders ──────────────────────────────────────────────────────
+        self.ego_encoder = nn.Sequential(
+            nn.Linear(EGO_STATE_DIM, hidden // 2),
+            nn.LayerNorm(hidden // 2),
+            nn.Tanh(),
+            nn.Linear(hidden // 2, ego_embed),
+            nn.LayerNorm(ego_embed),
+        )
+        self.ticket_encoder = nn.Sequential(
+            nn.Linear(TICKET_VECTOR_DIM, hidden // 2),
+            nn.LayerNorm(hidden // 2),
+            nn.ReLU(),
+            nn.Linear(hidden // 2, ticket_embed),
+            nn.LayerNorm(ticket_embed),
+        )
+
+        # ── Pass 1: perceive (ego queries tickets) ───────────────────────
+        self.attn_pass1 = nn.MultiheadAttention(
+            embed_dim=ego_embed, num_heads=n_heads,
+            dropout=0.0, batch_first=True,
+        )
+        self.norm1 = nn.LayerNorm(ego_embed)
+
+        # ── Reflection gate: fuse ego + pass1 context for second query ───
+        self.reflect_proj = nn.Sequential(
+            nn.Linear(ego_embed * 2, ego_embed),
+            nn.LayerNorm(ego_embed),
+            nn.Tanh(),
+        )
+
+        # ── Pass 2: plan (refined query re-attends to tickets) ───────────
+        self.attn_pass2 = nn.MultiheadAttention(
+            embed_dim=ego_embed, num_heads=n_heads,
+            dropout=0.0, batch_first=True,
+        )
+        self.norm2 = nn.LayerNorm(ego_embed)
+
+        # ── Fused representation ─────────────────────────────────────────
+        fused_dim = ego_embed + ego_embed  # ego + refined context
+
+        # ── Steer head (conservative, smooth output) ─────────────────────
+        self.steer_head = nn.Sequential(
+            nn.Linear(fused_dim, hidden // 2),
+            nn.LayerNorm(hidden // 2),
+            nn.Tanh(),
+            nn.Linear(hidden // 2, hidden // 4),
+            nn.Tanh(),
+            nn.Linear(hidden // 4, 1),
+            nn.Tanh(),
+        )
+
+        # ── Drive head (throttle + brake) ────────────────────────────────
+        self.drive_head = nn.Sequential(
+            nn.Linear(fused_dim, hidden // 2),
+            nn.LayerNorm(hidden // 2),
+            nn.Tanh(),
+            nn.Linear(hidden // 2, hidden // 4),
+            nn.Tanh(),
+            nn.Linear(hidden // 4, 2),
+            nn.Tanh(),
+        )
+
+        # ── Critic head ──────────────────────────────────────────────────
+        self.critic = nn.Sequential(
+            nn.Linear(fused_dim, hidden),
+            nn.LayerNorm(hidden),
+            nn.Tanh(),
+            nn.Linear(hidden, hidden // 2),
+            nn.Tanh(),
+            nn.Linear(hidden // 2, 1),
+        )
+
+        self._init_weights()
+
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.orthogonal_(m.weight, gain=1.0)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+        # Very small initial actions — start by doing almost nothing
+        nn.init.orthogonal_(self.steer_head[-2].weight, gain=0.01)
+        nn.init.orthogonal_(self.drive_head[-2].weight, gain=0.01)
+        # Critic starts near zero
+        nn.init.orthogonal_(self.critic[-1].weight, gain=0.1)
+
+    def _attend(self, attn_module, norm_module, query, tk_emb, is_padding, all_empty):
+        """Run one attention pass with NaN-safe masking."""
+        B = query.shape[0]
+        q = query if query.dim() == 3 else query.unsqueeze(1)
+
+        if all_empty.all():
+            return torch.zeros(B, self.ego_embed, device=query.device)
+
+        safe_mask = is_padding.clone()
+        safe_mask[all_empty, 0] = False
+        attn_out, _ = attn_module(
+            query=q, key=tk_emb, value=tk_emb,
+            key_padding_mask=safe_mask,
+        )
+        context = attn_out.squeeze(1)
+        context[all_empty] = 0.0
+        return norm_module(context)
+
+    def forward(self, obs: torch.Tensor):
+        B = obs.shape[0]
+
+        # Split observation
+        ego_raw = obs[:, :EGO_STATE_DIM]
+        tk_raw  = obs[:, EGO_STATE_DIM:].view(B, self.max_tickets, self.ticket_dim)
+
+        # Encode
+        ego_emb = self.ego_encoder(ego_raw)
+        tk_emb  = self.ticket_encoder(tk_raw)
+
+        # Padding mask
+        is_padding = (tk_raw.abs().sum(dim=-1) == 0)
+        all_empty  = is_padding.all(dim=-1)
+
+        # ── Pass 1: perceive ─────────────────────────────────────────────
+        ctx1 = self._attend(self.attn_pass1, self.norm1,
+                            ego_emb, tk_emb, is_padding, all_empty)
+
+        # ── Reflect: combine ego + initial context into refined query ────
+        reflected = self.reflect_proj(torch.cat([ego_emb, ctx1], dim=-1))
+
+        # ── Pass 2: plan (re-attend with richer query) ───────────────────
+        ctx2 = self._attend(self.attn_pass2, self.norm2,
+                            reflected, tk_emb, is_padding, all_empty)
+
+        # ── Fuse and decode ──────────────────────────────────────────────
+        fused = torch.cat([ego_emb, ctx2], dim=-1)
+
+        steer  = self.steer_head(fused)          # (B, 1)
+        drive  = self.drive_head(fused)           # (B, 2)
+        action = torch.cat([steer, drive], dim=-1)  # (B, 3)
+        value  = self.critic(fused)               # (B, 1)
+
+        return action, value
+
+    def get_attention_weights(self, obs: torch.Tensor) -> torch.Tensor:
+        """Returns pass-2 attention weights for interpretability."""
+        B = obs.shape[0]
+        ego_raw = obs[:, :EGO_STATE_DIM]
+        tk_raw  = obs[:, EGO_STATE_DIM:].view(B, self.max_tickets, self.ticket_dim)
+        ego_emb = self.ego_encoder(ego_raw)
+        tk_emb  = self.ticket_encoder(tk_raw)
+        is_padding = (tk_raw.abs().sum(dim=-1) == 0)
+        all_empty = is_padding.all(dim=-1)
+
+        # Pass 1
+        ctx1 = self._attend(self.attn_pass1, self.norm1,
+                            ego_emb, tk_emb, is_padding, all_empty)
+        reflected = self.reflect_proj(torch.cat([ego_emb, ctx1], dim=-1))
+
+        # Pass 2 — get weights
+        safe_mask = is_padding.clone()
+        safe_mask[all_empty, 0] = False
+        query = reflected.unsqueeze(1)
+        _, weights = self.attn_pass2(
+            query=query, key=tk_emb, value=tk_emb,
+            key_padding_mask=safe_mask,
+            need_weights=True, average_attn_weights=False,
+        )
+        weights[all_empty] = 0.0
+        return weights

requirements.txt

@@ -0,0 +1,7 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.5.1+cpu
+numpy>=1.24.0
+gradio>=4.44.0
+pydantic>=2.0.0
+requests>=2.31.0
+openenv-overflow-env @ git+https://huggingface.co/spaces/SteveDusty/overflow_env