Upload folder using huggingface_hub

inference.py

@@ -4,8 +4,10 @@ import os
 import sys
 import json
 import time
+import math
 import urllib.request
-from typing import List, Optional
+from collections import deque
+from typing import Deque, List, Optional, Tuple
 
 # Allow imports from repo root so both container-root and package contexts work
 _HERE = os.path.dirname(os.path.abspath(__file__))
@@ -18,88 +20,325 @@ from openai import OpenAI
 
 try:
     from traffic_control.client import TrafficControlEnv
-    from traffic_control.models import TrafficAction, TrafficObservation
+    from traffic_control.models import TrafficAction, TrafficObservation, TrafficState
 except ImportError:
     from client import TrafficControlEnv  # type: ignore
-    from models import TrafficAction, TrafficObservation  # type: ignore
+    from models import TrafficAction, TrafficObservation, TrafficState  # type: ignore
 
 # ---------------------------------------------------------------------------
-# Environment variables — MUST be injected by the hackathon validator
+# Environment variables — injected by hackathon validator
 # ---------------------------------------------------------------------------
 API_KEY      = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
 API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
 MODEL_NAME   = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
 SERVER_URL   = os.getenv("SERVER_URL", "http://localhost:7860")
 SEED         = 42
-MAX_TOKENS   = 64
+MAX_TOKENS   = 256   # enough for chain-of-thought + JSON
 TEMPERATURE  = 0.0
 
+# Minimum green hold steps (ensures stability_bonus in grading)
+MIN_GREEN_HOLD = 4
+
+# Task max_steps for score projection
+TASK_MAX_STEPS = {"basic_flow": 200, "emergency_priority": 300, "dynamic_scenarios": 400}
+
+
 # ---------------------------------------------------------------------------
-# LLM Prompt
+# Heuristic core — uses actual environment reward formula
 # ---------------------------------------------------------------------------
 
-SYSTEM_PROMPT = """You are an expert Autonomous Traffic Signal Controller.
+def _em_pressure(count: int, urgency: int) -> float:
+    """Emergency pressure = count × urgency^1.5 × 0.5 (matches env._compute_reward)."""
+    return count * (max(urgency, 1) ** 1.5) * 0.5 if count > 0 else 0.0
 
-PHASES:
-  0 = North-South Green
-  1 = East-West Green  
-  2 = All Red
 
-SCORING:
-  +0.2 per regular vehicle cleared
-  +10.0 per emergency vehicle cleared
-  -0.4 * urgency per step emergency waits
-  -0.5 for switching to empty queue
+def _dir_pressure(queue: int, em_count: int, urgency: int) -> float:
+    """Combined directional pressure: throughput value + weighted emergency penalty."""
+    return queue * 0.30 + _em_pressure(em_count, urgency) * 4.0
 
-OUTPUT: {"light_phase": 0, 1, or 2} only JSON, no other text."""
 
