prashant-9457/my-openenv-task
0
1"""2ICU Resource Allocation — OpenEnv Environment3==============================================4A real-world environment modelling a 20-bed ICU in a 500-bed Indian tertiary-5care hospital. An AI agent acts as the ICU charge-coordinator, deciding every630 minutes how to allocate beds, staff and equipment across an incoming stream7of critically ill patients.8 9Clinical grounding10------------------11- Patient severity measured by the SOFA score (Sequential Organ Failure12 Assessment, range 0-24), the gold standard triage tool used in ICUs globally.13- Nurse : patient ratios follow NABH (National Accreditation Board for Hospitals)14 guidelines — 1 : 2 for ICU.15- Bed turnover time (cleaning + preparation) modelled at 45-90 min, matching16 published Indian hospital data.17- Patient arrival follows a non-homogeneous Poisson process with higher rates18 during 08-12 h and 20-24 h (documented admission peaks).19- Equipment (ventilators, monitors, dialysis) tracked against real 500-bed20 tertiary-care inventories.21- Costs in INR, calibrated to CGHS package rates (Central Govt Health Scheme).22 23Action space (Discrete 7)24--------------------------250 HOLD – Observe; no allocation change this step.261 ADMIT_CRITICAL – Admit the highest-SOFA patient from the waiting queue.272 ADMIT_FIFO – Admit the longest-waiting patient from the queue.283 TRANSFER_OUT – Transfer the most stable current ICU patient to step-down.294 CALL_EXTRA_NURSE– Arrange an overtime nurse for this shift (₹1 200 premium).305 SPECIALIST_CONSULT – Request urgent specialist consult for the sickest31 current patient (₹3 500, reduces mortality risk).326 EXPEDITE_BED – Pay porter/housekeeping overtime to clean next bed faster33 (₹600, cuts turnover time by ~30 min).34 35Observation space (23 fields)36------------------------------37See _build_obs() for full description with units and ranges.38 39Reward (partial progress at every step, no sparse end-of-episode)40------41See _calculate_reward() for breakdown.42"""43 44import math45import random46from dataclasses import dataclass, field47from typing import Optional48 49 50# ─────────────────────────────────────────────────────────────────────────────51# Data structures52# ─────────────────────────────────────────────────────────────────────────────53 54@dataclass55class Patient:56 """Represents a single patient."""57 pid: int58 sofa: float # 0-24 (Sequential Organ Failure Assessment score)59 needs_ventilator: bool60 needs_dialysis: bool61 arrival_step: int # Step when patient arrived in queue62 admitted_step: Optional[int] = None63 los_steps: int = 0 # Expected length of stay in steps (each step=30 min)64 mortality_risk: float = 0.0 # 0-1 (derived from SOFA)65 66 @staticmethod67 def sofa_to_mortality(sofa: float) -> float:68 """69 Approximate ICU mortality from SOFA score.70 Based on: Ferreira et al., JAMA 2001 — SOFA score as predictor of ICU outcome.71 """72 # SOFA 0-6: ~10%, 7-9: ~21%, 10-12: ~33%, 13-14: ~50%, 15-24: ~95%73 breakpoints = [(6, 0.10), (9, 0.21), (12, 0.33), (14, 0.50), (24, 0.95)]74 for threshold, risk in breakpoints:75 if sofa <= threshold:76 return risk77 return 0.9578 79 def __post_init__(self):80 self.mortality_risk = self.sofa_to_mortality(self.sofa)81 82 83@dataclass84class Bed:85 """ICU bed state."""86 bed_id: int87 patient: Optional[Patient] = None88 turnover_steps_remaining: int = 0 # >0 means bed being cleaned89 90 @property91 def is_available(self) -> bool:92 return self.patient is None and self.turnover_steps_remaining == 093 94 @property95 def is_occupied(self) -> bool:96 return self.patient is not None97 98 @property99 def in_turnover(self) -> bool:100 return self.patient is None and self.turnover_steps_remaining > 0101 102 103# ─────────────────────────────────────────────────────────────────────────────104# Main environment105# ─────────────────────────────────────────────────────────────────────────────106 107class ICUEnv:108 """109 OpenEnv-compliant ICU Resource Allocation environment.110 111 Each step = 30 minutes of real time.112 One episode = 48 steps = 24 hours (one full ICU duty cycle).113 """114 115 # ── Hospital configuration (typical 500-bed tertiary care, India) ────────116 TOTAL_ICU_BEDS = 20117 TOTAL_VENTILATORS = 12118 TOTAL_DIALYSIS = 4119 TOTAL_MONITORS = 20 # 1 per bed120 121 # Staff baseline per shift122 BASE_NURSES_DAY = 10 # 1:2 ratio for 20 beds123 BASE_NURSES_NIGHT = 8 # slightly reduced night staffing124 BASE_DOCTORS = 2 # intensivists on call125 126 # Financial (INR, calibrated to CGHS 2023 rates)127 DAILY_BUDGET_INR = 150_000 # ₹1.5 lakh daily ICU operating budget128 ICU_BED_COST_STEP = 3_125 # ₹3 125 per bed per step (₹1.5L / 48 steps / ~1 bed)129 OVERTIME_NURSE_COST = 1_200 # Per shift overtime premium130 SPECIALIST_COST = 3_500 # Specialist consult fee131 EXPEDITE_BED_COST = 600 # Housekeeping overtime for fast bed prep132 133 # Clinical thresholds134 SAFE_NURSE_RATIO = 2.0 # Max patients per nurse (NABH standard)135 CRITICAL_SOFA = 11 # SOFA ≥ 11 → critical, time-sensitive136 TRANSFER_SOFA_MAX = 6 # SOFA ≤ 6 → eligible for step-down transfer137 MAX_QUEUE_WAIT_STEPS = 4 # >4 steps (2h) wait for critical → outcome worsens138 139 # Time-to-admission mortality penalty scaling140 # Every 30-min delay for critical patient increases mortality risk by ~3%141 DELAY_MORTALITY_INCREMENT = 0.03142 143 MAX_STEPS = 48144 145 def __init__(self, seed: int = 42):146 self.seed = seed147 self._rng = random.Random(seed)148 self._step = 0149 self._pid_counter = 0150 self.reset()151 152 # ─────────────────────────────────────────────────────────────────────153 # OpenEnv API154 # ─────────────────────────────────────────────────────────────────────155 156 def reset(self) -> dict:157 """Reset to beginning of a fresh 24-hour duty cycle."""158 self._rng = random.Random(self.seed)159 self._step = 0160 self._pid_counter = 0161 self._hour = 8.0 # Duty cycle starts at 08:00162 163 # Beds164 self._beds = [Bed(bed_id=i) for i in range(self.TOTAL_ICU_BEDS)]165 166 # Pre-populate ~60% bed occupancy at start of shift (realistic handover)167 initial_patients = int(self.TOTAL_ICU_BEDS * 0.60)168 for i in range(initial_patients):169 p = self._generate_patient(is_initial=True)170 self._beds[i].patient = p171 172 # Queues and tracking173 self._queue: list[Patient] = []174 self._discharged: list[Patient] = []175 self._deaths_in_queue: int = 0176 self._adverse_events: int = 0177 self._admissions_today: int = 0178 self._transfers_today: int = 0179 180 # Staff181 self._extra_nurses_called: int = 0182 self._specialist_consults: int = 0183 184 # Equipment185 self._ventilators_in_use: int = sum(186 1 for b in self._beds if b.is_occupied and b.patient.needs_ventilator187 )188 self._dialysis_in_use: int = sum(189 1 for b in self._beds if b.is_occupied and b.patient.needs_dialysis190 )191 192 # Budget193 self._budget_remaining = self.DAILY_BUDGET_INR194 self._cost_this_step = 0.0195 196 # Outcome tracking197 self._mortality_risks_avoided = 0.0198 self._total_sofa_admitted = 0.0199 self._wait_violations = 0200 201 # Generate initial queue (2-5 waiting patients at shift start)202 for _ in range(self._rng.randint(2, 5)):203 self._queue.append(self._generate_patient())204 205 return self._build_obs()206 207 def step(self, action: int) -> tuple[dict, float, bool, dict]:208 """209 Apply action, simulate 30 minutes, return (obs, reward, done, info).210 211 Actions:212 0 HOLD213 1 ADMIT_CRITICAL – admit highest-SOFA patient214 2 ADMIT_FIFO – admit longest-waiting patient215 3 TRANSFER_OUT – move most stable ICU patient to step-down216 4 CALL_EXTRA_NURSE217 5 SPECIALIST_CONSULT218 6 EXPEDITE_BED219 """220 if action not in range(7):221 action = 0222 223 self._cost_this_step = 0.0224 action_result = self._apply_action(action)225 226 # Simulate 30-minute time passage227 self._simulate_time_passage()228 229 # Advance clock230 self._step += 1231 self._hour = (8.0 + self._step * 0.5) % 24232 233 reward = self._calculate_reward(action)234 done = self._step >= self.MAX_STEPS235 236 obs = self._build_obs()237 info = {238 "action_result": action_result,239 "cost_this_step_inr": round(self._cost_this_step, 2),240 "deaths_in_queue": self._deaths_in_queue,241 "adverse_events": self._adverse_events,242 "admissions_today": self._admissions_today,243 "transfers_today": self._transfers_today,244 "wait_violations": self._wait_violations,245 "nurse_ratio": round(self._nurse_patient_ratio(), 2),246 }247 return obs, round(reward, 4), done, info248 249 def state(self) -> dict:250 """Return current observation without advancing time."""251 return self._build_obs()252 253 # ─────────────────────────────────────────────────────────────────────254 # Action handlers255 # ─────────────────────────────────────────────────────────────────────256 257 def _apply_action(self, action: int) -> str:258 if action == 0:259 return "HOLD: no allocation change"260 261 elif action == 1: # ADMIT_CRITICAL262 if not self._queue:263 return "ADMIT_CRITICAL: queue empty"264 bed = self._first_available_bed()265 if bed is None:266 return "ADMIT_CRITICAL: no available bed"267 # Admit highest-SOFA patient268 patient = max(self._queue, key=lambda p: p.sofa)269 self._queue.remove(patient)270 self._admit_patient(bed, patient)271 return f"ADMIT_CRITICAL: admitted P{patient.pid} (SOFA={patient.sofa:.1f}) to Bed {bed.bed_id}"272 273 elif action == 2: # ADMIT_FIFO274 if not self._queue:275 return "ADMIT_FIFO: queue empty"276 bed = self._first_available_bed()277 if bed is None:278 return "ADMIT_FIFO: no available bed"279 # Admit longest-waiting patient280 patient = min(self._queue, key=lambda p: p.arrival_step)281 self._queue.remove(patient)282 self._admit_patient(bed, patient)283 return f"ADMIT_FIFO: admitted P{patient.pid} (SOFA={patient.sofa:.1f}) to Bed {bed.bed_id}"284 285 elif action == 3: # TRANSFER_OUT286 candidates = [b for b in self._beds287 if b.is_occupied and b.patient.sofa <= self.TRANSFER_SOFA_MAX]288 if not candidates:289 return "TRANSFER_OUT: no stable patients eligible"290 # Transfer lowest-SOFA patient291 bed = min(candidates, key=lambda b: b.patient.sofa)292 patient = bed.patient293 bed.patient = None294 bed.turnover_steps_remaining = self._rng.randint(1, 3) # 30-90 min cleanup295 self._discharged.append(patient)296 self._transfers_today += 1297 # Reclaim equipment298 if patient.needs_ventilator:299 self._ventilators_in_use = max(0, self._ventilators_in_use - 1)300 if patient.needs_dialysis:301 self._dialysis_in_use = max(0, self._dialysis_in_use - 1)302 return f"TRANSFER_OUT: transferred P{patient.pid} (SOFA={patient.sofa:.1f}) to step-down"303 304 elif action == 4: # CALL_EXTRA_NURSE305 cost = self.OVERTIME_NURSE_COST306 if self._budget_remaining < cost:307 return "CALL_EXTRA_NURSE: insufficient budget"308 self._budget_remaining -= cost309 self._cost_this_step += cost310 self._extra_nurses_called += 1311 return f"CALL_EXTRA_NURSE: +1 nurse this shift (₹{cost})"312 313 elif action == 5: # SPECIALIST_CONSULT314 # Reduce mortality risk of sickest current patient315 occupied = [b for b in self._beds if b.is_occupied]316 if not occupied:317 return "SPECIALIST_CONSULT: no current patients"318 cost = self.SPECIALIST_COST319 if self._budget_remaining < cost:320 return "SPECIALIST_CONSULT: insufficient budget"321 sickest = max(occupied, key=lambda b: b.patient.mortality_risk)322 old_risk = sickest.patient.mortality_risk323 sickest.patient.mortality_risk = max(0.05, old_risk - 0.15)324 self._budget_remaining -= cost325 self._cost_this_step += cost326 self._specialist_consults += 1327 self._mortality_risks_avoided += (old_risk - sickest.patient.mortality_risk)328 return (f"SPECIALIST_CONSULT: P{sickest.patient.pid} mortality risk "329 f"{old_risk:.2f}→{sickest.patient.mortality_risk:.2f} (₹{cost})")330 331 elif action == 6: # EXPEDITE_BED332 turnover_beds = [b for b in self._beds if b.in_turnover]333 if not turnover_beds:334 return "EXPEDITE_BED: no beds in turnover"335 cost = self.EXPEDITE_BED_COST336 if self._budget_remaining < cost:337 return "EXPEDITE_BED: insufficient budget"338 # Reduce turnover time of the bed closest to ready339 target = min(turnover_beds, key=lambda b: b.turnover_steps_remaining)340 target.turnover_steps_remaining = max(0, target.turnover_steps_remaining - 1)341 self._budget_remaining -= cost342 self._cost_this_step += cost343 return f"EXPEDITE_BED: Bed {target.bed_id} ready sooner (₹{cost})"344 345 return "UNKNOWN action"346 347 # ─────────────────────────────────────────────────────────────────────348 # Simulation349 # ─────────────────────────────────────────────────────────────────────350 351 def _simulate_time_passage(self):352 """Advance simulation by 30 minutes."""353 # 1. Existing ICU patients: progress LOS, possibly deteriorate or improve354 for bed in self._beds:355 if not bed.is_occupied:356 continue357 p = bed.patient358 p.los_steps += 1359 360 # Natural deterioration/improvement (small random walk on SOFA)361 delta = self._rng.gauss(0, 0.4)362 p.sofa = max(0.0, min(24.0, p.sofa + delta))363 p.mortality_risk = Patient.sofa_to_mortality(p.sofa)364 365 # Check if patient ready for discharge (completed LOS)366 if p.los_steps >= p.admitted_step + self._rng.randint(4, 16):367 # Patient stable enough for general ward368 if p.sofa <= 8:369 bed.patient = None370 bed.turnover_steps_remaining = self._rng.randint(1, 3)371 self._discharged.append(p)372 if p.needs_ventilator:373 self._ventilators_in_use = max(0, self._ventilators_in_use - 1)374 if p.needs_dialysis:375 self._dialysis_in_use = max(0, self._dialysis_in_use - 1)376 377 # Adverse event if nurse ratio is unsafe378 if self._nurse_patient_ratio() > self.SAFE_NURSE_RATIO * 1.5:379 if self._rng.random() < 0.08: # 8% chance per step per patient380 self._adverse_events += 1381 p.sofa = min(24.0, p.sofa + 1.5)382 p.mortality_risk = Patient.sofa_to_mortality(p.sofa)383 384 # 2. Bed turnover countdown385 for bed in self._beds:386 if bed.in_turnover:387 bed.turnover_steps_remaining -= 1388 389 # 3. Waiting queue deterioration and deaths390 for p in list(self._queue):391 wait = self._step - p.arrival_step392 if wait >= self.MAX_QUEUE_WAIT_STEPS and p.sofa >= self.CRITICAL_SOFA:393 self._wait_violations += 1394 # Mortality risk worsens with delay395 p.mortality_risk = min(0.99, p.mortality_risk + self.DELAY_MORTALITY_INCREMENT)396 p.sofa = min(24.0, p.sofa + 0.5)397 # Patient may die in queue398 if p.mortality_risk > 0.90 and self._rng.random() < 0.15:399 self._queue.remove(p)400 self._deaths_in_queue += 1401 402 # 4. New arrivals (non-homogeneous Poisson, peaks at 08-12h and 20-24h)403 arrival_rate = self._arrival_rate_per_step()404 n_arrivals = self._rng.poisson_approx(arrival_rate)405 for _ in range(n_arrivals):406 self._queue.append(self._generate_patient())407 408 # 5. Deduct bed operating costs from budget409 occupied_count = sum(1 for b in self._beds if b.is_occupied)410 step_cost = occupied_count * (self.DAILY_BUDGET_INR / self.MAX_STEPS / self.TOTAL_ICU_BEDS)411 self._budget_remaining = max(0.0, self._budget_remaining - step_cost)412 413 def _arrival_rate_per_step(self) -> float:414 """415 Non-homogeneous Poisson arrival rate.416 Peak hours: 08-12 (post-morning rounds referrals) and 20-24 (evening emergencies).417 Based on: Arias-Verdú et al., Critical Care Medicine 2017.418 """419 h = self._hour420 base = 0.4421 if 8 <= h < 12:422 return base * 2.0423 elif 20 <= h < 24:424 return base * 1.8425 elif 14 <= h < 18:426 return base * 1.2427 elif 0 <= h < 6:428 return base * 0.5429 return base430 431 def _generate_patient(self, is_initial: bool = False) -> Patient:432 """Generate a patient with realistic SOFA distribution."""433 self._pid_counter += 1434 # SOFA distribution in Indian ICU referrals (based on published data)435 # ~20% critical (≥11), ~40% severe (7-10), ~40% moderate (0-6)436 r = self._rng.random()437 if r < 0.20:438 sofa = self._rng.uniform(11, 20) # Critical439 elif r < 0.60:440 sofa = self._rng.uniform(7, 11) # Severe441 else:442 sofa = self._rng.uniform(1, 7) # Moderate443 444 # Equipment needs correlate with severity445 needs_vent = sofa >= 9 and self._rng.random() < 0.55446 needs_dial = sofa >= 10 and self._rng.random() < 0.30447 448 expected_los = max(4, int(sofa * 1.5 + self._rng.gauss(0, 2)))449 450 return Patient(451 pid=self._pid_counter,452 sofa=round(sofa, 1),453 needs_ventilator=needs_vent,454 needs_dialysis=needs_dial,455 arrival_step=self._step if not is_initial else -self._rng.randint(2, 8),456 los_steps=0,457 admitted_step=0 if is_initial else None,458 )459 460 def _admit_patient(self, bed: Bed, patient: Patient):461 """Place a patient into a bed and update equipment counts."""462 bed.patient = patient463 patient.admitted_step = self._step464 self._admissions_today += 1465 466 if patient.needs_ventilator and self._ventilators_in_use < self.TOTAL_VENTILATORS:467 self._ventilators_in_use += 1468 elif patient.needs_ventilator:469 patient.needs_ventilator = False # Can't provide — document as constraint470 471 if patient.needs_dialysis and self._dialysis_in_use < self.TOTAL_DIALYSIS:472 self._dialysis_in_use += 1473 elif patient.needs_dialysis:474 patient.needs_dialysis = False475 476 # ─────────────────────────────────────────────────────────────────────477 # Reward478 # ─────────────────────────────────────────────────────────────────────479 480 def _calculate_reward(self, action: int) -> float:481 """482 Multi-objective reward with partial signals at every step.483 484 Components485 ----------486 +3.0 per critical patient admitted before 2-hour breach487 +1.0 for maintaining safe nurse:patient ratio488 -5.0 per patient death in queue this step489 -2.0 per adverse event this step490 -1.5 per critical patient waiting > 2 hours (ongoing)491 +0.5 per effective specialist consult (action=5 AND patient at risk)492 -0.3 budget overspend fraction (if budget depleted)493 -0.5 for HOLD when critical patient in queue AND bed available (missed opportunity)494 """495 reward = 0.0496 497 # Nurse ratio component498 ratio = self._nurse_patient_ratio()499 if ratio <= self.SAFE_NURSE_RATIO:500 reward += 1.0501 else:502 reward -= (ratio - self.SAFE_NURSE_RATIO) * 1.5503 504 # Critical patients in queue breaching wait time505 for p in self._queue:506 wait = self._step - p.arrival_step507 if p.sofa >= self.CRITICAL_SOFA and wait > self.MAX_QUEUE_WAIT_STEPS:508 reward -= 1.5509 510 # Deaths in queue penalise heavily511 # (deaths_in_queue is cumulative; reward on incremental change is handled512 # by tracking _last_deaths — approximated here as step-level signal)513 # We track last step deaths via adverse events counter delta514 # Simple: penalise for current step deaths via wait_violations increase515 breach_count = sum(516 1 for p in self._queue517 if p.sofa >= self.CRITICAL_SOFA and (self._step - p.arrival_step) >= self.MAX_QUEUE_WAIT_STEPS518 )519 reward -= breach_count * 0.5520 521 # Missed opportunity: HOLD when could have admitted critical patient522 if action == 0:523 has_critical_queue = any(p.sofa >= self.CRITICAL_SOFA for p in self._queue)524 has_free_bed = self._first_available_bed() is not None525 if has_critical_queue and has_free_bed:526 reward -= 0.5527 528 # Budget management529 budget_fraction = self._budget_remaining / self.DAILY_BUDGET_INR530 if budget_fraction <= 0:531 reward -= 0.3532 else:533 reward += budget_fraction * 0.2534 535 # Throughput bonus: reward for high admissions-to-capacity ratio536 throughput = self._admissions_today / max(1, self._step)537 reward += min(0.5, throughput * 0.5)538 539 # Equipment utilisation (reward efficient use, penalise over-saturation)540 vent_util = self._ventilators_in_use / self.TOTAL_VENTILATORS541 if vent_util > 0.95:542 reward -= 0.4 # Near capacity is dangerous543 544 return reward545 546 # ─────────────────────────────────────────────────────────────────────547 # Observation builder548 # ─────────────────────────────────────────────────────────────────────549 550 def _build_obs(self) -> dict:551 occupied = [b for b in self._beds if b.is_occupied]552 available = [b for b in self._beds if b.is_available]553 turnover = [b for b in self._beds if b.in_turnover]554 555 q_sofa = [p.sofa for p in self._queue]556 q_critical = [p for p in self._queue if p.sofa >= self.CRITICAL_SOFA]557 q_severe = [p for p in self._queue if 7 <= p.sofa < self.CRITICAL_SOFA]558 q_moderate = [p for p in self._queue if p.sofa < 7]559 560 current_sofa = [b.patient.sofa for b in occupied]561 avg_icu_sofa = sum(current_sofa) / max(1, len(current_sofa))562 avg_icu_mortality = sum(b.patient.mortality_risk for b in occupied) / max(1, len(occupied))563 564 longest_wait = 0565 if self._queue:566 longest_wait = self._step - min(p.arrival_step for p in self._queue)567 568 nurses = self._nurses_on_duty()569 570 return {571 # Bed status572 "beds_total": self.TOTAL_ICU_BEDS,573 "beds_occupied": len(occupied),574 "beds_available": len(available),575 "beds_in_turnover": len(turnover),576 577 # Queue status578 "queue_total": len(self._queue),579 "queue_critical": len(q_critical), # SOFA ≥ 11580 "queue_severe": len(q_severe), # SOFA 7-10581 "queue_moderate": len(q_moderate), # SOFA < 7582 "queue_max_wait_steps": longest_wait, # Steps since oldest arrival583 584 # Current ICU patient acuity585 "avg_icu_sofa": round(avg_icu_sofa, 2),586 "avg_icu_mortality_risk": round(avg_icu_mortality, 3),587 588 # Equipment589 "ventilators_available": self.TOTAL_VENTILATORS - self._ventilators_in_use,590 "ventilators_in_use": self._ventilators_in_use,591 "dialysis_available": self.TOTAL_DIALYSIS - self._dialysis_in_use,592 593 # Staff594 "nurses_on_duty": nurses,595 "nurse_patient_ratio": round(self._nurse_patient_ratio(), 2),596 "doctors_on_duty": self.BASE_DOCTORS,597 598 # Time599 "shift": self._current_shift(), # 0=day 1=evening 2=night600 "time_of_day": round(self._hour, 1),601 "step": self._step,602 603 # Finance604 "budget_remaining_inr": round(self._budget_remaining, 2),605 "budget_utilisation_pct": round((1 - self._budget_remaining / self.DAILY_BUDGET_INR) * 100, 1),606 607 # Cumulative outcomes608 "admissions_today": self._admissions_today,609 "transfers_today": self._transfers_today,610 "deaths_in_queue": self._deaths_in_queue,611 "adverse_events": self._adverse_events,612 "wait_violations": self._wait_violations,613 }614 615 # ─────────────────────────────────────────────────────────────────────616 # Helpers617 # ─────────────────────────────────────────────────────────────────────618 619 def _first_available_bed(self) -> Optional[Bed]:620 for b in self._beds:621 if b.is_available:622 return b623 return None624 625 def _nurses_on_duty(self) -> int:626 base = self.BASE_NURSES_DAY if self._current_shift() == 0 else self.BASE_NURSES_NIGHT627 return base + self._extra_nurses_called628 629 def _nurse_patient_ratio(self) -> float:630 occupied = sum(1 for b in self._beds if b.is_occupied)631 nurses = self._nurses_on_duty()632 if nurses == 0:633 return float("inf")634 return occupied / nurses635 636 def _current_shift(self) -> int:637 h = self._hour638 if 8 <= h < 16:639 return 0 # Day640 elif 16 <= h < 24:641 return 1 # Evening642 else:643 return 2 # Night644 645 class _RNG(random.Random):646 pass647 648 649# Monkey-patch a Poisson approximation onto the rng650def _poisson_approx(self, lam: float) -> int:651 """Approximate Poisson using sum of Bernoulli trials (works for small λ)."""652 n, p = 20, lam / 20653 return sum(1 for _ in range(n) if self.random() < p)654 655random.Random.poisson_approx = _poisson_approx656 657 658# ─────────────────────────────────────────────────────────────────────────────659# Quick smoke test660# ─────────────────────────────────────────────────────────────────────────────661if __name__ == "__main__":662 env = ICUEnv(seed=42)663 obs = env.reset()664 print("── INITIAL STATE ─────────────────────────────────────────")665 for k, v in obs.items():666 print(f" {k:35s}: {v}")667 668 print("\n── FIRST 6 STEPS ─────────────────────────────────────────")669 total_reward = 0670 for i in range(6):671 action = i % 7672 obs, reward, done, info = env.step(action)673 total_reward += reward674 print(f"Step {i+1} | act={action} | beds={obs['beds_occupied']}/20 "675 f"| queue={obs['queue_total']} (crit={obs['queue_critical']}) "676 f"| reward={reward:+.3f} | budget=₹{obs['budget_remaining_inr']:,.0f} "677 f"| {info['action_result'][:50]}")678 679 print(f"\nTotal reward so far: {total_reward:.3f}")680 print("State OK:", len(env.state()) == 27)681