""" Comparative Agent Arena Run multiple agents simultaneously to compare performance. Supports: - Rule-based agent - LLM agent (makes live API calls) - Random agent """ import os import json import textwrap import time import random import asyncio from typing import Dict, List, Any, Optional from dataclasses import dataclass, field from datetime import datetime from traffic_control.environment import TrafficControlEnvironment from traffic_control.models import TrafficAction, TrafficObservation from traffic_control.tasks import grade, GradeResult try: from openai import OpenAI _HAS_OPENAI = True except ImportError: _HAS_OPENAI = False @dataclass class AgentResult: """Result for a single agent run.""" agent_name: str agent_type: str task_id: str episode_id: str steps: int = 0 total_reward: float = 0.0 score: float = 0.0 metrics: Dict[str, Any] = field(default_factory=dict) grade_result: Optional[GradeResult] = None decision_times: List[float] = field(default_factory=list) @property def avg_decision_time_ms(self) -> float: if self.decision_times: return sum(self.decision_times) / len(self.decision_times) * 1000 return 0.0 class RuleBasedAgent: """Rule-based traffic controller (optimized for high scores).""" def __init__(self, name: str = "RuleBased"): self.name = name def decide(self, obs: TrafficObservation) -> int: """Return light phase based on rules.""" em_q = obs.emergency_queue em_u = obs.emergency_urgency q = obs.queue_lengths current = obs.current_phase time_in = obs.time_in_phase # Emergency prioritization ns_em_urgency = em_u[0] + em_u[1] + em_q[0] * 2 + em_q[1] * 2 ew_em_urgency = em_u[2] + em_u[3] + em_q[2] * 2 + em_q[3] * 2 if ns_em_urgency > 0 or ew_em_urgency > 0: return 0 if ns_em_urgency >= ew_em_urgency else 1 # Queue-based switching ns_total = q[0] + q[1] ew_total = q[2] + q[3] min_phase_time = min(3 + max(ns_total, ew_total) // 5, 8) if current == 0 and time_in < min_phase_time and ns_total > 0: return 0 if current == 1 and time_in < min_phase_time and ew_total > 0: return 1 if ns_total >= ew_total + 2: return 0 elif ew_total >= ns_total + 2: return 1 else: return current if current in (0, 1) else 0 class RandomAgent: """Random traffic controller for baseline comparison.""" def __init__(self, name: str = "Random", seed: int = 42): self.name = name self.rng = random.Random(seed) def decide(self, obs: TrafficObservation) -> int: """Return random light phase.""" return self.rng.choice([0, 1, 2]) class RoundRobinAgent: """Simple round-robin controller.""" def __init__(self, name: str = "RoundRobin", switch_interval: int = 5): self.name = name self.switch_interval = switch_interval self.step_count = 0 def decide(self, obs: TrafficObservation) -> int: """Alternate between phases.""" self.step_count += 1 phase_index = (self.step_count // self.switch_interval) % 2 return int(phase_index) # LLM System Prompt for arena LLM_SYSTEM_PROMPT = textwrap.dedent(""" You are an Autonomous Traffic Control AI managing a 4-way intersection. PHASES: 0 = North-South Green (N/S vehicles may pass) 1 = East-West Green (E/W vehicles may pass) 2 = All Red (no vehicles pass) DECISION RULES (apply in order): 1. EMERGENCY CHECK: If emergency vehicles are waiting, prioritize them. 2. MINIMUM PHASE TIME: Stay in current phase at least 3 steps if traffic present. 3. QUEUE BALANCE: Switch to direction with significantly more traffic. OUTPUT: Reply with exactly one JSON object: {"light_phase": <0, 1, or 2>} """) class LLM_Agent: """LLM-powered agent that makes dynamic API calls.""" def __init__( self, name: str = "LLM-Agent", api_base_url: Optional[str] = None, api_key: Optional[str] = None, model_name: str = "gpt-4.1-mini", ): self.name = name self.model_name = model_name self.api_calls_made = 0 # Initialize OpenAI client if credentials available if _HAS_OPENAI and (api_base_url or api_key): self.client = OpenAI( base_url=api_base_url or "https://api.openai.com/v1", api_key=api_key or "dummy-key", ) else: self.client = None def _build_prompt(self, obs: TrafficObservation) -> str: """Build the user prompt from observation.""" return ( f"Current phase: {obs.current_phase} (0=NS Green, 1=EW Green, 2=All Red)\n" f"Time in phase: {obs.time_in_phase} steps\n" f"\n" f"Queue lengths (N, S, E, W): {obs.queue_lengths}\n" f"Emergency queues (N, S, E, W): {obs.emergency_queue}\n" f"Emergency urgency (N, S, E, W): {obs.emergency_urgency}\n" f"\n" f"What light phase should be set? Respond with JSON: {{\"light_phase\": 0, 1, or 2}}" ) def decide(self, obs: TrafficObservation) -> int: """Make LLM API call to get decision.""" if not self.client: # Fallback to rule-based if no client return self._rule_fallback(obs) try: resp = self.client.chat.completions.create( model=self.model_name, messages=[ {"role": "system", "content": LLM_SYSTEM_PROMPT}, {"role": "user", "content": self._build_prompt(obs)}, ], temperature=0.0, max_tokens=32, stream=False, ) self.api_calls_made += 1 content = resp.choices[0].message.content.strip() # Parse JSON response try: data = json.loads(content) phase = int(data.get("light_phase", 0)) return max(0, min(2, phase)) # Clamp to valid range except (json.JSONDecodeError, ValueError, KeyError): # Fallback if parsing fails return self._rule_fallback(obs) except Exception as e: # Fallback on API error print(f"[LLM Agent] API error: {e}, using fallback") return self._rule_fallback(obs) def _rule_fallback(self, obs: TrafficObservation) -> int: """Rule-based fallback when LLM fails.""" em_q = obs.emergency_queue em_u = obs.emergency_urgency q = obs.queue_lengths current = obs.current_phase # Emergency prioritization ns_em = em_u[0] + em_u[1] + em_q[0] + em_q[1] ew_em = em_u[2] + em_u[3] + em_q[2] + em_q[3] if ns_em > 0 or ew_em > 0: return 0 if ns_em >= ew_em else 1 # Queue-based ns_total = q[0] + q[1] ew_total = q[2] + q[3] if ns_total > ew_total: return 0 elif ew_total > ns_total: return 1 else: return current if current in (0, 1) else 0 class Arena: """Run multiple agents and compare results.""" def __init__(self): self.results: List[AgentResult] = [] # Get LLM credentials from env (for arena LLM agent) api_base = os.environ.get("API_BASE_URL") api_key = os.environ.get("API_KEY") model = os.environ.get("MODEL_NAME", "gpt-4.1-mini") self.agents = { "llm": LLM_Agent("Dynamic LLM", api_base, api_key, model), "rule_based": RuleBasedAgent("Smart Rule-Based"), "random": RandomAgent("Random Baseline"), "round_robin": RoundRobinAgent("Round Robin"), } async def run_agent( self, agent_type: str, task_id: str, max_steps: int = 300, seed: int = 42, ) -> AgentResult: """Run a single agent episode.""" env = TrafficControlEnvironment(task_id=task_id) agent = self.agents.get(agent_type, self.agents["rule_based"]) obs = env.reset(seed=seed) episode_id = env._episode_id result = AgentResult( agent_name=agent.name, agent_type=agent_type, task_id=task_id, episode_id=episode_id, ) for step in range(max_steps): import time start_time = time.time() action_id = agent.decide(obs) action = TrafficAction(light_phase=action_id) decision_time = time.time() - start_time result.decision_times.append(decision_time) obs = env.step(action) result.steps = step + 1 result.total_reward += obs.reward or 0.0 if obs.done: break # Grade the result state = env.state result.grade_result = grade( task_id, total_vehicles_passed=state.total_vehicles_passed, total_emergency_passed=state.total_emergency_passed, total_waiting_time=state.total_waiting_time, total_collisions=state.total_collisions, total_emergency_delay=state.total_emergency_delay, total_phase_changes=state.total_phase_changes, step_count=result.steps, ) result.score = result.grade_result.score result.metrics = result.grade_result.metrics return result async def run_comparison( self, task_id: str, agents: Optional[List[str]] = None, runs_per_agent: int = 1, ) -> Dict[str, Any]: """Run multiple agents and compare.""" agents_to_run = agents or list(self.agents.keys()) all_results = [] for agent_type in agents_to_run: for run in range(runs_per_agent): seed = 42 + run result = await self.run_agent(agent_type, task_id, seed=seed) all_results.append(result) self.results.extend(all_results) # Aggregate results summary = self._aggregate_results(all_results) return summary def _aggregate_results(self, results: List[AgentResult]) -> Dict[str, Any]: """Aggregate results by agent type.""" by_agent: Dict[str, List[AgentResult]] = {} for r in results: if r.agent_type not in by_agent: by_agent[r.agent_type] = [] by_agent[r.agent_type].append(r) summary = { "timestamp": datetime.now().isoformat(), "total_runs": len(results), "agents": {}, "winner": None, } best_score = -1 best_agent = None for agent_type, agent_results in by_agent.items(): avg_score = sum(r.score for r in agent_results) / len(agent_results) avg_reward = sum(r.total_reward for r in agent_results) / len(agent_results) avg_steps = sum(r.steps for r in agent_results) / len(agent_results) avg_time = sum(r.avg_decision_time_ms for r in agent_results) / len(agent_results) summary["agents"][agent_type] = { "name": agent_results[0].agent_name, "runs": len(agent_results), "avg_score": round(avg_score, 4), "avg_total_reward": round(avg_reward, 2), "avg_steps": round(avg_steps, 1), "avg_decision_time_ms": round(avg_time, 2), "best_run": max(agent_results, key=lambda r: r.score).episode_id, } if avg_score > best_score: best_score = avg_score best_agent = agent_type summary["winner"] = best_agent summary["all_runs"] = [ { "agent": r.agent_type, "episode_id": r.episode_id, "score": r.score, "reward": round(r.total_reward, 2), "steps": r.steps, } for r in results ] return summary # Global arena instance _arena = Arena() def get_arena() -> Arena: """Get the global arena instance.""" return _arena