| import sys |
| import os |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..')) |
|
|
| import copy |
| from uuid import uuid4 |
| from typing import Optional, Tuple |
|
|
| from server.grader import TrajectoryStep |
| from openenv.core.env_server.types import State |
| from openenv.core.env_server.interfaces import Environment |
|
|
| from models import ( |
| CityState, |
| CityObservation, |
| ContainmentAction, |
| ) |
| from server.constants import ( |
| HOSPITAL_BREACH_POINT, |
| TASK_CONFIG, |
| INFECTION_THRESHOLD, |
| SAFE_THRESHOLD, |
| LOW_THRESHOLD, |
| ALLOCATE_REDUCTION, |
| RESTRICT_REDUCTION, |
| TREATMENT_REDUCTION, |
| RESOURCE_REPLENISH, |
| REWARD_INFECTION_PENALTY, |
| REWARD_HOSPITAL_BREACH, |
| REWARD_EARLY_CONTAINMENT, |
| REWARD_UNNECESSARY_RESTRICTION, |
| REWARD_CORRECT_PRIORITISATION, |
| ) |
| from server.utils import ( |
| build_observation, |
| compute_spread, |
| get_highest_infected_district, |
| all_districts_contained, |
| any_hospital_breached, |
| districts_above_threshold, |
| snapshot_infection_rates, |
| generate_episode_id, |
| ) |
| from server.tasks.registry import get_task |
| from server.grader import grade_trajectory |
|
|
| _last_grade: dict = {} |
|
|
|
|
| class EpidemicContainmentEnv(Environment): |
| """ |
| Cascade Containment — RL environment for epidemic response. |
| |
| The agent acts as a city health authority making sequential resource |
| allocation decisions under uncertainty and potentially delayed data. |
| |
| Interface: |
| reset(task_name) → CityObservation |
| step(action) → CityObservation |
| state → State (property) |
| """ |
|
|
| def __init__(self): |
| self._city: CityState = CityState() |
| self._state: State = State(episode_id=str(uuid4()), step_count=0) |
| self._task_name: str = "easy" |
| self._trajectory: list = [] |
|
|
| |
|
|
| def reset(self, task_name: str = "easy") -> CityObservation: |
| self._task_name = task_name |
| task = get_task(task_name) |
| self._city = task.build_initial_state() |
| self._state = State(episode_id=generate_episode_id(), step_count=0) |
| self._trajectory = [] |
|
|
| return build_observation( |
| state = self._city, |
| step_count = self._state.step_count, |
| reward = None, |
| message = ( |
| f"Episode started. Task: {task_name}. " |
| f"Districts: {len(self._city.districts)}, " |
| f"Steps: {self._city.max_steps} available." |
| ), |
| done = False, |
| ) |
|
|
| def step(self, action: ContainmentAction) -> CityObservation: |
| assert self._city is not None, "Call reset() before step()." |
| assert self._state is not None, "Call reset() before step()." |
|
|
| action, message = self._validate_action(action) |
|
|
| self._city.infection_history.append( |
| snapshot_infection_rates(self._city.districts) |
| ) |
|
|
| self._apply_action(action) |
|
|
| new_rates = compute_spread(self._city.districts) |
| for i, district in enumerate(self._city.districts): |
| district.true_infection_rate = new_rates[i] |
|
|
| self._update_hospital_capacity() |
|
|
| self._city.available_resources = min( |
| self._city.available_resources + RESOURCE_REPLENISH, |
| TASK_CONFIG[self._task_name]["resource_pool"], |
| ) |
|
|
| for district in self._city.districts: |
| district.deployed_resources = 0 |
|
|
| self._city.day += 1 |
| self._state.step_count += 1 |
|
|
| reward = self._compute_reward(action) |
|
|
| self._trajectory.append(TrajectoryStep( |
| step = self._state.step_count, |
| city_state = copy.deepcopy(self._city), |
| action = action, |
| reward = reward, |
| done = False, |
| )) |
|
|
| done, terminal_message = self._check_terminal() |
|
|
| if done and self._trajectory: |
| self._trajectory[-1].done = True |
| import server.environment as _self_module |
| result = grade_trajectory(self._trajectory, self._task_name) |
| _self_module._last_grade = { |
| "final_score": result.final_score, |
| "containment_score": result.containment_score, |
| "hospital_score": result.hospital_score, |
| "efficiency_score": result.efficiency_score, |
| "speed_score": result.speed_score, |
| "hospital_breached": result.hospital_breached, |
| "districts_contained": result.districts_contained, |
| "total_steps": result.total_steps, |
| "task_name": self._task_name, |
| } |
|
|
| return build_observation( |
| state = self._city, |
| step_count = self._state.step_count, |
| reward = reward, |
| message = terminal_message if terminal_message else message, |
| done = done, |
| ) |
|
|
| @property |
| def state(self) -> State: |
| return self._state |
|
|
| def get_trajectory(self) -> list: |
| return self._trajectory |
|
|
| |
|
|
| def _validate_action( |
| self, action: ContainmentAction |
| ) -> Tuple[ContainmentAction, str]: |
| """ |
| Invalid actions are replaced with a safe default rather than raising — |
| the episode must continue even when the LLM returns malformed output. |
| """ |
| valid_types = {"test", "restrict", "allocate"} |
| num_districts = len(self._city.districts) |
|
|
| if action.action_type not in valid_types: |
| return ( |
| ContainmentAction(action_type="allocate", district_id=0), |
| f"Invalid action_type '{action.action_type}'. Defaulted to allocate on district 0.", |
| ) |
|
|
| if not (0 <= action.district_id < num_districts): |
| safe_id = max(0, min(action.district_id, num_districts - 1)) |
| return ( |
| ContainmentAction(action_type=action.action_type, district_id=safe_id), |
| f"district_id {action.district_id} out of range. Clamped to {safe_id}.", |
| ) |
|
|
| if action.action_type in {"test", "allocate"} and self._city.available_resources <= 0: |
| return ( |
| ContainmentAction(action_type="restrict", district_id=action.district_id), |
| f"No resources left. Switched to restrict on district {action.district_id}.", |
| ) |
|
|
| return action, f"{action.action_type.capitalize()} on district {action.district_id}." |
|
|
| def _apply_action(self, action: ContainmentAction) -> None: |
| district = self._city.districts[action.district_id] |
|
|
| if action.action_type == "test": |
| district.days_since_tested = 0 |
| self._city.available_resources -= 1 |
|
|
| elif action.action_type == "restrict": |
| district.restriction_active = True |
| district.days_since_tested += 1 |
|
|
| elif action.action_type == "allocate": |
| district.deployed_resources += 1 |
| district.true_infection_rate = max(0.0, district.true_infection_rate - TREATMENT_REDUCTION) |
| self._city.available_resources -= 1 |
| district.days_since_tested += 1 |
|
|
| for d in self._city.districts: |
| if d.district_id != action.district_id: |
| d.days_since_tested += 1 |
|
|
| |
|
|
| def _update_hospital_capacity(self) -> None: |
| for district in self._city.districts: |
| if district.true_infection_rate > INFECTION_THRESHOLD: |
| excess = district.true_infection_rate - INFECTION_THRESHOLD |
| drain = round(excess * 0.25, 4) |
| district.hospital_capacity_remaining = max( |
| 0.0, |
| district.hospital_capacity_remaining - drain |
| ) |
| else: |
| district.hospital_capacity_remaining = min( |
| 1.0, |
| district.hospital_capacity_remaining + 0.02 |
| ) |
| if district.true_infection_rate < SAFE_THRESHOLD: |
| district.restriction_active = False |
|
|
| def _compute_reward(self, action: ContainmentAction) -> float: |
| reward = 0.0 |
|
|
| for district in districts_above_threshold(self._city.districts): |
| density_weight = max(0.5, district.population_density * len(self._city.districts)) |
| reward += REWARD_INFECTION_PENALTY * min(2.0, density_weight) |
|
|
| for district in self._city.districts: |
| if district.hospital_capacity_remaining <= HOSPITAL_BREACH_POINT: |
| reward += REWARD_HOSPITAL_BREACH |
|
|
| for district in self._city.districts: |
| if district.true_infection_rate < SAFE_THRESHOLD: |
| time_factor = 1 - (self._state.step_count / self._city.max_steps) |
| reward += REWARD_EARLY_CONTAINMENT * time_factor |
|
|
| if action.action_type == "restrict": |
| target = self._city.districts[action.district_id] |
| if target.true_infection_rate < LOW_THRESHOLD: |
| reward += REWARD_UNNECESSARY_RESTRICTION |
|
|
| if action.action_type == "allocate": |
| if action.district_id == get_highest_infected_district(self._city.districts): |
| reward += REWARD_CORRECT_PRIORITISATION |
|
|
| return round(reward, 4) |
|
|
| def _check_terminal(self) -> Tuple[bool, Optional[str]]: |
| if all_districts_contained(self._city.districts): |
| return True, "✓ Outbreak contained. All districts below safe threshold." |
|
|
| if any_hospital_breached(self._city.districts): |
| return True, "✗ Hospital capacity breached. Episode failed." |
|
|
| if self._state.step_count >= self._city.max_steps: |
| return True, f"Episode complete. {self._city.max_steps} steps reached." |
|
|
| return False, None |
|
|
