MedTriage / README.md
VP-TT
Initial commit: MedTriage OpenEnv environment
8df47e6
|
Raw
History Blame Contribute Delete
9.91 kB
metadata
title: MedTriage OpenEnv
emoji: πŸ₯
colorFrom: blue
colorTo: green
sdk: docker
app_port: 7860
tags:
  - openenv
  - medical
  - emergency-medicine
  - reinforcement-learning
  - clinical-decision-support
pinned: false
license: mit

πŸ₯ MedTriage OpenEnv

Emergency Department Triage & Clinical Decision Support
An OpenEnv-compliant reinforcement learning environment where AI agents act as emergency physicians.

OpenEnv Python 3.10+ License: MIT


🌍 Environment Description & Motivation

Emergency departments worldwide face critical challenges: physician shortages, overcrowding, and delayed triage lead to preventable deaths. AI agents capable of clinical reasoning could dramatically improve throughput and safety.

MedTriage-Env trains agents to master three core emergency medicine skills, each modelled on real clinical workflows used by certified emergency physicians:

  1. Vital Sign Triage β€” classify patient urgency using the internationally-standardized Emergency Severity Index (ESI v4)
  2. Differential Diagnosis β€” generate ranked differential diagnoses and identify life-threatening red flags from clinical presentations and lab data
  3. Treatment Safety β€” propose evidence-based treatment plans while avoiding drug allergies, contraindications, and nephrotoxic agents

All patient cases are synthetic but clinically realistic, designed by emergency medicine scenarios. All graders are fully deterministic β€” using rule-based algorithms, clinical reference ranges, and drug safety lookup tables.


🎯 Action Space

class MedTriageAction(BaseModel):
    # Task 1 β€” vital-triage
    esi_level: Optional[int]       # 1 (immediate) β†’ 5 (non-urgent)
    triage_reason: Optional[str]   # Brief clinical justification

    # Task 2 β€” differential-diagnosis
    diagnoses: Optional[List[str]]         # Ranked list (most likely first)
    red_flags: Optional[List[str]]         # Critical findings identified
    recommended_tests: Optional[List[str]] # Additional tests to order

    # Task 3 β€” treatment-safety
    diagnosis: Optional[str]        # Working diagnosis
    drug_name: Optional[str]        # Chosen drug or intervention
    dose_mg: Optional[float]        # Dose in mg (None for non-drug interventions)
    route: Optional[str]            # "IV" | "PO" | "IM" | "Inhaled" | "Other"
    rationale: Optional[str]        # Clinical reasoning including safety checks

πŸ‘οΈ Observation Space

class MedTriageObservation(BaseModel):
    patient: PatientInfo            # Full patient data (age, sex, CC, history, vitals, labs)
    current_task: str               # "vital-triage" | "differential-diagnosis" | "treatment-safety"
    task_instruction: str           # Natural-language instruction
    step_number: int
    max_steps: int                  # 30
    last_action_result: str         # Human-readable feedback
    last_action_error: Optional[str]
    cumulative_reward: float        # Running total
    progress: float                 # 0.0–1.0
    done: bool
    context: Dict[str, Any]        # Task-specific context (e.g., confirmed diagnosis for Task 3)

PatientInfo contains:

  • Demographics (age, sex)
  • Chief complaint & medical history
  • Vital signs (HR, BP, RR, Temp, SpO2, GCS, pain score)
  • Lab results with reference ranges and critical flags
  • Allergies & current medications
  • Known conditions

πŸ“‹ Tasks

Task 1: vital-triage β€” Easy

Property Value
Difficulty Easy
Max Steps 30
Score Range 0.0 – 1.0

Objective: Classify patient urgency using Emergency Severity Index (ESI v4):

  • ESI 1 (Immediate): Life-threatening β€” airway compromise, pulselessness, shock
  • ESI 2 (Emergent): High-risk, confused/lethargic, or severe pain
  • ESI 3 (Urgent): Stable but requires multiple diagnostic resources
  • ESI 4 (Less Urgent): Stable, requires one resource
  • ESI 5 (Non-Urgent): No resources needed

Grading:

  • Exact ESI match β†’ 1.0
  • Off by 1 β†’ 0.5
  • Off by 2 β†’ 0.2
  • Off by 3+ β†’ 0.0
  • Bonus +0.1 for mentioning key clinical features in rationale (capped at 1.0)

Task 2: differential-diagnosis β€” Medium

Property Value
Difficulty Medium
Max Steps 30
Score Range 0.0 – 1.0

Objective: Given symptoms, vitals, and lab results, generate:

  • Ranked differential diagnoses (most likely first, up to 3)
  • Critical red flags in the presentation
  • Recommended additional diagnostic tests

