""" Core simulation logic for the Autonomous Traffic Control Environment. Inherits from openenv.core.env_server.interfaces.Environment so this class is directly compatible with openenv-core's create_app() factory. Simulates a 4-way intersection with: - Poisson vehicle arrivals per approach - Emergency vehicles with urgency levels - Yellow-light transition state machine - Traffic-surge events (hard task) - Reward shaping for throughput, waiting time, emergency priority and safety """ from __future__ import annotations import math import random import uuid from dataclasses import dataclass from enum import IntEnum from typing import Dict, List, Optional, Set, Tuple from openenv.core.env_server.interfaces import Environment from models import ( TrafficAction, TrafficObservation, TrafficState, PHASE_NS_GREEN, PHASE_EW_GREEN, PHASE_ALL_RED, PHASE_NS_YELLOW, PHASE_EW_YELLOW, ) # --------------------------------------------------------------------------- # Internal enums # --------------------------------------------------------------------------- class VehicleType(IntEnum): CAR = 0 BUS = 1 EMERGENCY = 2 class Direction(IntEnum): NORTH = 0 SOUTH = 1 EAST = 2 WEST = 3 class LightPhase(IntEnum): NS_GREEN = PHASE_NS_GREEN EW_GREEN = PHASE_EW_GREEN ALL_RED = PHASE_ALL_RED NS_YELLOW = PHASE_NS_YELLOW EW_YELLOW = PHASE_EW_YELLOW # --------------------------------------------------------------------------- # Phase transition tables # --------------------------------------------------------------------------- PHASE_ALLOWS: Dict[LightPhase, Set[int]] = { LightPhase.NS_GREEN: {Direction.NORTH, Direction.SOUTH}, LightPhase.EW_GREEN: {Direction.EAST, Direction.WEST}, LightPhase.ALL_RED: set(), LightPhase.NS_YELLOW: {Direction.NORTH, Direction.SOUTH}, LightPhase.EW_YELLOW: {Direction.EAST, Direction.WEST}, } PHASE_FLOW_RATE: Dict[LightPhase, int] = { LightPhase.NS_GREEN: 3, LightPhase.EW_GREEN: 3, LightPhase.ALL_RED: 0, LightPhase.NS_YELLOW: 1, LightPhase.EW_YELLOW: 1, } YELLOW_DURATION = 2 # --------------------------------------------------------------------------- # Task configurations # --------------------------------------------------------------------------- TASK_CONFIGS: Dict[str, dict] = { "basic_flow": { "vehicle_arrival_rate": 0.4, "emergency_arrival_rate": 0.0, "emergency_urgency_range": (0, 0), "max_steps": 200, "max_queue_per_lane": 20, "surge_probability": 0.0, "surge_multiplier": 1.0, }, "emergency_priority": { "vehicle_arrival_rate": 0.5, "emergency_arrival_rate": 0.015, "emergency_urgency_range": (7, 10), "max_steps": 300, "max_queue_per_lane": 20, "surge_probability": 0.0, "surge_multiplier": 1.0, }, "dynamic_scenarios": { "vehicle_arrival_rate": 0.7, "emergency_arrival_rate": 0.035, "emergency_urgency_range": (8, 10), "max_steps": 400, "max_queue_per_lane": 30, "surge_probability": 0.04, "surge_multiplier": 3.0, }, } # --------------------------------------------------------------------------- # Internal vehicle dataclass # --------------------------------------------------------------------------- @dataclass class Vehicle: vehicle_type: VehicleType direction: Direction waiting_time: int = 0 urgency: int = 0 # --------------------------------------------------------------------------- # Environment – extends openenv-core Environment base class # --------------------------------------------------------------------------- class TrafficControlEnvironment(Environment): """ OpenEnv-compliant Autonomous Traffic Control environment. Inherits from openenv.core.env_server.interfaces.Environment, making it compatible with openenv-core's create_app() factory without any adapter. Methods ------- reset(seed, episode_id, **kwargs) -> TrafficObservation step(action) -> TrafficObservation state -> TrafficState (property) """ # Allow multiple concurrent WebSocket sessions (each session gets its own # env instance when max_concurrent_envs > 1 in create_app). SUPPORTS_CONCURRENT_SESSIONS: bool = True def __init__(self, task_id: str = "basic_flow") -> None: if task_id not in TASK_CONFIGS: raise ValueError( f"Unknown task_id '{task_id}'. " f"Valid options: {list(TASK_CONFIGS.keys())}" ) self.task_id = task_id self._cfg = TASK_CONFIGS[task_id] self._rng = random.Random() self._episode_id: str = "" self._step_count: int = 0 self._queues: List[List[Vehicle]] = [[] for _ in range(4)] self._current_phase: LightPhase = LightPhase.NS_GREEN self._time_in_phase: int = 0 self._pending_phase: Optional[int] = None self._total_vehicles_passed: int = 0 self._total_emergency_passed: int = 0 self._total_waiting_time: float = 0.0 self._total_emergency_delay: float = 0.0 self._total_collisions: int = 0 self._total_phase_changes: int = 0 # ------------------------------------------------------------------ # openenv-core Environment interface # ------------------------------------------------------------------ def reset( self, seed: Optional[int] = None, episode_id: Optional[str] = None, **kwargs, ) -> TrafficObservation: """Start a fresh episode – required by openenv-core Environment.""" self._rng = random.Random(seed) self._episode_id = episode_id or str(uuid.uuid4()) self._step_count = 0 self._queues = [[] for _ in range(4)] self._current_phase = LightPhase.NS_GREEN self._time_in_phase = 0 self._pending_phase = None self._total_vehicles_passed = 0 self._total_emergency_passed = 0 self._total_waiting_time = 0.0 self._total_emergency_delay = 0.0 self._total_collisions = 0 self._total_phase_changes = 0 return self._build_obs(0, 0, 0.0, False, 0.0, False) def step(self, action: TrafficAction) -> TrafficObservation: """Execute one simulation step – required by openenv-core Environment.""" self._step_count += 1 self._spawn_vehicles() phase_changed = self._apply_action(action) self._advance_phase() vehicles_passed, emergency_passed = self._flow_traffic() waiting_delta = self._tick_waiting_times() collision = self._check_collision() reward = self._compute_reward( vehicles_passed, emergency_passed, waiting_delta, collision, phase_changed ) self._total_vehicles_passed += vehicles_passed self._total_emergency_passed += emergency_passed self._total_waiting_time += waiting_delta if collision: self._total_collisions += 1 done = collision or self._step_count >= self._cfg["max_steps"] return self._build_obs(vehicles_passed, emergency_passed, waiting_delta, collision, reward, done) @property def state(self) -> TrafficState: """Return cumulative episode-level state – required by openenv-core Environment.""" return TrafficState( episode_id=self._episode_id, step_count=self._step_count, task_id=self.task_id, total_vehicles_passed=self._total_vehicles_passed, total_emergency_passed=self._total_emergency_passed, total_waiting_time=self._total_waiting_time, total_emergency_delay=self._total_emergency_delay, total_collisions=self._total_collisions, total_phase_changes=self._total_phase_changes, ) # ------------------------------------------------------------------ # Simulation internals # ------------------------------------------------------------------ def _spawn_vehicles(self) -> None: arr = self._cfg["vehicle_arrival_rate"] em = self._cfg["emergency_arrival_rate"] urg = self._cfg["emergency_urgency_range"] surge_p = self._cfg["surge_probability"] surge_m = self._cfg["surge_multiplier"] max_q = self._cfg["max_queue_per_lane"] surge_dir = -1 surge_extra = 0 if surge_p > 0.0 and self._rng.random() < surge_p: surge_dir = self._rng.randint(0, 3) surge_extra = max(0, int(self._rng.gauss(3, 1) * surge_m)) for d in range(4): n = self._poisson(arr) if d == surge_dir: n += surge_extra for _ in range(n): if len(self._queues[d]) < max_q: vt = VehicleType.BUS if self._rng.random() < 0.10 else VehicleType.CAR self._queues[d].append(Vehicle(vt, Direction(d))) if em > 0.0 and self._rng.random() < em: if len(self._queues[d]) < max_q: urgency = self._rng.randint(urg[0], urg[1]) self._queues[d].insert( 0, Vehicle(VehicleType.EMERGENCY, Direction(d), urgency=urgency), ) def _apply_action(self, action: TrafficAction) -> bool: req = action.light_phase if req not in (PHASE_NS_GREEN, PHASE_EW_GREEN, PHASE_ALL_RED): return False if self._current_phase in (LightPhase.NS_YELLOW, LightPhase.EW_YELLOW): return False current_base = int(self._current_phase) if current_base == req: return False self._total_phase_changes += 1 self._pending_phase = req if req == PHASE_ALL_RED: self._current_phase = LightPhase.ALL_RED self._time_in_phase = 0 self._pending_phase = None elif self._current_phase == LightPhase.NS_GREEN: self._current_phase = LightPhase.NS_YELLOW self._time_in_phase = 0 elif self._current_phase == LightPhase.EW_GREEN: self._current_phase = LightPhase.EW_YELLOW self._time_in_phase = 0 elif self._current_phase == LightPhase.ALL_RED: self._current_phase = LightPhase(req) self._time_in_phase = 0 self._pending_phase = None return True def _advance_phase(self) -> None: self._time_in_phase += 1 if self._current_phase in (LightPhase.NS_YELLOW, LightPhase.EW_YELLOW): if self._time_in_phase >= YELLOW_DURATION: target = self._pending_phase if self._pending_phase is not None else PHASE_ALL_RED self._current_phase = LightPhase(target) self._time_in_phase = 0 self._pending_phase = None def _flow_traffic(self) -> Tuple[int, int]: allowed = PHASE_ALLOWS[self._current_phase] flow_rate = PHASE_FLOW_RATE[self._current_phase] vehicles_passed = 0 emergency_passed = 0 for d in allowed: queue = self._queues[int(d)] passed_dir = 0 while queue and passed_dir < flow_rate: vehicle = queue.pop(0) passed_dir += 1 if vehicle.vehicle_type == VehicleType.EMERGENCY: emergency_passed += 1 else: vehicles_passed += 1 return vehicles_passed, emergency_passed def _tick_waiting_times(self) -> float: total = 0.0 for d in range(4): for v in self._queues[d]: v.waiting_time += 1 total += 1.0 if v.vehicle_type == VehicleType.EMERGENCY: self._total_emergency_delay += 1.0 return total def _check_collision(self) -> bool: total_queued = sum(len(q) for q in self._queues) if total_queued > 40 and self._time_in_phase > 20: return self._rng.random() < 0.04 return False def _compute_reward( self, vehicles_passed: int, emergency_passed: int, waiting_delta: float, collision: bool, phase_changed: bool, ) -> float: r = 0.0 r += vehicles_passed * 0.20 r += emergency_passed * 10.0 r -= waiting_delta * 0.05 for d in range(4): for v in self._queues[d]: if v.vehicle_type == VehicleType.EMERGENCY: r -= v.urgency * 0.4 if collision: r -= 200.0 if phase_changed: p = int(self._current_phase) if p == PHASE_NS_GREEN: if (len(self._queues[0]) + len(self._queues[1])) == 0: r -= 0.5 elif p == PHASE_EW_GREEN: if (len(self._queues[2]) + len(self._queues[3])) == 0: r -= 0.5 return r def _build_obs( self, vehicles_passed: int, emergency_passed: int, waiting_delta: float, collision: bool, reward: float, done: bool, ) -> TrafficObservation: queue_lengths = [] emergency_queue = [] emergency_urgency = [] for d in range(4): reg = sum(1 for v in self._queues[d] if v.vehicle_type != VehicleType.EMERGENCY) em = sum(1 for v in self._queues[d] if v.vehicle_type == VehicleType.EMERGENCY) max_u = max( (v.urgency for v in self._queues[d] if v.vehicle_type == VehicleType.EMERGENCY), default=0, ) queue_lengths.append(reg) emergency_queue.append(em) emergency_urgency.append(max_u) return TrafficObservation( current_phase=int(self._current_phase), time_in_phase=self._time_in_phase, queue_lengths=queue_lengths, emergency_queue=emergency_queue, emergency_urgency=emergency_urgency, vehicles_passed=vehicles_passed, emergency_passed=emergency_passed, total_waiting_time=waiting_delta, collision=collision, reward=reward, done=done, metadata={ "step_count": self._step_count, "task_id": self.task_id, }, ) def _poisson(self, lam: float) -> int: if lam <= 0.0: return 0 threshold = math.exp(-lam) k, p = 0, 1.0 while p > threshold: k += 1 p *= self._rng.random() return k - 1