import numpy as np from typing import Dict, List, Optional, Tuple, Any from dataclasses import dataclass, asdict from datetime import datetime import random from .models import Observation, Action, Reward, State, ServiceRequest class EmergencyResourceEnv: """ Emergency Resource Allocation Environment Grid-based city where agent allocates resources to service requests """ def __init__(self, grid_size: int = 5, max_time_steps: int = 50): self.grid_size = grid_size self.max_time_steps = max_time_steps # Environment state self.agent_pos: Tuple[int, int] = (0, 0) self.requests: List[ServiceRequest] = [] self.time_left: int = max_time_steps self.resources_left: int = 10 self.current_step: int = 0 self.total_reward: float = 0.0 self.done: bool = False self.task_config: Dict[str, Any] = {} def reset(self, task_config: Optional[Dict[str, Any]] = None) -> Observation: """Reset environment to initial state""" self.agent_pos = (0, 0) self.time_left = self.max_time_steps self.resources_left = 10 self.current_step = 0 self.total_reward = 0.0 self.done = False # Load task configuration self.task_config = task_config or {} # Generate service requests based on task self._generate_requests() return self._get_observation() def _generate_requests(self): """Generate service requests based on task difficulty""" task_difficulty = self.task_config.get('difficulty', 'medium') if task_difficulty == 'easy': num_requests = 3 critical_ratio = 0.33 # 1 critical, 2 normal elif task_difficulty == 'medium': num_requests = 5 critical_ratio = 0.4 # 2 critical, 3 normal else: # hard num_requests = 8 critical_ratio = 0.5 # 4 critical, 4 normal self.requests = [] for i in range(num_requests): priority = 'critical' if random.random() < critical_ratio else 'normal' # Ensure critical requests aren't too clustered if priority == 'critical' and len([r for r in self.requests if r.priority == 'critical']) >= num_requests * critical_ratio: priority = 'normal' # Generate position not overlapping with existing requests while True: pos = (random.randint(0, self.grid_size-1), random.randint(0, self.grid_size-1)) if not any(r.position == pos for r in self.requests): break self.requests.append(ServiceRequest( id=i, position=pos, priority=priority, allocated=False, created_at=self.current_step )) def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict]: """Execute action and return new state""" if self.done: return self._get_observation(), Reward(score=0.0), True, {"error": "Episode already done"} # Execute action action_type = action.action_type error = None try: if action_type == "move": self._move(action.direction) elif action_type == "allocate": reward_change = self._allocate_resources() else: error = f"Unknown action type: {action_type}" reward_change = -1.0 except Exception as e: error = str(e) reward_change = -1.0 # Apply step penalty step_penalty = -0.5 self.total_reward += step_penalty self.current_step += 1 self.time_left -= 1 # Check episode termination all_requests_served = all(r.allocated for r in self.requests) time_out = self.time_left <= 0 resources_depleted = self.resources_left <= 0 if all_requests_served: self.total_reward += 30 # Completion bonus self.done = True elif time_out: self.total_reward -= 20 # Time out penalty self.done = True elif resources_depleted: self.done = True elif self.current_step >= self.max_time_steps: self.done = True # Get final reward reward = Reward(score=max(0.0, min(1.0, (self.total_reward + 50) / 100))) return self._get_observation(), reward, self.done, {"error": error} def _move(self, direction: str): """Move agent in grid""" x, y = self.agent_pos if direction == "up" and x > 0: self.agent_pos = (x - 1, y) elif direction == "down" and x < self.grid_size - 1: self.agent_pos = (x + 1, y) elif direction == "left" and y > 0: self.agent_pos = (x, y - 1) elif direction == "right" and y < self.grid_size - 1: self.agent_pos = (x, y + 1) def _allocate_resources(self) -> float: """Allocate resources at current position""" reward_change = 0.0 # Check if any request at current position for request in self.requests: if not request.allocated and request.position == self.agent_pos: if self.resources_left > 0: request.allocated = True self.resources_left -= 1 if request.priority == "critical": reward_change = 20.0 else: reward_change = 10.0 self.total_reward += reward_change break return reward_change def _get_observation(self) -> Observation: """Get current observation""" # Create grid representation grid = np.zeros((self.grid_size, self.grid_size, 3)) # 3 channels: agent, critical, normal # Agent position grid[self.agent_pos[0], self.agent_pos[1], 0] = 1 # Requests for req in self.requests: if not req.allocated: channel = 1 if req.priority == "critical" else 2 grid[req.position[0], req.position[1], channel] = 1 return Observation( agent_position=self.agent_pos, grid=grid.tolist(), requests=[r for r in self.requests if not r.allocated], time_left=self.time_left, resources_left=self.resources_left, step=self.current_step, task_config=self.task_config ) def state(self) -> State: """Get current state""" return State( agent_position=self.agent_pos, requests=[asdict(r) for r in self.requests], time_left=self.time_left, resources_left=self.resources_left, current_step=self.current_step, done=self.done, total_reward=self.total_reward )