CoolFace
Apppublic

Tulsi2006/HOSPITAL_TRIAGE

sourceHugging Faceupdated 5mo agoView on Hugging Face
0likes
App README

Hospital Triage RL — Hackathon Upgrade

End-to-end RL training system built on the original Hospital Triage OpenEnv. Compatible with OpenEnv + TRL (GRPO) + Unsloth.

File Overview

FilePurpose
environment.pyUpgraded TriageEnv with curriculum learning, multi-component rewards, anti-hacking
action_parser.pyJSON action parser + validator + prompt builder
train_grpo.pyTRL GRPO training loop with rollout, reward fn, curriculum advancement
server.pyFastAPI server (upgraded, fully backward-compatible)
example_usage.pyRunnable examples for all features

Quickstart

bash
pip install -r requirements.txt

# Run examples (no GPU needed)
python example_usage.py

# Run demo mode
python train_grpo.py --demo --demo-level 1

# Start API server
python server.py

Curriculum Learning

python
from environment import TriageEnv

env = TriageEnv(level=1)   # 1 patient, abundant resources — learn triage basics
env = TriageEnv(level=2)   # 5 patients, limited resources — learn prioritisation
env = TriageEnv(level=3)   # 15 patients, mass casualty — full simulation

# Legacy named tasks still work
env = TriageEnv(task_name="easy_triage")

Multi-Component Reward

Every step() returns an info dict with individual reward signals:

python
result = env.step(action)
print(result.reward)                          # scalar for RL loop
print(result.info["reward_components"])        # dict for logging
# {
#   "triage_accuracy": +0.50,
#   "wait_time_penalty": -0.05,
#   "resource_util": +0.30,
#   "critical_penalty": 0.00,
#   "invalid_action_penalty": 0.00
# }
print(result.info["reward_explanation"])       # human-readable

Action Interface

The model should output JSON in this format:

json
{"action": "assign_triage", "patient_id": "P001", "priority": "CRITICAL"}
{"action": "admit_to_bed",  "patient_id": "P001", "resource_id": "icu_bed_1"}
{"action": "discharge",     "patient_id": "P003"}
{"action": "wait"}

The parser handles markdown fences, integer patient_id, and alias names (HIGH → urgent, etc.).

GRPO Training

bash
# Start from level 1, auto-advance to level 2 and 3 as agent improves
python train_grpo.py \
  --model unsloth/Qwen2.5-7B-Instruct \
  --level 1 \
  --total-steps 300 \
  --episodes-per-batch 4 \
  --output-dir ./triage_model

Anti-Reward-Hacking Safeguards

SafeguardImplementation
Critical patient discharge blockedstep() returns error + penalty if agent tries to discharge a CRITICAL patient who isn't admitted
Idle loop penaltyConsecutive WAIT actions accumulate increasing negative reward
Invalid action penaltyAny malformed or impossible action gives -0.15
Max step limitEpisode terminates at task_config["max_steps"]
Critical patient penaltyHard -0.3 × n_critical penalty for ignoring IMMEDIATE patients while idling or discharging

API Endpoints

EndpointMethodDescription
/resetPOSTStart new episode (level=1/2/3 or task_name=)
/stepPOSTTake action, receive reward breakdown
/stateGETFull internal state
/gradeGETFinal score [0, 1]
/reward_schemaGETReward component descriptions
/healthGETLiveness probe

Demo Mode

python
from train_grpo import run_demo

# Before training (random actions)
run_demo(model=None, level=1)

# After training
run_demo(model=trained_model, tokenizer=tokenizer, level=2)