--- title: Hospital ED Resource Allocator emoji: ๐Ÿฅ colorFrom: red colorTo: blue sdk: docker pinned: false app_port: 8000 tags: - openenv - reinforcement-learning - hospital - healthcare - triage base_path: /web --- # Hospital Emergency Department Resource Allocator **An OpenEnv reinforcement-learning environment for triaging patients and allocating beds, ICU slots, and ventilators in a public hospital emergency department.** Built for the **Meta PyTorch OpenEnv Hackathon (Round 1, April 2026)**. * **OpenEnv-compliant**: subclasses `openenv.core.Environment`, ships Pydantic `Action` / `Observation` / `State`, FastAPI app at `app:app`, drivable via `openenv.core.GenericEnvClient` over WebSocket. * **Gymnasium-native core**: same simulation backs the Gymnasium env (`HospitalEnv`) so all SB3 / sb3-contrib agents work unchanged. * **MaskablePPO agent that beats the heuristic on every scenario**, trained in ~90 s of CPU. * **5 stress scenarios**: `normal_day`, `surge`, `mass_casualty`, `night_shift`, `ventilator_crisis`. * **47 tests** including a real WebSocket round-trip via `GenericEnvClient` against the FastAPI app. ## TL;DR results | Agent | Overall | Survival | Crit save | Wait | Invalid | |-------------|---------|----------|-----------|------|---------| | Random | 33.36 | 30.9% | 11.2% | 6.54 | 57.5% | | Heuristic | 68.43 | 65.0% | 70.5% | 4.62 | 0.0% | | **MaskablePPO (50k)** | **78.85** | **80.6%** | **76.3%** | **2.27** | **0.0%** | ``` normal_day surge mass_casualty night_shift vent_crisis random 35.02 33.86 30.10 35.06 32.75 heuristic 83.46 54.84 50.23 90.25 63.39 ppo 89.54 68.38 61.67 90.81 83.88 ``` PPO trained for **50,000 timesteps in โ‰ˆ90 seconds** on a CPU MacBook, using `sb3-contrib.MaskablePPO` with the env's action mask. --- ## Why this matters Indian public hospitals routinely run at 150%+ capacity during surge events (COVID waves, heat waves, mass-casualty incidents). The decision of *which* patient gets the next ICU bed or ventilator is, literally, a life-or-death scheduling problem โ€” one that rewards both fast triage and disciplined long-horizon planning. This environment turns that problem into a reproducible RL benchmark. --- ## How it works ### Observation (`gym.spaces.Dict`) | Key | Shape | Meaning | |----------------------|------------|-----------------------------------------------------| | `bed_occupancy` | `(20,)` | Severity of patient in each general bed (0 = empty) | | `icu_occupancy` | `(5,)` | Severity of patient in each ICU bed | | `ventilator_status` | `(3,)` | 0 = free, 1 = in use | | `waiting_queue` | `(10, 3)` | `[severity, condition_id, waiting_time]` per slot | | `time_step` | `(1,)` | Current timestep (0โ€“100) | | `stats` | `(3,)` | `[total_treated, total_deaths, total_waiting]` | ### Action (`gym.spaces.Discrete(39)`) ``` 0 No-op 1-10 Assign waiting patient [idx] to a general bed 11-20 Assign waiting patient [idx] to an ICU bed 21-23 Use ventilator slot [idx] on the most severe ICU patient without one 24-28 Transfer general bed [idx] patient to an ICU bed 29-33 Early-discharge general bed [idx] patient 34-38 Discharge ICU bed [idx] patient ``` Invalid actions (e.g. assigning from an empty queue slot) are handled gracefully: they incur a small `-1` penalty and never crash. ### Action masking The env exposes a boolean validity mask via two channels: * `info["action_mask"]` โ€” present in every `reset` / `step` info dict. * `env.action_masks()` โ€” direct method, named to be auto-detected by `sb3-contrib.MaskablePPO` through wrapper chains. The mask never marks an invalid action as valid, and it never marks a clearly-valid action as invalid. Tested by running 400 masked-random steps and asserting **zero** invalid actions (`tests/test_env.py`). ### Reward (dense) | Event | Reward | |-------------------------------------------------|--------| | Successful discharge | `+5` | | Successful discharge of a critical patient | `+10` | | Patient death | `-15` | | Per-timestep per waiting **critical** patient | `-0.5` | | Per-timestep per waiting non-critical patient | `-0.1` | | Invalid action | `-1` | | Timestep with ICU utilization > 60% | `+0.1` | | End-of-episode (`+20 * survival_rate`) | โ‰ค `+20`| ### Episode ``` โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ 24-hour ED shift โ”€โ”€โ”€ 100 timesteps โ”‚ โ”‚ โ”‚ โ”‚ arrivals (Poisson) โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ โ–ผ โ”‚ โ”‚ [waiting queue] โ”€โ”€โ–บ agent action โ”€โ”€โ–บ [beds / ICU] โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ โ–ผ โ–ผ โ”‚ โ”‚ deterioration tick_treatment โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ โ–ผ โ–ผ โ”‚ โ”‚ death? discharged โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ ``` Episode ends at step 100 (truncated) or when `total_deaths >= max_deaths` (terminated). --- ## Setup ```bash # 1. Clone and enter the repo git clone cd hospital-resource-allocator # 2. Create a virtual environment and install dependencies python -m venv .venv source .venv/bin/activate pip install -r requirements.txt ``` Or with Docker: ```bash docker build -t hospital-env . docker run --rm hospital-env # runs evaluate.py docker run --rm hospital-env python demo.py # runs visual demo ``` --- ## Running it ### Visual demo ```bash python demo.py # heuristic agent, normal day python demo.py --agent random # random baseline python demo.py --scenario surge # COVID-like surge python demo.py --scenario mass_casualty # trauma burst ``` ### Grade an agent ```bash python evaluate.py --agent heuristic # prints JSON score to stdout python evaluate.py --agent random --episodes 10 python evaluate.py --agent ppo --output ppo_score.json ``` ### Side-by-side comparison ```bash python compare.py # random vs heuristic vs ppo python compare.py --agents random heuristic --episodes 10 python compare.py --output comparison.json ``` Sample output: ``` ======================================================================================== Hospital ED Resource Allocator โ€” agent comparison ======================================================================================== Agent Overall Survive Crit Wait Util Inval bar ---------------------------------------------------------------------------------------- random 33.36 30.9% 11.2% 6.54 10.6% 57.5% โ–ˆโ–ˆโ–ˆโ–ˆยทยทยทยทยทยทยทยท heuristic 68.43 65.0% 70.5% 4.62 29.3% 0.0% โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆยทยทยทยท ppo 78.85 80.6% 76.3% 2.27 36.3% 0.0% โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆยทยทยท ``` ### Train MaskablePPO ```bash pip install -r requirements.txt # pulls torch + sb3-contrib python -m agents.train_ppo --timesteps 50000 --save-path ppo_hospital python evaluate.py --agent ppo ``` 50k timesteps takes โ‰ˆ90 s on a CPU MacBook and reliably beats the heuristic. Use `--no-mask` to compare against vanilla PPO without action masking. ### OpenEnv-compliant FastAPI server (the hackathon entry point) The environment ships with a real **OpenEnv** wrapper at `app:app`, built via `openenv.core.create_app(...)`. This is the file the hackathon judges' OpenEnv test client (and any other `EnvClient` consumer) will talk to. ```bash uvicorn app:app --host 0.0.0.0 --port 8000 ``` Routes registered by the framework: | Method | Path | Purpose | |--------|-------------|--------------------------------------------------| | GET | /health | liveness probe | | GET | /metadata | environment name / description / version | | GET | /schema | JSON Schema for action / observation / state | | GET | /state | current episode state (Pydantic `HospitalState`) | | POST | /reset | start a new episode (`{"seed":0}`) | | POST | /step | apply an action (`{"action":{"action":0}}`) | | WS | /ws | full stateful session (used by `EnvClient`) | | WS | /mcp | MCP JSON-RPC tool-calling protocol | | GET | /docs | auto-generated FastAPI/OpenAPI docs | Drive it with the official OpenEnv client: ```python import asyncio from openenv.core import GenericEnvClient async def main(): client = GenericEnvClient(base_url="http://localhost:8000") await client.connect() result = await client.reset(seed=0) print("initial obs time_step:", result.observation["time_step"]) for _ in range(10): result = await client.step({"action": 0}) print("reward:", result.reward, "done:", result.done) state = await client.state() print("step_count:", state["step_count"], "queue_len:", state["queue_len"]) await client.disconnect() asyncio.run(main()) ``` Or via Docker (the default `CMD`): ```bash docker build -t hospital-env . docker run --rm -p 8000:8000 hospital-env docker run --rm -p 8000:8000 -e HOSPITAL_SCENARIO=surge hospital-env # surge config ``` There is also a minimal stdlib HTTP server at `server.py` (zero dependencies, custom JSON shape) โ€” useful as a fallback when `openenv-core` isn't installed, but **not** the contract the judges will speak. Use `app:app` for grading. ### Run the test suite ```bash pytest tests/ -v ``` 34 tests covering Gymnasium API compliance, action mask correctness, reward-hack regression, grader reproducibility, all 5 scenarios ร— 2 agents, and the HTTP server endpoints. --- ## Grader output `python evaluate.py --agent ppo` produces: ```json { "overall_score": 78.85, "survival_rate": 0.8060, "avg_wait_time": 2.27, "resource_utilization": 0.3631, "critical_patient_survival": 0.7625, "invalid_action_rate": 0.0, "scenario_scores": { "normal_day": 89.54, "surge": 68.38, "mass_casualty": 61.67, "night_shift": 90.81, "ventilator_crisis": 83.88 } } ``` ### Scoring formula ``` overall = 40 * survival_rate + 20 * critical_patient_survival + 20 * (1 - min(avg_wait_time / 20, 1)) + 10 * resource_utilization + 10 * (1 - invalid_action_rate) ``` Range: **0 โ€“ 100**. Reference scores on a fresh machine (5 episodes / scenario): | Agent | overall | survival | crit survival | invalid rate | |-------------|---------|----------|---------------|--------------| | Random | 33.36 | 0.31 | 0.11 | 0.58 | | Heuristic | 68.43 | 0.65 | 0.71 | 0.00 | | **MaskablePPO (50k)** | **78.85** | **0.81** | **0.76** | **0.00** | PPO trained for **50,000 timesteps in โ‰ˆ90 seconds** beats the heuristic on every scenario, with the biggest gains on the hardest ones: `surge` (+13), `mass_casualty` (+11), `ventilator_crisis` (+20). ### Per-scenario notes | Scenario | What it tests | |----------------------|-----------------------------------------------------------------| | `normal_day` | Steady-state triage, mild-leaning mix | | `surge` | COVID-style respiratory wave, sustained pressure | | `mass_casualty` | Sudden trauma burst, rapid-triage stress test | | `night_shift` | Sparse arrivals โ€” penalises wasted resources / panic vent use | | `ventilator_crisis` | 80% respiratory at sev โ‰ฅ 3, 3 vents are the binding constraint | ### Example mid-episode render `python demo.py --agent heuristic` will print a frame like this each timestep: ``` ======================================================================== HOSPITAL ED | Step 35 / 100 | Treated: 17 Deaths: 6 Queue: 10 ======================================================================== Gen Beds [. . . . . . . . . . . . . . . . . . . .] occ= 0.0% ICU Beds [4v . 4v 4 . ] occ=60.0% Vents [# # .] util=66.7% ------------------------------------------------------------------------ Waiting queue (10): #0 sev=2 cond=trauma wait=29 vent_needed= #1 sev=2 cond=infection wait=20 vent_needed= #2 sev=3 cond=trauma wait=17 vent_needed= #3 sev=2 cond=respiratory wait=14 vent_needed= #4 sev=3 cond=infection wait=13 vent_needed= ... ------------------------------------------------------------------------ Crit saved: 6/12 Invalid rate: 0.00% Ep reward: +40.40 ======================================================================== ``` Read it as: 0/20 general beds occupied, 3/5 ICU beds full (two on ventilators), 2/3 ventilators in use, queue of 10 with the longest waiter at 29 timesteps. This is a **mid-surge moment** where the agent has run out of ICU room and the queue is backing up โ€” exactly the kind of state where the reward function is pressuring the agent to discharge or transfer. --- ## Architecture ``` hospital-resource-allocator/ โ”œโ”€โ”€ hospital_env/ # Simulation core + both interfaces โ”‚ โ”œโ”€โ”€ patient.py # Patient & PatientGenerator (Poisson arrivals) โ”‚ โ”œโ”€โ”€ hospital.py # Hospital state (beds / ICU / vents / queue) โ”‚ โ”œโ”€โ”€ env.py # HospitalEnv(gym.Env) + action_masks() โ”‚ โ”œโ”€โ”€ renderer.py # ASCII rendering for demo.py โ”‚ โ”œโ”€โ”€ openenv_types.py # โ˜… Pydantic Action / Observation / State โ”‚ โ””โ”€โ”€ openenv_env.py # โ˜… HospitalOpenEnv(openenv.core.Environment) โ”‚ โ”œโ”€โ”€ grader/ # Programmatic scoring across scenarios โ”‚ โ”œโ”€โ”€ scenarios.py # 5 scenario configs โ”‚ โ””โ”€โ”€ grader.py # Per-episode rollout + composite score โ”‚ โ”œโ”€โ”€ agents/ # Baseline and trainable agents โ”‚ โ”œโ”€โ”€ random_agent.py # Uniform-random baseline (~33) โ”‚ โ”œโ”€โ”€ heuristic_agent.py # Rule-based triage (~68) โ”‚ โ””โ”€โ”€ train_ppo.py # MaskablePPO trainer + DictObs flattener โ”‚ โ”œโ”€โ”€ tests/ # pytest test suite (47 tests) โ”‚ โ”œโ”€โ”€ test_env.py # Gym API + mechanics + mask + reward-hack regression โ”‚ โ”œโ”€โ”€ test_grader.py # Grader correctness + reproducibility โ”‚ โ”œโ”€โ”€ test_scenarios.py # Smoke each scenario ร— random/heuristic โ”‚ โ”œโ”€โ”€ test_server.py # stdlib HTTP server endpoint smoke โ”‚ โ””โ”€โ”€ test_openenv.py # โ˜… HospitalOpenEnv unit + FastAPI TestClient round-trip โ”‚ โ”œโ”€โ”€ app.py # โ˜… OpenEnv FastAPI entry point (uvicorn app:app) โ”œโ”€โ”€ server.py # Minimal stdlib HTTP server (fallback, non-OpenEnv shape) โ”œโ”€โ”€ compare.py # Side-by-side agent comparison CLI โ”œโ”€โ”€ demo.py # Visual episode with ASCII rendering โ”œโ”€โ”€ evaluate.py # Grade an agent, output JSON score โ”œโ”€โ”€ Dockerfile # python:3.11-slim, default CMD = uvicorn app:app โ”œโ”€โ”€ requirements.txt # openenv-core, gymnasium, sb3, sb3-contrib, torch, pytest โ””โ”€โ”€ README.md # (this file) ``` ### Key design choices 0. **OpenEnv compliance via a thin adapter, not a rewrite.** The simulation core is a vanilla Gymnasium `HospitalEnv`. A separate `HospitalOpenEnv(openenv.core.Environment[โ€ฆ])` adapter wraps it, converts numpy obs to a Pydantic `HospitalObservation`, exposes a `state` property of type `HospitalState`, and is mounted as `app:app` via `openenv.core.create_app(...)`. This means: * The judges' OpenEnv test client speaks to a real `Environment` subclass, with proper Pydantic validation, JSON schemas at `/schema`, and stateful WebSocket sessions at `/ws`. * All the existing agents, grader, tests, demo and CLI keep using the Gymnasium env unchanged. * Verified end-to-end: a real `GenericEnvClient` round-trip (`reset` โ†’ 5 ร— `step` โ†’ `state`) returns correct `episode_id`, `step_count`, `queue_len`, cumulative `reward`. 1. **Dense reward shaping.** A death costs `-15` but a successful discharge only `+5`, so the agent cannot afford to lose patients; the per-timestep waiting penalty nudges it to act rather than idle. 2. **Early-discharge reward hack is closed.** Action `29-38` only awards `+5` / `+10` if `treatment_time_remaining โ‰ค 0` โ€” i.e. the patient is actually fully treated. Discharging early ("against medical advice") incurs a `-3` penalty scaled by remaining severity and **does not** count as a successful treatment. Without this, a trained agent could learn the loop *admit โ†’ immediately discharge โ†’ +5 free reward* (verified: a 40-step admit-then-discharge policy went from ~+200 reward to **โˆ’958**). 3. **Ventilator action is slot-indexed.** Action `21+v` "uses ventilator slot `v` on the most severe ICU patient who doesn't already have one." This is strictly more agent-friendly than a per-ICU-bed action because the observation exposes ventilator availability but not the ventilator-to-ICU-bed mapping. 4. **Critical-survivor tracking is monotone.** A patient is counted as "critical" the first tick their severity hits 4 (either on arrival or after deteriorating in the queue). This guarantees `critical_saved โ‰ค critical_total` even though patients can walk up the severity ladder. 5. **Per-episode deterministic seeding.** Both `HospitalEnv` and `Grader` derive episode seeds from a user-supplied base seed so that `evaluate.py` gives bit-identical output across runs. --- ## License MIT.