+def _compute_pressures(obs: TrafficObservation) -> Tuple[float, float]:
+    """Return (ns_pressure, ew_pressure)."""
+    ns_p = _dir_pressure(
+        obs.queue_lengths[0] + obs.queue_lengths[1],
+        obs.emergency_queue[0] + obs.emergency_queue[1],
+        max(obs.emergency_urgency[0], obs.emergency_urgency[1]),
+    )
+    ew_p = _dir_pressure(
+        obs.queue_lengths[2] + obs.queue_lengths[3],
+        obs.emergency_queue[2] + obs.emergency_queue[3],
+        max(obs.emergency_urgency[2], obs.emergency_urgency[3]),
+    )
+    return ns_p, ew_p
+
+
+def _dynamic_hold_time(obs: TrafficObservation) -> int:
+    """
+    Adaptive minimum hold: longer when current direction has more traffic
+    so we don't leave vehicles stranded mid-queue.
+    """
+    current = obs.current_phase
+    if current in (0, 3):   # NS_GREEN / NS_YELLOW
+        q = obs.queue_lengths[0] + obs.queue_lengths[1]
+    elif current in (1, 4): # EW_GREEN / EW_YELLOW
+        q = obs.queue_lengths[2] + obs.queue_lengths[3]
+    else:
+        return MIN_GREEN_HOLD
+    # Hold longer if queue is deep (drain rate ~3 veh/step)
+    return max(MIN_GREEN_HOLD, min(q // 3, 10))
+
+
+def _heuristic_phase(obs: TrafficObservation) -> int:
+    """
+    Mathematically optimal phase recommendation.
+    Priority order:
+      1. Critical emergency (urgency ≥ 8) — clear NOW
+      2. Moderate emergency (urgency 5-7) — prioritise unless other side is worse
+      3. Hysteresis — don't switch if hold time not reached
+      4. Queue pressure — switch to higher pressure direction
+      5. Default — hold current
+    """
+    ns_em  = obs.emergency_queue[0] + obs.emergency_queue[1]
+    ew_em  = obs.emergency_queue[2] + obs.emergency_queue[3]
+    ns_urg = max(obs.emergency_urgency[0], obs.emergency_urgency[1])
+    ew_urg = max(obs.emergency_urgency[2], obs.emergency_urgency[3])
+    cur    = obs.current_phase
+
+    # 1. Critical emergency
+    if ns_em > 0 and ns_urg >= 8 and ew_em > 0 and ew_urg >= 8:
+        return 2  # ALL_RED: both critical, momentary pause to avoid collision
+    if ns_em > 0 and ns_urg >= 8:
+        return 0
+    if ew_em > 0 and ew_urg >= 8:
+        return 1
+
+    # 2. Moderate emergency — switch if other side isn't also urgent
+    if ns_em > 0 and ns_urg >= 5:
+        if ew_em == 0 or ns_urg >= ew_urg:
+            return 0
+    if ew_em > 0 and ew_urg >= 5:
+        return 1
+
+    # 3. Hysteresis
+    hold = _dynamic_hold_time(obs)
+    if obs.time_in_phase < hold:
+        if cur in (0, 3): return 0
+        if cur in (1, 4): return 1
+
+    # 4. Queue pressure
+    ns_p, ew_p = _compute_pressures(obs)
+    if ns_p > ew_p * 1.3:
+        return 0
+    if ew_p > ns_p * 1.3:
+        return 1
+
+    # 5. Hold current
+    if cur in (0, 3): return 0
+    if cur in (1, 4): return 1
+    return 0
+
+
+# ---------------------------------------------------------------------------
+# Live grade projection (shows LLM how its decisions impact the final score)
+# ---------------------------------------------------------------------------
+
+def _project_score(task: str, state: Optional[TrafficState], step: int) -> str:
+    """Compute projected grading scores from current episode state."""
+    if state is None or step == 0:
+        return "(no data yet)"
+
+    s     = state
+    steps = max(s.step_count, 1)
+    max_s = TASK_MAX_STEPS.get(task, 300)
+
+    throughput_per_step = s.total_vehicles_passed / steps
+    em_rate             = s.total_emergency_passed / steps
+    avg_wait            = s.total_waiting_time / steps
+
+    if task == "basic_flow":
+        tput  = min(throughput_per_step / 1.8, 1.0)
+        eff   = 1.0 / (1.0 + avg_wait * 0.1)
+        sw    = s.total_phase_changes / steps
+        stab  = max(0.0, 0.05 * (1.0 - min(sw * 4, 1.0)))
+        proj  = tput * 0.6 + eff * 0.4 + stab
+        return (
+            f"throughput={tput:.2f}(×0.6) eff={eff:.2f}(×0.4) stability={stab:.3f} "
+            f"→ projected={proj:.3f}  "
+            f"[veh/step={throughput_per_step:.2f} target=1.8, switch_rate={sw:.2f} target<0.25]"
+        )
+
+    if task == "emergency_priority":
+        tput     = min(throughput_per_step / 1.5, 1.0)
+        em_score = min(em_rate / (1.0 / 20.0), 1.0)
+        if s.total_emergency_passed > 0:
+            delay = max(0.0, 1.0 - (s.total_emergency_delay / s.total_emergency_passed) / 12.0)
+        else:
+            delay = 0.5
+        eff  = 1.0 / (1.0 + avg_wait * 0.05)
+        proj = tput * 0.30 + em_score * 0.35 + delay * 0.20 + eff * 0.15
+        return (
+            f"tput={tput:.2f}(×0.30) em_rate={em_score:.2f}(×0.35) "
+            f"delay={delay:.2f}(×0.20) eff={eff:.2f}(×0.15) → projected={proj:.3f}  "
+            f"[em_cleared={s.total_emergency_passed} need≥{steps//20}]"
+        )
+
+    if task == "dynamic_scenarios":
+        tput     = min(throughput_per_step / 2.0, 1.0)
+        em_score = min(em_rate / (1.0 / 15.0), 1.0)
+        if s.total_emergency_passed > 0:
+            delay = max(0.0, 1.0 - (s.total_emergency_delay / s.total_emergency_passed) / 5.0)
+        else:
+            delay = 0.0
+        eff   = 1.0 / (1.0 + avg_wait * 0.08)
+        adapt = 1.0 / (1.0 + (s.total_phase_changes / steps) * 0.5)
+        proj  = tput * 0.25 + em_score * 0.30 + delay * 0.20 + eff * 0.15 + adapt * 0.10
+        return (
+            f"tput={tput:.2f}(×0.25) em={em_score:.2f}(×0.30) delay={delay:.2f}(×0.20) "
+            f"eff={eff:.2f}(×0.15) adapt={adapt:.2f}(×0.10) → projected={proj:.3f}"
+        )
+
+    return "(unknown task)"
+
+
+# ---------------------------------------------------------------------------
+# LLM prompt
+# ---------------------------------------------------------------------------
+
+SYSTEM_PROMPT = """You are an expert Autonomous Traffic Signal Controller optimising a 4-way intersection.
+
+PHASES:  0=NS_GREEN  1=EW_GREEN  2=ALL_RED
+
+REWARD FUNCTION (per step):
+  +0.30 × regular vehicles cleared
+  +12.0 × emergency vehicles cleared
+  -(urgency^1.5)×0.5 per waiting emergency vehicle (EVERY step it waits!)
+  -0.08 × total vehicles waiting
+  -0.5 to -2.0 for unnecessary phase switch (proportional to how empty the new direction is)
+  +0.05 stability bonus when traffic flows without switching
+
+GRADING WEIGHTS:
+  basic_flow:          throughput×0.60  efficiency×0.40  stability_bonus
+  emergency_priority:  throughput×0.30  em_rate×0.35     delay×0.20  efficiency×0.15
+  dynamic_scenarios:   throughput×0.25  em_rate×0.30     delay×0.20  efficiency×0.15  adaptability×0.10
+
+DECISION RULES (follow strictly):
+  1. Urgency ≥ 8 emergency → switch to that direction IMMEDIATELY
+  2. Urgency 5-7 emergency → switch unless other side is equally urgent
+  3. Hold current phase ≥ 4 steps before switching (preserves stability bonus)
+  4. Switch only when pressure ratio > 1.3× (avoid unnecessary switches)
+  5. ALL_RED only when both directions have simultaneous critical emergencies
+
+Think step by step, then output ONLY valid JSON on the last line: {"light_phase": 0}"""
+
+
+def _build_prompt(
+    obs: TrafficObservation,
+    step: int,
+    task: str,
+    history: Deque[str],
+    heuristic: int,
+    score_projection: str,
+) -> str:
+    ns_p, ew_p = _compute_pressures(obs)
+    ns_q  = obs.queue_lengths[0] + obs.queue_lengths[1]
+    ew_q  = obs.queue_lengths[2] + obs.queue_lengths[3]
+    ns_em = obs.emergency_queue[0] + obs.emergency_queue[1]
+    ew_em = obs.emergency_queue[2] + obs.emergency_queue[3]
+    ns_urg = max(obs.emergency_urgency[0], obs.emergency_urgency[1])
+    ew_urg = max(obs.emergency_urgency[2], obs.emergency_urgency[3])
+
+    phase_name = {0:"NS_GREEN",1:"EW_GREEN",2:"ALL_RED",3:"NS_YELLOW",4:"EW_YELLOW"}
+    hint_name  = {0:"NS_GREEN (0)",1:"EW_GREEN (1)",2:"ALL_RED (2)"}
+    trend_str  = f"[{obs.queue_trend[0]:+d},{obs.queue_trend[1]:+d},{obs.queue_trend[2]:+d},{obs.queue_trend[3]:+d}]"
+
+    history_str = "\n".join(history) if history else "  (episode start)"
 
-def _build_prompt(obs: TrafficObservation, step: int) -> str:
     return (
-        f"Step {step}, phase={obs.current_phase}, time={obs.time_in_phase}\n"
-        f"Queues N,S,E,W: {list(obs.queue_lengths)}\n"
-        f"Emergency N,S,E,W: {list(obs.emergency_queue)} urgency={list(obs.emergency_urgency)}\n"
-        f"What phase? Return only JSON: {{\"light_phase\": 0, 1, or 2}}"
+        f"TASK: {task}  |  Step {step}  |  Current phase: {phase_name.get(obs.current_phase,'?')} (held {obs.time_in_phase} steps)\n"
+        f"\n"
+        f"STATE:\n"
+        f"  NS: {ns_q} regular + {ns_em} emergency(urgency={ns_urg})  pressure={ns_p:.1f}\n"
+        f"  EW: {ew_q} regular + {ew_em} emergency(urgency={ew_urg})  pressure={ew_p:.1f}\n"
+        f"  Queues [N,S,E,W]: {list(obs.queue_lengths)}  trend={trend_str}\n"
+        f"  Avg wait: {obs.avg_wait_time:.1f} steps  |  Collision: {obs.collision}\n"
+        f"\n"
+        f"LIVE SCORE PROJECTION:\n  {score_projection}\n"
+        f"\n"
+        f"RECENT HISTORY (last {len(history)} steps):\n{history_str}\n"
+        f"\n"
+        f"Heuristic recommendation: {hint_name[heuristic]}\n"
+        f"Reason through the decision, then output JSON on the last line."
     )
 
 
+# ---------------------------------------------------------------------------
+# Parse + LLM call
+# ---------------------------------------------------------------------------
+
 def _sanitize(s: str) -> str:
-    """Remove characters that break log parsing."""
     return s.replace('"', "'").replace("\n", " ")
 
 
-def _parse_phase(raw: str) -> int:
-    """Extract phase from LLM response."""
-    try:
-        data = json.loads(raw)
-        phase = int(data.get("light_phase", data.get("phase", 0)))
-        return max(0, min(2, phase))
-    except Exception:
-        import re
-        m = re.search(r'\b([012])\b', raw)
-        return int(m.group(1)) if m else 0
-
+def _parse_phase(raw: str) -> Optional[int]:
+    """Extract phase from LLM chain-of-thought output (JSON on last line)."""
+    import re
+    # Try last non-empty line first (chain-of-thought ends with JSON)
+    lines = [l.strip() for l in raw.split("\n") if l.strip()]
+    for line in reversed(lines):
+        try:
+            data = json.loads(line)
+            p = int(data.get("light_phase", data.get("phase", -1)))
+            if 0 <= p <= 2:
+                return p
+        except Exception:
+            pass
+    # Fallback: regex anywhere
+    m = re.search(r'"light_phase"\s*:\s*([012])', raw)
+    if m:
+        return int(m.group(1))
+    m = re.search(r'\b([012])\b', raw)
+    if m:
+        return int(m.group(1))
+    return None
+
+
+def get_llm_action(
+    client: OpenAI,
+    obs: TrafficObservation,
+    step: int,
+    task: str,
+    history: Deque[str],
+    state: Optional[TrafficState],
+) -> TrafficAction:
+    heuristic  = _heuristic_phase(obs)
+    score_proj = _project_score(task, state, step)
 
-def get_llm_action(client: OpenAI, obs: TrafficObservation, step: int) -> TrafficAction:
-    """Call LLM proxy for a traffic phase decision."""
     resp = client.chat.completions.create(
         model=MODEL_NAME,
         messages=[
             {"role": "system", "content": SYSTEM_PROMPT},
-            {"role": "user",   "content": _build_prompt(obs, step)},
+            {"role": "user",   "content": _build_prompt(obs, step, task, history, heuristic, score_proj)},
         ],
         temperature=TEMPERATURE,
         max_tokens=MAX_TOKENS,
     )
-    raw = resp.choices[0].message.content.strip()
+    raw   = resp.choices[0].message.content.strip()
     phase = _parse_phase(raw)
+
+    # Fallback to heuristic if LLM output is unparseable or clearly wrong
+    if phase is None:
+        phase = heuristic
+
     return TrafficAction(light_phase=phase)
 
 
+# ---------------------------------------------------------------------------
+# Server health check
+# ---------------------------------------------------------------------------
+
 def _wait_for_server(url: str, timeout: int = 60) -> None:
-    """Block until the env server is healthy or timeout expires."""
     health_url = url.rstrip("/") + "/health"
-    deadline = time.time() + timeout
+    deadline   = time.time() + timeout
     while time.time() < deadline:
         try:
             with urllib.request.urlopen(health_url, timeout=3) as r:
@@ -108,22 +347,25 @@ def _wait_for_server(url: str, timeout: int = 60) -> None:
         except Exception:
             pass
         time.sleep(2)
-    # If the server never came up, log and continue anyway
     print(f"[WARN] Server not healthy after {timeout}s — proceeding anyway", flush=True)
 
 
+# ---------------------------------------------------------------------------
+# Episode runner
+# ---------------------------------------------------------------------------
+
 def run_task(task: str, client: OpenAI) -> dict:
-    """Run a single task episode."""
     print(f'[START] task={task} env=traffic_control model={MODEL_NAME}', flush=True)
 
-    rewards: List[float] = []
-    step = 0
-    last_error: Optional[str] = None
-    done = False
+    rewards:    List[float]        = []
+    history:    Deque[str]         = deque(maxlen=6)
+    step        = 0
+    last_error: Optional[str]      = None
+    done        = False
+    state:      Optional[TrafficState] = None
 
     try:
         with TrafficControlEnv(base_url=SERVER_URL).sync() as env:
-            # env.reset() returns a StepResult; unwrap the observation
             reset_result = env.reset(task_id=task, seed=SEED)
             obs: TrafficObservation = reset_result.observation
             done = reset_result.done
@@ -131,18 +373,32 @@ def run_task(task: str, client: OpenAI) -> dict:
             while not done:
                 step += 1
 
-                # Call the LLM proxy — this is the call the validator monitors
-                action = get_llm_action(client, obs, step)
+                # Refresh cumulative state every 10 steps for score projection
+                if step % 10 == 1:
+                    try:
+                        state = env.state()
+                    except Exception:
+                        pass
+
+                action     = get_llm_action(client, obs, step, task, history, state)
                 action_str = f"light_phase={action.light_phase}"
 
                 try:
-                    # env.step() returns a StepResult; unwrap the observation
-                    result = env.step(action)
-                    obs        = result.observation   # TrafficObservation
+                    result     = env.step(action)
+                    obs        = result.observation
                     reward_val = result.reward if result.reward is not None else 0.0
                     rewards.append(reward_val)
                     done       = result.done
                     last_error = None
+
+                    phase_name = {0:"NS",1:"EW",2:"AR",3:"NSy",4:"EWy"}
+                    history.append(
+                        f"  s{step}: →{action.light_phase}"
+                        f" clr={obs.vehicles_passed}r+{obs.emergency_passed}em"
+                        f" r={reward_val:+.1f}"
+                        f" now={phase_name.get(obs.current_phase,'?')}"
+                        f" queues={list(obs.queue_lengths)}"
+                    )
                 except Exception as exc:
                     reward_val = 0.0
                     done       = True
@@ -182,11 +438,8 @@ def run_task(task: str, client: OpenAI) -> dict:
 # ---------------------------------------------------------------------------
 
 def main():
-    """Main entry point."""
-    # Wait for the env server to be ready before starting inference
     _wait_for_server(SERVER_URL)
 
-    # Initialize OpenAI client with hackathon-injected proxy credentials
     client = OpenAI(
         base_url=API_BASE_URL,
         api_key=API_KEY,