Spaces:
Sleeping
Sleeping
| import uuid | |
| import logging | |
| from openenv.core.env_server.interfaces import Environment | |
| from openenv.core.env_server.types import State | |
| from models import TrafficAction, TrafficObservation | |
| from simulator import TrafficSimulator | |
| logger = logging.getLogger(__name__) | |
| REWARD_MIN_EXCLUSIVE = 0.01 | |
| REWARD_MAX_EXCLUSIVE = 0.99 | |
| class TrafficEnvironment(Environment): | |
| """ | |
| Traffic Signal RL Environment. | |
| The agent manages traffic light phases to optimize throughput and fairness. | |
| """ | |
| SUPPORTS_CONCURRENT_SESSIONS: bool = True | |
| MAX_STEPS = 20 | |
| DIFFICULTIES = ["easy", "medium", "hard"] | |
| def __init__(self): | |
| self._state = State(episode_id=str(uuid.uuid4()), step_count=0) | |
| self.current_task_idx = 0 | |
| self._init_simulator() | |
| def _init_simulator(self): | |
| diff = self.DIFFICULTIES[min(self.current_task_idx, len(self.DIFFICULTIES) - 1)] | |
| self.simulator = TrafficSimulator(difficulty=diff) | |
| self.cumulative_reward = 0.0 | |
| def _normalize_step_reward(self, raw_reward: float) -> float: | |
| # Map raw simulator rewards to the grading-friendly [0,1] band first, | |
| # then enforce strict exclusivity via epsilon margins. | |
| mapped = (raw_reward + 500.0) / 1000.0 | |
| return max(REWARD_MIN_EXCLUSIVE, min(REWARD_MAX_EXCLUSIVE, mapped)) | |
| def reset(self) -> TrafficObservation: | |
| self._state = State(episode_id=str(uuid.uuid4()), step_count=0) | |
| self._init_simulator() | |
| return self._get_observation(reward=0.0, done=False) | |
| def _get_observation(self, reward: float, done: bool) -> TrafficObservation: | |
| return TrafficObservation( | |
| queue_lengths=self.simulator.queues.copy(), | |
| waiting_times=self.simulator.wait_times.copy(), | |
| signal_phase=self.simulator.signal_phase, | |
| time_since_last_change=self.simulator.time_since_last_change, | |
| emergency_presence=self.simulator.emergency.copy(), | |
| pedestrian_requests=self.simulator.pedestrians.copy(), | |
| task_difficulty=self.simulator.difficulty, | |
| reward=reward, | |
| done=done, | |
| metadata={ | |
| "cumulative_reward": self.cumulative_reward, | |
| "step": self._state.step_count | |
| } | |
| ) | |
| def step(self, action: TrafficAction) -> TrafficObservation: | |
| # Check if we already hit done | |
| if self._state.step_count >= self.MAX_STEPS: | |
| # Reached end of current task, reset for the next task automatically or just halt? | |
| # Standard RL environments return done. Up to client to call reset(). | |
| return self._get_observation(reward=REWARD_MIN_EXCLUSIVE, done=True) | |
| raw_step_reward = self.simulator.step(action.action_type) | |
| step_reward = self._normalize_step_reward(raw_step_reward) | |
| self.cumulative_reward += raw_step_reward | |
| self._state.step_count += 1 | |
| done = (self._state.step_count >= self.MAX_STEPS) | |
| if done: | |
| # At the end of the episode, map raw cumulative reward to a 0.0 - 1.0 score. | |
| # Empirical mapping (since we don't know the exact max bounds without running an agent, | |
| # we'll frame it so anything >= 0 is mapped roughly into a normalized scale) | |
| # Easy task gets mostly positive rewards. Medium gets mixed. Hard gets heavily penalized. | |
| score = max(REWARD_MIN_EXCLUSIVE, min(REWARD_MAX_EXCLUSIVE, (self.cumulative_reward + 500) / 1000.0)) | |
| # The hackathon asks for graders to produce score between 0-1. | |
| # We overwrite the final step reward with this score for the client to read, | |
| # or the client computes based on cumulative reward. | |
| # To provide dense signal, step_reward is returned. The log script logs the 'score' independently. | |
| # Auto-advance task index so next reset gives next difficulty | |
| if self.current_task_idx < len(self.DIFFICULTIES) - 1: | |
| self.current_task_idx += 1 | |
| else: | |
| self.current_task_idx = 0 # Loop back to easy for endless play! | |
| return self._get_observation(reward=step_reward, done=done) | |
| def state(self) -> State: | |
| return self._state | |