README.md

---
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](https://github.com/meta-pytorch/OpenEnv)-compliant reinforcement learning environment where AI agents act as emergency physicians.

[![OpenEnv](https://img.shields.io/badge/OpenEnv-compatible-blue)](https://github.com/meta-pytorch/OpenEnv)
[![Python 3.10+](https://img.shields.io/badge/python-3.10+-brightgreen)](https://python.org)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow)](LICENSE)

---

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

```python
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

```python
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

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

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

### Run Server Locally

```bash
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

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

### Python Client Usage

```python
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

```python
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

```bash
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.