| title | Hospital ED Resource Allocator | |||||
|---|---|---|---|---|---|---|
| emoji | 🏥 | |||||
| colorFrom | red | |||||
| colorTo | blue | |||||
| sdk | docker | |||||
| pinned | false | |||||
| app_port | 8000 | |||||
| tags |
|
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 PydanticAction/Observation/State, FastAPI app atapp:app, drivable viaopenenv.core.GenericEnvClientover 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
GenericEnvClientagainst the FastAPI app.
| 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.
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.
| 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] |
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.
The env exposes a boolean validity mask via two channels:
info["action_mask"]— present in everyreset/stepinfo dict.env.action_masks()— direct method, named to be auto-detected bysb3-contrib.MaskablePPOthrough 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).
| 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 |
┌─────────────────────────────────────────────────────────────┐
│ 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).
# 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.txtOr 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 demopython 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 burstpython evaluate.py --agent heuristic # prints JSON score to stdout
python evaluate.py --agent random --episodes 10
python evaluate.py --agent ppo --output ppo_score.jsonpython compare.py # random vs heuristic vs ppo
python compare.py --agents random heuristic --episodes 10
python compare.py --output comparison.jsonSample 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% █████████···
pip install -r requirements.txt # pulls torch + sb3-contrib
python -m agents.train_ppo --timesteps 50000 --save-path ppo_hospital
python evaluate.py --agent ppo50k timesteps takes ≈90 s on a CPU MacBook and reliably beats the
heuristic. Use --no-mask to compare against vanilla PPO without
action masking.
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 8000Routes 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 configThere 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.
pytest tests/ -v34 tests covering Gymnasium API compliance, action mask correctness, reward-hack regression, grader reproducibility, all 5 scenarios × 2 agents, and the HTTP server endpoints.
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
}
}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).
| 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 |
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.
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)
-
OpenEnv compliance via a thin adapter, not a rewrite. The simulation core is a vanilla Gymnasium
HospitalEnv. A separateHospitalOpenEnv(openenv.core.Environment[…])adapter wraps it, converts numpy obs to a PydanticHospitalObservation, exposes astateproperty of typeHospitalState, and is mounted asapp:appviaopenenv.core.create_app(...). This means:- The judges' OpenEnv test client speaks to a real
Environmentsubclass, 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
GenericEnvClientround-trip (reset→ 5 ×step→state) returns correctepisode_id,step_count,queue_len, cumulativereward.
- The judges' OpenEnv test client speaks to a real
-
Dense reward shaping. A death costs
-15but 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. -
Early-discharge reward hack is closed. Action
29-38only awards+5/+10iftreatment_time_remaining ≤ 0— i.e. the patient is actually fully treated. Discharging early ("against medical advice") incurs a-3penalty 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). -
Ventilator action is slot-indexed. Action
21+v"uses ventilator slotvon 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. -
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_totaleven though patients can walk up the severity ladder. -
Per-episode deterministic seeding. Both
HospitalEnvandGraderderive episode seeds from a user-supplied base seed so thatevaluate.pygives bit-identical output across runs.
MIT.
