"""Inference Script - Indian Traffic Signal OpenEnv ================================================ MANDATORY environment variables (injected by the validator): API_BASE_URL The LiteLLM proxy endpoint. API_KEY Your API key for the proxy. MODEL_NAME The model identifier to use for inference. STDOUT FORMAT (exact - do not deviate): [START] task= env= model= [STEP] step= action= reward=<0.00> done= error= [END] success= steps= score=<0.000> rewards= """ import json import os import sys from typing import List, Optional from openai import OpenAI from env import IndianTrafficEnv from grader import grade_rollout from models import TrafficAction, TrafficState # ------------------------------------------------------------------- # MANDATORY: read from injected environment variables — no hardcoding. # The validator checks that all LLM calls flow through API_BASE_URL. # ------------------------------------------------------------------- API_BASE_URL: str = os.environ["API_BASE_URL"] # must be set by validator API_KEY: str = os.environ.get("API_KEY") or os.environ.get("HF_TOKEN", "") MODEL_NAME: str = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-7B-Instruct") BENCHMARK: str = os.environ.get("BENCHMARK", "indian-traffic-signal-openenv") SUCCESS_SCORE_THRESHOLD = 0.5 TASKS = ["single_intersection", "rush_hour", "emergency_priority"] VALID_ACTIONS = [a.value for a in TrafficAction] # Single shared client — always routed through the injected proxy URL. _client = OpenAI( base_url=API_BASE_URL, api_key=API_KEY, timeout=30.0, # generous timeout for proxy round-trips max_retries=1, ) # --------------------------------------------------------------------------- # Logging helpers — exact format required by the validator # --------------------------------------------------------------------------- def log_start(task: str, env: str, model: str) -> None: print(f"[START] task={task} env={env} model={model}", flush=True) def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None: error_val = error if error else "null" print( f"[STEP] step={step} action={action} reward={reward:.2f} " f"done={str(done).lower()} error={error_val}", flush=True, ) def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None: rewards_str = ",".join(f"{r:.2f}" for r in rewards) print( f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True, ) # --------------------------------------------------------------------------- # Fallback policy — used only if the LLM call itself raises an exception # --------------------------------------------------------------------------- def _fallback_action(task_name: str, state: TrafficState) -> str: """Deterministic fallback — mirrors the baseline policy.""" if state.emergency_vehicle.present: return TrafficAction.EMERGENCY_OVERRIDE.value if state.pedestrian_count >= 16 and state.pedestrian_wait_time > 18: return TrafficAction.PEDESTRIAN_CROSS.value if state.time_since_last_phase_switch < 3 and state.current_signal_phase in ( TrafficAction.NS_GREEN, TrafficAction.EW_GREEN, TrafficAction.LEFT_PRIORITY, ): return TrafficAction.EXTEND_GREEN.value ns = state.lane_queues["N"].total + state.lane_queues["S"].total ew = state.lane_queues["E"].total + state.lane_queues["W"].total if task_name == "emergency_priority": if abs(ns - ew) >= 10: return TrafficAction.NS_GREEN.value if ns > ew else TrafficAction.EW_GREEN.value return TrafficAction.NS_GREEN.value if abs(ns - ew) >= 12: return TrafficAction.NS_GREEN.value if ns > ew else TrafficAction.EW_GREEN.value cycle = (state.tick // 8) % 4 return [ TrafficAction.NS_GREEN.value, TrafficAction.EW_GREEN.value, TrafficAction.LEFT_PRIORITY.value, TrafficAction.PEDESTRIAN_CROSS.value, ][cycle] # --------------------------------------------------------------------------- # LLM call — ALWAYS goes through the injected proxy (API_BASE_URL / _client) # --------------------------------------------------------------------------- def get_action_from_llm(state: TrafficState, task_name: str) -> str: """Call the LLM via the injected proxy to choose a signal action.""" preferred = _fallback_action(task_name, state) state_summary = { "tick": state.tick, "current_phase": state.current_signal_phase.value, "time_since_switch": state.time_since_last_phase_switch, "emergency": state.emergency_vehicle.model_dump(), "pedestrian_count": state.pedestrian_count, "pedestrian_wait": round(state.pedestrian_wait_time, 2), "rain_level": round(state.rain_level, 3), "lane_queues": { lane: {"total": q.total, **q.model_dump()} for lane, q in state.lane_queues.items() }, } system_prompt = ( "You are an AI traffic controller managing an Indian urban intersection. " f"Task: {task_name}. " f"Choose exactly one action from: {', '.join(VALID_ACTIONS)}. " "Analyse the state and pick the best signal phase. " f"Suggested action: {preferred}. " "Reply with only the action name — no explanation, no punctuation." ) user_prompt = f"Intersection state: {json.dumps(state_summary)}" # This call MUST reach the proxy — do not wrap in a silent broad except. response = _client.chat.completions.create( model=MODEL_NAME, messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ], temperature=0.0, max_tokens=16, stream=False, ) action = (response.choices[0].message.content or "").strip().upper() return action if action in VALID_ACTIONS else preferred # --------------------------------------------------------------------------- # Main inference loop # --------------------------------------------------------------------------- def run_inference() -> None: if not API_KEY: print("Warning: API_KEY / HF_TOKEN not set.", file=sys.stderr, flush=True) for task_name in TASKS: env = IndianTrafficEnv(task_id=task_name) env.reset(seed=42, task_id=task_name) rewards: List[float] = [] steps_taken = 0 success = False score = 0.001 done = False log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME) try: step = 1 while not done: state = env.get_state() error: Optional[str] = None try: action_str = get_action_from_llm(state, task_name) except Exception as exc: # LLM call failed — log it, use fallback, keep running action_str = _fallback_action(task_name, state) error = f"llm_error:{type(exc).__name__}" try: traffic_action = TrafficAction(action_str) except ValueError: traffic_action = TrafficAction.ALL_RED action_str = TrafficAction.ALL_RED.value try: _, reward, done, _ = env.step(traffic_action) except Exception as exc: reward = 0.0 done = True error = str(exc) rewards.append(reward) steps_taken = step log_step(step=step, action=action_str, reward=reward, done=done, error=error) step += 1 grader_result = grade_rollout(task_id=task_name, seed=42) score = float(grader_result.score) success = score >= SUCCESS_SCORE_THRESHOLD except Exception as exc: print(f"Fatal error in task {task_name}: {exc}", file=sys.stderr, flush=True) success = False score = 0.001 finally: log_end(success=success, steps=steps_taken, score=score, rewards=rewards) if __name__ == "__main__": run_inference()