Spaces:
Sleeping
Sleeping
| 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, | |
| } | |
| 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 | |
| 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, | |
| } | |