import random from dataclasses import dataclass, field from typing import Dict, Iterable, Optional, Tuple from models import EmergencyVehicle, LANES, VEHICLE_TYPES, LaneQueues, TrafficAction, TrafficState from tasks import get_task LANE_GROUPS = { TrafficAction.NS_GREEN: ("N", "S"), TrafficAction.EW_GREEN: ("E", "W"), TrafficAction.LEFT_PRIORITY: ("N", "E"), } VEHICLE_CLEARANCE_UNITS = { "cars": 1.0, "bikes": 0.45, "autos": 0.75, "buses": 2.2, "trucks": 2.5, } @dataclass class TrafficMetrics: total_vehicles_cleared: int = 0 total_wait_observations: float = 0.0 wait_samples: int = 0 max_queue_length: int = 0 emergency_seen: int = 0 emergency_cleared_fast: int = 0 unsafe_switches: int = 0 full_clearances: int = 0 @dataclass class StepDiagnostics: vehicles_cleared: int = 0 unsafe_switch_penalty: int = 0 emergency_clear_bonus: int = 0 raw_reward: float = 0.0 normalized_reward: float = 0.0 invalid_action: bool = False notes: list = field(default_factory=list) class IndianTrafficEnv: """Deterministic mixed-traffic signal environment for one Indian urban intersection. The class is intentionally network-ready: lane state and flow calculations are kept per-lane, so a future MultiIntersectionEnv can compose several instances. """ def __init__(self, task_id: str = "single_intersection"): self.task_id = task_id self.task = get_task(task_id) self.rng = random.Random() self.seed_value = 42 self.tick = 0 self.current_phase = TrafficAction.ALL_RED self.time_since_switch = 0 self.queues: Dict[str, Dict[str, int]] = {} self.waiting_time: Dict[str, float] = {} self.last_inflow: Dict[str, Dict[str, int]] = {} self.pedestrian_count = 0 self.pedestrian_wait_time = 0.0 self.emergency = EmergencyVehicle(present=False) self.rain_level = float(self.task.constraints["rain_level"]) self.driver_aggression = 0.4 self.random_blockage_probability = 0.02 self.peak_hour_multiplier = float(self.task.constraints["initial_peak_multiplier"]) self.metrics = TrafficMetrics() self.done = False self.reset(self.seed_value, task_id) def reset(self, seed: Optional[int] = None, task_id: Optional[str] = None) -> TrafficState: if task_id is not None: self.task_id = task_id self.task = get_task(task_id) self.seed_value = 42 if seed is None else int(seed) self.rng = random.Random(self.seed_value) self.tick = 0 self.current_phase = TrafficAction.ALL_RED self.time_since_switch = 0 self.rain_level = float(self.task.constraints["rain_level"]) self.driver_aggression = self.rng.uniform(0.22, 0.72) self.random_blockage_probability = self.rng.uniform(0.01, 0.05) self.peak_hour_multiplier = float(self.task.constraints["initial_peak_multiplier"]) self.queues = {lane: self._sample_initial_lane(lane) for lane in LANES} self.waiting_time = {lane: 0.0 for lane in LANES} self.last_inflow = {lane: {vehicle: 0 for vehicle in VEHICLE_TYPES} for lane in LANES} self.pedestrian_count = self.rng.randint(2, 14) self.pedestrian_wait_time = 0.0 self.emergency = EmergencyVehicle(present=False) self.metrics = TrafficMetrics() self.done = False return self.get_state() def step(self, action: TrafficAction) -> Tuple[TrafficState, float, bool, Dict[str, object]]: if self.done: return self.get_state(), 0.0, True, {"message": "Episode already complete."} diagnostics = StepDiagnostics() try: action = TrafficAction(action) except ValueError: action = TrafficAction.ALL_RED diagnostics.invalid_action = True diagnostics.notes.append("Invalid action converted to ALL_RED.") effective_action = self._apply_action_constraints(action, diagnostics) self._maybe_switch_phase(effective_action) self.last_inflow = self._generate_inflow() self._add_inflow(self.last_inflow) self._maybe_spawn_emergency() diagnostics.vehicles_cleared = self._clear_traffic(effective_action) emergency_cleared = self._update_emergency(effective_action) diagnostics.emergency_clear_bonus = 1 if emergency_cleared else 0 self._update_pedestrians(effective_action) self._update_waiting_times() queue_length = self._total_queue_length() total_wait = sum(self.waiting_time.values()) self.metrics.total_vehicles_cleared += diagnostics.vehicles_cleared self.metrics.total_wait_observations += total_wait self.metrics.wait_samples += 1 self.metrics.max_queue_length = max(self.metrics.max_queue_length, queue_length) if queue_length == 0: self.metrics.full_clearances += 1 reward = self._calculate_reward( vehicles_cleared=diagnostics.vehicles_cleared, total_waiting_time=total_wait, queue_length=queue_length, pedestrian_wait_time=self.pedestrian_wait_time, unsafe_switch_penalty=diagnostics.unsafe_switch_penalty, emergency_clear_bonus=diagnostics.emergency_clear_bonus, diagnostics=diagnostics, ) self.tick += 1 self.time_since_switch += 1 self.done = self._is_done(queue_length) return self.get_state(), reward, self.done, self._info(diagnostics) def get_state(self) -> TrafficState: return TrafficState( tick=self.tick, lane_queues={lane: LaneQueues(**self.queues[lane]) for lane in LANES}, lane_waiting_time={lane: round(self.waiting_time[lane], 3) for lane in LANES}, current_signal_phase=self.current_phase, time_since_last_phase_switch=self.time_since_switch, pedestrian_count=self.pedestrian_count, pedestrian_wait_time=round(self.pedestrian_wait_time, 3), emergency_vehicle=self.emergency, rain_level=round(self.rain_level, 3), random_traffic_inflow={lane: LaneQueues(**self.last_inflow[lane]) for lane in LANES}, ) def _sample_initial_lane(self, lane: str) -> Dict[str, int]: arterial_bonus = 3 if lane in ("N", "S") else 1 return { "cars": self.rng.randint(3 + arterial_bonus, 8 + arterial_bonus), "bikes": self.rng.randint(6 + arterial_bonus, 14 + arterial_bonus), "autos": self.rng.randint(2, 6), "buses": self.rng.randint(0, 2), "trucks": self.rng.randint(0, 2), } def _generate_inflow(self) -> Dict[str, Dict[str, int]]: inflow = {} rain_slowdown = 1.0 + self.rain_level * 0.25 for lane in LANES: arterial = 1.25 if lane in ("N", "S") else 0.95 base = self.peak_hour_multiplier * arterial * rain_slowdown inflow[lane] = { "cars": self._bounded_arrivals(base, 2), "bikes": self._bounded_arrivals(base * 1.8, 4), "autos": self._bounded_arrivals(base * 0.85, 2), "buses": 1 if self.rng.random() < 0.08 * base else 0, "trucks": 1 if self.rng.random() < 0.06 * base else 0, } if self.tick % 40 == 0 and self.tick > 0: self.peak_hour_multiplier = max(0.8, self.peak_hour_multiplier * 0.94) return inflow def _bounded_arrivals(self, intensity: float, cap: int) -> int: arrivals = int(intensity) fractional = intensity - arrivals if self.rng.random() < fractional: arrivals += 1 if self.rng.random() < 0.12 * self.peak_hour_multiplier: arrivals += 1 return max(0, min(cap, arrivals)) def _add_inflow(self, inflow: Dict[str, Dict[str, int]]) -> None: for lane in LANES: for vehicle in VEHICLE_TYPES: self.queues[lane][vehicle] += inflow[lane][vehicle] def _apply_action_constraints(self, action: TrafficAction, diagnostics: StepDiagnostics) -> TrafficAction: min_green = int(self.task.constraints["min_green_time"]) switching = action != self.current_phase and action not in (TrafficAction.EXTEND_GREEN, TrafficAction.EMERGENCY_OVERRIDE) green_to_green = self.current_phase in (TrafficAction.NS_GREEN, TrafficAction.EW_GREEN, TrafficAction.LEFT_PRIORITY) wants_green = action in (TrafficAction.NS_GREEN, TrafficAction.EW_GREEN, TrafficAction.LEFT_PRIORITY) if switching and green_to_green and wants_green and self.time_since_switch < min_green: diagnostics.unsafe_switch_penalty += 1 diagnostics.notes.append("Minimum green-time violation.") if action == TrafficAction.PEDESTRIAN_CROSS and self.pedestrian_count == 0: diagnostics.unsafe_switch_penalty += 1 diagnostics.notes.append("Pedestrian phase requested without demand.") if action == TrafficAction.EMERGENCY_OVERRIDE and not self.emergency.present: diagnostics.unsafe_switch_penalty += 1 diagnostics.notes.append("Emergency override requested without emergency vehicle.") if action == TrafficAction.EXTEND_GREEN: if self.current_phase in (TrafficAction.NS_GREEN, TrafficAction.EW_GREEN, TrafficAction.LEFT_PRIORITY): return self.current_phase diagnostics.notes.append("EXTEND_GREEN from non-green phase converted to ALL_RED.") return TrafficAction.ALL_RED return action def _maybe_switch_phase(self, action: TrafficAction) -> None: if action != self.current_phase: self.current_phase = action self.time_since_switch = 0 def _clear_traffic(self, action: TrafficAction) -> int: if action == TrafficAction.ALL_RED or action == TrafficAction.PEDESTRIAN_CROSS: return 0 lanes = self._active_lanes(action) if not lanes: return 0 cleared = 0 rain_factor = 1.0 - self.rain_level * 0.32 blockage_factor = 0.45 if self.rng.random() < self.random_blockage_probability else 1.0 aggression_bonus = 1.0 + self.driver_aggression * 0.18 capacity_units = 8.5 * rain_factor * blockage_factor * aggression_bonus if action == TrafficAction.LEFT_PRIORITY: capacity_units *= 0.65 if action == TrafficAction.EMERGENCY_OVERRIDE: capacity_units *= 1.15 for lane in lanes: remaining_units = capacity_units for vehicle in ("bikes", "autos", "cars", "buses", "trucks"): cleared_count, remaining_units = self._clear_vehicle_type(lane, vehicle, remaining_units) cleared += cleared_count if self.queues[lane] and sum(self.queues[lane].values()) == 0: self.waiting_time[lane] = 0.0 return cleared def _clear_vehicle_type(self, lane: str, vehicle: str, remaining_units: float) -> Tuple[int, float]: unit = VEHICLE_CLEARANCE_UNITS[vehicle] possible = int(remaining_units // unit) count = min(self.queues[lane][vehicle], possible) self.queues[lane][vehicle] -= count return count, remaining_units - count * unit def _active_lanes(self, action: TrafficAction) -> Iterable[str]: if action == TrafficAction.EMERGENCY_OVERRIDE and self.emergency.present and self.emergency.lane: return (self.emergency.lane,) return LANE_GROUPS.get(action, ()) def _maybe_spawn_emergency(self) -> None: if self.emergency.present: return if self.rng.random() < float(self.task.constraints["emergency_rate"]): self.emergency = EmergencyVehicle( present=True, lane=self.rng.choice(LANES), type=self.rng.choice(["ambulance", "fire_truck"]), wait_time=0.0, ) self.metrics.emergency_seen += 1 def _update_emergency(self, action: TrafficAction) -> bool: if not self.emergency.present or not self.emergency.lane: return False active = set(self._active_lanes(action)) if self.emergency.lane in active: if action == TrafficAction.EMERGENCY_OVERRIDE or self.rng.random() < 0.55: if self.emergency.wait_time <= 8: self.metrics.emergency_cleared_fast += 1 self.emergency = EmergencyVehicle(present=False) return True self.emergency.wait_time += 1.0 return False def _update_pedestrians(self, action: TrafficAction) -> None: arrivals = 1 if self.rng.random() < 0.35 else 0 if self.rng.random() < 0.08 * self.peak_hour_multiplier: arrivals += self.rng.randint(1, 3) self.pedestrian_count += arrivals if action == TrafficAction.PEDESTRIAN_CROSS: crossed = min(self.pedestrian_count, 18) self.pedestrian_count -= crossed if self.pedestrian_count == 0: self.pedestrian_wait_time = 0.0 return if self.pedestrian_count: unsafe_crossing_pressure = self.driver_aggression * self.pedestrian_wait_time / 60.0 if self.rng.random() < unsafe_crossing_pressure: self.pedestrian_count = max(0, self.pedestrian_count - 1) self.pedestrian_wait_time += self.pedestrian_count * 0.55 def _update_waiting_times(self) -> None: active = set(self._active_lanes(self.current_phase)) for lane in LANES: queue = sum(self.queues[lane].values()) if queue == 0: self.waiting_time[lane] = 0.0 continue pressure = queue * (1.0 + self.rain_level * 0.3) if lane in active: pressure *= 0.35 self.waiting_time[lane] += pressure def _calculate_reward( self, vehicles_cleared: int, total_waiting_time: float, queue_length: int, pedestrian_wait_time: float, unsafe_switch_penalty: int, emergency_clear_bonus: int, diagnostics: StepDiagnostics, ) -> float: weights = self.task.reward_weights raw = ( vehicles_cleared * weights["vehicles_cleared"] + total_waiting_time * weights["total_waiting_time"] + queue_length * weights["queue_length"] + pedestrian_wait_time * weights["pedestrian_wait_time"] + unsafe_switch_penalty * weights["unsafe_switch_penalty"] + emergency_clear_bonus * weights["emergency_clear_bonus"] ) if queue_length == 0: raw += 10.0 diagnostics.notes.append("Full clearance milestone.") if diagnostics.invalid_action: raw -= 8.0 diagnostics.raw_reward = raw normalized = max(0.0, min(1.0, (raw + 180.0) / 260.0)) diagnostics.normalized_reward = normalized if unsafe_switch_penalty: self.metrics.unsafe_switches += unsafe_switch_penalty return normalized def _total_queue_length(self) -> int: return sum(sum(queue.values()) for queue in self.queues.values()) def _is_done(self, queue_length: int) -> bool: if self.tick + 1 >= int(self.task.constraints["max_steps"]): return True if self.metrics.total_vehicles_cleared >= int(self.task.termination["target_cleared"]): return True if queue_length >= int(self.task.constraints["max_queue_before_failure"]): return True max_emergency_wait = self.task.termination.get("max_emergency_wait") if max_emergency_wait and self.emergency.present and self.emergency.wait_time > float(max_emergency_wait): return True return False def _info(self, diagnostics: StepDiagnostics) -> Dict[str, object]: total_wait = sum(self.waiting_time.values()) return { "task_id": self.task_id, "vehicles_cleared": diagnostics.vehicles_cleared, "total_vehicles_cleared": self.metrics.total_vehicles_cleared, "total_waiting_time": round(total_wait, 3), "queue_length": self._total_queue_length(), "unsafe_switch_penalty": diagnostics.unsafe_switch_penalty, "emergency_clear_bonus": diagnostics.emergency_clear_bonus, "raw_reward": round(diagnostics.raw_reward, 3), "normalized_reward": round(diagnostics.normalized_reward, 3), "hidden_dynamics": { "driver_aggression": round(self.driver_aggression, 3), "random_blockage_probability": round(self.random_blockage_probability, 3), "peak_hour_multiplier": round(self.peak_hour_multiplier, 3), }, "notes": diagnostics.notes, }