feat: continuous PPO training + live web dashboard

- Replace server/app.py with FastAPI continuous training loop
- PPO training runs in background thread (no manual stepping)
- Live dashboard: 2D road canvas on left, reward chart + episode log on right
- SSE + polling for real-time updates (0.5s heartbeat)
- Capped reward mode: min(reward, 2.0) per step
- Uncapped reward mode: base + token_count * 0.001 (frontier LLM scaling)
- Toggle capped/uncapped via UI button or POST /api/mode
- Curriculum stage display with pip indicators
- PPO losses (pg/vf/entropy) shown live
- Dockerfile cleaned: removed openenv-core dep to avoid websockets conflict

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Dockerfile

@@ -2,21 +2,33 @@ FROM python:3.11-slim
 
 WORKDIR /app
 
-# Install system deps
 RUN apt-get update && \
     apt-get install -y --no-install-recommends git curl && \
     rm -rf /var/lib/apt/lists/*
 
-# Copy environment code into a proper package directory
+# Copy the package into /app/overflow_env
 COPY . /app/overflow_env
 
-# Install dependencies via pip using requirements.txt
-RUN pip install --no-cache-dir -r /app/overflow_env/server/requirements.txt
+# Install runtime dependencies (no openenv-core to avoid websockets conflict)
+RUN pip install --no-cache-dir \
+    --extra-index-url https://download.pytorch.org/whl/cpu \
+    "fastapi>=0.115.0" \
+    "pydantic>=2.0.0" \
+    "uvicorn[standard]>=0.24.0" \
+    "requests>=2.31.0" \
+    "torch==2.5.1+cpu" \
+    "numpy>=1.24.0" \
+    "gymnasium>=0.29.0" \
+    "matplotlib>=3.8.0" \
+    "pillow==10.4.0"
 
-HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+# Make overflow_env importable as a top-level package
+ENV PYTHONPATH=/app
+ENV ENABLE_WEB_INTERFACE=true
+
+HEALTHCHECK --interval=30s --timeout=10s --start-period=30s --retries=3 \
     CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1
 
 EXPOSE 8000
 
-ENV ENABLE_WEB_INTERFACE=true
 CMD ["uvicorn", "overflow_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]

README.md

@@ -1,5 +1,5 @@
 ---
-title: Overflow Environment Server
+title: Overflow OpenENV
 emoji: 🚗
 colorFrom: red
 colorTo: yellow

__init__.py

@@ -1,6 +1,10 @@
 """Overflow Environment — Autonomous vehicle fleet oversight for OpenEnv."""
 
-from .client import OverflowEnv
+try:
+    from .client import OverflowEnv
+except ImportError:
+    OverflowEnv = None  # openenv-core not installed; training-only mode
+
 from .models import (
     CarStateData,
     LaneOccupancyData,

models.py

@@ -12,7 +12,16 @@ from typing import Any, Dict, List, Optional
 
 from pydantic import BaseModel, Field
 
-from openenv.core.env_server.types import Action, Observation, State
+try:
+    from openenv.core.env_server.types import Action, Observation, State
+except ImportError:
+    class Action(BaseModel): pass
+    class Observation(BaseModel):
+        done: bool = False
+        reward: float = 0.0
+    class State(BaseModel):
+        episode_id: str = ""
+        step_count: int = 0
 
 # ── Structured sub-models (frontend-compatible) ─────────────────────────
 

policies/__init__.py

@@ -0,0 +1,5 @@
+from .base_policy import BasePolicy
+from .flat_mlp_policy import FlatMLPPolicy
+from .ticket_attention_policy import TicketAttentionPolicy
+
+__all__ = ["BasePolicy", "FlatMLPPolicy", "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

pyproject.toml

@@ -13,6 +13,14 @@ dependencies = [
     "pydantic>=2.0.0",
     "uvicorn[standard]>=0.24.0",
     "requests>=2.31.0",
+    "torch>=2.10.0",
+    "numpy>=2.2.6",
+    "pillow>=12.1.1",
+    "gymnasium>=1.2.3",
+    "matplotlib>=3.10.8",
+    "gradio>=6.9.0",
+    "uvloop>=0.22.1",
+    "pyaudioop",
 ]
 
 [project.optional-dependencies]

requirements.txt

@@ -0,0 +1,9 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.5.1+cpu
+numpy>=1.24.0
+pillow==10.4.0
+matplotlib>=3.8.0
+pydantic>=2.0.0
+requests>=2.31.0
+gymnasium>=0.29.0
+pyaudioop>=0.1.0

server/app.py

@@ -1,44 +1,882 @@
 """
-FastAPI application for the Overflow Environment.
+Overflow OpenENV — Continuous PPO Training + Live Web Dashboard
 
-Exposes the OverflowEnvironment over HTTP and WebSocket endpoints.
+Runs PPO training in a background thread and exposes:
+  GET /          → live HTML dashboard (road animation + reward charts)
+  GET /api/state → JSON snapshot of current training state
+  GET /api/stream → SSE stream of state updates
+  POST /api/mode  → switch reward mode {capped, uncapped}
 
-Usage:
-    uvicorn server.app:app --reload --host 0.0.0.0 --port 8000
+Reward modes:
+  capped   — standard shaped reward, capped at REWARD_CAP per step
+  uncapped — base reward + token_bonus (scales with LLM reasoning token count)
+             frontier models that produce more reasoning tokens earn more reward
 """
 
-import inspect
+from __future__ import annotations
 
-from openenv.core.env_server.http_server import create_app
+import asyncio
+import json
+import math
+import sys
+import threading
+import time
+from collections import deque
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any, Deque, Dict, List, Optional
 
-from ..models import OverflowAction, OverflowObservation
-from .overflow_environment import OverflowEnvironment
+import numpy as np
+import torch
 
+# ── Absolute-import fallback ──────────────────────────────────────────────────
+try:
+    from ..training.overflow_gym_env import OverflowGymEnv, _obs_to_vector, _action_to_decision
+    from ..training.curriculum import CurriculumManager, STAGES
+    from ..training.reward import compute_reward, compute_episode_bonus, W_COLLISION
+    from ..training.ppo_trainer import PPOTrainer, RolloutBuffer
+    from ..policies.flat_mlp_policy import FlatMLPPolicy
+    from ..policies.ticket_attention_policy import TicketAttentionPolicy
+    from ..policies.policy_spec import OBS_DIM
+except ImportError:
+    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+    from training.overflow_gym_env import OverflowGymEnv, _obs_to_vector, _action_to_decision
+    from training.curriculum import CurriculumManager, STAGES
+    from training.reward import compute_reward, compute_episode_bonus, W_COLLISION
+    from training.ppo_trainer import PPOTrainer, RolloutBuffer
+    from policies.flat_mlp_policy import FlatMLPPolicy
+    from policies.ticket_attention_policy import TicketAttentionPolicy
+    from policies.policy_spec import OBS_DIM
 
-def _create_overflow_app():
-    """Build app across create_app variants that may expect a factory or an instance."""
+from fastapi import FastAPI
+from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
+from pydantic import BaseModel
+
+# ── Reward mode config ────────────────────────────────────────────────────────
+
+REWARD_CAP     = 2.0          # per-step ceiling in capped mode
+TOKEN_SCALE    = 0.001        # uncapped: reward += tokens * TOKEN_SCALE
+MAX_TOKEN_BONUS = 5.0         # uncapped mode bonus ceiling per step
+
+# ── Shared training state (updated by training thread, read by API) ───────────
+
+@dataclass
+class CarSnapshot:
+    car_id: int
+    x: float
+    y: float
+    lane: int
+    speed: float
+
+@dataclass
+class EpisodeRecord:
+    episode: int
+    steps: int
+    reward: float
+    outcome: str   # "crash" | "goal" | "timeout"
+    stage: int
+    reward_mode: str
+
+@dataclass
+class TrainingState:
+    # Current frame
+    cars: List[CarSnapshot] = field(default_factory=list)
+    ego_x: float = 0.0
+    ego_lane: int = 2
+    # Live metrics
+    total_steps: int = 0
+    n_updates: int = 0
+    n_episodes: int = 0
+    episode_reward: float = 0.0
+    episode_steps: int = 0
+    mean_reward_100: float = 0.0
+    mean_ep_len: float = 0.0
+    stage: int = 1
+    stage_name: str = "Survival"
+    reward_mode: str = "capped"
+    steps_per_sec: float = 0.0
+    # History (capped at 500 for the chart)
+    reward_history: List[float] = field(default_factory=list)
+    episode_history: List[Dict] = field(default_factory=list)
+    # PPO update metrics
+    last_pg_loss: float = 0.0
+    last_vf_loss: float = 0.0
+    last_entropy: float = 0.0
+    running: bool = True
+    error: Optional[str] = None
+
+
+_state = TrainingState()
+_state_lock = threading.Lock()
+_sse_queue: Deque[str] = deque(maxlen=50)
+
+# ── Reward mode switch (thread-safe) ─────────────────────────────────────────
+
+def get_reward_mode() -> str:
+    with _state_lock:
+        return _state.reward_mode
+
+def set_reward_mode(mode: str) -> None:
+    with _state_lock:
+        _state.reward_mode = mode
+
+def apply_reward_mode(base_reward: float, token_count: int = 0) -> float:
+    mode = get_reward_mode()
+    if mode == "uncapped":
+        bonus = min(token_count * TOKEN_SCALE, MAX_TOKEN_BONUS)
+        return base_reward + bonus
+    else:
+        return min(base_reward, REWARD_CAP) if base_reward > 0 else base_reward
+
+
+# ── Training thread ───────────────────────────────────────────────────────────
+
+def _push_sse(data: dict) -> None:
+    _sse_queue.append(json.dumps(data))
+
+def _snapshot_cars(overflow_obs) -> List[CarSnapshot]:
+    snaps = []
+    if not overflow_obs or not overflow_obs.cars:
+        return snaps
+    for c in overflow_obs.cars:
+        snaps.append(CarSnapshot(
+            car_id=c.carId,
+            x=c.position.x,
+            y=c.position.y if hasattr(c.position, "y") else (c.lane - 2) * 3.7,
+            lane=c.lane,
+            speed=c.speed,
+        ))
+    return snaps
+
+def _training_loop() -> None:
+    global _state
     try:
-        first_param = next(iter(inspect.signature(create_app).parameters.values()))
-        annotation_text = str(first_param.annotation)
-    except (StopIteration, TypeError, ValueError):
-        annotation_text = "typing.Callable"
+        # Build policy + env + curriculum
+        policy = TicketAttentionPolicy(obs_dim=OBS_DIM)
+        env = OverflowGymEnv()
+        curriculum = CurriculumManager()
+
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        policy.to(device)
+
+        optimizer = torch.optim.Adam(policy.parameters(), lr=3e-4, eps=1e-5)
+
+        # PPO hyperparams
+        GAMMA      = 0.99
+        GAE_LAMBDA = 0.95
+        CLIP       = 0.2
+        ENT_COEF   = 0.02
+        VF_COEF    = 0.5
+        MAX_GRAD   = 0.5
+        N_STEPS    = 512
+        BATCH_SIZE = 128
+        N_EPOCHS   = 6
+
+        buf = RolloutBuffer(N_STEPS, OBS_DIM, device)
+
+        ep_rewards: Deque[float] = deque(maxlen=100)
+        ep_lengths: Deque[int]   = deque(maxlen=100)
+
+        obs, _ = env.reset()
+        ep_reward = 0.0
+        ep_steps  = 0
+        total_steps = 0
+        n_updates = 0
+        n_episodes = 0
+        t0 = time.time()
+
+        while True:
+            buf.reset()
+            policy.eval()
+
+            # ── Collect rollout ──────────────────────────────────────────────
+            for _ in range(N_STEPS):
+                curriculum.step(env._sim_time)
+
+                obs_t = torch.as_tensor(obs, dtype=torch.float32, device=device)
+                with torch.no_grad():
+                    act_mean, val = policy(obs_t.unsqueeze(0))
+                act_mean = act_mean.squeeze(0)
+                val      = val.squeeze(0)
+
+                dist   = torch.distributions.Normal(act_mean, torch.ones_like(act_mean) * 0.3)
+                action = dist.sample().clamp(-1, 1)
+                logp   = dist.log_prob(action).sum()
+
+                next_obs, base_reward, term, trunc, info = env.step(action.cpu().numpy())
+
+                # Apply reward mode
+                reward = apply_reward_mode(base_reward)
+
+                buf.add(obs, action.cpu().numpy(), reward, float(val), float(logp), float(term or trunc))
+
+                obs        = next_obs
+                ep_reward += reward
+                ep_steps  += 1
+                total_steps += 1
+
+                # Update live car positions from OverflowEnvironment._cars
+                cars = []
+                overflow_env = env._env  # OverflowEnvironment instance
+                if hasattr(overflow_env, "_cars"):
+                    for c in overflow_env._cars:
+                        cars.append(CarSnapshot(
+                            car_id=c.car_id,
+                            x=c.position,
+                            y=(c.lane - 2) * 3.7,
+                            lane=c.lane,
+                            speed=c.speed,
+                        ))
+
+                with _state_lock:
+                    _state.total_steps = total_steps
+                    _state.episode_reward = ep_reward
+                    _state.episode_steps = ep_steps
+                    _state.steps_per_sec = total_steps / max(time.time() - t0, 1.0)
+                    if cars:
+                        _state.cars = cars
+                        ego = next((c for c in cars if c.car_id == 0), None)
+                        if ego:
+                            _state.ego_x = ego.x
+                            _state.ego_lane = ego.lane
+
+                if term or trunc:
+                    bonus = compute_episode_bonus(
+                        total_steps=ep_steps,
+                        survived=not info.get("collision", False),
+                    )
+                    ep_reward += bonus
+                    ep_rewards.append(ep_reward)
+                    ep_lengths.append(ep_steps)
+                    n_episodes += 1
+
+                    advanced = curriculum.record_episode_reward(ep_reward)
+                    outcome = (
+                        "crash" if info.get("collision") else
+                        ("goal" if info.get("goal_reached") else "timeout")
+                    )
+
+                    ep_rec = {
+                        "episode":     n_episodes,
+                        "steps":       ep_steps,
+                        "reward":      round(ep_reward, 3),
+                        "outcome":     outcome,
+                        "stage":       curriculum.current_stage,
+                        "reward_mode": get_reward_mode(),
+                    }
+
+                    with _state_lock:
+                        _state.n_episodes = n_episodes
+                        _state.stage = curriculum.current_stage
+                        _state.stage_name = curriculum.config.name
+                        _state.mean_reward_100 = float(np.mean(ep_rewards))
+                        _state.mean_ep_len = float(np.mean(ep_lengths))
+                        _state.reward_history.append(round(ep_reward, 3))
+                        if len(_state.reward_history) > 500:
+                            _state.reward_history = _state.reward_history[-500:]
+                        _state.episode_history.append(ep_rec)
+                        if len(_state.episode_history) > 200:
+                            _state.episode_history = _state.episode_history[-200:]
+
+                    _push_sse({"type": "episode", "data": ep_rec})
+
+                    obs, _ = env.reset()
+                    ep_reward = 0.0
+                    ep_steps  = 0
+
+            # ── PPO update ───────────────────────────────────────────────────
+            with torch.no_grad():
+                obs_t = torch.as_tensor(obs, dtype=torch.float32, device=device)
+                _, last_val = policy(obs_t.unsqueeze(0))
+            buf.compute_returns(float(last_val), GAMMA, GAE_LAMBDA)
+
+            policy.train()
+
+            all_obs      = buf.obs
+            all_acts     = buf.acts
+            old_logp     = buf.logp
+            adv          = buf.ret - buf.val
+            adv          = (adv - adv.mean()) / (adv.std() + 1e-8)
+            ret          = buf.ret
+            old_val      = buf.val
+            indices      = torch.randperm(N_STEPS, device=device)
+
+            pg_losses, vf_losses, entropies = [], [], []
+
+            for _ in range(N_EPOCHS):
+                for start in range(0, N_STEPS, BATCH_SIZE):
+                    idx = indices[start: start + BATCH_SIZE]
+                    act_mean, val = policy(all_obs[idx])
+                    val = val.squeeze(-1)
+
+                    dist     = torch.distributions.Normal(act_mean, torch.ones_like(act_mean) * 0.3)
+                    logp     = dist.log_prob(all_acts[idx]).sum(dim=-1)
+                    entropy  = dist.entropy().sum(dim=-1).mean()
+
+                    ratio    = torch.exp(logp - old_logp[idx])
+                    pg_loss  = torch.max(-adv[idx] * ratio, -adv[idx] * ratio.clamp(1 - CLIP, 1 + CLIP)).mean()
+
+                    val_clip = old_val[idx] + (val - old_val[idx]).clamp(-CLIP, CLIP)
+                    vf_loss  = 0.5 * torch.max((val - ret[idx]) ** 2, (val_clip - ret[idx]) ** 2).mean()
+
+                    loss = pg_loss + VF_COEF * vf_loss - ENT_COEF * entropy
+                    optimizer.zero_grad()
+                    loss.backward()
+                    torch.nn.utils.clip_grad_norm_(policy.parameters(), MAX_GRAD)
+                    optimizer.step()
+
+                    pg_losses.append(float(pg_loss))
+                    vf_losses.append(float(vf_loss))
+                    entropies.append(float(entropy))
+
+            n_updates += 1
+            with _state_lock:
+                _state.n_updates     = n_updates
+                _state.last_pg_loss  = round(float(np.mean(pg_losses)), 5)
+                _state.last_vf_loss  = round(float(np.mean(vf_losses)), 5)
+                _state.last_entropy  = round(float(np.mean(entropies)), 5)
+
+            _push_sse({
+                "type": "update",
+                "data": {
+                    "n_updates":   n_updates,
+                    "total_steps": total_steps,
+                    "mean_reward": round(float(np.mean(ep_rewards)) if ep_rewards else 0.0, 3),
+                    "stage":       curriculum.current_stage,
+                    "pg_loss":     round(float(np.mean(pg_losses)), 5),
+                    "vf_loss":     round(float(np.mean(vf_losses)), 5),
+                    "entropy":     round(float(np.mean(entropies)), 5),
+                }
+            })
+
+    except Exception as exc:
+        import traceback
+        tb = traceback.format_exc()
+        with _state_lock:
+            _state.running = False
+            _state.error = f"{exc}\n\n{tb}"
+        print(f"[Training] ERROR: {exc}\n{tb}", flush=True)
+
 
-    expects_instance = (
-        "Environment" in annotation_text and "Callable" not in annotation_text
-    )
-    env_arg = OverflowEnvironment() if expects_instance else OverflowEnvironment
-    return create_app(
-        env_arg, OverflowAction, OverflowObservation, env_name="overflow_env"
-    )
+# ── FastAPI app ───────────────────────────────────────────────────────────────
 
+app = FastAPI(title="Overflow OpenENV")
 
-app = _create_overflow_app()
+@app.on_event("startup")
+def _start_training():
+    t = threading.Thread(target=_training_loop, daemon=True)
+    t.start()
+
+
+@app.get("/health")
+def health():
+    return {"status": "ok"}
+
+
+class ModeRequest(BaseModel):
+    mode: str   # "capped" or "uncapped"
+
+@app.post("/api/mode")
+def set_mode(req: ModeRequest):
+    if req.mode not in ("capped", "uncapped"):
+        return JSONResponse({"error": "mode must be 'capped' or 'uncapped'"}, status_code=400)
+    set_reward_mode(req.mode)
+    return {"mode": req.mode}
+
+
+@app.get("/api/state")
+def get_state():
+    with _state_lock:
+        s = _state
+        return {
+            "total_steps":    s.total_steps,
+            "n_updates":      s.n_updates,
+            "n_episodes":     s.n_episodes,
+            "episode_reward": round(s.episode_reward, 3),
+            "episode_steps":  s.episode_steps,
+            "mean_reward":    round(s.mean_reward_100, 3),
+            "mean_ep_len":    round(s.mean_ep_len, 1),
+            "stage":          s.stage,
+            "stage_name":     s.stage_name,
+            "reward_mode":    s.reward_mode,
+            "steps_per_sec":  round(s.steps_per_sec, 1),
+            "pg_loss":        s.last_pg_loss,
+            "vf_loss":        s.last_vf_loss,
+            "entropy":        s.last_entropy,
+            "reward_history": s.reward_history[-200:],
+            "episode_history": s.episode_history[-50:],
+            "cars":           [asdict(c) for c in s.cars],
+            "ego_x":          s.ego_x,
+            "ego_lane":       s.ego_lane,
+            "running":        s.running,
+            "error":          s.error,
+        }
+
+
+@app.get("/api/stream")
+async def sse_stream():
+    async def generator():
+        last_idx = len(_sse_queue)
+        while True:
+            current = list(_sse_queue)
+            for msg in current[last_idx:]:
+                yield f"data: {msg}\n\n"
+            last_idx = len(current)
+            # Also send a heartbeat state snapshot every 2s
+            with _state_lock:
+                s = _state
+                snap = {
+                    "type": "tick",
+                    "data": {
+                        "total_steps":    s.total_steps,
+                        "episode_reward": round(s.episode_reward, 3),
+                        "episode_steps":  s.episode_steps,
+                        "stage":          s.stage,
+                        "stage_name":     s.stage_name,
+                        "reward_mode":    s.reward_mode,
+                        "cars":           [asdict(c) for c in s.cars],
+                        "ego_x":          s.ego_x,
+                    }
+                }
+            yield f"data: {json.dumps(snap)}\n\n"
+            await asyncio.sleep(0.5)
+
+    return StreamingResponse(generator(), media_type="text/event-stream")
+
+
+# ── HTML Dashboard ────────────────────────────────────────────────────────────
+
+DASHBOARD_HTML = r"""<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<meta name="viewport" content="width=device-width, initial-scale=1.0">
+<title>Overflow OpenENV — Live Training</title>
+<style>
+  * { box-sizing: border-box; margin: 0; padding: 0; }
+  body { background: #0a0a0f; color: #e0e0e0; font-family: 'Courier New', monospace; height: 100vh; display: flex; flex-direction: column; }
+  header { background: #12121a; border-bottom: 1px solid #2a2a40; padding: 10px 20px; display: flex; align-items: center; gap: 16px; }
+  header h1 { font-size: 16px; color: #7eb8ff; letter-spacing: 2px; }
+  .badge { padding: 3px 10px; border-radius: 12px; font-size: 11px; font-weight: bold; }
+  .badge-running { background: #1a3a1a; color: #4caf50; border: 1px solid #4caf50; }
+  .badge-error { background: #3a1a1a; color: #f44336; border: 1px solid #f44336; }
+  .mode-toggle { margin-left: auto; display: flex; gap: 8px; align-items: center; }
+  .mode-btn { padding: 4px 14px; border-radius: 8px; border: 1px solid #444; background: #1a1a2a; color: #aaa; cursor: pointer; font-size: 12px; transition: all 0.2s; }
+  .mode-btn.active { background: #1a3a5a; color: #7eb8ff; border-color: #7eb8ff; }
+  main { flex: 1; display: flex; gap: 0; overflow: hidden; }
+  .left-panel { flex: 0 0 55%; display: flex; flex-direction: column; padding: 16px; gap: 12px; border-right: 1px solid #2a2a40; overflow: hidden; }
+  .right-panel { flex: 1; display: flex; flex-direction: column; padding: 16px; gap: 12px; overflow: hidden; }
+  .panel-title { font-size: 11px; color: #667; letter-spacing: 2px; text-transform: uppercase; margin-bottom: 4px; }
+  canvas { border-radius: 8px; }
+  #road-canvas { width: 100%; height: 200px; background: #111118; border: 1px solid #2a2a40; border-radius: 8px; }
+  .metrics-grid { display: grid; grid-template-columns: repeat(4, 1fr); gap: 8px; }
+  .metric-card { background: #12121a; border: 1px solid #2a2a40; border-radius: 8px; padding: 10px; }
+  .metric-label { font-size: 10px; color: #667; text-transform: uppercase; letter-spacing: 1px; }
+  .metric-value { font-size: 18px; color: #7eb8ff; font-weight: bold; margin-top: 2px; }
+  .metric-sub { font-size: 10px; color: #556; margin-top: 2px; }
+  .stage-bar { background: #12121a; border: 1px solid #2a2a40; border-radius: 8px; padding: 12px; }
+  .stage-stages { display: flex; gap: 6px; margin-top: 8px; }
+  .stage-pip { flex: 1; height: 6px; border-radius: 3px; background: #2a2a40; transition: background 0.5s; }
+  .stage-pip.active { background: #7eb8ff; }
+  .stage-pip.done { background: #4caf50; }
+  #reward-canvas { width: 100%; flex: 1; min-height: 160px; background: #12121a; border: 1px solid #2a2a40; border-radius: 8px; }
+  .episode-table { flex: 1; overflow-y: auto; background: #12121a; border: 1px solid #2a2a40; border-radius: 8px; min-height: 0; }
+  table { width: 100%; border-collapse: collapse; font-size: 11px; }
+  th { padding: 6px 10px; color: #667; text-align: left; border-bottom: 1px solid #2a2a40; position: sticky; top: 0; background: #1a1a28; }
+  td { padding: 5px 10px; border-bottom: 1px solid #1a1a28; }
+  tr:hover td { background: #1a1a2a; }
+  .outcome-crash { color: #f44336; }
+  .outcome-goal  { color: #4caf50; }
+  .outcome-timeout { color: #ff9800; }
+  .ppo-row { display: flex; gap: 8px; flex-wrap: wrap; }
+  .ppo-stat { background: #12121a; border: 1px solid #2a2a40; border-radius: 6px; padding: 6px 12px; font-size: 11px; }
+  .ppo-stat span { color: #7eb8ff; }
+  .error-box { background: #2a0a0a; border: 1px solid #f44336; border-radius: 8px; padding: 12px; font-size: 11px; color: #f44336; white-space: pre-wrap; overflow-y: auto; max-height: 200px; }
+</style>
+</head>
+<body>
+<header>
+  <h1>OVERFLOW OPENENV</h1>
+  <span id="status-badge" class="badge badge-running">TRAINING</span>
+  <div class="mode-toggle">
+    <span style="font-size:11px;color:#667">REWARD MODE:</span>
+    <button class="mode-btn active" id="btn-capped" onclick="setMode('capped')">CAPPED</button>
+    <button class="mode-btn" id="btn-uncapped" onclick="setMode('uncapped')">UNCAPPED (LLM)</button>
+  </div>
+</header>
+<main>
+  <!-- LEFT: road + metrics + stage -->
+  <div class="left-panel">
+    <div>
+      <div class="panel-title">Road View — Live</div>
+      <canvas id="road-canvas"></canvas>
+    </div>
+    <div class="metrics-grid">
+      <div class="metric-card">
+        <div class="metric-label">Total Steps</div>
+        <div class="metric-value" id="m-steps">0</div>
+        <div class="metric-sub" id="m-sps">0 sps</div>
+      </div>
+      <div class="metric-card">
+        <div class="metric-label">Episodes</div>
+        <div class="metric-value" id="m-eps">0</div>
+        <div class="metric-sub" id="m-eplen">avg len: 0</div>
+      </div>
+      <div class="metric-card">
+        <div class="metric-label">Mean Reward (100)</div>
+        <div class="metric-value" id="m-reward">0.00</div>
+        <div class="metric-sub" id="m-curr-r">ep: 0.00</div>
+      </div>
+      <div class="metric-card">
+        <div class="metric-label">PPO Updates</div>
+        <div class="metric-value" id="m-updates">0</div>
+        <div class="metric-sub">policy gradient</div>
+      </div>
+    </div>
+    <div class="stage-bar">
+      <div class="panel-title">Curriculum Stage — <span id="stage-name">Survival</span></div>
+      <div class="stage-stages">
+        <div class="stage-pip active" id="pip-1"></div>
+        <div class="stage-pip" id="pip-2"></div>
+        <div class="stage-pip" id="pip-3"></div>
+        <div class="stage-pip" id="pip-4"></div>
+      </div>
+    </div>
+    <div>
+      <div class="panel-title">PPO Losses</div>
+      <div class="ppo-row">
+        <div class="ppo-stat">Policy: <span id="pg-loss">—</span></div>
+        <div class="ppo-stat">Value: <span id="vf-loss">—</span></div>
+        <div class="ppo-stat">Entropy: <span id="entropy">—</span></div>
+      </div>
+    </div>
+    <div id="error-section" style="display:none">
+      <div class="panel-title" style="color:#f44336">Error</div>
+      <div class="error-box" id="error-text"></div>
+    </div>
+  </div>
+
+  <!-- RIGHT: reward chart + episode table -->
+  <div class="right-panel">
+    <div style="flex:0 0 auto">
+      <div class="panel-title">Reward History</div>
+    </div>
+    <canvas id="reward-canvas"></canvas>
+    <div style="flex:0 0 auto">
+      <div class="panel-title">Episode Log</div>
+    </div>
+    <div class="episode-table">
+      <table>
+        <thead><tr>
+          <th>#</th><th>Steps</th><th>Reward</th><th>Outcome</th><th>Stage</th><th>Mode</th>
+        </tr></thead>
+        <tbody id="ep-tbody"></tbody>
+      </table>
+    </div>
+  </div>
+</main>
+
+<script>
+// ── State ──────────────────────────────────────────────────────────────────
+let state = { cars: [], reward_history: [], episode_history: [], stage: 1, reward_mode: 'capped' };
+
+// ── Road canvas ────────────────────────────────────────────────────────────
+const roadCanvas = document.getElementById('road-canvas');
+const roadCtx    = roadCanvas.getContext('2d');
+const N_LANES    = 3;
+const LANE_H     = 40;
+const ROAD_PAD   = 20;
+
+function resizeRoad() {
+  roadCanvas.width  = roadCanvas.offsetWidth;
+  roadCanvas.height = roadCanvas.offsetHeight;
+}
+resizeRoad();
+window.addEventListener('resize', resizeRoad);
+
+function laneY(lane) {
+  // lane 1..3, top = lane 1
+  const totalH = N_LANES * LANE_H;
+  const offsetY = (roadCanvas.height - totalH) / 2;
+  return offsetY + (lane - 1) * LANE_H + LANE_H / 2;
+}
+
+function carX(x, egoX) {
+  // Center ego at 30% of canvas
+  const w = roadCanvas.width;
+  const scale = 0.8;   // pixels per unit
+  return w * 0.3 + (x - egoX) * scale;
+}
+
+function drawRoad() {
+  const w = roadCanvas.width, h = roadCanvas.height;
+  roadCtx.clearRect(0, 0, w, h);
+
+  // Road background
+  const totalH = N_LANES * LANE_H;
+  const offsetY = (h - totalH) / 2;
+  roadCtx.fillStyle = '#1a1a28';
+  roadCtx.fillRect(0, offsetY, w, totalH);
+
+  // Lane dividers
+  roadCtx.setLineDash([20, 15]);
+  roadCtx.strokeStyle = '#3a3a50';
+  roadCtx.lineWidth = 1;
+  for (let i = 1; i < N_LANES; i++) {
+    const y = offsetY + i * LANE_H;
+    roadCtx.beginPath();
+    roadCtx.moveTo(0, y);
+    roadCtx.lineTo(w, y);
+    roadCtx.stroke();
+  }
+  roadCtx.setLineDash([]);
+
+  // Road edges
+  roadCtx.strokeStyle = '#5a5a80';
+  roadCtx.lineWidth = 2;
+  roadCtx.beginPath(); roadCtx.moveTo(0, offsetY); roadCtx.lineTo(w, offsetY); roadCtx.stroke();
+  roadCtx.beginPath(); roadCtx.moveTo(0, offsetY + totalH); roadCtx.lineTo(w, offsetY + totalH); roadCtx.stroke();
+
+  // Cars
+  const egoX = state.ego_x || 0;
+  for (const car of state.cars || []) {
+    const cx = carX(car.x, egoX);
+    if (cx < -50 || cx > w + 50) continue;
+    const cy = laneY(car.lane);
+    const isEgo = car.car_id === 0;
+
+    // Car body
+    roadCtx.save();
+    roadCtx.translate(cx, cy);
+    const cw = 36, ch = 18;
+    roadCtx.fillStyle = isEgo ? '#2a5a9a' : '#3a2a1a';
+    roadCtx.strokeStyle = isEgo ? '#7eb8ff' : '#ff9800';
+    roadCtx.lineWidth = isEgo ? 2 : 1;
+    roadCtx.beginPath();
+    roadCtx.roundRect(-cw/2, -ch/2, cw, ch, 4);
+    roadCtx.fill();
+    roadCtx.stroke();
+
+    // Label
+    roadCtx.fillStyle = isEgo ? '#7eb8ff' : '#ff9800';
+    roadCtx.font = isEgo ? 'bold 9px Courier New' : '8px Courier New';
+    roadCtx.textAlign = 'center';
+    roadCtx.textBaseline = 'middle';
+    roadCtx.fillText(isEgo ? 'EGO' : `C${car.car_id}`, 0, 0);
+    roadCtx.restore();
+  }
+}
+
+// ── Reward chart ────────────────────────────────────────────────────────────
+const rwCanvas = document.getElementById('reward-canvas');
+const rwCtx    = rwCanvas.getContext('2d');
+
+function resizeRw() {
+  rwCanvas.width  = rwCanvas.offsetWidth;
+  rwCanvas.height = rwCanvas.offsetHeight;
+}
+resizeRw();
+window.addEventListener('resize', resizeRw);
+
+function drawRewardChart() {
+  const w = rwCanvas.width, h = rwCanvas.height;
+  rwCtx.clearRect(0, 0, w, h);
+
+  const hist = state.reward_history || [];
+  if (hist.length < 2) {
+    rwCtx.fillStyle = '#667';
+    rwCtx.font = '12px Courier New';
+    rwCtx.textAlign = 'center';
+    rwCtx.fillText('Waiting for episodes...', w/2, h/2);
+    return;
+  }
+
+  const pad = { top: 10, right: 10, bottom: 30, left: 50 };
+  const pw = w - pad.left - pad.right;
+  const ph = h - pad.top - pad.bottom;
+
+  const minR = Math.min(...hist);
+  const maxR = Math.max(...hist);
+  const rangeR = maxR - minR || 1;
+
+  const xScale = pw / (hist.length - 1);
+  const yScale = ph / rangeR;
+
+  // Grid
+  rwCtx.strokeStyle = '#1e1e2e';
+  rwCtx.lineWidth = 1;
+  for (let i = 0; i <= 4; i++) {
+    const y = pad.top + ph * (i / 4);
+    rwCtx.beginPath(); rwCtx.moveTo(pad.left, y); rwCtx.lineTo(pad.left + pw, y); rwCtx.stroke();
+    const val = maxR - rangeR * (i / 4);
+    rwCtx.fillStyle = '#556';
+    rwCtx.font = '9px Courier New';
+    rwCtx.textAlign = 'right';
+    rwCtx.fillText(val.toFixed(1), pad.left - 4, y + 3);
+  }
+
+  // Zero line
+  if (minR < 0 && maxR > 0) {
+    const zy = pad.top + (maxR / rangeR) * ph;
+    rwCtx.strokeStyle = '#3a3a50';
+    rwCtx.lineWidth = 1;
+    rwCtx.setLineDash([4, 4]);
+    rwCtx.beginPath(); rwCtx.moveTo(pad.left, zy); rwCtx.lineTo(pad.left + pw, zy); rwCtx.stroke();
+    rwCtx.setLineDash([]);
+  }
+
+  // Moving average (window=10)
+  const MA = 10;
+  const ma = hist.map((_, i) => {
+    const sl = hist.slice(Math.max(0, i - MA + 1), i + 1);
+    return sl.reduce((a, b) => a + b, 0) / sl.length;
+  });
+
+  // Raw line
+  rwCtx.strokeStyle = 'rgba(126,184,255,0.25)';
+  rwCtx.lineWidth = 1;
+  rwCtx.beginPath();
+  hist.forEach((v, i) => {
+    const x = pad.left + i * xScale;
+    const y = pad.top + (maxR - v) * yScale;
+    i === 0 ? rwCtx.moveTo(x, y) : rwCtx.lineTo(x, y);
+  });
+  rwCtx.stroke();
+
+  // MA line
+  rwCtx.strokeStyle = '#7eb8ff';
+  rwCtx.lineWidth = 2;
+  rwCtx.beginPath();
+  ma.forEach((v, i) => {
+    const x = pad.left + i * xScale;
+    const y = pad.top + (maxR - v) * yScale;
+    i === 0 ? rwCtx.moveTo(x, y) : rwCtx.lineTo(x, y);
+  });
+  rwCtx.stroke();
+
+  // X axis label
+  rwCtx.fillStyle = '#556';
+  rwCtx.font = '9px Courier New';
+  rwCtx.textAlign = 'center';
+  rwCtx.fillText(`Episodes (${hist.length})`, pad.left + pw / 2, h - 6);
+}
+
+// ── Episode table ───────────────────────────────────────────────────────────
+function updateEpisodeTable() {
+  const tbody = document.getElementById('ep-tbody');
+  const episodes = (state.episode_history || []).slice().reverse().slice(0, 50);
+  tbody.innerHTML = episodes.map(ep => `
+    <tr>
+      <td>${ep.episode}</td>
+      <td>${ep.steps}</td>
+      <td style="color:${ep.reward >= 0 ? '#4caf50' : '#f44336'}">${ep.reward.toFixed(2)}</td>
+      <td class="outcome-${ep.outcome}">${ep.outcome.toUpperCase()}</td>
+      <td>${ep.stage}</td>
+      <td style="color:#ff9800">${ep.mode || ep.reward_mode || 'capped'}</td>
+    </tr>
+  `).join('');
+}
+
+// ── Metric update ───────────────────────────────────────────────────────────
+function updateMetrics(s) {
+  document.getElementById('m-steps').textContent = s.total_steps?.toLocaleString() || '0';
+  document.getElementById('m-sps').textContent   = `${(s.steps_per_sec||0).toFixed(0)} sps`;
+  document.getElementById('m-eps').textContent    = s.n_episodes || '0';
+  document.getElementById('m-eplen').textContent  = `avg len: ${(s.mean_ep_len||0).toFixed(0)}`;
+  document.getElementById('m-reward').textContent = (s.mean_reward||0).toFixed(2);
+  document.getElementById('m-curr-r').textContent = `ep: ${(s.episode_reward||0).toFixed(2)}`;
+  document.getElementById('m-updates').textContent = s.n_updates || '0';
+  document.getElementById('pg-loss').textContent  = s.pg_loss ?? '—';
+  document.getElementById('vf-loss').textContent  = s.vf_loss ?? '—';
+  document.getElementById('entropy').textContent  = s.entropy ?? '—';
+  document.getElementById('stage-name').textContent = s.stage_name || 'Survival';
+
+  const stage = s.stage || 1;
+  for (let i = 1; i <= 4; i++) {
+    const pip = document.getElementById(`pip-${i}`);
+    pip.className = 'stage-pip' + (i < stage ? ' done' : (i === stage ? ' active' : ''));
+  }
+
+  // Mode buttons
+  const mode = s.reward_mode || 'capped';
+  document.getElementById('btn-capped').className   = 'mode-btn' + (mode === 'capped' ? ' active' : '');
+  document.getElementById('btn-uncapped').className = 'mode-btn' + (mode === 'uncapped' ? ' active' : '');
+
+  if (s.error) {
+    document.getElementById('error-section').style.display = 'block';
+    document.getElementById('error-text').textContent = s.error;
+    document.getElementById('status-badge').textContent = 'ERROR';
+    document.getElementById('status-badge').className = 'badge badge-error';
+  }
+}
+
+// ── Mode switch ─────────────────────────────────────────────────────────────
+function setMode(mode) {
+  fetch('/api/mode', {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify({ mode })
+  }).then(r => r.json()).then(d => {
+    state.reward_mode = d.mode;
+    updateMetrics(state);
+  });
+}
+
+// ── Poll /api/state every 2s (SSE as supplement) ───────────────────────────
+function poll() {
+  fetch('/api/state').then(r => r.json()).then(s => {
+    Object.assign(state, s);
+    drawRoad();
+    drawRewardChart();
+    updateMetrics(s);
+    updateEpisodeTable();
+  }).catch(() => {});
+}
+setInterval(poll, 2000);
+poll();
+
+// ── SSE for fast episode events ─────────────────────────────────────────────
+const evtSrc = new EventSource('/api/stream');
+evtSrc.onmessage = (e) => {
+  try {
+    const msg = JSON.parse(e.data);
+    if (msg.type === 'episode') {
+      if (!state.episode_history) state.episode_history = [];
+      state.episode_history.push(msg.data);
+      if (!state.reward_history) state.reward_history = [];
+      state.reward_history.push(msg.data.reward);
+      updateEpisodeTable();
+      drawRewardChart();
+    } else if (msg.type === 'tick') {
+      Object.assign(state, msg.data);
+      drawRoad();
+    } else if (msg.type === 'update') {
+      Object.assign(state, msg.data);
+      updateMetrics(state);
+    }
+  } catch(err) {}
+};
+
+// Render loop for smooth road animation
+function renderLoop() {
+  drawRoad();
+  requestAnimationFrame(renderLoop);
+}
+renderLoop();
+</script>
+</body>
+</html>
+"""
+
+
+@app.get("/", response_class=HTMLResponse)
+@app.get("/web", response_class=HTMLResponse)
+def dashboard():
+    return HTMLResponse(content=DASHBOARD_HTML)
 
 
 def main():
-    """Entry point for direct execution via uv run or python -m."""
     import uvicorn
-
     uvicorn.run(app, host="0.0.0.0", port=8000)
 
 

server/overflow_environment.py

@@ -15,18 +15,33 @@ from dataclasses import dataclass, field
 from typing import Any, List, Optional
 from uuid import uuid4
 
-from openenv.core.env_server.interfaces import Environment
-from openenv.core.env_server.types import State
-
-from ..models import (
-    CarStateData,
-    LaneOccupancyData,
-    OverflowAction,
-    OverflowObservation,
-    OverflowState,
-    Position,
-    ProximityData,
-)
+try:
+    from openenv.core.env_server.interfaces import Environment
+    from openenv.core.env_server.types import State
+except ImportError:
+    class Environment:  # stub for training-only mode
+        pass
+    class State:
+        pass
+
+try:
+    from ..models import (
+        CarStateData, LaneOccupancyData, OverflowAction,
+        OverflowObservation, OverflowState, Position, ProximityData,
+    )
+    from ..policies.flat_mlp_policy import FlatMLPPolicy
+    from ..policies.ticket_attention_policy import TicketAttentionPolicy
+    from ..policies.policy_spec import OBS_DIM
+    from .policy_adapter import overflow_obs_to_policy_obs, policy_action_to_decision
+except ImportError:
+    from models import (
+        CarStateData, LaneOccupancyData, OverflowAction,
+        OverflowObservation, OverflowState, Position, ProximityData,
+    )
+    from policies.flat_mlp_policy import FlatMLPPolicy
+    from policies.ticket_attention_policy import TicketAttentionPolicy
+    from policies.policy_spec import OBS_DIM
+    from server.policy_adapter import overflow_obs_to_policy_obs, policy_action_to_decision
 
 # --- Constants ---
 NUM_LANES = 3
@@ -253,6 +268,11 @@ class OverflowEnvironment(Environment):
         self._cars: List[Car] = []
         self._rng = random.Random()
         self._done = False
+        self._last_obs: Optional[OverflowObservation] = None
+        self._policies = {
+            "flat_mlp":         FlatMLPPolicy(obs_dim=OBS_DIM),
+            "ticket_attention":  TicketAttentionPolicy(obs_dim=OBS_DIM),
+        }
 
     def _build_observation(
         self,
@@ -332,7 +352,8 @@ class OverflowEnvironment(Environment):
                 )
             )
 
-        return self._build_observation(incident_report="", reward=0.0)
+        self._last_obs = self._build_observation(incident_report="", reward=0.0)
+        return self._last_obs
 
     def step(
         self,
@@ -347,6 +368,18 @@ class OverflowEnvironment(Environment):
                 reward=0.0,
             )
 
+        # Policy intercept: decision="policy:flat_mlp" or "policy:ticket_attention"
+        if action.decision.startswith("policy:") and self._last_obs is not None:
+            policy_name = action.decision.split(":", 1)[1].lower()
+            if policy_name in self._policies:
+                obs_vec = overflow_obs_to_policy_obs(self._last_obs)
+                act_vec = self._policies[policy_name].predict(obs_vec)
+                decision, reasoning = policy_action_to_decision(act_vec)
+                action = OverflowAction(
+                    decision=decision,
+                    reasoning=f"[{policy_name}] {reasoning}",
+                )
+
         self._state.step_count += 1
         reward = 0.0
         incidents = []
@@ -451,11 +484,12 @@ class OverflowEnvironment(Environment):
             "\n".join(incidents) if incidents else "Observer: No incidents this step."
         )
 
-        return self._build_observation(
+        self._last_obs = self._build_observation(
             incident_report=incident_report,
             reward=reward,
             proximities=proximity_list,
         )
+        return self._last_obs
 
     @property
     def state(self) -> OverflowState:

server/policy_adapter.py

@@ -0,0 +1,80 @@
+"""
+Adapter between OverflowObservation (2D road grid) and the OpenENV policy
+observation format (ego state + ticket matrix).
+
+Nearby cars are converted to collision_risk tickets so TicketAttentionPolicy
+can reason about them using the same mechanism it was designed for.
+"""
+
+from __future__ import annotations
+
+import math
+import numpy as np
+
+try:
+    from ..policies.policy_spec import build_obs, build_ticket_vector, OBS_DIM
+except ImportError:
+    from policies.policy_spec import build_obs, build_ticket_vector, OBS_DIM
+
+
+def overflow_obs_to_policy_obs(obs) -> np.ndarray:
+    """OverflowObservation → 603-dim numpy vector for our policies."""
+    cars = 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_speed_ms = ego.speed / 4.5            # OverflowEnv speed units → m/s
+    ego_x        = ego.position.x
+    ego_y        = (ego.lane - 2) * 3.7      # lane → lateral metres
+
+    ticket_vectors = []
+    for car in cars:
+        if car.carId == 0:
+            continue
+        rel_x    = car.position.x - ego.position.x
+        rel_y    = (car.lane - ego.lane) * 3.7
+        car_spd  = car.speed / 4.5
+        distance = math.sqrt(rel_x ** 2 + rel_y ** 2)
+        if distance > 80:
+            continue
+        closing  = max(ego_speed_ms - car_spd * math.copysign(1, max(rel_x, 0.01)), 0.1)
+        ttc      = min(distance / closing, 30.0)
+        severity = 1.0 if distance < 8 else (0.75 if distance < 15 else 0.5)
+        ticket_vectors.append(build_ticket_vector(
+            severity_weight=severity, ttl=5.0,
+            pos_x=rel_x, pos_y=rel_y, pos_z=0.0,
+            vel_x=car_spd, vel_y=0.0, vel_z=0.0,
+            heading=0.0,
+            size_length=4.0, size_width=2.0, size_height=1.5,
+            distance=distance, time_to_collision=ttc,
+            bearing=math.atan2(rel_y, max(rel_x, 0.01)),
+            ticket_type="collision_risk", entity_type="vehicle", confidence=1.0,
+        ))
+
+    tv = np.array(ticket_vectors, dtype=np.float32) if ticket_vectors else None
+    return build_obs(
+        ego_x=ego_x, ego_y=ego_y, ego_z=0.0,
+        ego_vx=ego_speed_ms, ego_vy=0.0,
+        heading=0.0, speed=ego_speed_ms,
+        steer=0.0, throttle=0.5, brake=0.0,
+        ticket_vectors=tv,
+    )
+
+
+def policy_action_to_decision(action_vec: np.ndarray) -> tuple[str, str]:
+    """Continuous [steer, throttle, brake] → (text decision, reasoning)."""
+    steer, throttle, brake = float(action_vec[0]), float(action_vec[1]), float(action_vec[2])
+    if abs(steer) > 0.35:
+        decision  = "lane_change_left" if steer < 0 else "lane_change_right"
+        reasoning = f"steer={steer:.2f}: lateral avoidance"
+    elif brake > 0.25:
+        decision  = "brake"
+        reasoning = f"brake={brake:.2f}: closing gap"
+    elif throttle > 0.20:
+        decision  = "accelerate"
+        reasoning = f"throttle={throttle:.2f}: clear ahead"
+    else:
+        decision  = "maintain"
+        reasoning = f"s={steer:.2f} t={throttle:.2f} b={brake:.2f}: holding course"
+    return decision, reasoning

training/curriculum.py

@@ -0,0 +1,99 @@
+"""
+CurriculumManager — ported from openenv/training/curriculum.py.
+
+Same 4-stage progression and same reward thresholds. Adapted for
+OverflowEnvironment: no ticket injection (the env has its own scripted
+NPCs), stages instead control training logging and advancement criteria.
+
+Stage 1  No extra pressure.  Goal: learn basic speed + lane keeping.
+Stage 2  Standard traffic.   Goal: survive without crashing.
+Stage 3  Evaluate more.      Goal: consistent goal-reaching.
+Stage 4  Full evaluation.    Goal: high mean reward over long window.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List
+
+
+@dataclass
+class StageConfig:
+    stage:             int
+    name:              str
+    description:       str
+    advance_threshold: float   # mean episode reward to advance
+    advance_window:    int     # consecutive episodes required
+
+
+STAGES: List[StageConfig] = [
+    StageConfig(
+        stage=1, name="Survival",
+        description="Learn basic speed control and lane keeping.",
+        advance_threshold=50.0, advance_window=8,
+    ),
+    StageConfig(
+        stage=2, name="Crash Avoidance",
+        description="Navigate traffic without colliding.",
+        advance_threshold=120.0, advance_window=15,
+    ),
+    StageConfig(
+        stage=3, name="Goal Reaching",
+        description="Consistently reach the goal position.",
+        advance_threshold=200.0, advance_window=15,
+    ),
+    StageConfig(
+        stage=4, name="Mastery",
+        description="High reward, smooth driving, minimal near-misses.",
+        advance_threshold=280.0, advance_window=15,
+    ),
+]
+
+
+class CurriculumManager:
+    """
+    Tracks stage progression based on episode rewards.
+    Same API as openenv CurriculumManager — PPOTrainer calls it unchanged.
+    """
+
+    def __init__(self, seed: int = 0):
+        self._stage_idx    = 0
+        self._rewards: List[float] = []
+        self._auto_advance = True
+
+    @property
+    def current_stage(self) -> int:
+        return STAGES[self._stage_idx].stage
+
+    @property
+    def config(self) -> StageConfig:
+        return STAGES[self._stage_idx]
+
+    def step(self, sim_time: float) -> list:
+        """No ticket injection in OverflowEnvironment — always returns []."""
+        return []
+
+    def record_episode_reward(self, reward: float) -> bool:
+        """Record episode reward and advance stage if threshold met."""
+        self._rewards.append(reward)
+        cfg    = self.config
+        window = self._rewards[-cfg.advance_window:]
+
+        if (
+            self._auto_advance
+            and len(window) >= cfg.advance_window
+            and sum(window) / len(window) >= cfg.advance_threshold
+            and self._stage_idx < len(STAGES) - 1
+        ):
+            self._stage_idx += 1
+            self._rewards = []
+            print(f"[Curriculum] Advanced to Stage {self.current_stage}: {self.config.name}")
+            return True
+        return False
+
+    def force_stage(self, stage: int) -> None:
+        idx = stage - 1
+        if 0 <= idx < len(STAGES):
+            self._stage_idx = idx
+            self._rewards   = []
+            print(f"[Curriculum] Forced to Stage {stage}: {self.config.name}")

training/overflow_gym_env.py

@@ -0,0 +1,170 @@
+"""
+Gymnasium wrapper around OverflowEnvironment.
+
+Bridges the gap between OverflowEnvironment (text actions, structured obs)
+and our PPO trainer (continuous actions, numeric obs vector).
+
+Observation: 603-dim float32 vector (same layout as CarEnv3D — ego state +
+             collision-risk ticket matrix built from nearby cars)
+
+Action:      [steer, throttle, brake] all in [-1, 1]
+             → mapped to text decision for OverflowEnvironment
+
+This makes OverflowEnvironment a drop-in replacement for CarEnv3D so that
+FlatMLPPolicy and TicketAttentionPolicy train with the exact same PPO loop.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Dict, Optional, Tuple
+
+import numpy as np
+import gymnasium as gym
+from gymnasium import spaces
+
+from ..server.overflow_environment import OverflowEnvironment
+from ..models import OverflowAction
+from ..policies.policy_spec import (
+    build_obs, build_ticket_vector, OBS_DIM,
+)
+from .reward import compute_reward
+
+
+# ── Action mapping ────────────────────────────────────────────────────────────
+
+def _action_to_decision(action: np.ndarray) -> str:
+    steer, throttle, brake = float(action[0]), float(action[1]), float(action[2])
+    if abs(steer) > 0.35:
+        return "lane_change_left" if steer < 0 else "lane_change_right"
+    if brake > 0.25:
+        return "brake"
+    if throttle > 0.20:
+        return "accelerate"
+    return "maintain"
+
+
+# ── Observation extraction ────────────────────────────────────────────────────
+
+def _obs_to_vector(overflow_obs) -> np.ndarray:
+    """OverflowObservation → 603-dim numpy vector matching policy_spec layout."""
+    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_speed_ms = ego.speed / 4.5
+    ego_x        = ego.position.x
+    ego_y        = (ego.lane - 2) * 3.7
+
+    ticket_vectors = []
+    for car in cars:
+        if car.carId == 0:
+            continue
+        rel_x    = car.position.x - ego.position.x
+        rel_y    = (car.lane - ego.lane) * 3.7
+        car_spd  = car.speed / 4.5
+        distance = math.sqrt(rel_x ** 2 + rel_y ** 2)
+        if distance > 80:
+            continue
+        closing  = max(ego_speed_ms - car_spd * math.copysign(1, max(rel_x, 0.01)), 0.1)
+        ttc      = min(distance / closing, 30.0)
+        severity = 1.0 if distance < 8 else (0.75 if distance < 15 else 0.5)
+        ticket_vectors.append(build_ticket_vector(
+            severity_weight=severity, ttl=5.0,
+            pos_x=rel_x, pos_y=rel_y, pos_z=0.0,
+            vel_x=car_spd, vel_y=0.0, vel_z=0.0,
+            heading=0.0,
+            size_length=4.0, size_width=2.0, size_height=1.5,
+            distance=distance, time_to_collision=ttc,
+            bearing=math.atan2(rel_y, max(rel_x, 0.01)),
+            ticket_type="collision_risk", entity_type="vehicle", confidence=1.0,
+        ))
+
+    tv = np.array(ticket_vectors, dtype=np.float32) if ticket_vectors else None
+    return build_obs(
+        ego_x=ego_x, ego_y=ego_y, ego_z=0.0,
+        ego_vx=ego_speed_ms, ego_vy=0.0,
+        heading=0.0, speed=ego_speed_ms,
+        steer=0.0, throttle=0.5, brake=0.0,
+        ticket_vectors=tv,
+    )
+
+
+# ── Gymnasium wrapper ─────────────────────────────────────────────────────────
+
+class OverflowGymEnv(gym.Env):
+    """
+    Gymnasium-compatible wrapper around OverflowEnvironment.
+
+    Provides the same interface as CarEnv3D so PPOTrainer works unchanged.
+    """
+
+    metadata = {"render_modes": []}
+
+    def __init__(self):
+        super().__init__()
+        self._env = OverflowEnvironment()
+        self._last_overflow_obs = None
+        self._prev_action = np.zeros(3, dtype=np.float32)
+        self._sim_time = 0.0         # incremented each step (mirrors CarEnv3D._sim_time)
+        self._step_dt  = 0.1         # seconds per step
+
+        self.observation_space = spaces.Box(
+            low=-1.0, high=1.0, shape=(OBS_DIM,), dtype=np.float32
+        )
+        self.action_space = spaces.Box(
+            low=-1.0, high=1.0, shape=(3,), dtype=np.float32
+        )
+
+    def reset(
+        self,
+        seed: Optional[int] = None,
+        options: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[np.ndarray, Dict]:
+        super().reset(seed=seed)
+        self._last_overflow_obs = self._env.reset(seed=seed)
+        self._prev_action = np.zeros(3, dtype=np.float32)
+        self._sim_time = 0.0
+        return _obs_to_vector(self._last_overflow_obs), {}
+
+    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict]:
+        decision = _action_to_decision(action)
+        overflow_action = OverflowAction(decision=decision, reasoning="")
+        overflow_obs = self._env.step(overflow_action)
+        self._last_overflow_obs = overflow_obs
+        self._sim_time += self._step_dt
+
+        obs_vec = _obs_to_vector(overflow_obs)
+
+        # Extract signals for reward shaping
+        ego = next((c for c in overflow_obs.cars if c.carId == 0), None)
+        ego_speed_ms = (ego.speed / 4.5) if ego else 0.0
+        ego_y        = ((ego.lane - 2) * 3.7) if ego else 0.0
+
+        collision    = any("CRASH" in p for p in (overflow_obs.incident_report or "").split("\n")
+                           if "Car 0" in p)
+        goal_reached = overflow_obs.done and not collision
+
+        reward = compute_reward(
+            ego_speed    = ego_speed_ms,
+            ego_y        = ego_y,
+            action       = action,
+            prev_action  = self._prev_action,
+            collision    = collision,
+            goal_reached = goal_reached,
+            near_miss    = "NEAR MISS" in (overflow_obs.incident_report or ""),
+            raw_reward   = overflow_obs.reward or 0.0,
+        )
+
+        self._prev_action = action.copy()
+
+        terminated = overflow_obs.done
+        truncated  = False
+        info: Dict[str, Any] = {
+            "collision":    collision,
+            "goal_reached": goal_reached,
+            "incident":     overflow_obs.incident_report,
+        }
+
+        return obs_vec, reward, terminated, truncated, info

training/ppo_trainer.py

@@ -0,0 +1,329 @@
+"""
+PPO trainer — ported directly from openenv/training/ppo_trainer.py.
+
+Same algorithm, same hyperparameters, same GAE implementation.
+Only change: uses OverflowGymEnv instead of CarEnv3D.
+
+Usage:
+    from overflow_env.training.ppo_trainer import run_training
+    run_training(policy_type="attention", total_steps=2_000_000)
+"""
+
+from __future__ import annotations
+
+import time
+from collections import deque
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+from .overflow_gym_env import OverflowGymEnv
+from .curriculum import CurriculumManager
+from .reward import compute_episode_bonus
+from ..policies.base_policy import BasePolicy
+from ..policies.policy_spec import OBS_DIM
+
+
+# ── Rollout buffer ─────────────────────────────────────────────────────────────
+# Identical to openenv/training/ppo_trainer.py
+
+class RolloutBuffer:
+    def __init__(self, n_steps: int, obs_dim: int, device: torch.device):
+        self.n    = n_steps
+        self.obs  = torch.zeros(n_steps, obs_dim, device=device)
+        self.acts = torch.zeros(n_steps, 3,        device=device)
+        self.rew  = torch.zeros(n_steps,            device=device)
+        self.val  = torch.zeros(n_steps,            device=device)
+        self.logp = torch.zeros(n_steps,            device=device)
+        self.done = torch.zeros(n_steps,            device=device)
+        self.ptr  = 0
+
+    def add(self, obs, act, rew, val, logp, done):
+        i = self.ptr
+        self.obs[i]  = torch.as_tensor(obs, dtype=torch.float32)
+        self.acts[i] = torch.as_tensor(act, dtype=torch.float32)
+        self.rew[i]  = float(rew)
+        self.val[i]  = float(val)
+        self.logp[i] = float(logp)
+        self.done[i] = float(done)
+        self.ptr += 1
+
+    def full(self) -> bool:
+        return self.ptr >= self.n
+
+    def reset(self):
+        self.ptr = 0
+
+    def compute_returns(self, last_val: float, gamma: float, gae_lambda: float):
+        """Generalized Advantage Estimation — identical to openenv."""
+        adv = torch.zeros_like(self.rew)
+        gae = 0.0
+        for t in reversed(range(self.n)):
+            next_val = last_val if t == self.n - 1 else float(self.val[t + 1])
+            delta    = self.rew[t] + gamma * next_val * (1 - self.done[t]) - self.val[t]
+            gae      = delta + gamma * gae_lambda * (1 - self.done[t]) * gae
+            adv[t]   = gae
+        self.ret = adv + self.val
+
+
+# ── PPO Trainer ────────────────────────────────────────────────────────────────
+
+class PPOTrainer:
+    """
+    Identical to openenv PPOTrainer — same hyperparameters, same PPO update.
+    Environment is OverflowGymEnv instead of CarEnv3D.
+    """
+
+    def __init__(
+        self,
+        policy:        BasePolicy,
+        env:           OverflowGymEnv,
+        curriculum:    Optional[CurriculumManager] = None,
+        # PPO hyperparameters — same defaults as openenv
+        lr:            float = 3e-4,
+        gamma:         float = 0.99,
+        gae_lambda:    float = 0.95,
+        clip_range:    float = 0.2,
+        clip_range_vf: float = 0.2,
+        ent_coef:      float = 0.02,
+        vf_coef:       float = 0.5,
+        max_grad_norm: float = 0.5,
+        n_steps:       int   = 2048,
+        batch_size:    int   = 256,
+        n_epochs:      int   = 10,
+        save_dir:      str   = "checkpoints",
+        log_interval:  int   = 10,
+        device:        str   = "auto",
+    ):
+        self.policy     = policy
+        self.env        = env
+        self.curriculum = curriculum or CurriculumManager()
+        self.gamma      = gamma
+        self.gae_lambda = gae_lambda
+        self.clip       = clip_range
+        self.clip_vf    = clip_range_vf
+        self.ent_coef   = ent_coef
+        self.vf_coef    = vf_coef
+        self.max_grad   = max_grad_norm
+        self.n_steps    = n_steps
+        self.batch_size = batch_size
+        self.n_epochs   = n_epochs
+        self.log_every  = log_interval
+        self.save_dir   = Path(save_dir)
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+
+        if device == "auto":
+            device = "cuda" if torch.cuda.is_available() else \
+                     "mps"  if torch.backends.mps.is_available() else "cpu"
+        self.device = torch.device(device)
+        self.policy.to(self.device)
+
+        self.optimizer = optim.Adam(policy.parameters(), lr=lr, eps=1e-5)
+        self.scheduler = optim.lr_scheduler.LinearLR(
+            self.optimizer, start_factor=1.0, end_factor=0.1, total_iters=500,
+        )
+
+        self.buffer = RolloutBuffer(n_steps, OBS_DIM, self.device)
+
+        self.ep_rewards  = deque(maxlen=100)
+        self.ep_lengths  = deque(maxlen=100)
+        self.total_steps = 0
+        self.n_updates   = 0
+
+    # ── Main training loop ─────────────────────────────────────────────────────
+
+    def train(self, total_steps: int = 2_000_000) -> None:
+        print(f"\n{'='*70}", flush=True)
+        print(f"  OpenENV PPO Training — policy={self.policy.__class__.__name__}", flush=True)
+        print(f"  total_steps={total_steps}  n_steps={self.n_steps}  lr={self.optimizer.param_groups[0]['lr']:.0e}", flush=True)
+        print(f"  gamma={self.gamma}  gae_lambda={self.gae_lambda}  clip={self.clip}  ent_coef={self.ent_coef}", flush=True)
+        print(f"{'='*70}\n", flush=True)
+
+        obs, _ = self.env.reset()
+        ep_reward = 0.0
+        ep_steps  = 0
+        t0 = time.time()
+
+        while self.total_steps < total_steps:
+            self.buffer.reset()
+            self.policy.eval()
+
+            # ── Collect rollout ──────────────────────────────────────────────
+            for _ in range(self.n_steps):
+                # Curriculum step (returns [] for OverflowEnv — kept for API compat)
+                self.curriculum.step(self.env._sim_time)
+
+                obs_t = torch.as_tensor(obs, dtype=torch.float32, device=self.device)
+                with torch.no_grad():
+                    act_mean, val = self.policy(obs_t.unsqueeze(0))
+                act_mean = act_mean.squeeze(0)
+                val      = val.squeeze(0)
+
+                dist   = torch.distributions.Normal(act_mean, torch.ones_like(act_mean) * 0.3)
+                action = dist.sample().clamp(-1, 1)
+                logp   = dist.log_prob(action).sum()
+
+                next_obs, reward, term, trunc, info = self.env.step(action.cpu().numpy())
+
+                self.buffer.add(
+                    obs, action.cpu().numpy(), reward,
+                    float(val), float(logp), float(term or trunc),
+                )
+
+                obs        = next_obs
+                ep_reward += reward
+                ep_steps  += 1
+                self.total_steps += 1
+
+                if term or trunc:
+                    bonus = compute_episode_bonus(
+                        total_steps=ep_steps,
+                        survived=not info.get("collision", False),
+                    )
+                    ep_reward += bonus
+                    self.ep_rewards.append(ep_reward)
+                    self.ep_lengths.append(ep_steps)
+                    advanced = self.curriculum.record_episode_reward(ep_reward)
+
+                    outcome = "CRASH" if info.get("collision") else ("GOAL" if info.get("goal_reached") else "timeout")
+                    print(
+                        f"  ep#{len(self.ep_rewards):>4d} | "
+                        f"steps={ep_steps:>3d} | "
+                        f"reward={ep_reward:>8.2f} | "
+                        f"outcome={outcome:<8} | "
+                        f"stage={self.curriculum.current_stage} | "
+                        f"total_steps={self.total_steps}",
+                        flush=True,
+                    )
+
+                    obs, _ = self.env.reset()
+                    ep_reward = 0.0
+                    ep_steps  = 0
+
+            # ── PPO update ───────────────────────────────────────────────────
+            with torch.no_grad():
+                obs_t = torch.as_tensor(obs, dtype=torch.float32, device=self.device)
+                _, last_val = self.policy(obs_t.unsqueeze(0))
+            self.buffer.compute_returns(float(last_val), self.gamma, self.gae_lambda)
+
+            self.policy.train()
+            self._ppo_update()
+            self.n_updates += 1
+            self.scheduler.step()
+
+            elapsed = time.time() - t0
+            sps     = self.total_steps / max(elapsed, 1)
+            mean_r  = np.mean(self.ep_rewards) if self.ep_rewards else 0.0
+            mean_l  = np.mean(self.ep_lengths) if self.ep_lengths else 0.0
+            print(
+                f"\n[PPO update #{self.n_updates}] "
+                f"step={self.total_steps}  "
+                f"mean_reward={mean_r:.2f}  "
+                f"mean_ep_len={mean_l:.0f}  "
+                f"stage={self.curriculum.current_stage}  "
+                f"sps={sps:.0f}\n",
+                flush=True,
+            )
+
+            # ── Checkpoint ───────────────────────────────────────────────────
+            if self.n_updates % 50 == 0:
+                ckpt = self.save_dir / f"policy_step{self.total_steps}_stage{self.curriculum.current_stage}.pt"
+                torch.save({
+                    "step":   self.total_steps,
+                    "stage":  self.curriculum.current_stage,
+                    "policy": self.policy.state_dict(),
+                    "optim":  self.optimizer.state_dict(),
+                }, ckpt)
+                print(f"[PPO] Saved checkpoint → {ckpt}")
+
+    # ── PPO update pass — identical to openenv ─────────────────────────────────
+
+    def _ppo_update(self):
+        obs      = self.buffer.obs
+        acts     = self.buffer.acts
+        old_logp = self.buffer.logp
+        adv      = self.buffer.ret - self.buffer.val
+        adv      = (adv - adv.mean()) / (adv.std() + 1e-8)
+        ret      = self.buffer.ret
+        old_val  = self.buffer.val
+
+        indices = torch.randperm(self.n_steps, device=self.device)
+
+        for _ in range(self.n_epochs):
+            for start in range(0, self.n_steps, self.batch_size):
+                idx = indices[start: start + self.batch_size]
+
+                act_mean, val = self.policy(obs[idx])
+                val = val.squeeze(-1)
+
+                dist    = torch.distributions.Normal(act_mean, torch.ones_like(act_mean) * 0.3)
+                logp    = dist.log_prob(acts[idx]).sum(dim=-1)
+                entropy = dist.entropy().sum(dim=-1).mean()
+
+                ratio    = torch.exp(logp - old_logp[idx])
+                pg_loss1 = -adv[idx] * ratio
+                pg_loss2 = -adv[idx] * ratio.clamp(1 - self.clip, 1 + self.clip)
+                pg_loss  = torch.max(pg_loss1, pg_loss2).mean()
+
+                val_unclipped = (val - ret[idx]) ** 2
+                val_clipped   = (
+                    old_val[idx]
+                    + (val - old_val[idx]).clamp(-self.clip_vf, self.clip_vf)
+                    - ret[idx]
+                ) ** 2
+                vf_loss = 0.5 * torch.max(val_unclipped, val_clipped).mean()
+
+                loss = pg_loss + self.vf_coef * vf_loss - self.ent_coef * entropy
+
+                self.optimizer.zero_grad()
+                loss.backward()
+                nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad)
+                self.optimizer.step()
+
+
+# ── Entry point ────────────────────────────────────────────────────────────────
+
+def run_training(
+    policy_type: str = "attention",
+    total_steps: int = 2_000_000,
+    start_stage: int = 1,
+    checkpoint:  Optional[str] = None,
+    device:      str = "auto",
+) -> None:
+    from ..policies.ticket_attention_policy import TicketAttentionPolicy
+    from ..policies.flat_mlp_policy         import FlatMLPPolicy
+
+    policy_map = {
+        "attention": lambda: TicketAttentionPolicy(obs_dim=OBS_DIM),
+        "mlp":       lambda: FlatMLPPolicy(obs_dim=OBS_DIM),
+    }
+    policy = policy_map[policy_type]()
+
+    if checkpoint:
+        ckpt = torch.load(checkpoint, map_location="cpu")
+        policy.load_state_dict(ckpt["policy"])
+        print(f"[PPO] Loaded checkpoint from {checkpoint}")
+
+    env = OverflowGymEnv()
+    cm  = CurriculumManager()
+    if start_stage > 1:
+        cm.force_stage(start_stage)
+
+    trainer = PPOTrainer(policy=policy, env=env, curriculum=cm, device=device, n_steps=512)
+    trainer.train(total_steps=total_steps)
+
+
+if __name__ == "__main__":
+    import argparse
+    p = argparse.ArgumentParser()
+    p.add_argument("--policy",     default="attention", choices=["attention", "mlp"])
+    p.add_argument("--steps",      default=2_000_000,  type=int)
+    p.add_argument("--stage",      default=1,           type=int)
+    p.add_argument("--checkpoint", default=None)
+    p.add_argument("--device",     default="auto")
+    args = p.parse_args()
+    run_training(args.policy, args.steps, args.stage, args.checkpoint, args.device)

training/reward.py

@@ -0,0 +1,94 @@
+"""
+Reward shaping for OverflowEnvironment — ported from openenv/training/reward.py.
+
+Same core principle: BASE + THREAT_RESPONSE with clear gradient direction.
+Adapted to OverflowEnvironment's signals (no EventTicket objects — uses
+collision/near-miss flags and raw reward from the environment).
+
+  BASE:           survival + speed + lane  ~+0.4/step
+  COLLISION:      -50  (terminal)
+  NEAR MISS:      -0.8 per event
+  GOAL REACHED:   +5.0 (terminal bonus)
+  SMOOTH DRIVING: small bonus when no threats
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+# ── Same weights as openenv/training/reward.py ────────────────────────────────
+
+W_ALIVE         =  0.40
+W_SPEED         =  0.10
+W_LANE          =  0.15
+W_SMOOTH        =  0.03
+TARGET_SPEED    = 11.0    # m/s (~40 km/h)
+TARGET_SPEED_TOL =  3.0
+
+W_COLLISION     = -50.0
+W_NEAR_MISS     =  -0.8
+W_GOAL          =   5.0
+W_SURVIVE_BONUS =   5.0
+
+ROAD_HALF_WIDTH = 3.7 * 1.5   # ~2.5 lanes worth of tolerance
+
+
+def compute_reward(
+    ego_speed:    float,
+    ego_y:        float,
+    action:       np.ndarray,
+    prev_action:  np.ndarray,
+    collision:    bool,
+    goal_reached: bool,
+    near_miss:    bool,
+    raw_reward:   float,       # OverflowEnvironment's built-in reward (used as baseline)
+) -> float:
+    """
+    Shaped reward. Mirrors openenv reward structure:
+      - collision → large terminal penalty
+      - base survival + speed + lane keeping
+      - near-miss penalty
+      - goal bonus
+      - smooth driving bonus when clear
+    """
+    if collision:
+        return W_COLLISION
+
+    reward = 0.0
+
+    # 1. Survival
+    reward += W_ALIVE
+
+    # 2. Speed maintenance (same formula as openenv)
+    speed_err = abs(ego_speed - TARGET_SPEED)
+    if speed_err < TARGET_SPEED_TOL:
+        reward += W_SPEED * (1.0 - speed_err / TARGET_SPEED_TOL)
+    else:
+        reward -= 0.03 * min(speed_err - TARGET_SPEED_TOL, 5.0)
+
+    # 3. Lane keeping
+    norm_y = abs(ego_y) / ROAD_HALF_WIDTH
+    reward += W_LANE * max(0.0, 1.0 - norm_y ** 2)
+
+    # 4. Near miss penalty
+    if near_miss:
+        reward += W_NEAR_MISS
+
+    # 5. Goal bonus
+    if goal_reached:
+        reward += W_GOAL
+
+    # 6. Smooth driving
+    action_delta = np.abs(action - prev_action).sum()
+    reward += W_SMOOTH * max(0.0, 1.0 - action_delta * 3.0)
+
+    return float(reward)
+
+
+def compute_episode_bonus(total_steps: int, survived: bool) -> float:
+    """End-of-episode bonus — same as openenv."""
+    if not survived:
+        return 0.0
+    bonus  = W_SURVIVE_BONUS
+    bonus += min(total_steps, 500) * 0.02   # longevity reward
+    return float(bonus)