Spaces:
Sleeping
Sleeping
| 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 <this-repo> | |
| 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. | |