README.md

metadata

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

# 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:

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

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

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

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.

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:

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):

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

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:

{
  "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.