Grading (weighted score):

  • Primary diagnosis correct (top of list matches expected): 50%
  • Differential overlap (recall over top-3 correct diagnoses): 25%
  • Red flags identified (recall over ground-truth flags): 25%

Task 3: treatment-safety β€” Hard

Property Value
Difficulty Hard
Max Steps 30
Score Range 0.0 – 1.0

Objective: Given a confirmed working diagnosis, propose a safe, evidence-based treatment plan. Must avoid:

  • Drug allergies documented in patient record
  • Condition-specific contraindications (e.g., beta-blockers in asthma)
  • Nephrotoxic agents in CKD patients
  • Drugs known to cause harm in specific diagnoses

Grading:

  • Contraindicated drug selected β†’ 0.0 immediately
  • Correct drug: base 0.6
    • Dose within therapeutic range: +0.2
    • Correct route: +0.1
    • Rationale contains relevant clinical keywords: +0.1
  • Unknown drug (not contraindicated, not in list): 0.15 (partial)

πŸ”§ Setup & Usage

Prerequisites

  • Python 3.10+
  • Docker (for containerized deployment)
  • OpenAI-compatible API key (for inference)

Local Installation

git clone <your-repo-url>
cd medtriage-env

pip install -e .
# or
pip install -e ".[dev]"

Run Server Locally

uvicorn medtriage_env.server.app:app --host 0.0.0.0 --port 7860

Then open http://localhost:7860/docs for the interactive API docs.

Docker Build & Run

docker build -t medtriage-env .
docker run -p 7860:7860 medtriage-env

Python Client Usage

import asyncio
from medtriage_env import MedTriageEnv, MedTriageAction

async def main():
    async with MedTriageEnv(base_url="http://localhost:7860", task="vital-triage") as env:
        result = await env.reset()
        print(result.observation.patient.chief_complaint)

        action = MedTriageAction(esi_level=2, triage_reason="Hypotension and hypoxia")
        result = await env.step(action)
        print(f"Reward: {result.reward}, Done: {result.done}")

asyncio.run(main())

Sync Usage

from medtriage_env import MedTriageEnv, MedTriageAction

with MedTriageEnv(base_url="http://localhost:7860").sync() as env:
    result = env.reset(task="differential-diagnosis")
    action = MedTriageAction(
        diagnoses=["STEMI", "Aortic dissection", "NSTEMI"],
        red_flags=["ST elevation", "hypotension", "diaphoresis"],
        recommended_tests=["12-lead ECG", "Troponin", "CXR"]
    )
    result = env.step(action)
    print(f"Score: {result.reward:.2f}")

πŸš€ Running Inference

export API_BASE_URL="https://router.huggingface.co/v1"
export MODEL_NAME="Qwen/Qwen2.5-72B-Instruct"
export HF_TOKEN="hf_your_token_here"
export MEDTRIAGE_URL="https://your-space.hf.space"

python inference.py

πŸ“Š Baseline Scores

Scores obtained with Qwen/Qwen2.5-72B-Instruct via HuggingFace Inference Router:

Task Score Notes
vital-triage ~0.75 ESI classification is straightforward with good clinical context
differential-diagnosis ~0.60 Primary diagnosis often correct; red flag recall varies
treatment-safety ~0.55 Drug safety requires careful allergy/contraindication checking
Average ~0.63

Scores are reproducible across runs given the same model and fixed temperature=0.2.


πŸ—οΈ API Endpoints

Method Endpoint Description
GET /health Health check
GET /tasks List available tasks with metadata
POST /reset Start new episode ({"task": "vital-triage"})
POST /step Execute an action
GET /state Get current episode state
GET /docs Interactive Swagger UI

πŸ“ Project Structure

medtriage-env/
β”œβ”€β”€ inference.py                    # Baseline inference script (root)
β”œβ”€β”€ Dockerfile                      # Production container
β”œβ”€β”€ openenv.yaml                    # OpenEnv manifest
β”œβ”€β”€ pyproject.toml                  # Package config
β”œβ”€β”€ README.md                       # This file
β”œβ”€β”€ .dockerignore
└── medtriage_env/
    β”œβ”€β”€ __init__.py                 # Public API exports
    β”œβ”€β”€ models.py                   # Pydantic models (Action, Obs, State)
    β”œβ”€β”€ client.py                   # HTTP client (async + sync)
    └── server/
        β”œβ”€β”€ app.py                  # FastAPI application
        β”œβ”€β”€ medical_environment.py  # Core environment logic
        β”œβ”€β”€ graders.py              # Deterministic graders
        β”œβ”€β”€ patient_data.py         # Synthetic patient dataset
        └── requirements.txt

⚠️ Disclaimer

All patient cases are entirely fictional and intended solely for AI training and evaluation purposes. This environment does not constitute medical advice and should not be used for actual clinical decision-making.