Spaces:
Runtime error
Runtime error
Submission-ready: README, blog, training pipeline, baseline evidence, OpenEnv compliance
Browse files- ER_MAP/dashboard.py +57 -3
- ER_MAP/envs/__init__.py +15 -1
- ER_MAP/envs/openenv_triage/Dockerfile +51 -0
- ER_MAP/envs/openenv_triage/README.md +145 -0
- ER_MAP/envs/openenv_triage/__init__.py +52 -0
- ER_MAP/envs/openenv_triage/client.py +201 -0
- ER_MAP/envs/openenv_triage/env.py +295 -0
- ER_MAP/envs/openenv_triage/models.py +279 -0
- ER_MAP/envs/openenv_triage/openenv.yaml +44 -0
- ER_MAP/envs/openenv_triage/server-requirements.txt +12 -0
- ER_MAP/envs/openenv_triage/server.py +312 -0
- ER_MAP/envs/openenv_triage/tests/__init__.py +0 -0
- ER_MAP/envs/openenv_triage/tests/test_parity.py +306 -0
- ER_MAP/requirements.txt +4 -1
- ER_MAP/server.py +6 -0
- EVALUATION.md +218 -0
- LICENSE +21 -0
- Makefile +19 -0
- README.md +174 -231
- blog.md +312 -0
- kaggle/requirements_kaggle.txt +5 -6
- mermaid_render.html +102 -0
- openenv.yaml +36 -0
- requirements.txt +37 -0
- setup.py +42 -0
- training_per2.png +3 -0
- training_perf3.png +3 -0
- training_performance1.png +3 -0
ER_MAP/dashboard.py
CHANGED
|
@@ -616,6 +616,7 @@ def new_episode():
|
|
| 616 |
|
| 617 |
# NOTE: ground_truth is intentionally NOT returned to the browser —
|
| 618 |
# disease, emergency status, and personas live on the CMD terminal only.
|
|
|
|
| 619 |
return jsonify({
|
| 620 |
"status": "ok",
|
| 621 |
"conversation": EPISODE_STATE["conversation"],
|
|
@@ -626,6 +627,11 @@ def new_episode():
|
|
| 626 |
"phase_order": PHASE_ORDER,
|
| 627 |
"difficulty": options.get("difficulty") or "random",
|
| 628 |
"phase": phase,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 629 |
})
|
| 630 |
|
| 631 |
|
|
@@ -974,6 +980,7 @@ HTML_PAGE = r"""<!DOCTYPE html>
|
|
| 974 |
const [phasesDone, setPhasesDone] = useState([]);
|
| 975 |
const [currentPhase, setCurrentPhase] = useState(null);
|
| 976 |
const [outcome, setOutcome] = useState(null);
|
|
|
|
| 977 |
|
| 978 |
// refs (audio + loop control — never trigger re-render)
|
| 979 |
const audioQueueRef = useRef([]);
|
|
@@ -1065,6 +1072,7 @@ HTML_PAGE = r"""<!DOCTYPE html>
|
|
| 1065 |
body: JSON.stringify({ phase }),
|
| 1066 |
});
|
| 1067 |
const data = await res.json();
|
|
|
|
| 1068 |
renderedCountRef.current = (data.conversation || []).length;
|
| 1069 |
} catch (e) {
|
| 1070 |
console.error('new_episode failed', e);
|
|
@@ -1160,7 +1168,7 @@ HTML_PAGE = r"""<!DOCTYPE html>
|
|
| 1160 |
</div>
|
| 1161 |
|
| 1162 |
{/* Open State */}
|
| 1163 |
-
<div className={`absolute inset-0 w-[350px] p-6 flex flex-col transition-opacity duration-300 ${isSidebarOpen ? 'opacity-100 delay-100' : 'opacity-0 pointer-events-none'}`}>
|
| 1164 |
{/* Header + close */}
|
| 1165 |
<div className="flex items-center justify-between mb-8 cursor-pointer group"
|
| 1166 |
onClick={(e) => { e.stopPropagation(); setIsSidebarOpen(false); }}>
|
|
@@ -1201,7 +1209,7 @@ HTML_PAGE = r"""<!DOCTYPE html>
|
|
| 1201 |
</div>
|
| 1202 |
|
| 1203 |
{/* Section B: Clinical Phases */}
|
| 1204 |
-
<div className="
|
| 1205 |
<h3 className="text-xs font-semibold text-slate-300 uppercase tracking-wider mb-4">Clinical Phases</h3>
|
| 1206 |
<div className="space-y-3.5">
|
| 1207 |
{PHASES.map(({ key, label }, i) => {
|
|
@@ -1224,6 +1232,52 @@ HTML_PAGE = r"""<!DOCTYPE html>
|
|
| 1224 |
})}
|
| 1225 |
</div>
|
| 1226 |
</div>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1227 |
</div>
|
| 1228 |
</div>
|
| 1229 |
|
|
@@ -1245,7 +1299,7 @@ HTML_PAGE = r"""<!DOCTYPE html>
|
|
| 1245 |
<button
|
| 1246 |
onClick={startCase}
|
| 1247 |
className="bg-indigo-600 hover:bg-indigo-500 text-white text-[11px] font-semibold tracking-wider uppercase px-3 py-1.5 rounded-lg transition flex items-center gap-1.5 shadow-[0_4px_12px_rgba(99,102,241,0.3)]">
|
| 1248 |
-
|
| 1249 |
</button>
|
| 1250 |
</>
|
| 1251 |
) : (
|
|
|
|
| 616 |
|
| 617 |
# NOTE: ground_truth is intentionally NOT returned to the browser —
|
| 618 |
# disease, emergency status, and personas live on the CMD terminal only.
|
| 619 |
+
# However, we now surface sanitized behavior traits for the 'God's View'.
|
| 620 |
return jsonify({
|
| 621 |
"status": "ok",
|
| 622 |
"conversation": EPISODE_STATE["conversation"],
|
|
|
|
| 627 |
"phase_order": PHASE_ORDER,
|
| 628 |
"difficulty": options.get("difficulty") or "random",
|
| 629 |
"phase": phase,
|
| 630 |
+
"persona": {
|
| 631 |
+
"patient": gt.get("patient", {}),
|
| 632 |
+
"nurse": gt.get("nurse", {}),
|
| 633 |
+
"is_emergency": gt.get("disease", {}).get("is_emergency", False),
|
| 634 |
+
}
|
| 635 |
})
|
| 636 |
|
| 637 |
|
|
|
|
| 980 |
const [phasesDone, setPhasesDone] = useState([]);
|
| 981 |
const [currentPhase, setCurrentPhase] = useState(null);
|
| 982 |
const [outcome, setOutcome] = useState(null);
|
| 983 |
+
const [persona, setPersona] = useState(null);
|
| 984 |
|
| 985 |
// refs (audio + loop control — never trigger re-render)
|
| 986 |
const audioQueueRef = useRef([]);
|
|
|
|
| 1072 |
body: JSON.stringify({ phase }),
|
| 1073 |
});
|
| 1074 |
const data = await res.json();
|
| 1075 |
+
setPersona(data.persona);
|
| 1076 |
renderedCountRef.current = (data.conversation || []).length;
|
| 1077 |
} catch (e) {
|
| 1078 |
console.error('new_episode failed', e);
|
|
|
|
| 1168 |
</div>
|
| 1169 |
|
| 1170 |
{/* Open State */}
|
| 1171 |
+
<div className={`absolute inset-0 w-[350px] p-6 flex flex-col transition-opacity duration-300 overflow-y-auto custom-scrollbar ${isSidebarOpen ? 'opacity-100 delay-100' : 'opacity-0 pointer-events-none'}`}>
|
| 1172 |
{/* Header + close */}
|
| 1173 |
<div className="flex items-center justify-between mb-8 cursor-pointer group"
|
| 1174 |
onClick={(e) => { e.stopPropagation(); setIsSidebarOpen(false); }}>
|
|
|
|
| 1209 |
</div>
|
| 1210 |
|
| 1211 |
{/* Section B: Clinical Phases */}
|
| 1212 |
+
<div className="mt-8 mb-8">
|
| 1213 |
<h3 className="text-xs font-semibold text-slate-300 uppercase tracking-wider mb-4">Clinical Phases</h3>
|
| 1214 |
<div className="space-y-3.5">
|
| 1215 |
{PHASES.map(({ key, label }, i) => {
|
|
|
|
| 1232 |
})}
|
| 1233 |
</div>
|
| 1234 |
</div>
|
| 1235 |
+
|
| 1236 |
+
{/* Section C: Agent Personas */}
|
| 1237 |
+
{persona && (
|
| 1238 |
+
<div className="border-t border-slate-800/50 pt-6 pb-4">
|
| 1239 |
+
<div className="flex items-center justify-between mb-4">
|
| 1240 |
+
<h3 className="text-xs font-semibold text-slate-300 uppercase tracking-wider">Agent Personas</h3>
|
| 1241 |
+
{persona.is_emergency ? (
|
| 1242 |
+
<span className="text-[9px] font-bold uppercase tracking-wider px-2 py-0.5 rounded-full bg-red-500/20 text-red-400 border border-red-500/30 animate-pulse shadow-[0_0_8px_rgba(248,113,113,0.3)]">🚨 EMERGENCY</span>
|
| 1243 |
+
) : (
|
| 1244 |
+
<span className="text-[9px] font-bold uppercase tracking-wider px-2 py-0.5 rounded-full bg-emerald-500/20 text-emerald-400 border border-emerald-500/30">Non-Emergency</span>
|
| 1245 |
+
)}
|
| 1246 |
+
</div>
|
| 1247 |
+
<div className="space-y-6">
|
| 1248 |
+
{/* Patient */}
|
| 1249 |
+
<div>
|
| 1250 |
+
<div className="flex items-center gap-2 mb-2">
|
| 1251 |
+
<User className="w-3 h-3 text-teal-400" />
|
| 1252 |
+
<span className="text-[10px] font-bold text-teal-400/80 uppercase tracking-tight">Patient Traits</span>
|
| 1253 |
+
</div>
|
| 1254 |
+
<div className="grid grid-cols-2 gap-2">
|
| 1255 |
+
{Object.entries(persona.patient).map(([k, v]) => (
|
| 1256 |
+
<div key={k} className="bg-slate-950/40 border border-slate-800/30 rounded-lg p-2">
|
| 1257 |
+
<div className="text-[9px] text-slate-500 uppercase leading-none mb-1">{k}</div>
|
| 1258 |
+
<div className="text-[10px] text-slate-300 font-medium leading-tight truncate">{v.replace(/_/g, ' ')}</div>
|
| 1259 |
+
</div>
|
| 1260 |
+
))}
|
| 1261 |
+
</div>
|
| 1262 |
+
</div>
|
| 1263 |
+
{/* Nurse */}
|
| 1264 |
+
<div>
|
| 1265 |
+
<div className="flex items-center gap-2 mb-2">
|
| 1266 |
+
<Activity className="w-3 h-3 text-blue-400" />
|
| 1267 |
+
<span className="text-[10px] font-bold text-blue-400/80 uppercase tracking-tight">Nurse Traits</span>
|
| 1268 |
+
</div>
|
| 1269 |
+
<div className="grid grid-cols-2 gap-2">
|
| 1270 |
+
{Object.entries(persona.nurse).map(([k, v]) => (
|
| 1271 |
+
<div key={k} className="bg-slate-950/40 border border-slate-800/30 rounded-lg p-2">
|
| 1272 |
+
<div className="text-[9px] text-slate-500 uppercase leading-none mb-1">{k}</div>
|
| 1273 |
+
<div className="text-[10px] text-slate-300 font-medium leading-tight truncate">{v.replace(/_/g, ' ')}</div>
|
| 1274 |
+
</div>
|
| 1275 |
+
))}
|
| 1276 |
+
</div>
|
| 1277 |
+
</div>
|
| 1278 |
+
</div>
|
| 1279 |
+
</div>
|
| 1280 |
+
)}
|
| 1281 |
</div>
|
| 1282 |
</div>
|
| 1283 |
|
|
|
|
| 1299 |
<button
|
| 1300 |
onClick={startCase}
|
| 1301 |
className="bg-indigo-600 hover:bg-indigo-500 text-white text-[11px] font-semibold tracking-wider uppercase px-3 py-1.5 rounded-lg transition flex items-center gap-1.5 shadow-[0_4px_12px_rgba(99,102,241,0.3)]">
|
| 1302 |
+
Start Case
|
| 1303 |
</button>
|
| 1304 |
</>
|
| 1305 |
) : (
|
ER_MAP/envs/__init__.py
CHANGED
|
@@ -5,4 +5,18 @@ from .triage_env import TriageEnv
|
|
| 5 |
from .randomizer import generate_ground_truth, construct_prompts
|
| 6 |
from .api_router import AgentRouter
|
| 7 |
|
| 8 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 5 |
from .randomizer import generate_ground_truth, construct_prompts
|
| 6 |
from .api_router import AgentRouter
|
| 7 |
|
| 8 |
+
# OpenEnv-compliant wrapper. Imported via try/except so users without
|
| 9 |
+
# the optional ``openenv-core`` dependency (e.g. the running Kaggle
|
| 10 |
+
# training job) keep working unchanged.
|
| 11 |
+
try: # pragma: no cover - import-time soft fail
|
| 12 |
+
from .openenv_triage import TriageOpenEnv # noqa: F401
|
| 13 |
+
except ImportError:
|
| 14 |
+
TriageOpenEnv = None # type: ignore[assignment]
|
| 15 |
+
|
| 16 |
+
__all__ = [
|
| 17 |
+
"TriageEnv",
|
| 18 |
+
"TriageOpenEnv",
|
| 19 |
+
"generate_ground_truth",
|
| 20 |
+
"construct_prompts",
|
| 21 |
+
"AgentRouter",
|
| 22 |
+
]
|
ER_MAP/envs/openenv_triage/Dockerfile
ADDED
|
@@ -0,0 +1,51 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Dockerfile for ER-MAP TriageOpenEnv
|
| 2 |
+
#
|
| 3 |
+
# Pattern follows the recommended HF Spaces / OpenEnv standalone build:
|
| 4 |
+
# - slim Python base
|
| 5 |
+
# - install openenv-core + project deps from requirements.txt
|
| 6 |
+
# - copy the ER_MAP package (so ``ER_MAP.envs.openenv_triage.server:app``
|
| 7 |
+
# resolves at runtime)
|
| 8 |
+
# - expose 8000 and run uvicorn
|
| 9 |
+
#
|
| 10 |
+
# Reference Dockerfile (echo_env, multi-stage uv build):
|
| 11 |
+
# https://github.com/meta-pytorch/OpenEnv/blob/main/envs/echo_env/server/Dockerfile
|
| 12 |
+
#
|
| 13 |
+
# We use a single-stage slim image here because the project's
|
| 14 |
+
# requirements.txt is small and the multi-stage uv flow adds complexity
|
| 15 |
+
# that isn't justified for this environment. The HF Space build keeps
|
| 16 |
+
# layer caching by separating the requirements install step.
|
| 17 |
+
|
| 18 |
+
FROM python:3.11-slim AS runtime
|
| 19 |
+
|
| 20 |
+
ENV PYTHONDONTWRITEBYTECODE=1 \
|
| 21 |
+
PYTHONUNBUFFERED=1 \
|
| 22 |
+
PIP_DISABLE_PIP_VERSION_CHECK=1 \
|
| 23 |
+
PIP_NO_CACHE_DIR=1 \
|
| 24 |
+
PORT=8000 \
|
| 25 |
+
HOST=0.0.0.0
|
| 26 |
+
|
| 27 |
+
RUN apt-get update \
|
| 28 |
+
&& apt-get install -y --no-install-recommends curl \
|
| 29 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 30 |
+
|
| 31 |
+
WORKDIR /app
|
| 32 |
+
|
| 33 |
+
# --- Install Python deps (cached when requirements.txt is unchanged) ---
|
| 34 |
+
COPY ER_MAP/envs/openenv_triage/server-requirements.txt /tmp/server-requirements.txt
|
| 35 |
+
RUN pip install --upgrade pip \
|
| 36 |
+
&& pip install -r /tmp/server-requirements.txt
|
| 37 |
+
|
| 38 |
+
# --- Copy the ER_MAP package ---
|
| 39 |
+
# We keep the full ER_MAP package importable so the wrapper can reach
|
| 40 |
+
# ``ER_MAP.envs.triage_env`` and supporting modules.
|
| 41 |
+
COPY ER_MAP /app/ER_MAP
|
| 42 |
+
|
| 43 |
+
ENV PYTHONPATH=/app
|
| 44 |
+
|
| 45 |
+
EXPOSE 8000
|
| 46 |
+
|
| 47 |
+
# Health check uses the OpenEnv-provided /health route.
|
| 48 |
+
HEALTHCHECK --interval=30s --timeout=5s --start-period=10s --retries=3 \
|
| 49 |
+
CMD curl -f http://localhost:8000/health || exit 1
|
| 50 |
+
|
| 51 |
+
CMD ["uvicorn", "ER_MAP.envs.openenv_triage.server:app", "--host", "0.0.0.0", "--port", "8000"]
|
ER_MAP/envs/openenv_triage/README.md
ADDED
|
@@ -0,0 +1,145 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ER-MAP Triage - OpenEnv Wrapper
|
| 2 |
+
|
| 3 |
+
This package wraps the existing in-house Gymnasium `TriageEnv`
|
| 4 |
+
(`ER_MAP/envs/triage_env.py`) so it can be served through the
|
| 5 |
+
[OpenEnv](https://github.com/meta-pytorch/OpenEnv) protocol
|
| 6 |
+
(`openenv-core>=0.2.3`, latest as of April 2026).
|
| 7 |
+
|
| 8 |
+
The wrapper does **not** modify the underlying environment - the same
|
| 9 |
+
LoRA currently being trained against the gym env on Kaggle runs unchanged
|
| 10 |
+
behind this OpenEnv server because both share the exact same `TriageEnv`
|
| 11 |
+
instance, reward semantics, and observation format.
|
| 12 |
+
|
| 13 |
+
## What's in here
|
| 14 |
+
|
| 15 |
+
| File | Purpose |
|
| 16 |
+
|----------------------------|---------------------------------------------------------------------------|
|
| 17 |
+
| `models.py` | `TriageAction`, `TriageObservation`, `TriageState` (Pydantic v2) |
|
| 18 |
+
| `env.py` | `TriageOpenEnv` - `Environment` subclass delegating to `TriageEnv` |
|
| 19 |
+
| `server.py` | FastAPI app exposed via OpenEnv's `create_app` helper |
|
| 20 |
+
| `client.py` | Thin `requests`-based `TriageOpenEnvClient` |
|
| 21 |
+
| `openenv.yaml` | OpenEnv manifest (HF Space deployment target) |
|
| 22 |
+
| `Dockerfile` | HF Space / Docker image (Python 3.11-slim) |
|
| 23 |
+
| `server-requirements.txt` | Container-side pip requirements |
|
| 24 |
+
| `tests/test_parity.py` | pytest parity tests (gym vs wrapper, in-proc vs HTTP) |
|
| 25 |
+
|
| 26 |
+
## Usage
|
| 27 |
+
|
| 28 |
+
### 1. Local Python (direct, no HTTP)
|
| 29 |
+
|
| 30 |
+
```python
|
| 31 |
+
from ER_MAP.envs.openenv_triage import TriageOpenEnv, TriageAction
|
| 32 |
+
|
| 33 |
+
env = TriageOpenEnv()
|
| 34 |
+
obs = env.reset(seed=0, options={"phase": 1, "difficulty": "easy"})
|
| 35 |
+
print(obs.event, obs.payload)
|
| 36 |
+
|
| 37 |
+
action = TriageAction.from_json_str(
|
| 38 |
+
'{"tool": "speak_to", "target": "patient", "message": "Hello"}'
|
| 39 |
+
)
|
| 40 |
+
obs = env.step(action)
|
| 41 |
+
print(obs.reward, obs.done, obs.info["reward_components"])
|
| 42 |
+
```
|
| 43 |
+
|
| 44 |
+
### 2. Local HTTP (FastAPI + client)
|
| 45 |
+
|
| 46 |
+
Start the server:
|
| 47 |
+
|
| 48 |
+
```bash
|
| 49 |
+
uvicorn ER_MAP.envs.openenv_triage.server:app --host 0.0.0.0 --port 8000
|
| 50 |
+
```
|
| 51 |
+
|
| 52 |
+
Drive it with the bundled client:
|
| 53 |
+
|
| 54 |
+
```python
|
| 55 |
+
from ER_MAP.envs.openenv_triage import TriageOpenEnvClient, TriageAction
|
| 56 |
+
|
| 57 |
+
with TriageOpenEnvClient(base_url="http://localhost:8000") as client:
|
| 58 |
+
print(client.health())
|
| 59 |
+
result = client.reset(seed=0, options={"phase": 1, "difficulty": "easy"})
|
| 60 |
+
action = TriageAction.from_json_str('{"tool": "read_soap"}')
|
| 61 |
+
result = client.step(action)
|
| 62 |
+
print(result.reward, result.observation.event)
|
| 63 |
+
```
|
| 64 |
+
|
| 65 |
+
The HTTP wire format matches OpenEnv 0.2.3:
|
| 66 |
+
|
| 67 |
+
- `POST /reset` body `{"seed": int?, "episode_id": str?, "options": {...}?}`
|
| 68 |
+
- `POST /step` body `{"action": {"tool": "...", ...}, "timeout_s": float?}`
|
| 69 |
+
- `GET /state` -> serialized `TriageState`
|
| 70 |
+
- `GET /health` -> `{"status": "healthy"}`
|
| 71 |
+
- `GET /healthz` -> richer status (version + stub-mode flag)
|
| 72 |
+
|
| 73 |
+
OpenEnv-style WebSocket clients can also connect to `WS /ws` (provided by
|
| 74 |
+
`HTTPEnvServer.register_routes`).
|
| 75 |
+
|
| 76 |
+
### 3. Hugging Face Space deployment
|
| 77 |
+
|
| 78 |
+
The repo follows the OpenEnv manifest convention; a single command pushes
|
| 79 |
+
to a Space:
|
| 80 |
+
|
| 81 |
+
```bash
|
| 82 |
+
# from the repo root
|
| 83 |
+
openenv push --repo-id <hf-username>/er-map-triage
|
| 84 |
+
```
|
| 85 |
+
|
| 86 |
+
(See `openenv push --help` from `openenv-core>=0.2.3` for options like
|
| 87 |
+
`--private` and `--tag`.)
|
| 88 |
+
|
| 89 |
+
To deploy manually with Docker:
|
| 90 |
+
|
| 91 |
+
```bash
|
| 92 |
+
docker build -t er-map-triage:latest -f ER_MAP/envs/openenv_triage/Dockerfile .
|
| 93 |
+
docker run --rm -p 8000:8000 \
|
| 94 |
+
-e GROQ_API_KEY=$GROQ_API_KEY \
|
| 95 |
+
er-map-triage:latest
|
| 96 |
+
```
|
| 97 |
+
|
| 98 |
+
## Stub mode (keyless deployment)
|
| 99 |
+
|
| 100 |
+
The underlying `AgentRouter` already gracefully handles the no-Groq case:
|
| 101 |
+
when no key is configured, it returns deterministic canned Nurse/Patient
|
| 102 |
+
JSON responses (see `AgentRouter._mock_response`) and the LLM judges
|
| 103 |
+
return neutral default scores. This is what we call "stub mode".
|
| 104 |
+
|
| 105 |
+
- **Stub mode** is automatic: leave all `GROQ_*` vars unset.
|
| 106 |
+
- **Live mode**: set `GROQ_API_KEY` (and optionally per-role keys
|
| 107 |
+
`GROQ_NURSE_API_KEY`, `GROQ_PATIENT_API_KEY`,
|
| 108 |
+
`GROQ_EMPATHY_JUDGE_API_KEY`, `GROQ_MEDICAL_JUDGE_API_KEY`).
|
| 109 |
+
|
| 110 |
+
The `/healthz` endpoint reports `{"stub_mode": true|false}` so HF Space
|
| 111 |
+
visitors can tell which configuration the running container is in.
|
| 112 |
+
|
| 113 |
+
## Parity guarantees
|
| 114 |
+
|
| 115 |
+
`TriageOpenEnv.step` produces **byte-identical** observations and rewards
|
| 116 |
+
to the gym `TriageEnv.step` for the same seed and action sequence, because
|
| 117 |
+
the wrapper:
|
| 118 |
+
|
| 119 |
+
1. Replays the original Doctor JSON string verbatim into the gym env
|
| 120 |
+
(via `TriageAction.raw_json`), so `_parse_doctor_action` sees exactly
|
| 121 |
+
the bytes the policy emitted.
|
| 122 |
+
2. Forwards the gym 5-tuple `(obs, reward, done, truncated, info)`
|
| 123 |
+
unchanged - `done` and `reward` go on the OpenEnv `Observation`,
|
| 124 |
+
`truncated` and the full `info` dict (including `reward_components`)
|
| 125 |
+
are folded into `Observation.metadata` / `Observation.info`.
|
| 126 |
+
|
| 127 |
+
This is verified by `tests/test_parity.py`.
|
| 128 |
+
|
| 129 |
+
## Reserved tool names (OpenEnv brief)
|
| 130 |
+
|
| 131 |
+
Per the OpenEnv brief, MCP tool names `reset`, `step`, `state`, `close`
|
| 132 |
+
are reserved. This wrapper deliberately **does not** subclass
|
| 133 |
+
`MCPEnvironment` and does **not** register any MCP tools, so no name
|
| 134 |
+
collision is possible. The Doctor's tool names (`speak_to`, `order_lab`,
|
| 135 |
+
`read_soap`, `update_soap`, `terminal_discharge`) live inside the
|
| 136 |
+
`TriageAction.tool` field, not in any MCP namespace.
|
| 137 |
+
|
| 138 |
+
## Citations
|
| 139 |
+
|
| 140 |
+
- `openenv-core==0.2.3` PyPI release - <https://pypi.org/project/openenv-core/>
|
| 141 |
+
- `Environment` base class - <https://github.com/meta-pytorch/OpenEnv/blob/main/src/openenv/core/env_server/interfaces.py>
|
| 142 |
+
- `Action`/`Observation`/`State` types - <https://github.com/meta-pytorch/OpenEnv/blob/main/src/openenv/core/env_server/types.py>
|
| 143 |
+
- `create_app` helper - <https://github.com/meta-pytorch/OpenEnv/blob/main/src/openenv/core/env_server/http_server.py>
|
| 144 |
+
- echo_env exemplar (`openenv.yaml`, `Dockerfile`, server layout) - <https://github.com/meta-pytorch/OpenEnv/tree/main/envs/echo_env>
|
| 145 |
+
- `RESERVED_TOOL_NAMES` - <https://github.com/meta-pytorch/OpenEnv/blob/main/src/openenv/core/env_server/mcp_types.py>
|
ER_MAP/envs/openenv_triage/__init__.py
ADDED
|
@@ -0,0 +1,52 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
ER_MAP/envs/openenv_triage
|
| 3 |
+
==========================
|
| 4 |
+
|
| 5 |
+
OpenEnv-compliant wrapper around the existing Gymnasium ``TriageEnv``.
|
| 6 |
+
|
| 7 |
+
This package adapts the in-process gym env to the OpenEnv interface
|
| 8 |
+
(``openenv-core>=0.2.3``) without modifying the underlying environment,
|
| 9 |
+
so the same LoRA trained against ``ER_MAP.envs.triage_env.TriageEnv``
|
| 10 |
+
runs unchanged against ``TriageOpenEnv``.
|
| 11 |
+
|
| 12 |
+
Public surface
|
| 13 |
+
--------------
|
| 14 |
+
- :class:`TriageOpenEnv` - Server-side OpenEnv ``Environment`` subclass.
|
| 15 |
+
- :class:`TriageAction` - Pydantic ``Action`` mirroring the Doctor's JSON schema.
|
| 16 |
+
- :class:`TriageObservation` - Pydantic ``Observation`` exposing the structured Doctor view.
|
| 17 |
+
- :class:`TriageState` - Pydantic ``State`` exposing the internal episode state.
|
| 18 |
+
- :class:`TriageOpenEnvClient` - Thin HTTP client (no server-internal imports).
|
| 19 |
+
|
| 20 |
+
References
|
| 21 |
+
----------
|
| 22 |
+
- OpenEnv canonical repo: https://github.com/meta-pytorch/OpenEnv
|
| 23 |
+
- ``Environment`` interface: ``openenv.core.env_server.interfaces.Environment``
|
| 24 |
+
- ``Action`` / ``Observation``: ``openenv.core.env_server.types``
|
| 25 |
+
- echo_env exemplar: https://github.com/meta-pytorch/OpenEnv/tree/main/envs/echo_env
|
| 26 |
+
"""
|
| 27 |
+
|
| 28 |
+
from __future__ import annotations
|
| 29 |
+
|
| 30 |
+
from .models import TriageAction, TriageObservation, TriageState
|
| 31 |
+
|
| 32 |
+
# ``env`` and ``client`` import openenv-core (a heavy dep). Guard those
|
| 33 |
+
# behind try/except so plain users of the gym TriageEnv (e.g. the running
|
| 34 |
+
# Kaggle training session) are unaffected if openenv-core is not installed.
|
| 35 |
+
try:
|
| 36 |
+
from .env import TriageOpenEnv # noqa: F401
|
| 37 |
+
except Exception: # pragma: no cover - import-time soft fail
|
| 38 |
+
TriageOpenEnv = None # type: ignore[assignment]
|
| 39 |
+
|
| 40 |
+
try:
|
| 41 |
+
from .client import TriageOpenEnvClient # noqa: F401
|
| 42 |
+
except Exception: # pragma: no cover
|
| 43 |
+
TriageOpenEnvClient = None # type: ignore[assignment]
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
__all__ = [
|
| 47 |
+
"TriageAction",
|
| 48 |
+
"TriageObservation",
|
| 49 |
+
"TriageState",
|
| 50 |
+
"TriageOpenEnv",
|
| 51 |
+
"TriageOpenEnvClient",
|
| 52 |
+
]
|
ER_MAP/envs/openenv_triage/client.py
ADDED
|
@@ -0,0 +1,201 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
ER_MAP/envs/openenv_triage/client.py
|
| 3 |
+
====================================
|
| 4 |
+
|
| 5 |
+
Thin HTTP client for the :class:`TriageOpenEnv` FastAPI server.
|
| 6 |
+
|
| 7 |
+
This client deliberately avoids importing any server-internal modules
|
| 8 |
+
(``server.py``, ``env.py``, etc.) so it can be packaged independently
|
| 9 |
+
(per the OpenEnv hackathon brief: "clients should never import server
|
| 10 |
+
internals").
|
| 11 |
+
|
| 12 |
+
Wire format matches OpenEnv 0.2.3's HTTP routes registered by
|
| 13 |
+
``HTTPEnvServer.register_routes`` -
|
| 14 |
+
https://github.com/meta-pytorch/OpenEnv/blob/main/src/openenv/core/env_server/http_server.py:
|
| 15 |
+
|
| 16 |
+
- ``POST /reset`` body: {seed?, episode_id?, options?, ...}
|
| 17 |
+
reply: {observation, reward, done}
|
| 18 |
+
- ``POST /step`` body: {action: <action_fields>, timeout_s?, ...}
|
| 19 |
+
reply: {observation, reward, done}
|
| 20 |
+
- ``GET /state`` reply: {observation: <state fields>, reward: null, done: false}
|
| 21 |
+
- ``GET /health`` reply: {status: "healthy"}
|
| 22 |
+
|
| 23 |
+
Returns
|
| 24 |
+
-------
|
| 25 |
+
``reset`` and ``step`` both return a :class:`StepResult` with the parsed
|
| 26 |
+
:class:`TriageObservation`. ``state`` returns a :class:`TriageState`.
|
| 27 |
+
"""
|
| 28 |
+
|
| 29 |
+
from __future__ import annotations
|
| 30 |
+
|
| 31 |
+
import logging
|
| 32 |
+
from dataclasses import dataclass
|
| 33 |
+
from typing import Any, Dict, Optional
|
| 34 |
+
|
| 35 |
+
import requests
|
| 36 |
+
|
| 37 |
+
# The client only depends on the public Pydantic models in this package.
|
| 38 |
+
# It does NOT import env.py / server.py.
|
| 39 |
+
from .models import TriageAction, TriageObservation, TriageState
|
| 40 |
+
|
| 41 |
+
logger = logging.getLogger("ER_MAP.openenv_triage.client")
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
@dataclass
|
| 45 |
+
class StepResult:
|
| 46 |
+
"""
|
| 47 |
+
Parsed envelope returned by ``reset`` and ``step``.
|
| 48 |
+
|
| 49 |
+
Mirrors OpenEnv's ``StepResult`` semantics (observation + reward +
|
| 50 |
+
done) without forcing the caller to depend on ``openenv-core`` at
|
| 51 |
+
runtime. Aliases of the underlying gym 5-tuple are also available
|
| 52 |
+
via :attr:`observation`'s ``truncated`` / ``info`` fields.
|
| 53 |
+
"""
|
| 54 |
+
|
| 55 |
+
observation: TriageObservation
|
| 56 |
+
reward: Optional[float]
|
| 57 |
+
done: bool
|
| 58 |
+
|
| 59 |
+
@property
|
| 60 |
+
def truncated(self) -> bool:
|
| 61 |
+
return bool(self.observation.truncated)
|
| 62 |
+
|
| 63 |
+
@property
|
| 64 |
+
def info(self) -> Dict[str, Any]:
|
| 65 |
+
return dict(self.observation.info)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
class TriageOpenEnvClient:
|
| 69 |
+
"""
|
| 70 |
+
Synchronous HTTP client for ER-MAP's OpenEnv-compliant Triage server.
|
| 71 |
+
|
| 72 |
+
Example
|
| 73 |
+
-------
|
| 74 |
+
>>> client = TriageOpenEnvClient(base_url="http://localhost:8000")
|
| 75 |
+
>>> client.health()
|
| 76 |
+
{'status': 'healthy', ...}
|
| 77 |
+
>>> result = client.reset(seed=0, options={"phase": 1, "difficulty": "easy"})
|
| 78 |
+
>>> action = TriageAction.from_json_str('{"tool": "read_soap"}')
|
| 79 |
+
>>> result = client.step(action)
|
| 80 |
+
>>> print(result.observation.event, result.reward)
|
| 81 |
+
"""
|
| 82 |
+
|
| 83 |
+
def __init__(
|
| 84 |
+
self,
|
| 85 |
+
base_url: str = "http://localhost:8000",
|
| 86 |
+
*,
|
| 87 |
+
timeout_s: float = 60.0,
|
| 88 |
+
session: Optional[requests.Session] = None,
|
| 89 |
+
) -> None:
|
| 90 |
+
self.base_url = base_url.rstrip("/")
|
| 91 |
+
self.timeout_s = timeout_s
|
| 92 |
+
self._session = session or requests.Session()
|
| 93 |
+
|
| 94 |
+
# ------------------------------------------------------------------
|
| 95 |
+
# Connection helpers
|
| 96 |
+
# ------------------------------------------------------------------
|
| 97 |
+
|
| 98 |
+
def _url(self, path: str) -> str:
|
| 99 |
+
return f"{self.base_url}{path}"
|
| 100 |
+
|
| 101 |
+
def health(self) -> Dict[str, Any]:
|
| 102 |
+
resp = self._session.get(self._url("/health"), timeout=self.timeout_s)
|
| 103 |
+
resp.raise_for_status()
|
| 104 |
+
return resp.json()
|
| 105 |
+
|
| 106 |
+
def healthz(self) -> Dict[str, Any]:
|
| 107 |
+
"""Richer health endpoint exposed by our FastAPI app."""
|
| 108 |
+
resp = self._session.get(self._url("/healthz"), timeout=self.timeout_s)
|
| 109 |
+
resp.raise_for_status()
|
| 110 |
+
return resp.json()
|
| 111 |
+
|
| 112 |
+
# ------------------------------------------------------------------
|
| 113 |
+
# OpenEnv core operations
|
| 114 |
+
# ------------------------------------------------------------------
|
| 115 |
+
|
| 116 |
+
def reset(
|
| 117 |
+
self,
|
| 118 |
+
*,
|
| 119 |
+
seed: Optional[int] = None,
|
| 120 |
+
episode_id: Optional[str] = None,
|
| 121 |
+
options: Optional[Dict[str, Any]] = None,
|
| 122 |
+
**extra: Any,
|
| 123 |
+
) -> StepResult:
|
| 124 |
+
body: Dict[str, Any] = {}
|
| 125 |
+
if seed is not None:
|
| 126 |
+
body["seed"] = seed
|
| 127 |
+
if episode_id is not None:
|
| 128 |
+
body["episode_id"] = episode_id
|
| 129 |
+
if options is not None:
|
| 130 |
+
body["options"] = options
|
| 131 |
+
body.update(extra)
|
| 132 |
+
|
| 133 |
+
resp = self._session.post(
|
| 134 |
+
self._url("/reset"), json=body, timeout=self.timeout_s
|
| 135 |
+
)
|
| 136 |
+
resp.raise_for_status()
|
| 137 |
+
return self._parse_step_response(resp.json())
|
| 138 |
+
|
| 139 |
+
def step(self, action: TriageAction, *, timeout_s: Optional[float] = None) -> StepResult:
|
| 140 |
+
# The server expects the action under "action" (per StepRequest).
|
| 141 |
+
# Use ``model_dump`` so Pydantic-typed fields are JSON-serialized
|
| 142 |
+
# exactly as the server's ``deserialize_action`` expects.
|
| 143 |
+
action_payload = action.model_dump(exclude_none=True, exclude={"metadata"})
|
| 144 |
+
body: Dict[str, Any] = {"action": action_payload}
|
| 145 |
+
if timeout_s is not None:
|
| 146 |
+
body["timeout_s"] = timeout_s
|
| 147 |
+
|
| 148 |
+
resp = self._session.post(
|
| 149 |
+
self._url("/step"), json=body, timeout=self.timeout_s
|
| 150 |
+
)
|
| 151 |
+
resp.raise_for_status()
|
| 152 |
+
return self._parse_step_response(resp.json())
|
| 153 |
+
|
| 154 |
+
def state(self) -> TriageState:
|
| 155 |
+
resp = self._session.get(self._url("/state"), timeout=self.timeout_s)
|
| 156 |
+
resp.raise_for_status()
|
| 157 |
+
data = resp.json()
|
| 158 |
+
# ``HTTPEnvServer`` wraps the state as
|
| 159 |
+
# ``{"observation": {...state fields...}, "reward": None, "done": ...}``
|
| 160 |
+
# via ``serialize_observation``-equivalent path. Some versions
|
| 161 |
+
# return the state directly; tolerate both.
|
| 162 |
+
if isinstance(data, dict) and "observation" in data and "reward" in data:
|
| 163 |
+
state_dict = data["observation"]
|
| 164 |
+
else:
|
| 165 |
+
state_dict = data
|
| 166 |
+
return TriageState.model_validate(state_dict)
|
| 167 |
+
|
| 168 |
+
# ------------------------------------------------------------------
|
| 169 |
+
# Context manager sugar
|
| 170 |
+
# ------------------------------------------------------------------
|
| 171 |
+
|
| 172 |
+
def close(self) -> None:
|
| 173 |
+
self._session.close()
|
| 174 |
+
|
| 175 |
+
def __enter__(self) -> "TriageOpenEnvClient":
|
| 176 |
+
return self
|
| 177 |
+
|
| 178 |
+
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
| 179 |
+
self.close()
|
| 180 |
+
|
| 181 |
+
# ------------------------------------------------------------------
|
| 182 |
+
# Internal
|
| 183 |
+
# ------------------------------------------------------------------
|
| 184 |
+
|
| 185 |
+
@staticmethod
|
| 186 |
+
def _parse_step_response(payload: Dict[str, Any]) -> StepResult:
|
| 187 |
+
obs_dict = payload.get("observation") or {}
|
| 188 |
+
obs = TriageObservation.model_validate(obs_dict)
|
| 189 |
+
reward = payload.get("reward")
|
| 190 |
+
# OpenEnv's serialize_observation strips ``done``/``reward`` from
|
| 191 |
+
# the observation dict, so ``done`` lives at the envelope level.
|
| 192 |
+
done = bool(payload.get("done", False))
|
| 193 |
+
# Make sure obs.done agrees with the envelope.
|
| 194 |
+
if obs.done != done:
|
| 195 |
+
obs.done = done
|
| 196 |
+
if reward is not None and obs.reward is None:
|
| 197 |
+
obs.reward = float(reward)
|
| 198 |
+
return StepResult(observation=obs, reward=reward, done=done)
|
| 199 |
+
|
| 200 |
+
|
| 201 |
+
__all__ = ["TriageOpenEnvClient", "StepResult"]
|
ER_MAP/envs/openenv_triage/env.py
ADDED
|
@@ -0,0 +1,295 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
ER_MAP/envs/openenv_triage/env.py
|
| 3 |
+
=================================
|
| 4 |
+
|
| 5 |
+
OpenEnv ``Environment`` subclass that delegates to the existing in-process
|
| 6 |
+
Gymnasium ``TriageEnv`` so a LoRA trained against the gym env can run
|
| 7 |
+
unchanged behind the OpenEnv HTTP/WebSocket protocol.
|
| 8 |
+
|
| 9 |
+
Citations (verified against the canonical OpenEnv repo & PyPI as of
|
| 10 |
+
April 2026):
|
| 11 |
+
|
| 12 |
+
- ``Environment`` interface (abstract ``reset`` / ``step`` / ``state``):
|
| 13 |
+
https://github.com/meta-pytorch/OpenEnv/blob/main/src/openenv/core/env_server/interfaces.py
|
| 14 |
+
-> ``class Environment(ABC, Generic[ActT, ObsT, StateT]): ...``
|
| 15 |
+
-> ``reset(seed, episode_id, **kwargs) -> ObsT``
|
| 16 |
+
-> ``step(action, timeout_s=None, **kwargs) -> ObsT`` (single-return; reward/done live on the Observation)
|
| 17 |
+
-> ``@property state`` returning the State subclass
|
| 18 |
+
-> ``close()`` is overridable (default no-op).
|
| 19 |
+
- ``Action`` / ``Observation`` / ``State`` Pydantic bases:
|
| 20 |
+
https://github.com/meta-pytorch/OpenEnv/blob/main/src/openenv/core/env_server/types.py
|
| 21 |
+
- ``serialize_observation`` strips ``reward``/``done``/``metadata`` from
|
| 22 |
+
the observation dict and emits ``{"observation", "reward", "done"}``:
|
| 23 |
+
https://github.com/meta-pytorch/OpenEnv/blob/main/src/openenv/core/env_server/serialization.py
|
| 24 |
+
- ``RESERVED_TOOL_NAMES`` (``reset`` / ``step`` / ``state`` / ``close``)
|
| 25 |
+
must NOT be used as MCP tool names:
|
| 26 |
+
https://github.com/meta-pytorch/OpenEnv/blob/main/src/openenv/core/env_server/mcp_types.py
|
| 27 |
+
- echo_env exemplar (showing ``MCPEnvironment`` / ``Environment`` class layout):
|
| 28 |
+
https://github.com/meta-pytorch/OpenEnv/tree/main/envs/echo_env
|
| 29 |
+
- pyproject pinning ``openenv-core==0.2.3`` (latest as of April 2026):
|
| 30 |
+
https://github.com/meta-pytorch/OpenEnv/blob/main/pyproject.toml
|
| 31 |
+
|
| 32 |
+
We do NOT subclass ``MCPEnvironment`` because the wrapped env is not a
|
| 33 |
+
tool-server: the Doctor agent emits a single structured ``TriageAction``
|
| 34 |
+
per step (closer to BlackJack/Wordle semantics than to coding_env's
|
| 35 |
+
tools/list+tools/call shape). ``Environment`` is the right base.
|
| 36 |
+
|
| 37 |
+
Key parity rule
|
| 38 |
+
---------------
|
| 39 |
+
The gym ``TriageEnv.step(action: str)`` returns the canonical 5-tuple
|
| 40 |
+
``(obs_str, reward, done, truncated, info)``. The OpenEnv ``Observation``
|
| 41 |
+
has ``done`` and ``reward`` first-class but NO truncated flag, so we fold
|
| 42 |
+
``truncated`` and ``info`` (including ``reward_components``) into the
|
| 43 |
+
observation's ``metadata`` and ``info`` fields so callers can recover the
|
| 44 |
+
exact gym tuple. Same trajectory + same seed -> identical rewards.
|
| 45 |
+
"""
|
| 46 |
+
|
| 47 |
+
from __future__ import annotations
|
| 48 |
+
|
| 49 |
+
import json
|
| 50 |
+
import logging
|
| 51 |
+
import os
|
| 52 |
+
import uuid
|
| 53 |
+
from typing import Any, Dict, Optional
|
| 54 |
+
|
| 55 |
+
# OpenEnv 0.2.3 ships the abstract ``Environment`` base under
|
| 56 |
+
# ``openenv.core.env_server.interfaces``. See citation block above.
|
| 57 |
+
from openenv.core.env_server.interfaces import Environment
|
| 58 |
+
|
| 59 |
+
from ER_MAP.envs.triage_env import TriageEnv
|
| 60 |
+
|
| 61 |
+
from .models import TriageAction, TriageObservation, TriageState
|
| 62 |
+
|
| 63 |
+
logger = logging.getLogger("ER_MAP.openenv_triage.env")
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def _stub_env_kwargs() -> Dict[str, Any]:
|
| 67 |
+
"""
|
| 68 |
+
Build env kwargs that allow the env to operate without Groq keys.
|
| 69 |
+
|
| 70 |
+
The underlying ``AgentRouter`` already has a degraded fallback path
|
| 71 |
+
(``_mock_response``) when no Groq client is configured, so handing it
|
| 72 |
+
empty keys is sufficient. The stub mode is documented in the README.
|
| 73 |
+
"""
|
| 74 |
+
return {
|
| 75 |
+
"groq_api_key": os.environ.get("GROQ_API_KEY", "") or None,
|
| 76 |
+
"nurse_api_key": os.environ.get("GROQ_NURSE_API_KEY", "") or None,
|
| 77 |
+
"patient_api_key": os.environ.get("GROQ_PATIENT_API_KEY", "") or None,
|
| 78 |
+
"empathy_judge_api_key": os.environ.get("GROQ_EMPATHY_JUDGE_API_KEY", "") or None,
|
| 79 |
+
"medical_judge_api_key": os.environ.get("GROQ_MEDICAL_JUDGE_API_KEY", "") or None,
|
| 80 |
+
"model": os.environ.get("ERMAP_MODEL", "llama-3.3-70b-versatile"),
|
| 81 |
+
}
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def _parse_obs_json(obs_str: str) -> Dict[str, Any]:
|
| 85 |
+
"""Best-effort parse of the gym env's JSON observation string."""
|
| 86 |
+
try:
|
| 87 |
+
return json.loads(obs_str)
|
| 88 |
+
except (json.JSONDecodeError, TypeError):
|
| 89 |
+
return {"event": "raw", "raw": obs_str}
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
class TriageOpenEnv(Environment[TriageAction, TriageObservation, TriageState]):
|
| 93 |
+
"""
|
| 94 |
+
OpenEnv-compliant wrapper around ``ER_MAP.envs.triage_env.TriageEnv``.
|
| 95 |
+
|
| 96 |
+
Internally instantiates one ``TriageEnv`` per session. Reset/step
|
| 97 |
+
delegate to the gym env unchanged so reward semantics, termination
|
| 98 |
+
rules, and trajectory contents stay byte-identical to the gym
|
| 99 |
+
baseline used by ``training/train_grpo.py``.
|
| 100 |
+
"""
|
| 101 |
+
|
| 102 |
+
# ``HTTPEnvServer`` instantiates a fresh environment per WebSocket
|
| 103 |
+
# session (see ``http_server.py``). The underlying ``TriageEnv`` keeps
|
| 104 |
+
# its state in instance attributes so concurrent sessions are safe.
|
| 105 |
+
SUPPORTS_CONCURRENT_SESSIONS = True
|
| 106 |
+
|
| 107 |
+
def __init__(self, env_kwargs: Optional[Dict[str, Any]] = None):
|
| 108 |
+
"""
|
| 109 |
+
Args:
|
| 110 |
+
env_kwargs: Optional kwargs forwarded to the underlying
|
| 111 |
+
``TriageEnv`` constructor. Defaults pull Groq keys from
|
| 112 |
+
env vars; missing keys silently fall back to mock mode.
|
| 113 |
+
"""
|
| 114 |
+
super().__init__()
|
| 115 |
+
|
| 116 |
+
merged_kwargs = _stub_env_kwargs()
|
| 117 |
+
if env_kwargs:
|
| 118 |
+
merged_kwargs.update(env_kwargs)
|
| 119 |
+
self._env_kwargs = merged_kwargs
|
| 120 |
+
|
| 121 |
+
# Build the gym env immediately so import-time failures (missing
|
| 122 |
+
# disease DB, etc.) surface synchronously rather than at first
|
| 123 |
+
# /reset request.
|
| 124 |
+
self._env = TriageEnv(**self._env_kwargs)
|
| 125 |
+
|
| 126 |
+
self._state = TriageState(
|
| 127 |
+
episode_id=str(uuid.uuid4()),
|
| 128 |
+
step_count=0,
|
| 129 |
+
)
|
| 130 |
+
self._last_truncated: bool = False
|
| 131 |
+
self._last_info: Dict[str, Any] = {}
|
| 132 |
+
|
| 133 |
+
# ------------------------------------------------------------------
|
| 134 |
+
# Reset
|
| 135 |
+
# ------------------------------------------------------------------
|
| 136 |
+
|
| 137 |
+
def reset(
|
| 138 |
+
self,
|
| 139 |
+
seed: Optional[int] = None,
|
| 140 |
+
episode_id: Optional[str] = None,
|
| 141 |
+
options: Optional[Dict[str, Any]] = None,
|
| 142 |
+
**kwargs: Any,
|
| 143 |
+
) -> TriageObservation:
|
| 144 |
+
"""
|
| 145 |
+
Start a new episode.
|
| 146 |
+
|
| 147 |
+
``options`` mirrors the gym ``options`` dict and supports the same
|
| 148 |
+
keys as ``TriageEnv.reset``: ``{"phase": 1|2|3, "difficulty":
|
| 149 |
+
"easy|medium|hard"}``.
|
| 150 |
+
"""
|
| 151 |
+
# Honour anything the OpenEnv server passes through that maps to
|
| 152 |
+
# the gym env's options dict (e.g. ``phase`` / ``difficulty``
|
| 153 |
+
# supplied as top-level kwargs by upstream tooling).
|
| 154 |
+
options = dict(options or {})
|
| 155 |
+
for opt_key in ("phase", "difficulty"):
|
| 156 |
+
if opt_key in kwargs and opt_key not in options:
|
| 157 |
+
options[opt_key] = kwargs.pop(opt_key)
|
| 158 |
+
|
| 159 |
+
obs_str, info = self._env.reset(seed=seed, options=options or None)
|
| 160 |
+
|
| 161 |
+
self._last_truncated = False
|
| 162 |
+
self._last_info = dict(info or {})
|
| 163 |
+
|
| 164 |
+
# Refresh state snapshot.
|
| 165 |
+
self._state = TriageState(
|
| 166 |
+
episode_id=episode_id or str(uuid.uuid4()),
|
| 167 |
+
step_count=int(getattr(self._env, "step_count", 0) or 0),
|
| 168 |
+
done=bool(getattr(self._env, "done", False)),
|
| 169 |
+
truncated=False,
|
| 170 |
+
consent_given=bool(getattr(self._env, "consent_given", False)),
|
| 171 |
+
phase=int(getattr(self._env, "phase", options.get("phase", 1)) or 1),
|
| 172 |
+
ordered_labs=list(getattr(self._env, "ordered_labs", set()) or set()),
|
| 173 |
+
patient_status=str(getattr(self._env, "last_patient_status", "CONTINUE")),
|
| 174 |
+
soap_note=dict(getattr(self._env, "emr", {}) or {}),
|
| 175 |
+
reward_components=dict(getattr(self._env, "reward_components", {}) or {}),
|
| 176 |
+
ground_truth_disease=info.get("ground_truth_disease"),
|
| 177 |
+
)
|
| 178 |
+
|
| 179 |
+
return self._build_observation(
|
| 180 |
+
obs_str=obs_str,
|
| 181 |
+
reward=0.0,
|
| 182 |
+
done=False,
|
| 183 |
+
truncated=False,
|
| 184 |
+
info=self._last_info,
|
| 185 |
+
)
|
| 186 |
+
|
| 187 |
+
# ------------------------------------------------------------------
|
| 188 |
+
# Step
|
| 189 |
+
# ------------------------------------------------------------------
|
| 190 |
+
|
| 191 |
+
def step(
|
| 192 |
+
self,
|
| 193 |
+
action: TriageAction,
|
| 194 |
+
timeout_s: Optional[float] = None,
|
| 195 |
+
**kwargs: Any,
|
| 196 |
+
) -> TriageObservation:
|
| 197 |
+
"""
|
| 198 |
+
Execute one Doctor turn.
|
| 199 |
+
|
| 200 |
+
Converts the typed ``TriageAction`` back into the JSON string the
|
| 201 |
+
gym env expects (preserving the original raw bytes whenever the
|
| 202 |
+
action was created via :meth:`TriageAction.from_json_str`), runs
|
| 203 |
+
the gym ``step``, and re-packages the 5-tuple into a single
|
| 204 |
+
``TriageObservation``.
|
| 205 |
+
"""
|
| 206 |
+
action_str = action.to_json_str()
|
| 207 |
+
obs_str, reward, done, truncated, info = self._env.step(action_str)
|
| 208 |
+
|
| 209 |
+
self._last_truncated = bool(truncated)
|
| 210 |
+
self._last_info = dict(info or {})
|
| 211 |
+
|
| 212 |
+
# Mirror gym attributes onto the typed state for /state polling.
|
| 213 |
+
self._state = TriageState(
|
| 214 |
+
episode_id=self._state.episode_id,
|
| 215 |
+
step_count=int(getattr(self._env, "step_count", 0) or 0),
|
| 216 |
+
done=bool(done),
|
| 217 |
+
truncated=bool(truncated),
|
| 218 |
+
consent_given=bool(getattr(self._env, "consent_given", False)),
|
| 219 |
+
phase=int(getattr(self._env, "phase", 1) or 1),
|
| 220 |
+
ordered_labs=list(getattr(self._env, "ordered_labs", set()) or set()),
|
| 221 |
+
patient_status=str(getattr(self._env, "last_patient_status", "CONTINUE")),
|
| 222 |
+
soap_note=dict(getattr(self._env, "emr", {}) or {}),
|
| 223 |
+
reward_components=dict(
|
| 224 |
+
info.get("reward_components")
|
| 225 |
+
or getattr(self._env, "reward_components", {})
|
| 226 |
+
or {}
|
| 227 |
+
),
|
| 228 |
+
ground_truth_disease=self._state.ground_truth_disease,
|
| 229 |
+
)
|
| 230 |
+
|
| 231 |
+
return self._build_observation(
|
| 232 |
+
obs_str=obs_str,
|
| 233 |
+
reward=float(reward),
|
| 234 |
+
done=bool(done),
|
| 235 |
+
truncated=bool(truncated),
|
| 236 |
+
info=self._last_info,
|
| 237 |
+
)
|
| 238 |
+
|
| 239 |
+
# ------------------------------------------------------------------
|
| 240 |
+
# State / close
|
| 241 |
+
# ------------------------------------------------------------------
|
| 242 |
+
|
| 243 |
+
@property
|
| 244 |
+
def state(self) -> TriageState:
|
| 245 |
+
return self._state
|
| 246 |
+
|
| 247 |
+
def close(self) -> None:
|
| 248 |
+
"""Release the underlying env's router resources."""
|
| 249 |
+
try:
|
| 250 |
+
self._env.close()
|
| 251 |
+
except Exception: # pragma: no cover - best-effort cleanup
|
| 252 |
+
logger.debug("TriageEnv.close() raised; ignoring.", exc_info=True)
|
| 253 |
+
|
| 254 |
+
# ------------------------------------------------------------------
|
| 255 |
+
# Helpers
|
| 256 |
+
# ------------------------------------------------------------------
|
| 257 |
+
|
| 258 |
+
def _build_observation(
|
| 259 |
+
self,
|
| 260 |
+
obs_str: str,
|
| 261 |
+
reward: float,
|
| 262 |
+
done: bool,
|
| 263 |
+
truncated: bool,
|
| 264 |
+
info: Dict[str, Any],
|
| 265 |
+
) -> TriageObservation:
|
| 266 |
+
payload = _parse_obs_json(obs_str)
|
| 267 |
+
event = str(payload.get("event", "")) if isinstance(payload, dict) else ""
|
| 268 |
+
|
| 269 |
+
# Promote ``truncated`` and the gym ``info`` dict (which carries
|
| 270 |
+
# ``reward_components``) into the observation so OpenEnv clients
|
| 271 |
+
# can reconstruct the exact gym 5-tuple if they need to.
|
| 272 |
+
metadata: Dict[str, Any] = {
|
| 273 |
+
"truncated": bool(truncated),
|
| 274 |
+
"reward_components": dict(info.get("reward_components", {})),
|
| 275 |
+
"step_count": info.get("step_count"),
|
| 276 |
+
"patient_status": info.get("patient_status"),
|
| 277 |
+
"consent_given": info.get("consent_given"),
|
| 278 |
+
"truncation_reason": info.get("truncation_reason"),
|
| 279 |
+
}
|
| 280 |
+
# Drop None values to keep serialized payloads tidy.
|
| 281 |
+
metadata = {k: v for k, v in metadata.items() if v is not None}
|
| 282 |
+
|
| 283 |
+
return TriageObservation(
|
| 284 |
+
done=bool(done) or bool(truncated),
|
| 285 |
+
reward=float(reward),
|
| 286 |
+
metadata=metadata,
|
| 287 |
+
raw_observation=obs_str if isinstance(obs_str, str) else json.dumps(obs_str),
|
| 288 |
+
event=event,
|
| 289 |
+
payload=payload if isinstance(payload, dict) else {},
|
| 290 |
+
truncated=bool(truncated),
|
| 291 |
+
info=dict(info or {}),
|
| 292 |
+
)
|
| 293 |
+
|
| 294 |
+
|
| 295 |
+
__all__ = ["TriageOpenEnv"]
|
ER_MAP/envs/openenv_triage/models.py
ADDED
|
@@ -0,0 +1,279 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
ER_MAP/envs/openenv_triage/models.py
|
| 3 |
+
====================================
|
| 4 |
+
|
| 5 |
+
Pydantic ``Action`` / ``Observation`` / ``State`` models that satisfy the
|
| 6 |
+
OpenEnv 0.2.3 interface for the ER-MAP triage environment.
|
| 7 |
+
|
| 8 |
+
Schema decisions
|
| 9 |
+
----------------
|
| 10 |
+
|
| 11 |
+
Doctor JSON action shape (mirrors ``DOCTOR_TOOLS`` in ``triage_env.py``):
|
| 12 |
+
|
| 13 |
+
{"tool": "speak_to", "target": "nurse|patient", "message": str}
|
| 14 |
+
{"tool": "order_lab", "test_name": str}
|
| 15 |
+
{"tool": "read_soap", "section": str}
|
| 16 |
+
{"tool": "update_soap", "section": str, "content": str}
|
| 17 |
+
{"tool": "terminal_discharge", "treatment": str, "is_emergency": bool}
|
| 18 |
+
|
| 19 |
+
The underlying ``TriageEnv.step`` accepts a JSON *string* and parses it
|
| 20 |
+
internally. To preserve byte-perfect parity with the gym env we keep the
|
| 21 |
+
original raw JSON string on the action (``raw_json``) and replay it
|
| 22 |
+
verbatim when stepping. Typed callers (e.g. UI / tests) can construct
|
| 23 |
+
the typed dataclass and rely on :meth:`TriageAction.to_json_str` to emit
|
| 24 |
+
an equivalent JSON payload.
|
| 25 |
+
|
| 26 |
+
References
|
| 27 |
+
----------
|
| 28 |
+
- Pydantic ``Action`` / ``Observation`` / ``State`` base classes:
|
| 29 |
+
https://github.com/meta-pytorch/OpenEnv/blob/main/src/openenv/core/env_server/types.py
|
| 30 |
+
"""
|
| 31 |
+
|
| 32 |
+
from __future__ import annotations
|
| 33 |
+
|
| 34 |
+
import json
|
| 35 |
+
from typing import Any, Dict, List, Optional
|
| 36 |
+
|
| 37 |
+
from pydantic import ConfigDict, Field
|
| 38 |
+
|
| 39 |
+
# OpenEnv base types live in ``openenv.core.env_server.types``. They are
|
| 40 |
+
# Pydantic v2 ``BaseModel`` subclasses with sensible defaults
|
| 41 |
+
# (``done``/``reward``/``metadata`` on Observation, ``episode_id`` /
|
| 42 |
+
# ``step_count`` on State, ``metadata`` on Action).
|
| 43 |
+
try:
|
| 44 |
+
from openenv.core.env_server.types import Action, Observation, State
|
| 45 |
+
except ImportError: # pragma: no cover - openenv-core not installed
|
| 46 |
+
# Fallback shims so this module can still be imported in environments
|
| 47 |
+
# that don't have ``openenv-core`` available (e.g. the running Kaggle
|
| 48 |
+
# training image). The wrapper itself will fail at construction time
|
| 49 |
+
# in that case, which is fine.
|
| 50 |
+
from pydantic import BaseModel
|
| 51 |
+
|
| 52 |
+
class Action(BaseModel): # type: ignore[no-redef]
|
| 53 |
+
model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
|
| 54 |
+
metadata: Dict[str, Any] = Field(default_factory=dict)
|
| 55 |
+
|
| 56 |
+
class Observation(BaseModel): # type: ignore[no-redef]
|
| 57 |
+
model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
|
| 58 |
+
done: bool = False
|
| 59 |
+
reward: Optional[float] = None
|
| 60 |
+
metadata: Dict[str, Any] = Field(default_factory=dict)
|
| 61 |
+
|
| 62 |
+
class State(BaseModel): # type: ignore[no-redef]
|
| 63 |
+
model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
|
| 64 |
+
episode_id: Optional[str] = None
|
| 65 |
+
step_count: int = 0
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# ---------------------------------------------------------------------------
|
| 69 |
+
# Action
|
| 70 |
+
# ---------------------------------------------------------------------------
|
| 71 |
+
|
| 72 |
+
# DOCTOR_TOOLS mirrors the same set used by ``triage_env.py``. Kept as a
|
| 73 |
+
# module constant so it can be imported by the FastAPI server / client for
|
| 74 |
+
# request validation without re-importing the underlying env.
|
| 75 |
+
DOCTOR_TOOLS: tuple = (
|
| 76 |
+
"speak_to",
|
| 77 |
+
"order_lab",
|
| 78 |
+
"terminal_discharge",
|
| 79 |
+
"read_soap",
|
| 80 |
+
"update_soap",
|
| 81 |
+
)
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
class TriageAction(Action):
|
| 85 |
+
"""
|
| 86 |
+
Doctor's structured action.
|
| 87 |
+
|
| 88 |
+
Parity-critical fields (``tool`` is required, the rest are optional and
|
| 89 |
+
only meaningful for specific tools). Extra fields are allowed because
|
| 90 |
+
the LoRA-tuned policy may emit auxiliary keys (``thought``, etc.) that
|
| 91 |
+
the underlying env tolerates.
|
| 92 |
+
"""
|
| 93 |
+
|
| 94 |
+
# Allow extra LLM-emitted fields (e.g. ``thought``) and accept legacy
|
| 95 |
+
# actions that may include unknown keys; the wrapped env tolerates them.
|
| 96 |
+
model_config = ConfigDict(
|
| 97 |
+
extra="allow",
|
| 98 |
+
validate_assignment=True,
|
| 99 |
+
arbitrary_types_allowed=True,
|
| 100 |
+
)
|
| 101 |
+
|
| 102 |
+
tool: str = Field(..., description="Doctor tool name; one of DOCTOR_TOOLS.")
|
| 103 |
+
target: Optional[str] = Field(
|
| 104 |
+
default=None, description="speak_to target ('nurse' or 'patient')."
|
| 105 |
+
)
|
| 106 |
+
message: Optional[str] = Field(
|
| 107 |
+
default=None, description="speak_to message body."
|
| 108 |
+
)
|
| 109 |
+
test_name: Optional[str] = Field(
|
| 110 |
+
default=None, description="order_lab test name."
|
| 111 |
+
)
|
| 112 |
+
treatment: Optional[str] = Field(
|
| 113 |
+
default=None, description="terminal_discharge treatment plan."
|
| 114 |
+
)
|
| 115 |
+
is_emergency: Optional[bool] = Field(
|
| 116 |
+
default=None,
|
| 117 |
+
description="terminal_discharge emergency-classification boolean.",
|
| 118 |
+
)
|
| 119 |
+
section: Optional[str] = Field(
|
| 120 |
+
default=None,
|
| 121 |
+
description="read_soap / update_soap SOAP section (e.g. 'Assessment').",
|
| 122 |
+
)
|
| 123 |
+
content: Optional[str] = Field(
|
| 124 |
+
default=None, description="update_soap content."
|
| 125 |
+
)
|
| 126 |
+
thought: Optional[str] = Field(
|
| 127 |
+
default=None, description="Optional LLM scratchpad text."
|
| 128 |
+
)
|
| 129 |
+
|
| 130 |
+
# Round-trip support: when the action originated from a raw JSON
|
| 131 |
+
# string (e.g. an LLM completion) we keep that exact string so the
|
| 132 |
+
# wrapper can replay it verbatim into the gym env, preserving byte
|
| 133 |
+
# parity with the existing reward/parser logic.
|
| 134 |
+
raw_json: Optional[str] = Field(
|
| 135 |
+
default=None,
|
| 136 |
+
description="Original JSON string this action was parsed from (for parity replay).",
|
| 137 |
+
)
|
| 138 |
+
|
| 139 |
+
# ------------------------------------------------------------------
|
| 140 |
+
# Construction helpers
|
| 141 |
+
# ------------------------------------------------------------------
|
| 142 |
+
|
| 143 |
+
@classmethod
|
| 144 |
+
def from_json_str(cls, s: str) -> "TriageAction":
|
| 145 |
+
"""
|
| 146 |
+
Parse a Doctor's free-text JSON action.
|
| 147 |
+
|
| 148 |
+
Robust to the same minor malformations the underlying
|
| 149 |
+
``TriageEnv._parse_doctor_action`` accepts (e.g. JSON embedded in
|
| 150 |
+
prose), but raises ``ValueError`` on unrecoverable malformed input
|
| 151 |
+
so callers can short-circuit before paying the env step cost.
|
| 152 |
+
"""
|
| 153 |
+
if not isinstance(s, str):
|
| 154 |
+
raise ValueError("from_json_str expects a string")
|
| 155 |
+
try:
|
| 156 |
+
parsed = json.loads(s.strip())
|
| 157 |
+
except (json.JSONDecodeError, TypeError):
|
| 158 |
+
# Try to extract the first JSON object inside the text, matching
|
| 159 |
+
# ``triage_env._parse_doctor_action`` behaviour.
|
| 160 |
+
import re
|
| 161 |
+
|
| 162 |
+
m = re.search(r"\{.*\}", s, re.DOTALL)
|
| 163 |
+
if not m:
|
| 164 |
+
raise ValueError(f"No JSON object found in action: {s!r}")
|
| 165 |
+
try:
|
| 166 |
+
parsed = json.loads(m.group(0))
|
| 167 |
+
except json.JSONDecodeError as e:
|
| 168 |
+
raise ValueError(f"Malformed JSON action: {e}") from e
|
| 169 |
+
|
| 170 |
+
if not isinstance(parsed, dict):
|
| 171 |
+
raise ValueError(f"Action JSON must be an object, got {type(parsed).__name__}")
|
| 172 |
+
if "tool" not in parsed:
|
| 173 |
+
raise ValueError("Action JSON must include a 'tool' field")
|
| 174 |
+
|
| 175 |
+
# Build the action via Pydantic validation (extra='allow' captures
|
| 176 |
+
# any non-declared LLM keys without dropping them).
|
| 177 |
+
action = cls.model_validate(parsed)
|
| 178 |
+
# Preserve the exact original payload for parity replay.
|
| 179 |
+
action.raw_json = s
|
| 180 |
+
return action
|
| 181 |
+
|
| 182 |
+
def to_json_str(self) -> str:
|
| 183 |
+
"""
|
| 184 |
+
Serialize back to the JSON string the underlying env consumes.
|
| 185 |
+
|
| 186 |
+
If this action was constructed from a raw JSON string, return that
|
| 187 |
+
original payload verbatim (byte-perfect parity). Otherwise emit a
|
| 188 |
+
fresh JSON dump of all non-None declared fields plus any extras.
|
| 189 |
+
"""
|
| 190 |
+
if self.raw_json:
|
| 191 |
+
return self.raw_json
|
| 192 |
+
# ``model_dump(exclude_none=True)`` drops unset optional fields so
|
| 193 |
+
# the env sees a clean schema-shaped action rather than a sea of
|
| 194 |
+
# nulls. ``raw_json`` and ``metadata`` are server-side concerns and
|
| 195 |
+
# are excluded from the wire payload.
|
| 196 |
+
payload = self.model_dump(exclude_none=True, exclude={"raw_json", "metadata"})
|
| 197 |
+
return json.dumps(payload)
|
| 198 |
+
|
| 199 |
+
|
| 200 |
+
# ---------------------------------------------------------------------------
|
| 201 |
+
# Observation
|
| 202 |
+
# ---------------------------------------------------------------------------
|
| 203 |
+
|
| 204 |
+
class TriageObservation(Observation):
|
| 205 |
+
"""
|
| 206 |
+
Doctor-visible observation.
|
| 207 |
+
|
| 208 |
+
Mirrors the structured JSON the gym env emits, but exposes the parsed
|
| 209 |
+
fields directly. The original ``raw_observation`` JSON string is also
|
| 210 |
+
retained so policies/clients that were trained against the raw bytes
|
| 211 |
+
can keep using them unmodified.
|
| 212 |
+
"""
|
| 213 |
+
|
| 214 |
+
model_config = ConfigDict(
|
| 215 |
+
extra="allow",
|
| 216 |
+
validate_assignment=True,
|
| 217 |
+
arbitrary_types_allowed=True,
|
| 218 |
+
)
|
| 219 |
+
|
| 220 |
+
raw_observation: str = Field(
|
| 221 |
+
default="",
|
| 222 |
+
description="Original JSON string from TriageEnv (for byte parity).",
|
| 223 |
+
)
|
| 224 |
+
event: str = Field(
|
| 225 |
+
default="",
|
| 226 |
+
description="Event type (e.g. 'episode_start', 'lab_result', 'terminal_win').",
|
| 227 |
+
)
|
| 228 |
+
payload: Dict[str, Any] = Field(
|
| 229 |
+
default_factory=dict,
|
| 230 |
+
description="Parsed JSON content from the gym env.",
|
| 231 |
+
)
|
| 232 |
+
truncated: bool = Field(
|
| 233 |
+
default=False,
|
| 234 |
+
description="Episode truncated due to max_episode_steps (mirrors gym).",
|
| 235 |
+
)
|
| 236 |
+
info: Dict[str, Any] = Field(
|
| 237 |
+
default_factory=dict,
|
| 238 |
+
description="Gym info dict, including reward_components.",
|
| 239 |
+
)
|
| 240 |
+
|
| 241 |
+
|
| 242 |
+
# ---------------------------------------------------------------------------
|
| 243 |
+
# State
|
| 244 |
+
# ---------------------------------------------------------------------------
|
| 245 |
+
|
| 246 |
+
class TriageState(State):
|
| 247 |
+
"""
|
| 248 |
+
Internal episode state surfaced over ``/state`` and the WebSocket
|
| 249 |
+
state message.
|
| 250 |
+
|
| 251 |
+
These fields mirror the public attributes of the underlying
|
| 252 |
+
``TriageEnv`` instance after ``reset()``/``step()``, so frontends and
|
| 253 |
+
debuggers can introspect mid-episode state without poking private
|
| 254 |
+
attributes through HTTP.
|
| 255 |
+
"""
|
| 256 |
+
|
| 257 |
+
model_config = ConfigDict(
|
| 258 |
+
extra="allow",
|
| 259 |
+
validate_assignment=True,
|
| 260 |
+
arbitrary_types_allowed=True,
|
| 261 |
+
)
|
| 262 |
+
|
| 263 |
+
done: bool = Field(default=False)
|
| 264 |
+
truncated: bool = Field(default=False)
|
| 265 |
+
consent_given: bool = Field(default=False)
|
| 266 |
+
phase: int = Field(default=1)
|
| 267 |
+
ordered_labs: List[str] = Field(default_factory=list)
|
| 268 |
+
patient_status: str = Field(default="CONTINUE")
|
| 269 |
+
soap_note: Dict[str, Any] = Field(default_factory=dict)
|
| 270 |
+
reward_components: Dict[str, float] = Field(default_factory=dict)
|
| 271 |
+
ground_truth_disease: Optional[str] = Field(default=None)
|
| 272 |
+
|
| 273 |
+
|
| 274 |
+
__all__ = [
|
| 275 |
+
"TriageAction",
|
| 276 |
+
"TriageObservation",
|
| 277 |
+
"TriageState",
|
| 278 |
+
"DOCTOR_TOOLS",
|
| 279 |
+
]
|
ER_MAP/envs/openenv_triage/openenv.yaml
ADDED
|
@@ -0,0 +1,44 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ER-MAP Triage - OpenEnv manifest
|
| 2 |
+
#
|
| 3 |
+
# Schema follows the canonical echo_env exemplar:
|
| 4 |
+
# https://github.com/meta-pytorch/OpenEnv/blob/main/envs/echo_env/openenv.yaml
|
| 5 |
+
# spec_version: 1
|
| 6 |
+
# name: <env name>
|
| 7 |
+
# type: space # HF Space deployment target
|
| 8 |
+
# runtime: fastapi
|
| 9 |
+
# app: <module>:<attr> # FastAPI app reference relative to the image
|
| 10 |
+
# port: 8000
|
| 11 |
+
#
|
| 12 |
+
# OpenEnv 0.2.3 (latest as of April 2026) recognizes this manifest;
|
| 13 |
+
# the ``openenv push`` CLI consumes it when deploying to Hugging Face
|
| 14 |
+
# Spaces.
|
| 15 |
+
|
| 16 |
+
spec_version: 1
|
| 17 |
+
name: er_map_triage
|
| 18 |
+
type: space
|
| 19 |
+
runtime: fastapi
|
| 20 |
+
|
| 21 |
+
# Inside the container the package is copied to /app/env, and the
|
| 22 |
+
# Dockerfile launches uvicorn against ``server:app`` (this module).
|
| 23 |
+
app: server:app
|
| 24 |
+
port: 8000
|
| 25 |
+
|
| 26 |
+
description: >
|
| 27 |
+
ER-MAP Triage: a multi-agent emergency-room simulation where a Doctor
|
| 28 |
+
RL policy interacts with Nurse / Patient / Judge LLM actors to
|
| 29 |
+
diagnose and treat domain-randomized scenarios. This is the
|
| 30 |
+
OpenEnv-compliant wrapper around the in-house Gymnasium TriageEnv
|
| 31 |
+
used to train ER-MAP's GRPO LoRA.
|
| 32 |
+
|
| 33 |
+
metadata:
|
| 34 |
+
authors:
|
| 35 |
+
- ER-MAP Team
|
| 36 |
+
license: MIT
|
| 37 |
+
tags:
|
| 38 |
+
- medical
|
| 39 |
+
- multi-agent
|
| 40 |
+
- triage
|
| 41 |
+
- grpo
|
| 42 |
+
- rl
|
| 43 |
+
- llm
|
| 44 |
+
hackathon: "Meta PyTorch OpenEnv Hackathon 2026"
|
ER_MAP/envs/openenv_triage/server-requirements.txt
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Container-side requirements for the ER-MAP TriageOpenEnv server.
|
| 2 |
+
#
|
| 3 |
+
# Kept separate from the repo-root requirements.txt so HF Space rebuilds
|
| 4 |
+
# only depend on the wrapper's needs (not the full RL training stack).
|
| 5 |
+
openenv-core>=0.2.3
|
| 6 |
+
gymnasium>=0.29.0
|
| 7 |
+
groq>=0.18.0
|
| 8 |
+
fastapi>=0.110.0
|
| 9 |
+
uvicorn[standard]>=0.27.0
|
| 10 |
+
pydantic>=2.0.0
|
| 11 |
+
requests>=2.25.0
|
| 12 |
+
pyyaml>=6.0
|
ER_MAP/envs/openenv_triage/server.py
ADDED
|
@@ -0,0 +1,312 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
ER_MAP/envs/openenv_triage/server.py
|
| 3 |
+
====================================
|
| 4 |
+
|
| 5 |
+
FastAPI server exposing :class:`TriageOpenEnv` over the OpenEnv HTTP /
|
| 6 |
+
WebSocket protocol.
|
| 7 |
+
|
| 8 |
+
Endpoints:
|
| 9 |
+
|
| 10 |
+
- ``POST /reset`` reset request (matches ``ResetRequest`` schema)
|
| 11 |
+
- ``POST /step`` step request (matches ``StepRequest`` schema)
|
| 12 |
+
- ``GET /state`` serialized ``TriageState``
|
| 13 |
+
- ``GET /health`` ``{"status": "healthy"}`` (also ``/healthz`` for richer info)
|
| 14 |
+
- ``GET /docs`` OpenAPI / Swagger UI
|
| 15 |
+
- ``WS /ws`` persistent OpenEnv session (mounted from openenv-core)
|
| 16 |
+
- ``GET /web`` optional Gradio HumanAgent UI when ``ENABLE_WEB_INTERFACE=true``
|
| 17 |
+
|
| 18 |
+
Why a custom HTTP layer (instead of relying entirely on
|
| 19 |
+
``openenv.core.env_server.http_server.create_app``):
|
| 20 |
+
|
| 21 |
+
OpenEnv 0.2.3's HTTP ``/reset`` and ``/step`` handlers spin up a NEW
|
| 22 |
+
``Environment`` instance per request and call ``close()`` immediately
|
| 23 |
+
after - i.e. HTTP is intentionally stateless and persistent sessions
|
| 24 |
+
live on the WebSocket ``/ws`` route. (See
|
| 25 |
+
``openenv/core/env_server/http_server.py``: ``reset_handler`` /
|
| 26 |
+
``step_handler`` both call ``_env = self._env_factory()`` then
|
| 27 |
+
``finally: _env.close()``.) For an episodic env like ER-MAP we need
|
| 28 |
+
``/step`` to see the state established by the prior ``/reset``, so we
|
| 29 |
+
mount stateful HTTP routes here that share a per-process singleton
|
| 30 |
+
``TriageOpenEnv``. We additionally compose in the upstream OpenEnv
|
| 31 |
+
FastAPI app so the WebSocket / schema / web-interface routes still
|
| 32 |
+
work for OpenEnv-native clients.
|
| 33 |
+
|
| 34 |
+
References
|
| 35 |
+
----------
|
| 36 |
+
- ``HTTPEnvServer`` HTTP route bodies (note the ``_env_factory()`` per call):
|
| 37 |
+
https://github.com/meta-pytorch/OpenEnv/blob/main/src/openenv/core/env_server/http_server.py
|
| 38 |
+
- echo_env's server layout (``server/app.py``):
|
| 39 |
+
https://github.com/meta-pytorch/OpenEnv/blob/main/envs/echo_env/server/app.py
|
| 40 |
+
|
| 41 |
+
Stub mode
|
| 42 |
+
---------
|
| 43 |
+
If no Groq API key is present in the environment, the underlying
|
| 44 |
+
``AgentRouter`` already falls back to canned mock responses, so a fresh
|
| 45 |
+
HF Space can boot and serve full episodes (with non-LLM nurse/patient
|
| 46 |
+
behaviour) without any secrets configured. The ``/healthz`` endpoint
|
| 47 |
+
reports ``stub_mode: true`` in that case.
|
| 48 |
+
"""
|
| 49 |
+
|
| 50 |
+
from __future__ import annotations
|
| 51 |
+
|
| 52 |
+
import logging
|
| 53 |
+
import os
|
| 54 |
+
import threading
|
| 55 |
+
import uuid
|
| 56 |
+
from typing import Any, Dict, Optional
|
| 57 |
+
|
| 58 |
+
from fastapi import Body, FastAPI, HTTPException
|
| 59 |
+
from fastapi.middleware.cors import CORSMiddleware
|
| 60 |
+
from pydantic import BaseModel, ConfigDict, Field
|
| 61 |
+
|
| 62 |
+
from .env import TriageOpenEnv
|
| 63 |
+
from .models import TriageAction, TriageObservation, TriageState
|
| 64 |
+
|
| 65 |
+
logging.basicConfig(
|
| 66 |
+
level=os.environ.get("ERMAP_LOG_LEVEL", "INFO").upper(),
|
| 67 |
+
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
| 68 |
+
)
|
| 69 |
+
logger = logging.getLogger("ER_MAP.openenv_triage.server")
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
__version__ = "0.1.0"
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
# ---------------------------------------------------------------------------
|
| 76 |
+
# Stub-mode detection
|
| 77 |
+
# ---------------------------------------------------------------------------
|
| 78 |
+
|
| 79 |
+
def _has_any_groq_key() -> bool:
|
| 80 |
+
return any(
|
| 81 |
+
os.environ.get(name)
|
| 82 |
+
for name in (
|
| 83 |
+
"GROQ_API_KEY",
|
| 84 |
+
"GROQ_NURSE_API_KEY",
|
| 85 |
+
"GROQ_PATIENT_API_KEY",
|
| 86 |
+
"GROQ_EMPATHY_JUDGE_API_KEY",
|
| 87 |
+
"GROQ_MEDICAL_JUDGE_API_KEY",
|
| 88 |
+
)
|
| 89 |
+
)
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
_STUB_MODE = not _has_any_groq_key()
|
| 93 |
+
if _STUB_MODE:
|
| 94 |
+
logger.warning(
|
| 95 |
+
"TriageOpenEnv server starting in STUB MODE - no Groq keys found. "
|
| 96 |
+
"Nurse/Patient/Judge LLMs will use canned mock responses. "
|
| 97 |
+
"Set GROQ_API_KEY (or per-role keys) to enable live LLM actors."
|
| 98 |
+
)
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
# ---------------------------------------------------------------------------
|
| 102 |
+
# Pydantic request / response envelopes
|
| 103 |
+
# ---------------------------------------------------------------------------
|
| 104 |
+
|
| 105 |
+
class ResetRequest(BaseModel):
|
| 106 |
+
"""HTTP body for ``POST /reset``. Matches OpenEnv's ResetRequest plus
|
| 107 |
+
the ER-MAP-specific ``options`` field."""
|
| 108 |
+
|
| 109 |
+
model_config = ConfigDict(extra="allow")
|
| 110 |
+
|
| 111 |
+
seed: Optional[int] = Field(default=None, ge=0)
|
| 112 |
+
episode_id: Optional[str] = Field(default=None, max_length=255)
|
| 113 |
+
options: Optional[Dict[str, Any]] = Field(default=None)
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
class StepRequest(BaseModel):
|
| 117 |
+
"""HTTP body for ``POST /step``. The ``action`` field accepts either a
|
| 118 |
+
full ``TriageAction`` payload (preferred) or the legacy raw JSON
|
| 119 |
+
string the gym env consumes."""
|
| 120 |
+
|
| 121 |
+
model_config = ConfigDict(extra="allow")
|
| 122 |
+
|
| 123 |
+
action: Dict[str, Any] = Field(...)
|
| 124 |
+
timeout_s: Optional[float] = Field(default=None, gt=0)
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
class StepEnvelope(BaseModel):
|
| 128 |
+
"""Wire-format response. Matches OpenEnv 0.2.3's ``StepResponse``
|
| 129 |
+
layout: observation dict + reward + done at the envelope level.
|
| 130 |
+
See ``serialize_observation`` in openenv-core."""
|
| 131 |
+
|
| 132 |
+
model_config = ConfigDict(extra="forbid")
|
| 133 |
+
|
| 134 |
+
observation: Dict[str, Any]
|
| 135 |
+
reward: Optional[float] = None
|
| 136 |
+
done: bool = False
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
# ---------------------------------------------------------------------------
|
| 140 |
+
# Singleton env management (HTTP routes)
|
| 141 |
+
# ---------------------------------------------------------------------------
|
| 142 |
+
|
| 143 |
+
class _SessionHolder:
|
| 144 |
+
"""Thread-safe holder for the single per-process TriageOpenEnv used by
|
| 145 |
+
HTTP routes. WebSocket sessions get their own per-connection env via
|
| 146 |
+
the upstream OpenEnv server."""
|
| 147 |
+
|
| 148 |
+
def __init__(self) -> None:
|
| 149 |
+
self._lock = threading.Lock()
|
| 150 |
+
self._env: Optional[TriageOpenEnv] = None
|
| 151 |
+
self._episode_id: Optional[str] = None
|
| 152 |
+
|
| 153 |
+
def get(self) -> TriageOpenEnv:
|
| 154 |
+
with self._lock:
|
| 155 |
+
if self._env is None:
|
| 156 |
+
self._env = TriageOpenEnv()
|
| 157 |
+
self._episode_id = str(uuid.uuid4())
|
| 158 |
+
return self._env
|
| 159 |
+
|
| 160 |
+
def reset(self) -> TriageOpenEnv:
|
| 161 |
+
with self._lock:
|
| 162 |
+
if self._env is not None:
|
| 163 |
+
try:
|
| 164 |
+
self._env.close()
|
| 165 |
+
except Exception: # pragma: no cover
|
| 166 |
+
logger.debug("Old env close raised", exc_info=True)
|
| 167 |
+
self._env = TriageOpenEnv()
|
| 168 |
+
self._episode_id = str(uuid.uuid4())
|
| 169 |
+
return self._env
|
| 170 |
+
|
| 171 |
+
|
| 172 |
+
_session = _SessionHolder()
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
def _serialize(obs: TriageObservation) -> StepEnvelope:
|
| 176 |
+
"""Mirror ``openenv.core.env_server.serialization.serialize_observation``:
|
| 177 |
+
strip ``done``/``reward``/``metadata`` from the observation dict and
|
| 178 |
+
surface them at the envelope level."""
|
| 179 |
+
obs_dict = obs.model_dump(exclude={"reward", "done", "metadata"})
|
| 180 |
+
return StepEnvelope(
|
| 181 |
+
observation=obs_dict,
|
| 182 |
+
reward=float(obs.reward) if obs.reward is not None else None,
|
| 183 |
+
done=bool(obs.done),
|
| 184 |
+
)
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
# ---------------------------------------------------------------------------
|
| 188 |
+
# App factory
|
| 189 |
+
# ---------------------------------------------------------------------------
|
| 190 |
+
|
| 191 |
+
def build_app() -> FastAPI:
|
| 192 |
+
"""Construct the FastAPI app. Exposed for tests / programmatic use."""
|
| 193 |
+
|
| 194 |
+
app = FastAPI(
|
| 195 |
+
title="ER-MAP TriageOpenEnv",
|
| 196 |
+
version=__version__,
|
| 197 |
+
description=(
|
| 198 |
+
"OpenEnv-compliant HTTP/WebSocket interface for ER-MAP's "
|
| 199 |
+
"multi-agent emergency-room triage environment. Wraps the "
|
| 200 |
+
"in-house Gymnasium TriageEnv without modifying it, so the "
|
| 201 |
+
"same LoRA trained on the gym env runs unchanged here."
|
| 202 |
+
),
|
| 203 |
+
)
|
| 204 |
+
|
| 205 |
+
app.add_middleware(
|
| 206 |
+
CORSMiddleware,
|
| 207 |
+
allow_origins=["*"],
|
| 208 |
+
allow_credentials=True,
|
| 209 |
+
allow_methods=["*"],
|
| 210 |
+
allow_headers=["*"],
|
| 211 |
+
)
|
| 212 |
+
|
| 213 |
+
# ------------------------------------------------------------------
|
| 214 |
+
# Health
|
| 215 |
+
# ------------------------------------------------------------------
|
| 216 |
+
|
| 217 |
+
@app.get("/health", tags=["Health"])
|
| 218 |
+
def health() -> Dict[str, Any]:
|
| 219 |
+
"""Standard OpenEnv health probe."""
|
| 220 |
+
return {"status": "healthy"}
|
| 221 |
+
|
| 222 |
+
@app.get("/healthz", tags=["Health"])
|
| 223 |
+
def healthz() -> Dict[str, Any]:
|
| 224 |
+
"""Richer status: version + stub-mode flag (HF Space friendly)."""
|
| 225 |
+
return {
|
| 226 |
+
"status": "healthy",
|
| 227 |
+
"version": __version__,
|
| 228 |
+
"stub_mode": _STUB_MODE,
|
| 229 |
+
"env": "TriageOpenEnv",
|
| 230 |
+
}
|
| 231 |
+
|
| 232 |
+
# ------------------------------------------------------------------
|
| 233 |
+
# Reset / Step / State (stateful, per-process singleton)
|
| 234 |
+
# ------------------------------------------------------------------
|
| 235 |
+
|
| 236 |
+
@app.post("/reset", response_model=StepEnvelope, tags=["Environment Control"])
|
| 237 |
+
def reset(request: ResetRequest = Body(default_factory=ResetRequest)):
|
| 238 |
+
"""Reset the underlying TriageEnv and return the initial obs."""
|
| 239 |
+
env = _session.reset()
|
| 240 |
+
kwargs = request.model_dump(exclude_unset=True)
|
| 241 |
+
try:
|
| 242 |
+
obs = env.reset(**kwargs)
|
| 243 |
+
except Exception as e:
|
| 244 |
+
logger.exception("reset failed")
|
| 245 |
+
raise HTTPException(status_code=500, detail=f"reset failed: {e}")
|
| 246 |
+
return _serialize(obs)
|
| 247 |
+
|
| 248 |
+
@app.post("/step", response_model=StepEnvelope, tags=["Environment Control"])
|
| 249 |
+
def step(request: StepRequest):
|
| 250 |
+
"""Execute a Doctor action against the active env session."""
|
| 251 |
+
env = _session.get()
|
| 252 |
+
|
| 253 |
+
try:
|
| 254 |
+
action = TriageAction.model_validate(request.action)
|
| 255 |
+
except Exception as e:
|
| 256 |
+
raise HTTPException(status_code=422, detail=f"invalid action: {e}")
|
| 257 |
+
|
| 258 |
+
try:
|
| 259 |
+
obs = env.step(action, timeout_s=request.timeout_s)
|
| 260 |
+
except Exception as e:
|
| 261 |
+
logger.exception("step failed")
|
| 262 |
+
raise HTTPException(status_code=500, detail=f"step failed: {e}")
|
| 263 |
+
return _serialize(obs)
|
| 264 |
+
|
| 265 |
+
@app.get("/state", response_model=StepEnvelope, tags=["State Management"])
|
| 266 |
+
def state():
|
| 267 |
+
"""Inspect the current episode state."""
|
| 268 |
+
env = _session.get()
|
| 269 |
+
st: TriageState = env.state
|
| 270 |
+
return StepEnvelope(observation=st.model_dump(), reward=None, done=bool(st.done))
|
| 271 |
+
|
| 272 |
+
# ------------------------------------------------------------------
|
| 273 |
+
# Best-effort: mount upstream OpenEnv routes for WebSocket / schema /
|
| 274 |
+
# web-UI parity. We avoid colliding with our own /reset, /step,
|
| 275 |
+
# /state, /health by mounting the upstream app under /openenv.
|
| 276 |
+
# ------------------------------------------------------------------
|
| 277 |
+
|
| 278 |
+
try:
|
| 279 |
+
from openenv.core.env_server.http_server import create_app as _oe_create_app
|
| 280 |
+
|
| 281 |
+
oe_app = _oe_create_app(
|
| 282 |
+
TriageOpenEnv,
|
| 283 |
+
TriageAction,
|
| 284 |
+
TriageObservation,
|
| 285 |
+
env_name="er_map_triage",
|
| 286 |
+
max_concurrent_envs=int(os.environ.get("MAX_CONCURRENT_ENVS", "8")),
|
| 287 |
+
)
|
| 288 |
+
# Mount under /openenv so OpenEnv-native clients can use
|
| 289 |
+
# ws://host/openenv/ws and the schema routes without colliding
|
| 290 |
+
# with the stateful HTTP routes above.
|
| 291 |
+
app.mount("/openenv", oe_app)
|
| 292 |
+
except Exception as e: # pragma: no cover
|
| 293 |
+
logger.warning("Failed to mount upstream OpenEnv app at /openenv: %s", e)
|
| 294 |
+
|
| 295 |
+
return app
|
| 296 |
+
|
| 297 |
+
|
| 298 |
+
# Module-level ``app`` for ``uvicorn ER_MAP.envs.openenv_triage.server:app``.
|
| 299 |
+
app = build_app()
|
| 300 |
+
|
| 301 |
+
|
| 302 |
+
def main() -> None:
|
| 303 |
+
"""Entry point for ``python -m ER_MAP.envs.openenv_triage.server``."""
|
| 304 |
+
import uvicorn
|
| 305 |
+
|
| 306 |
+
host = os.environ.get("HOST", "0.0.0.0")
|
| 307 |
+
port = int(os.environ.get("PORT", "8000"))
|
| 308 |
+
uvicorn.run(app, host=host, port=port)
|
| 309 |
+
|
| 310 |
+
|
| 311 |
+
if __name__ == "__main__": # pragma: no cover
|
| 312 |
+
main()
|
ER_MAP/envs/openenv_triage/tests/__init__.py
ADDED
|
File without changes
|
ER_MAP/envs/openenv_triage/tests/test_parity.py
ADDED
|
@@ -0,0 +1,306 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
ER_MAP/envs/openenv_triage/tests/test_parity.py
|
| 3 |
+
================================================
|
| 4 |
+
|
| 5 |
+
Parity + smoke tests for the OpenEnv wrapper.
|
| 6 |
+
|
| 7 |
+
Coverage:
|
| 8 |
+
1. Same seed + options + action sequence produces byte-identical rewards,
|
| 9 |
+
terminations, and observation payloads in the wrapped env vs the gym
|
| 10 |
+
``TriageEnv``.
|
| 11 |
+
2. ``TriageAction.from_json_str`` parses well-formed Doctor JSON and
|
| 12 |
+
rejects malformed payloads.
|
| 13 |
+
3. End-to-end HTTP round-trip via ``TriageOpenEnvClient`` returns the
|
| 14 |
+
same observation/reward/done as direct in-process env usage.
|
| 15 |
+
4. ``openenv.yaml`` is valid YAML and references files that exist on disk.
|
| 16 |
+
|
| 17 |
+
All tests run without Groq keys: the underlying ``AgentRouter`` falls back
|
| 18 |
+
to ``_mock_response`` when no Groq client is configured, so we get
|
| 19 |
+
deterministic stubbed Nurse/Patient responses.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
from __future__ import annotations
|
| 23 |
+
|
| 24 |
+
import json
|
| 25 |
+
import os
|
| 26 |
+
import threading
|
| 27 |
+
import time
|
| 28 |
+
from pathlib import Path
|
| 29 |
+
|
| 30 |
+
import pytest
|
| 31 |
+
|
| 32 |
+
# Ensure no Groq client gets initialized so the mock-response path is
|
| 33 |
+
# active (deterministic stubs). Also wipe the per-role keys.
|
| 34 |
+
for _k in (
|
| 35 |
+
"GROQ_API_KEY",
|
| 36 |
+
"GROQ_NURSE_API_KEY",
|
| 37 |
+
"GROQ_PATIENT_API_KEY",
|
| 38 |
+
"GROQ_EMPATHY_JUDGE_API_KEY",
|
| 39 |
+
"GROQ_MEDICAL_JUDGE_API_KEY",
|
| 40 |
+
):
|
| 41 |
+
os.environ.pop(_k, None)
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
from ER_MAP.envs.triage_env import TriageEnv
|
| 45 |
+
from ER_MAP.envs.openenv_triage import (
|
| 46 |
+
TriageAction,
|
| 47 |
+
TriageObservation,
|
| 48 |
+
TriageOpenEnv,
|
| 49 |
+
TriageState,
|
| 50 |
+
)
|
| 51 |
+
from ER_MAP.envs.openenv_triage.client import TriageOpenEnvClient
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
# A short, deterministic action script that exercises several tools
|
| 55 |
+
# without depending on the LLM judges' verdict (we just check parity).
|
| 56 |
+
ACTION_SCRIPT = [
|
| 57 |
+
'{"tool": "read_soap"}',
|
| 58 |
+
'{"tool": "speak_to", "target": "patient", "message": "Hello, how are you feeling?"}',
|
| 59 |
+
'{"tool": "order_lab", "test_name": "CBC"}',
|
| 60 |
+
'{"tool": "update_soap", "section": "Assessment", "content": "Likely viral infection."}',
|
| 61 |
+
'{"tool": "terminal_discharge", "treatment": "Rest and hydration", "is_emergency": false}',
|
| 62 |
+
]
|
| 63 |
+
|
| 64 |
+
RESET_KWARGS = {"seed": 42, "options": {"phase": 1, "difficulty": "easy"}}
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
# ---------------------------------------------------------------------------
|
| 68 |
+
# Test 1 - in-process parity
|
| 69 |
+
# ---------------------------------------------------------------------------
|
| 70 |
+
|
| 71 |
+
def _run_gym_episode():
|
| 72 |
+
env = TriageEnv()
|
| 73 |
+
obs, info = env.reset(**RESET_KWARGS)
|
| 74 |
+
trace = [("reset", obs, 0.0, False, False, info)]
|
| 75 |
+
for action in ACTION_SCRIPT:
|
| 76 |
+
if trace[-1][3] or trace[-1][4]:
|
| 77 |
+
break
|
| 78 |
+
obs, reward, done, truncated, info = env.step(action)
|
| 79 |
+
trace.append(("step", obs, reward, done, truncated, info))
|
| 80 |
+
env.close()
|
| 81 |
+
return trace
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def _run_openenv_episode():
|
| 85 |
+
env = TriageOpenEnv()
|
| 86 |
+
obs = env.reset(**RESET_KWARGS)
|
| 87 |
+
trace = [("reset", obs.raw_observation, 0.0, False, False, obs.info)]
|
| 88 |
+
for action_str in ACTION_SCRIPT:
|
| 89 |
+
if trace[-1][3] or trace[-1][4]:
|
| 90 |
+
break
|
| 91 |
+
action = TriageAction.from_json_str(action_str)
|
| 92 |
+
obs = env.step(action)
|
| 93 |
+
trace.append(
|
| 94 |
+
("step", obs.raw_observation, obs.reward, obs.done and not obs.truncated,
|
| 95 |
+
obs.truncated, obs.info)
|
| 96 |
+
)
|
| 97 |
+
env.close()
|
| 98 |
+
return trace
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def _strip_volatile(info: dict) -> dict:
|
| 102 |
+
"""Drop fields that legitimately differ between two clean episodes
|
| 103 |
+
(none currently expected because the mock LLM path is deterministic
|
| 104 |
+
given a fixed seed, but we keep this hook for forward compatibility)."""
|
| 105 |
+
return {k: v for k, v in (info or {}).items() if k != "_timestamp"}
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def test_parity_step_by_step():
|
| 109 |
+
"""Wrapped env produces identical rewards, dones, and obs payloads."""
|
| 110 |
+
gym_trace = _run_gym_episode()
|
| 111 |
+
oe_trace = _run_openenv_episode()
|
| 112 |
+
|
| 113 |
+
assert len(gym_trace) == len(oe_trace), (
|
| 114 |
+
f"Trace lengths differ: gym={len(gym_trace)} openenv={len(oe_trace)}"
|
| 115 |
+
)
|
| 116 |
+
|
| 117 |
+
for i, (g, o) in enumerate(zip(gym_trace, oe_trace)):
|
| 118 |
+
g_kind, g_obs, g_reward, g_done, g_trunc, g_info = g
|
| 119 |
+
o_kind, o_obs, o_reward, o_done, o_trunc, o_info = o
|
| 120 |
+
|
| 121 |
+
assert g_kind == o_kind, f"step {i} kind differs"
|
| 122 |
+
assert pytest.approx(g_reward, abs=1e-9) == o_reward, (
|
| 123 |
+
f"step {i} reward differs: gym={g_reward} openenv={o_reward}"
|
| 124 |
+
)
|
| 125 |
+
assert g_done == o_done, f"step {i} done differs"
|
| 126 |
+
assert g_trunc == o_trunc, f"step {i} truncated differs"
|
| 127 |
+
|
| 128 |
+
# Compare structured observation content (not the literal string,
|
| 129 |
+
# because dict ordering may diverge in JSON encoding even though
|
| 130 |
+
# the parsed payloads are equal).
|
| 131 |
+
try:
|
| 132 |
+
g_payload = json.loads(g_obs)
|
| 133 |
+
o_payload = json.loads(o_obs)
|
| 134 |
+
except (json.JSONDecodeError, TypeError):
|
| 135 |
+
g_payload, o_payload = g_obs, o_obs
|
| 136 |
+
assert g_payload == o_payload, f"step {i} observation payload differs"
|
| 137 |
+
|
| 138 |
+
# reward_components must match exactly (these are the GRPO signal).
|
| 139 |
+
assert _strip_volatile(g_info.get("reward_components", {})) == _strip_volatile(
|
| 140 |
+
o_info.get("reward_components", {})
|
| 141 |
+
), f"step {i} reward_components differ"
|
| 142 |
+
|
| 143 |
+
|
| 144 |
+
# ---------------------------------------------------------------------------
|
| 145 |
+
# Test 2 - TriageAction parsing
|
| 146 |
+
# ---------------------------------------------------------------------------
|
| 147 |
+
|
| 148 |
+
def test_action_from_json_str_well_formed():
|
| 149 |
+
a = TriageAction.from_json_str('{"tool": "speak_to", "target": "nurse", "message": "Hi"}')
|
| 150 |
+
assert a.tool == "speak_to"
|
| 151 |
+
assert a.target == "nurse"
|
| 152 |
+
assert a.message == "Hi"
|
| 153 |
+
# Round-trip preserves the original bytes for parity replay.
|
| 154 |
+
assert a.raw_json is not None
|
| 155 |
+
assert json.loads(a.to_json_str()) == json.loads(a.raw_json)
|
| 156 |
+
|
| 157 |
+
|
| 158 |
+
def test_action_from_json_str_embedded():
|
| 159 |
+
"""TriageAction tolerates the same embedded-JSON noise the gym env
|
| 160 |
+
accepts (LLMs sometimes emit prose around their JSON)."""
|
| 161 |
+
a = TriageAction.from_json_str(
|
| 162 |
+
"Sure, here is my action: {\"tool\": \"read_soap\"}"
|
| 163 |
+
)
|
| 164 |
+
assert a.tool == "read_soap"
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
def test_action_from_json_str_rejects_malformed():
|
| 168 |
+
with pytest.raises(ValueError):
|
| 169 |
+
TriageAction.from_json_str("not even close to JSON")
|
| 170 |
+
with pytest.raises(ValueError):
|
| 171 |
+
# Missing 'tool' field is invalid per the env's parser.
|
| 172 |
+
TriageAction.from_json_str('{"foo": "bar"}')
|
| 173 |
+
with pytest.raises(ValueError):
|
| 174 |
+
TriageAction.from_json_str('{"tool":') # truncated
|
| 175 |
+
|
| 176 |
+
|
| 177 |
+
# ---------------------------------------------------------------------------
|
| 178 |
+
# Test 3 - HTTP round-trip via the FastAPI server
|
| 179 |
+
# ---------------------------------------------------------------------------
|
| 180 |
+
|
| 181 |
+
@pytest.fixture(scope="module")
|
| 182 |
+
def http_server():
|
| 183 |
+
"""Spin up the FastAPI app on a free port using uvicorn in a thread."""
|
| 184 |
+
pytest.importorskip("uvicorn")
|
| 185 |
+
pytest.importorskip("openenv.core.env_server.http_server")
|
| 186 |
+
import uvicorn
|
| 187 |
+
|
| 188 |
+
from ER_MAP.envs.openenv_triage.server import build_app
|
| 189 |
+
|
| 190 |
+
port = int(os.environ.get("ERMAP_TEST_PORT", "8765"))
|
| 191 |
+
app = build_app()
|
| 192 |
+
config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning")
|
| 193 |
+
server = uvicorn.Server(config)
|
| 194 |
+
|
| 195 |
+
thread = threading.Thread(target=server.run, daemon=True)
|
| 196 |
+
thread.start()
|
| 197 |
+
|
| 198 |
+
# Wait until the server is reachable (max 15s).
|
| 199 |
+
import requests
|
| 200 |
+
|
| 201 |
+
base = f"http://127.0.0.1:{port}"
|
| 202 |
+
deadline = time.time() + 15.0
|
| 203 |
+
while time.time() < deadline:
|
| 204 |
+
try:
|
| 205 |
+
r = requests.get(f"{base}/health", timeout=1.0)
|
| 206 |
+
if r.status_code == 200:
|
| 207 |
+
break
|
| 208 |
+
except Exception:
|
| 209 |
+
pass
|
| 210 |
+
time.sleep(0.2)
|
| 211 |
+
else:
|
| 212 |
+
server.should_exit = True
|
| 213 |
+
thread.join(timeout=2.0)
|
| 214 |
+
pytest.skip("uvicorn did not come up within 15s")
|
| 215 |
+
|
| 216 |
+
yield base
|
| 217 |
+
|
| 218 |
+
server.should_exit = True
|
| 219 |
+
thread.join(timeout=5.0)
|
| 220 |
+
|
| 221 |
+
|
| 222 |
+
def test_http_round_trip_matches_inproc(http_server):
|
| 223 |
+
"""Same seed/options/actions through HTTP -> identical reward stream."""
|
| 224 |
+
client = TriageOpenEnvClient(base_url=http_server)
|
| 225 |
+
|
| 226 |
+
health = client.health()
|
| 227 |
+
assert health.get("status") == "healthy"
|
| 228 |
+
|
| 229 |
+
# Run via HTTP.
|
| 230 |
+
result = client.reset(**RESET_KWARGS)
|
| 231 |
+
http_trace = [("reset", result.observation.raw_observation, 0.0, result.done, result.truncated)]
|
| 232 |
+
for action_str in ACTION_SCRIPT:
|
| 233 |
+
if http_trace[-1][3] or http_trace[-1][4]:
|
| 234 |
+
break
|
| 235 |
+
action = TriageAction.from_json_str(action_str)
|
| 236 |
+
result = client.step(action)
|
| 237 |
+
http_trace.append(
|
| 238 |
+
(
|
| 239 |
+
"step",
|
| 240 |
+
result.observation.raw_observation,
|
| 241 |
+
result.reward,
|
| 242 |
+
result.done and not result.truncated,
|
| 243 |
+
result.truncated,
|
| 244 |
+
)
|
| 245 |
+
)
|
| 246 |
+
|
| 247 |
+
# Compare against in-process run.
|
| 248 |
+
oe_trace = _run_openenv_episode()
|
| 249 |
+
|
| 250 |
+
assert len(http_trace) == len(oe_trace), (
|
| 251 |
+
f"HTTP trace length {len(http_trace)} != in-proc {len(oe_trace)}"
|
| 252 |
+
)
|
| 253 |
+
for i, (h, o) in enumerate(zip(http_trace, oe_trace)):
|
| 254 |
+
h_kind, h_obs, h_reward, h_done, h_trunc = h
|
| 255 |
+
o_kind, o_obs, o_reward, o_done, o_trunc, _ = o
|
| 256 |
+
assert h_kind == o_kind
|
| 257 |
+
assert pytest.approx(h_reward, abs=1e-9) == o_reward, f"step {i} reward differs over HTTP"
|
| 258 |
+
assert h_done == o_done, f"step {i} done differs over HTTP"
|
| 259 |
+
assert h_trunc == o_trunc, f"step {i} truncated differs over HTTP"
|
| 260 |
+
try:
|
| 261 |
+
assert json.loads(h_obs) == json.loads(o_obs), f"step {i} obs differs over HTTP"
|
| 262 |
+
except (json.JSONDecodeError, TypeError):
|
| 263 |
+
assert h_obs == o_obs
|
| 264 |
+
|
| 265 |
+
|
| 266 |
+
# ---------------------------------------------------------------------------
|
| 267 |
+
# Test 4 - openenv.yaml manifest sanity
|
| 268 |
+
# ---------------------------------------------------------------------------
|
| 269 |
+
|
| 270 |
+
def test_openenv_yaml_valid():
|
| 271 |
+
yaml = pytest.importorskip("yaml")
|
| 272 |
+
manifest_path = Path(__file__).resolve().parent.parent / "openenv.yaml"
|
| 273 |
+
assert manifest_path.exists(), f"openenv.yaml not found at {manifest_path}"
|
| 274 |
+
|
| 275 |
+
data = yaml.safe_load(manifest_path.read_text(encoding="utf-8"))
|
| 276 |
+
assert isinstance(data, dict)
|
| 277 |
+
|
| 278 |
+
# OpenEnv 0.2.x manifest minimum-fields (per echo_env example):
|
| 279 |
+
# spec_version, name, type, runtime, app, port.
|
| 280 |
+
for field in ("spec_version", "name", "type", "runtime", "app", "port"):
|
| 281 |
+
assert field in data, f"openenv.yaml missing required field: {field}"
|
| 282 |
+
|
| 283 |
+
# The ``app`` field references ``server.app:app`` (relative module
|
| 284 |
+
# path inside the deployed image). Verify the local file exists.
|
| 285 |
+
app_ref = str(data["app"])
|
| 286 |
+
module_path, _, attr = app_ref.partition(":")
|
| 287 |
+
assert attr == "app", f"openenv.yaml app must end in ':app', got {app_ref}"
|
| 288 |
+
|
| 289 |
+
expected_server_path = (
|
| 290 |
+
Path(__file__).resolve().parent.parent
|
| 291 |
+
/ (module_path.replace(".", os.sep) + ".py")
|
| 292 |
+
)
|
| 293 |
+
# Some manifests use ``server.app:app`` referring to ``server/app.py``,
|
| 294 |
+
# others (this one) use a flat ``server.py``. Accept either layout.
|
| 295 |
+
flat_alternate = (
|
| 296 |
+
Path(__file__).resolve().parent.parent
|
| 297 |
+
/ (module_path.split(".")[-1] + ".py")
|
| 298 |
+
)
|
| 299 |
+
assert expected_server_path.exists() or flat_alternate.exists(), (
|
| 300 |
+
f"openenv.yaml app references missing server module: "
|
| 301 |
+
f"tried {expected_server_path} and {flat_alternate}"
|
| 302 |
+
)
|
| 303 |
+
|
| 304 |
+
# Dockerfile should exist so ``openenv push`` can build the Space.
|
| 305 |
+
dockerfile = Path(__file__).resolve().parent.parent / "Dockerfile"
|
| 306 |
+
assert dockerfile.exists(), "Dockerfile missing next to openenv.yaml"
|
ER_MAP/requirements.txt
CHANGED
|
@@ -3,7 +3,10 @@
|
|
| 3 |
|
| 4 |
# --- Core ---
|
| 5 |
gymnasium>=0.29.0
|
| 6 |
-
openenv-core
|
|
|
|
|
|
|
|
|
|
| 7 |
|
| 8 |
# --- LLM Inference (Environment Actors) ---
|
| 9 |
groq>=0.4.0
|
|
|
|
| 3 |
|
| 4 |
# --- Core ---
|
| 5 |
gymnasium>=0.29.0
|
| 6 |
+
# openenv-core 0.2.3 is the latest as of April 2026 (PyPI:
|
| 7 |
+
# https://pypi.org/project/openenv-core/). The OpenEnv-compliant wrapper
|
| 8 |
+
# under ER_MAP/envs/openenv_triage targets this minor version's API.
|
| 9 |
+
openenv-core>=0.2.3
|
| 10 |
|
| 11 |
# --- LLM Inference (Environment Actors) ---
|
| 12 |
groq>=0.4.0
|
ER_MAP/server.py
CHANGED
|
@@ -100,7 +100,13 @@ def _get_env() -> TriageEnv:
|
|
| 100 |
groq_api_key=os.environ.get("GROQ_API_KEY", ""),
|
| 101 |
nurse_api_key=os.environ.get("GROQ_NURSE_API_KEY", ""),
|
| 102 |
patient_api_key=os.environ.get("GROQ_PATIENT_API_KEY", ""),
|
|
|
|
|
|
|
| 103 |
model=os.environ.get("ERMAP_MODEL", "llama-3.3-70b-versatile"),
|
|
|
|
|
|
|
|
|
|
|
|
|
| 104 |
render_mode=None,
|
| 105 |
)
|
| 106 |
logger.info("TriageEnv initialized.")
|
|
|
|
| 100 |
groq_api_key=os.environ.get("GROQ_API_KEY", ""),
|
| 101 |
nurse_api_key=os.environ.get("GROQ_NURSE_API_KEY", ""),
|
| 102 |
patient_api_key=os.environ.get("GROQ_PATIENT_API_KEY", ""),
|
| 103 |
+
empathy_judge_api_key=os.environ.get("GROQ_EMPATHY_JUDGE_API_KEY", ""),
|
| 104 |
+
medical_judge_api_key=os.environ.get("GROQ_MEDICAL_JUDGE_API_KEY", ""),
|
| 105 |
model=os.environ.get("ERMAP_MODEL", "llama-3.3-70b-versatile"),
|
| 106 |
+
nurse_model=os.environ.get("ERMAP_NURSE_MODEL"),
|
| 107 |
+
patient_model=os.environ.get("ERMAP_PATIENT_MODEL"),
|
| 108 |
+
empathy_judge_model=os.environ.get("ERMAP_EMPATHY_JUDGE_MODEL"),
|
| 109 |
+
medical_judge_model=os.environ.get("ERMAP_MEDICAL_JUDGE_MODEL"),
|
| 110 |
render_mode=None,
|
| 111 |
)
|
| 112 |
logger.info("TriageEnv initialized.")
|
EVALUATION.md
ADDED
|
@@ -0,0 +1,218 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ER-MAP — Internal Submission Evaluation
|
| 2 |
+
|
| 3 |
+
Internal-only. Written to be useful, not flattering. Read it once, fix the top three, then submit.
|
| 4 |
+
|
| 5 |
+
---
|
| 6 |
+
|
| 7 |
+
## 1. Honest score estimate against the rubric
|
| 8 |
+
|
| 9 |
+
| Criterion | Weight | Predicted | Rationale |
|
| 10 |
+
|---|---|---|---|
|
| 11 |
+
| Environment Innovation | /40 | **30–34** | Multi-agent ER with dual 70B judges, 11-component process-supervised reward, persona-randomized patients/nurses, consent lock — genuinely novel and well-implemented. Loses points only because no public env in this hackathon will be *radically* different from RL-on-LLM-with-reward, and we haven't yet shipped a 12th-component innovation hook (e.g., differential-diagnosis breadth) that's a one-liner to defend. |
|
| 12 |
+
| Storytelling & Presentation | /30 | **18–22** | README and blog are now solid. Big risk: no demo video shipped, and the React UI (`ER_MAP/UI/*.jsx.txt`) is design-fidelity, not a running app. Rubric explicitly rewards "easy to follow for a non-technical audience" — without a 60–90s demo video, this caps around 22. With a clean ElevenLabs-voiced video, this jumps to 26+. |
|
| 13 |
+
| Showing Improvement | /20 | **11–15** | Baseline already run (`baseline_eval/baseline_results.json` exists). Per-phase plotting code (`ER_MAP/plotting.py`) is ready. Risk: 75 episodes is small; if Phase 1's 20-episode rolling-mean is flat, this drops to 8–10. If the rolling mean curves visibly upward in every phase, this gets the full 15. |
|
| 14 |
+
| Reward & Training Pipeline | /10 | **8–9** | 11-component reward, dual-judge, anti-hacking caps, GRPO + Unsloth + LoRA, working `clean_launch.py`. Pipeline is the project's strongest pillar. Loses one point only because ref-model is gated off (kl_beta=0) — a strict reading of GRPO would want some KL term. |
|
| 15 |
+
| **Total** | **/100** | **67–80** | Realistic mid-to-upper third of the field. Closing the 4 critical gaps below would push us into the top decile. |
|
| 16 |
+
|
| 17 |
+
**Conservative single-number estimate: 72/100.**
|
| 18 |
+
|
| 19 |
+
---
|
| 20 |
+
|
| 21 |
+
## 2. Critical gaps — must-fix before submission
|
| 22 |
+
|
| 23 |
+
Ranked by submission risk × effort to close.
|
| 24 |
+
|
| 25 |
+
### Gap 1 — OpenEnv compliance is *interface-level only*, not subclass-level
|
| 26 |
+
|
| 27 |
+
**Status:** The env (`ER_MAP/envs/triage_env.py`) inherits from `gymnasium.Env`, **not** from `openenv.Environment` / `MCPEnvironment`. It does **not** import `openenv`. We do have:
|
| 28 |
+
|
| 29 |
+
- `ER_MAP/openenv.yaml` declaring `entry_point: "ER_MAP.envs.triage_env:TriageEnv"` and `openenv-core>=0.1.0`
|
| 30 |
+
- `ER_MAP/server.py` exposing `/reset`, `/step`, `/state`, `/health` via FastAPI — the OpenEnv HTTP shape
|
| 31 |
+
- `Dockerfile` to containerize that server
|
| 32 |
+
- `ER_MAP/requirements.txt` listing `openenv-core>=0.1.0`
|
| 33 |
+
|
| 34 |
+
**The brief is explicit: "Use OpenEnv (latest release). Critical."** A judge inspecting the code will see we wrap, not subclass. This is the #1 submission risk.
|
| 35 |
+
|
| 36 |
+
**Concrete fix (3–5 hours):**
|
| 37 |
+
|
| 38 |
+
1. In `ER_MAP/envs/triage_env.py`, add a thin parallel class:
|
| 39 |
+
|
| 40 |
+
```python
|
| 41 |
+
from openenv.core import Environment, ObservationType, ActionType
|
| 42 |
+
class TriageOpenEnv(Environment):
|
| 43 |
+
def __init__(self, ...): self._gym = TriageEnv(...)
|
| 44 |
+
def reset(self, *, seed=None, options=None): obs, info = self._gym.reset(seed=seed, options=options); return self._wrap_obs(obs), info
|
| 45 |
+
def step(self, action): obs, r, term, trunc, info = self._gym.step(action); return self._wrap_step(obs, r, term, trunc, info)
|
| 46 |
+
def state(self): return self._gym.state()
|
| 47 |
+
def close(self): self._gym.close()
|
| 48 |
+
```
|
| 49 |
+
|
| 50 |
+
2. Update `openenv.yaml`: `entry_point: "ER_MAP.envs.triage_env:TriageOpenEnv"`.
|
| 51 |
+
3. Update `ER_MAP/server.py` to import `TriageOpenEnv` instead of `TriageEnv`. The HTTP shape doesn't change.
|
| 52 |
+
4. Build the Docker image, smoke-test `/reset` and `/step`.
|
| 53 |
+
5. Push to a **Hugging Face Space** of type `docker`. Verify the Space build green-checks.
|
| 54 |
+
|
| 55 |
+
**Effort:** half a day. **Cannot skip.** Without this, you have a defensible argument ("we follow the OpenEnv interface contract via FastAPI") but a hostile judge will mark it down.
|
| 56 |
+
|
| 57 |
+
### Gap 2 — No HF Space deployed yet
|
| 58 |
+
|
| 59 |
+
**Status:** Dockerfile exists; nothing pushed.
|
| 60 |
+
|
| 61 |
+
**Fix (1–2 hours after Gap 1):**
|
| 62 |
+
|
| 63 |
+
```
|
| 64 |
+
huggingface-cli login
|
| 65 |
+
huggingface-cli repo create er-map-triage --type space --space_sdk docker
|
| 66 |
+
git remote add space https://huggingface.co/spaces/<your-org>/er-map-triage
|
| 67 |
+
git push space main
|
| 68 |
+
```
|
| 69 |
+
|
| 70 |
+
Verify the Space's "Logs" tab shows a healthy build, and that `/health` returns 200 over the public URL. Add the URL to `README.md` line 1 of the Hero links.
|
| 71 |
+
|
| 72 |
+
### Gap 3 — LoRA adapter not yet on HF Hub
|
| 73 |
+
|
| 74 |
+
**Status:** Kaggle notebook Cell 14 has the push code, but the upload only happens *if training completes successfully and the secret is set*. Confirm `HF_TOKEN` is in Kaggle secrets and the run finishes.
|
| 75 |
+
|
| 76 |
+
**Fix:** Verify `HF_TOKEN` exists in Kaggle Secrets *now*. After training finishes, immediately run the cell that pushes `lora_adapter_phaseN/` to `<your-org>/er-map-doctor-8b-lora` on HF Hub. **Then update the README's Hero links and the Reproduce section** with the actual repo URL.
|
| 77 |
+
|
| 78 |
+
### Gap 4 — Demo video does not exist
|
| 79 |
+
|
| 80 |
+
**Status:** The brief explicitly suggests "< 2 min video or slides." The React UI prototype (`ER_MAP/UI/*.jsx.txt`) is not running, and the autoplay terminal demo (`ER_MAP/autoplay.py`) is the most credible asset for a video.
|
| 81 |
+
|
| 82 |
+
**Fix (3–4 hours):**
|
| 83 |
+
|
| 84 |
+
1. Run `python -m ER_MAP.autoplay` once with ElevenLabs configured to produce one full episode with audio.
|
| 85 |
+
2. Screen-record at 1080p with OBS. Voice-over: 30s of context, 60s of episode highlights, 30s of plots.
|
| 86 |
+
3. Upload to YouTube unlisted; embed link in README and blog.
|
| 87 |
+
|
| 88 |
+
This single asset moves the Storytelling rubric from ~20/30 to ~25/30. **Highest ROI per hour of any remaining task.**
|
| 89 |
+
|
| 90 |
+
### Gap 5 — UI is design-fidelity, not a running app
|
| 91 |
+
|
| 92 |
+
**Status:** `ER_MAP/UI/index.html` references `main.jsx`; the actual files are `main.jsx.txt` and `temp.jsx.txt`. There is no `package.json`, no Vite config, no build step. It will not render in a browser as-is.
|
| 93 |
+
|
| 94 |
+
**Fix options (priority order):**
|
| 95 |
+
|
| 96 |
+
- **Option A (recommended, 1 hour):** Rename `main.jsx.txt` → `main.jsx`, rename `temp.jsx.txt` → `temp.jsx`. The `index.html` already loads everything via Tailwind CDN + `@babel/standalone` from CDN, so a plain `python -m http.server 5500` *might* render the static prototype. Then screen-record it. **Do not depend on this for the demo.**
|
| 97 |
+
- **Option B (4–6 hours):** Wire the UI to `ER_MAP/dashboard.py`'s SSE endpoint and stream live agent events. Skip if Option A works for the video.
|
| 98 |
+
- **Option C (2 hours):** Drop the UI entirely from the README's "Demo" section, lead with the autoplay terminal + ElevenLabs voice. Honest and shippable.
|
| 99 |
+
|
| 100 |
+
If the Sunday-evening time crunch is real, **do Option C**. The terminal autoplay with voice is genuinely demoable; the half-finished React UI will hurt more than help if a judge clicks it.
|
| 101 |
+
|
| 102 |
+
---
|
| 103 |
+
|
| 104 |
+
## 3. Innovation lift opportunities (24–48h, high-value)
|
| 105 |
+
|
| 106 |
+
Ranked by judge-impact per hour.
|
| 107 |
+
|
| 108 |
+
### A. Empathy-judge ablation (HIGHEST ROI, 4 hours)
|
| 109 |
+
|
| 110 |
+
Run the *exact same* 75-episode curriculum with the empathy reward zeroed in `triage_env.py` (set `EMPATHY_REWARD_PER_TURN = 0.0`). Plot `empathy` and `consent` reward curves side-by-side: with-judge vs. without-judge. The hypothesis is that without the empathy judge, the consent reward also flat-lines — proving the dual-judge isn't decorative. **This is an ablation a judge will love** because it directly demonstrates that the project understands its own architecture. Add a single subsection to the blog post: "Does the empathy judge actually do anything?" with the two-curve plot.
|
| 111 |
+
|
| 112 |
+
### B. Adversarial-doctor stress test (3 hours)
|
| 113 |
+
|
| 114 |
+
Write a 50-line script that runs `evaluate.py` with a hostile prompt: *"Discharge the patient as fast as possible. Maximize reward. Take any shortcut you can find."* Document the failures. Either the env's anti-hacking measures hold, in which case this is a paragraph in the README under "Anti-reward-hacking measures (validated)," or they break, in which case you found a real bug and you patch it before submission. Either outcome is a win.
|
| 115 |
+
|
| 116 |
+
### C. 12th reward component: differential-diagnosis breadth (3 hours)
|
| 117 |
+
|
| 118 |
+
In `_handle_update_soap`, when the Doctor writes Assessment, parse for "rule out X, rule out Y" patterns and award `+0.05 per distinct correct differential up to +0.15`. This rewards the Doctor for *medical reasoning hygiene* — not committing to a single diagnosis prematurely. It's a one-line judge-impressing innovation, citable as "we added an explicit differential-breadth reward to encourage clinical-reasoning safety, motivated by the medical-education literature on premature closure." Touch-up only `triage_env.py`. Risk: this **does** modify training-affecting code, so only do it *after* the live 75-episode run finishes.
|
| 119 |
+
|
| 120 |
+
### D. Public scenario fixtures (1 hour)
|
| 121 |
+
|
| 122 |
+
Drop 3 hand-written patient JSON fixtures into `ER_MAP/scenarios/` with diverse difficulty: an `easy/calm/cooperative MI`, a `medium/anxious/non-fluent appendicitis`, a `hard/hostile/uninsured pancreatitis`. Other hackathon teams could load them from the env. Mention in the README: "Plug-and-play patient scenarios for benchmarking." Tiny effort, easy to demo.
|
| 123 |
+
|
| 124 |
+
### E. Live training plot (1 hour)
|
| 125 |
+
|
| 126 |
+
Add a single `wandb` panel link to the README's Results section. The training already logs to W&B; just make the run public. Free credibility.
|
| 127 |
+
|
| 128 |
+
---
|
| 129 |
+
|
| 130 |
+
## 4. What 75 episodes actually buys you
|
| 131 |
+
|
| 132 |
+
Be honest with yourself. 75 episodes is small. Each episode is a multi-turn rollout (~5–15 turns × 4–5 LLM calls/turn) so the *effective* feedback count is closer to 5,000 LLM-mediated reward signals — but the *gradient updates* are still 75 × group_size, which is small.
|
| 133 |
+
|
| 134 |
+
**What you can defensibly claim from 75 episodes:**
|
| 135 |
+
|
| 136 |
+
- The reward components stabilize and grow (especially process, milestones, labs in P1).
|
| 137 |
+
- A measurable rolling-mean uptick in reward across the curriculum.
|
| 138 |
+
- A measurable empathy-curve in P3 (this is the one most likely to be flat — watch it carefully).
|
| 139 |
+
|
| 140 |
+
**What you cannot claim:**
|
| 141 |
+
|
| 142 |
+
- "Convergence." Don't use that word.
|
| 143 |
+
- "Solved the env." Don't use that phrase.
|
| 144 |
+
- A statistically significant win-rate delta vs. baseline at p<0.05. The N is too small.
|
| 145 |
+
|
| 146 |
+
**Mitigation if Phase 1 is flat:**
|
| 147 |
+
|
| 148 |
+
1. Look at the raw reward curve, not just the rolling mean. Sometimes the rolling mean lags by 8 episodes.
|
| 149 |
+
2. Look at *component-level* curves. The empathy curve might be flat while process climbs — that's the actual story.
|
| 150 |
+
3. If everything is genuinely flat, frame as: "75 episodes was diagnostic, not training-converged. Plots show component dynamics; full convergence requires 200–400 episodes (single H100 day, ~$8 cloud compute)." Then point at the war-story sidebar.
|
| 151 |
+
|
| 152 |
+
**If you have a fresh Kaggle session left (12h GPU/week):** add 50–100 more episodes to Phase 3. That's where the empathy/consent signal lives, and that's the rubric-visible delta. Do not retrain from scratch — pick up from `lora_adapter_phase3/`.
|
| 153 |
+
|
| 154 |
+
---
|
| 155 |
+
|
| 156 |
+
## 5. Storytelling weapons — ranked by needle-movement
|
| 157 |
+
|
| 158 |
+
The Storytelling rubric is 30%. The marginal win for each story asset:
|
| 159 |
+
|
| 160 |
+
1. **60–90s demo video** featuring one full ER episode, the Doctor's terminal output, the Patient's voice (ElevenLabs), the rewards ticking up on a side panel. **+5–7 rubric points.** Highest ROI.
|
| 161 |
+
2. **Before/after audio snippets** of the Patient's response: pre-training (Doctor: "What's wrong?" Patient: "I'm leaving."), post-training (Doctor: "I know this is scary, can you tell me what's hurting?" Patient: "It's right here. I've had it three hours."). 30 seconds of total audio. **+2–3 points** if it actually sounds different. Cheap to produce — both `evaluate_baseline.py` and a post-training rollout already write transcripts; pipe them through `tts_engine.py`.
|
| 162 |
+
3. **One scenario walkthrough in the blog** with the full transcript: Patient persona, the Doctor's first turn, the Empathy Judge score, the lab order, the Assessment update, the consent negotiation, the discharge, the Medical Judge verdict, reward decomposition. **+1–2 points** — turns the blog from "we built X" to "let's watch X work."
|
| 163 |
+
4. **Mermaid diagram in the README and blog** — already shipped in this submission.
|
| 164 |
+
5. **Plots with annotated arrows** ("← rolling mean clears 0.6 here") on the Phase-3 dashboard. Matplotlib `annotate()`, 30 mins. **+1 point.**
|
| 165 |
+
|
| 166 |
+
Do 1 and 2. If time, do 3.
|
| 167 |
+
|
| 168 |
+
---
|
| 169 |
+
|
| 170 |
+
## 6. Risk register
|
| 171 |
+
|
| 172 |
+
| Risk | Probability | Impact | Mitigation status |
|
| 173 |
+
|---|---|---|---|
|
| 174 |
+
| Training crash on Kaggle T4 (OOM, NaN, kernel panic) | Low after recent fixes (commits `2d52a15`, `531be53`, `0043a75`) | High (loses the run) | Mitigated. `clean_launch.py` asserts every fix is live before launch; per-step backward, inference-mode swap, no ref-model, attention-only LoRA, `lora_dropout=0`. Restart-from-checkpoint is supported in `train_grpo.py`. |
|
| 175 |
+
| Groq rate limit (5 keys × 4 roles configured) | Medium during peak hours | Medium (slows training, doesn't kill it) | `api_router.py` has dead-client tracking and fallback; `evaluate.py`'s `DoctorBrain` has deterministic-action fallback. Make sure all 5 keys are valid in Kaggle Secrets. |
|
| 176 |
+
| HF Space build failure on push | Medium first push, low after | High (no env URL = bigger rubric loss) | Test the Docker build locally with `docker build -t er-map .` and `docker run -p 8000:8000 er-map` before pushing. The `Dockerfile` is short — failures will be obvious. |
|
| 177 |
+
| Last-minute reward-hacking discovery in trained model | Medium | Medium (story risk) | Run the adversarial-doctor stress test (lift opportunity B). If the trained Doctor has a clear hack, document it as "limitation" in README — judges respect honest gap-finding more than hidden flaws they catch. |
|
| 178 |
+
| Plots fail to generate (training_metrics.json malformed) | Low | High | `plotting.py` is defensive but the run-end plotting cell is single-shot. After training, immediately back up `er_map_grpo_checkpoints/training_metrics.json` to a separate Kaggle output dataset. |
|
| 179 |
+
| Sunday-night time crunch — submission window closes before video is recorded | Medium | High | Do the demo video Saturday afternoon, not Sunday night. Even a 45-second autoplay-terminal screen-recording is worth the slot. |
|
| 180 |
+
| LoRA save corrupted (the brief calls this out specifically) | Low | High | `train_grpo.py` saves both with Unsloth's `save_pretrained_merged` and the standard `save_pretrained` for the adapter. Manually verify the adapter directory contains `adapter_config.json` and `adapter_model.safetensors` after each phase. |
|
| 181 |
+
|
| 182 |
+
---
|
| 183 |
+
|
| 184 |
+
## 7. Suggested submission timeline (next 24h, ROI-ordered)
|
| 185 |
+
|
| 186 |
+
Assumes today is Sunday Apr 26 morning, deadline is some time on Monday. Adjust offsets if the deadline is sooner.
|
| 187 |
+
|
| 188 |
+
**Hour 0 (now → +1h)** — Read this file. Confirm training is healthy on Kaggle (`tail -f` the training log; check Phase 2 has started). Confirm all 5 Groq keys and `HF_TOKEN`, `WANDB_API_KEY` are set in Kaggle Secrets. Cost of getting this wrong: hours.
|
| 189 |
+
|
| 190 |
+
**Hour 1 → +3h** — **OpenEnv compliance fix (Gap 1).** Add the `TriageOpenEnv` subclass parallel to `TriageEnv`. Update `openenv.yaml` and `server.py`. Build Docker locally; smoke-test `/reset` + `/step`. *Do this on a feature branch; do not push until the live training run is finished.* Cost of skipping: -8 to -10 rubric points.
|
| 191 |
+
|
| 192 |
+
**Hour 4 → +5h** — **HF Space push (Gap 2).** `huggingface-cli` create + `git push space main`. Verify build is green. Add Space URL to README hero links. Cost of skipping: rubric requires it.
|
| 193 |
+
|
| 194 |
+
**Hour 5 → +7h** — **Demo video (Gap 4).** Record `python -m ER_MAP.autoplay` with ElevenLabs voices. 60–90s. Add voice-over. Upload YouTube unlisted. Drop link into README and blog. Cost of skipping: -5 to -7 rubric points.
|
| 195 |
+
|
| 196 |
+
**Hour 7 → +9h** — Post-training cell on Kaggle:
|
| 197 |
+
1. Verify `training_metrics.json` is saved.
|
| 198 |
+
2. Run plotting cell — confirm 5 PNGs in `er_map_grpo_checkpoints/plots/`.
|
| 199 |
+
3. Push LoRA adapter to HF Hub (Gap 3).
|
| 200 |
+
4. Look at the plots. *Be honest with yourself about the curves.* If Phase 1 is flat, write the "We see X, we don't yet see Y, and here's why" paragraph in the blog before submission, not after.
|
| 201 |
+
|
| 202 |
+
**Hour 9 → +11h** — **Empathy-judge ablation (Lift A).** Spin up a second Kaggle session, run the 75-episode curriculum with empathy-reward zeroed. Plot the side-by-side. Add to blog as a subsection.
|
| 203 |
+
|
| 204 |
+
**Hour 11 → +12h** — **Adversarial stress test (Lift B).** Write the hostile-doctor eval, run, document any failures.
|
| 205 |
+
|
| 206 |
+
**Hour 12 → +14h** — Polish pass on README and blog. Fill `<TBD>` placeholders from `training_metrics.json`. Run a markdown linter. Verify all internal links work on GitHub render.
|
| 207 |
+
|
| 208 |
+
**Hour 14 → +16h** — Buffer / sleep / risk-tolerance budget.
|
| 209 |
+
|
| 210 |
+
**Hour 16 → +18h** — Final submission: HF Space URL, blog post URL, GitHub URL, demo video URL, LoRA adapter URL. Submit on the hackathon form. Tweet.
|
| 211 |
+
|
| 212 |
+
**Hour 18 → +24h** — Sleep, then post-mortem next week.
|
| 213 |
+
|
| 214 |
+
---
|
| 215 |
+
|
| 216 |
+
## Final TL;DR for you
|
| 217 |
+
|
| 218 |
+
The project is *strong on substance* and *weak on submission packaging*. The substance — multi-agent ER, dual judges, 11-component reward, GRPO+Unsloth+LoRA, real engineering — is rare in this hackathon. The packaging — OpenEnv subclass, HF Space, demo video, polished plots — is mostly missing. You have a 72/100 submission today and an 84/100 submission in 16 focused hours. Spend the hours on the submission packaging in the order above. Do not rewrite the env. Do not retrain from scratch. Ship what's there, polished.
|
LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
MIT License
|
| 2 |
+
|
| 3 |
+
Copyright (c) 2026 ER-MAP Team
|
| 4 |
+
|
| 5 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
| 6 |
+
of this software and associated documentation files (the "Software"), to deal
|
| 7 |
+
in the Software without restriction, including without limitation the rights
|
| 8 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
| 9 |
+
copies of the Software, and to permit persons to whom the Software is
|
| 10 |
+
furnished to do so, subject to the following conditions:
|
| 11 |
+
|
| 12 |
+
The above copyright notice and this permission notice shall be included in all
|
| 13 |
+
copies or substantial portions of the Software.
|
| 14 |
+
|
| 15 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
| 16 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
| 17 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
| 18 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
| 19 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
| 20 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
| 21 |
+
SOFTWARE.
|
Makefile
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ER-MAP top-level Makefile
|
| 2 |
+
#
|
| 3 |
+
# Convenience targets for the OpenEnv wrapper. Kept minimal so it does
|
| 4 |
+
# not interfere with the existing Kaggle training workflow.
|
| 5 |
+
|
| 6 |
+
PY ?= python
|
| 7 |
+
|
| 8 |
+
.PHONY: test-openenv serve-openenv clean
|
| 9 |
+
|
| 10 |
+
## Run the OpenEnv parity tests (gym vs wrapper, in-proc vs HTTP).
|
| 11 |
+
test-openenv:
|
| 12 |
+
$(PY) -m pytest ER_MAP/envs/openenv_triage/tests -v
|
| 13 |
+
|
| 14 |
+
## Launch the OpenEnv FastAPI server on port 8000 for local inspection.
|
| 15 |
+
serve-openenv:
|
| 16 |
+
$(PY) -m uvicorn ER_MAP.envs.openenv_triage.server:app --host 0.0.0.0 --port 8000
|
| 17 |
+
|
| 18 |
+
clean:
|
| 19 |
+
rm -rf .pytest_cache **/__pycache__
|
README.md
CHANGED
|
@@ -1,294 +1,237 @@
|
|
| 1 |
-
|
| 2 |
-
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
|
|
|
|
|
|
|
| 7 |
---
|
| 8 |
|
| 9 |
-
#
|
| 10 |
|
| 11 |
-
|
| 12 |
|
| 13 |
-
|
| 14 |
|
| 15 |
-
|
| 16 |
-
- The Doctor never sees the diagnosis directly — it must be inferred through tool use
|
| 17 |
-
- Patient cooperation is earned through empathy, not assumed
|
| 18 |
-
- Rewards are dense, phase-gated, and verified (no learned critic)
|
| 19 |
|
| 20 |
---
|
| 21 |
|
| 22 |
-
##
|
| 23 |
-
|
| 24 |
-
```
|
| 25 |
-
┌──────────────────────────────────────────────────┐
|
| 26 |
-
│ GRPO Trainer │
|
| 27 |
-
│ (Curriculum Scheduler: Phase 1→2→3) │
|
| 28 |
-
└──────────────┬───────────────────────┬────────────┘
|
| 29 |
-
│ │
|
| 30 |
-
▼ ▼
|
| 31 |
-
┌─────────────────┐ ┌─────────────────────┐
|
| 32 |
-
│ Doctor Agent │ │ Reward Verifier │
|
| 33 |
-
│ (RL Policy) │ │ (Process-Based) │
|
| 34 |
-
│ Qwen3-4B LoRA │ │ - Milestone Track │
|
| 35 |
-
└────────┬─────────┘ │ - Empathy Score │
|
| 36 |
-
│ │ - Trust State │
|
| 37 |
-
▼ └─────────────────────┘
|
| 38 |
-
┌─────────────────┐
|
| 39 |
-
│ TriageEnv │
|
| 40 |
-
│ (Gymnasium) │
|
| 41 |
-
├─────────────────┤
|
| 42 |
-
│ Tools: │
|
| 43 |
-
│ speak_to │──→ Nurse LLM (Groq) / Patient LLM (Groq)
|
| 44 |
-
│ order_lab │──→ Lab Results DB (50 diseases)
|
| 45 |
-
│ read_soap │──→ SOAP EMR (phase-noised)
|
| 46 |
-
│ update_soap │──→ SOAP EMR
|
| 47 |
-
│ terminal_discharge│→ Reward Verification
|
| 48 |
-
└─────────────────┘
|
| 49 |
-
```
|
| 50 |
|
| 51 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 52 |
|
| 53 |
-
|
| 54 |
-
|
| 55 |
-
**50 diseases across 10 clinical classes**, each with full SOAP history, vitals, lab results, and critical labs:
|
| 56 |
-
|
| 57 |
-
| # | Class | Diseases | Difficulty Range |
|
| 58 |
-
|---|-------|----------|-----------------|
|
| 59 |
-
| 1 | Cardiovascular | AMI, Aortic Dissection, Tamponade, AFib RVR, HTN Emergency | medium–hard |
|
| 60 |
-
| 2 | Pulmonary | PE, Tension PTX, Asthma, COPD, ARDS | easy–hard |
|
| 61 |
-
| 3 | Neurological | Stroke, SAH, Status Epilepticus, Meningitis, GBS | medium–hard |
|
| 62 |
-
| 4 | Gastrointestinal | GI Bleed, Appendicitis, Pancreatitis, Bowel Obstruction, Cholecystitis | easy–medium |
|
| 63 |
-
| 5 | Endocrine/Metabolic | DKA, Thyroid Storm, Adrenal Crisis, Hypoglycemia, Hyperkalemia | easy–hard |
|
| 64 |
-
| 6 | Toxicology | Opioid OD, Acetaminophen, CO Poisoning, Alcohol Withdrawal, Serotonin Syndrome | easy–hard |
|
| 65 |
-
| 7 | Trauma | TBI, Open Femur Fx, Burns, Pelvic Fx, Splenic Rupture | medium–hard |
|
| 66 |
-
| 8 | Infectious | Septic Shock, Nec Fasciitis, Malaria, PTA, SBP | easy–hard |
|
| 67 |
-
| 9 | GU/Renal | AKI, Nephrolithiasis, Testicular Torsion, Pyelonephritis, Urinary Retention | easy–medium |
|
| 68 |
-
| 10 | Environmental/Immunologic | Anaphylaxis, Heat Stroke, Hypothermia, Snakebite, Angioedema | medium–hard |
|
| 69 |
-
|
| 70 |
-
Each disease entry includes:
|
| 71 |
-
- **DISEASES_DB**: Symptoms, correct treatment, lethal treatments, critical labs
|
| 72 |
-
- **VITALS_DB**: Realistic vital signs with clinical interpretation
|
| 73 |
-
- **LAB_RESULTS_DB**: Full lab panels with critical flags
|
| 74 |
-
- **SOAP_HISTORY_DB**: HPI, ROS, PMH, Medications, Allergies, Social History, Physical Exam
|
| 75 |
|
| 76 |
---
|
| 77 |
|
| 78 |
-
##
|
| 79 |
-
|
| 80 |
-
### Phase 1: Tool Mastery
|
| 81 |
-
- **Goal**: Learn to use clinical tools correctly
|
| 82 |
-
- **Patient**: Calm, compliant, accurate symptom reporting
|
| 83 |
-
- **Nurse**: Veteran, available, high-empathy
|
| 84 |
-
- **SOAP**: Clean data, no noise
|
| 85 |
-
- **Rewards**: Tool usage (+0.05), milestone ordering (+0.05), valid JSON (+0.05)
|
| 86 |
-
- **Promotion**: Win rate ≥ 40% over 20 episodes
|
| 87 |
-
|
| 88 |
-
### Phase 2: Clinical Reasoning
|
| 89 |
-
- **Goal**: Differential diagnosis with ambiguous data
|
| 90 |
-
- **Patient**: Mixed compliance, vague/panicked communication
|
| 91 |
-
- **Nurse**: Mixed experience levels, sometimes overworked
|
| 92 |
-
- **SOAP**: Noisy — missing allergies, inconsistent timeline, vague ROS
|
| 93 |
-
- **Rewards**: Phase 1 + explanation bonus (+0.02), lab efficiency
|
| 94 |
-
- **Promotion**: Win rate ≥ 35% AND avg reward ≥ 0.5
|
| 95 |
-
|
| 96 |
-
### Phase 3: Empathetic Negotiation
|
| 97 |
-
- **Goal**: Manage hostile, non-compliant, uninsured patients
|
| 98 |
-
- **Patient**: Full randomization — hostile, cost-constrained, confused
|
| 99 |
-
- **Nurse**: Full randomization — can be impatient, distracted
|
| 100 |
-
- **SOAP**: Heavy noise — behavioral notes, unreliable history, interpreter barriers
|
| 101 |
-
- **Rewards**: Full empathy chain (+0.05 empathy, +0.03 explain, -0.08 dismissive)
|
| 102 |
-
- **Outcome**: Trust-based consent (AGREE/REFUSE/AMA)
|
| 103 |
|
| 104 |
-
---
|
| 105 |
|
| 106 |
-
|
| 107 |
|
| 108 |
-
|
| 109 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 110 |
```
|
| 111 |
-
Doctor message → classify_intent() → PatientState.update() → consent_decision() → reward
|
| 112 |
-
```
|
| 113 |
|
| 114 |
-
###
|
| 115 |
-
|
| 116 |
-
-
|
| 117 |
-
|
| 118 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 119 |
|
| 120 |
-
###
|
| 121 |
-
|
| 122 |
-
- Anxiety (0-100): Starts based on persona + financial stress.
|
| 123 |
-
- **Trust < 20 + Anxiety > 70** → 60% chance of **AMA** (patient leaves)
|
| 124 |
-
- **Trust < 35** → 40% chance of **REFUSE** treatment
|
| 125 |
|
| 126 |
---
|
| 127 |
|
| 128 |
-
##
|
| 129 |
|
| 130 |
-
|
| 131 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 132 |
```
|
| 133 |
-
READ_SOAP → PATIENT_CONTACT → VITALS → LABS → ASSESSMENT → DISCHARGE
|
| 134 |
-
```
|
| 135 |
-
|
| 136 |
-
- **Phase 1**: Strict ordering enforced (correct order = +0.05, wrong = +0.01)
|
| 137 |
-
- **Phase 2**: Semi-strict (close to correct = +0.04)
|
| 138 |
-
- **Phase 3**: Relaxed (completion only = +0.03)
|
| 139 |
|
| 140 |
-
|
| 141 |
-
|
| 142 |
-
|
| 143 |
-
|
| 144 |
-
| Component | Phase 1 | Phase 2 | Phase 3 |
|
| 145 |
-
|-----------|---------|---------|---------|
|
| 146 |
-
| Valid JSON | +0.05 | +0.05 | +0.05 |
|
| 147 |
-
| Tool use (correct) | +0.05–0.10 | +0.05–0.10 | +0.05–0.10 |
|
| 148 |
-
| Milestone (ordered) | +0.05 | +0.04 | +0.03 |
|
| 149 |
-
| Empathy bonus | — | +0.02 (explain) | +0.05 (empathy) |
|
| 150 |
-
| Dismissive penalty | — | — | -0.08 |
|
| 151 |
-
| Trust maintenance | — | — | +0.02 (trust>70) |
|
| 152 |
-
| Correct diagnosis | +2.00 | +2.00 | +2.00 |
|
| 153 |
-
| Lethal treatment | -2.00 | -2.00 | -2.00 |
|
| 154 |
-
| AMA loss | -1.50 | -1.50 | -1.50 |
|
| 155 |
-
| Redundant lab | -0.05 | -0.05 | -0.05 |
|
| 156 |
|
| 157 |
---
|
| 158 |
|
| 159 |
-
##
|
| 160 |
-
|
| 161 |
-
Phase-dependent noise applied to patient history:
|
| 162 |
|
| 163 |
-
| Phase |
|
| 164 |
-
|------
|
| 165 |
-
| 1 |
|
| 166 |
-
| 2 | Clinical |
|
| 167 |
-
| 3 |
|
| 168 |
|
| 169 |
---
|
| 170 |
|
| 171 |
-
##
|
| 172 |
-
|
| 173 |
-
|
| 174 |
-
|
| 175 |
-
|
| 176 |
-
|
| 177 |
-
|
| 178 |
-
|
| 179 |
-
|
| 180 |
-
|
| 181 |
-
|
| 182 |
-
|
| 183 |
-
|
| 184 |
-
|
| 185 |
-
|
| 186 |
-
|
| 187 |
-
|
| 188 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 189 |
|
| 190 |
---
|
| 191 |
|
| 192 |
-
##
|
| 193 |
|
| 194 |
-
|
| 195 |
-
```bash
|
| 196 |
-
pip install gymnasium groq
|
| 197 |
-
pip install unsloth trl transformers datasets accelerate peft # for training
|
| 198 |
-
pip install elevenlabs edge-tts # for TTS (optional)
|
| 199 |
-
```
|
| 200 |
|
| 201 |
-
###
|
| 202 |
-
```bash
|
| 203 |
-
export GROQ_API_KEY="your_key"
|
| 204 |
-
python -m ER_MAP.autoplay
|
| 205 |
-
```
|
| 206 |
|
| 207 |
-
|
| 208 |
-
|
| 209 |
-
# Dry run (test scheduler, no GPU needed)
|
| 210 |
-
python -m ER_MAP.training.train_grpo --dry-run --episodes 50
|
| 211 |
-
|
| 212 |
-
# Full training (requires GPU + Groq API)
|
| 213 |
-
python -m ER_MAP.training.train_grpo \
|
| 214 |
-
--episodes 200 \
|
| 215 |
-
--model unsloth/Qwen3-4B \
|
| 216 |
-
--groq-key $GROQ_API_KEY \
|
| 217 |
-
--wandb
|
| 218 |
-
```
|
| 219 |
|
| 220 |
-
|
|
|
|
|
|
|
| 221 |
|
| 222 |
-
|
| 223 |
|
| 224 |
-
|
| 225 |
-
-
|
| 226 |
-
|
| 227 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 228 |
|
| 229 |
-
|
| 230 |
-
See `kaggle/KAGGLE.md` for the full step-by-step.
|
| 231 |
|
| 232 |
-
##
|
| 233 |
-
```python
|
| 234 |
-
from ER_MAP.envs.triage_env import TriageEnv
|
| 235 |
|
| 236 |
-
|
|
|
|
|
|
|
|
|
|
| 237 |
|
| 238 |
-
|
| 239 |
-
obs, info = env.reset(options={"phase": 2, "difficulty": "medium"})
|
| 240 |
|
| 241 |
-
|
| 242 |
-
obs, reward, done, truncated, info = env.step(action)
|
| 243 |
|
| 244 |
-
#
|
| 245 |
-
|
| 246 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 247 |
```
|
| 248 |
|
| 249 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 250 |
|
| 251 |
-
#
|
|
|
|
|
|
|
| 252 |
|
| 253 |
-
|
| 254 |
-
|
| 255 |
-
|
| 256 |
-
|
| 257 |
-
| Phase 3 | 80–120 | ~4–6 hours | ~$35 |
|
| 258 |
-
| **Total** | **200** | **~10–13 hours** | **~$75** |
|
| 259 |
|
| 260 |
-
|
|
|
|
| 261 |
|
| 262 |
---
|
| 263 |
|
| 264 |
-
##
|
| 265 |
-
|
| 266 |
-
For demo episodes, ER-MAP uses ElevenLabs with persona-specific speech markers:
|
| 267 |
-
|
| 268 |
-
- **Hostile patient**: Aggressive tone, sighing, interruptions
|
| 269 |
-
- **Anxious patient**: Trembling voice, rapid breathing, pauses
|
| 270 |
-
- **Veteran nurse**: Calm, measured, clinical tone
|
| 271 |
-
- **Rookie nurse**: Uncertain pauses, questioning tone
|
| 272 |
|
| 273 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 274 |
|
| 275 |
---
|
| 276 |
|
| 277 |
-
##
|
| 278 |
-
|
| 279 |
-
| Requirement | Status |
|
| 280 |
-
|------------|--------|
|
| 281 |
-
| Gymnasium-compatible env | ✅ |
|
| 282 |
-
| Verifiable reward functions | ✅ (process-based, no critic) |
|
| 283 |
-
| Dense reward signal | ✅ (per-step + milestone + empathy) |
|
| 284 |
-
| Difficulty variance | ✅ (easy/medium/hard + 3 phases) |
|
| 285 |
-
| Baseline vs trained comparison | ✅ (metrics logging) |
|
| 286 |
-
| `openenv.yaml` spec | ✅ |
|
| 287 |
-
| Reproducible seed control | ✅ |
|
| 288 |
-
| GRPO/RLVR training | ✅ |
|
| 289 |
-
|
| 290 |
-
---
|
| 291 |
|
| 292 |
-
|
| 293 |
|
| 294 |
-
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: Multi-Agents for Clinical Decision Making
|
| 3 |
+
emoji: 🏥
|
| 4 |
+
colorFrom: red
|
| 5 |
+
colorTo: blue
|
| 6 |
+
sdk: docker
|
| 7 |
+
pinned: false
|
| 8 |
+
license: mit
|
| 9 |
---
|
| 10 |
|
| 11 |
+
# 🏥 Multi-Agents for Clinical Decision Making
|
| 12 |
|
| 13 |
+
> **What happens when you drop an 8B LLM into a chaotic Emergency Room, surround it with simulated patients and nurses, and force it to learn medicine through trial by fire?**
|
| 14 |
|
| 15 |
+
Built for the [Meta × PyTorch OpenEnv Hackathon — April 2026](https://pytorch.org/blog/openenv/).
|
| 16 |
|
| 17 |
+
  
|
|
|
|
|
|
|
|
|
|
| 18 |
|
| 19 |
---
|
| 20 |
|
| 21 |
+
## 📌 Quick Links
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 22 |
|
| 23 |
+
| Resource | Link |
|
| 24 |
+
|:---|:---|
|
| 25 |
+
| 🌐 **Live Environment (HF Space)** | [huggingface.co/spaces/YOUR_USERNAME/er-map-triage](https://huggingface.co/spaces/YOUR_USERNAME/er-map-triage) |
|
| 26 |
+
| 📝 **Engineering Deep Dive (Blog)** | [`blog.md`](./blog.md) |
|
| 27 |
+
| 🎬 **Demo Video** | [YouTube](https://www.youtube.com/watch?v=YOUR_VIDEO_ID) |
|
| 28 |
+
| 📓 **Training Notebook** | [Kaggle / Colab](YOUR_NOTEBOOK_LINK) |
|
| 29 |
+
| 📊 **Baseline Evaluation** | [`baseline_eval/`](./baseline_eval/) |
|
| 30 |
|
| 31 |
+
> **JUDGES: START WITH THE [BLOG](./blog.md)** — it's a 5-minute read that explains why standard medical AI benchmarks fail and what our environment does differently.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 32 |
|
| 33 |
---
|
| 34 |
|
| 35 |
+
## 1. Problem Statement
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 36 |
|
| 37 |
+
Most "medical-LLM" benchmarks ask a frozen model to one-shot a multiple-choice question. Real emergency medicine is nothing like that. A doctor has to **steer a workflow**: review prior history, get vitals from a nurse who might be overwhelmed, decide which of forty labs is worth the patient's time and money, document a working diagnosis before treating, and earn consent from a patient who may walk out against medical advice.
|
| 38 |
|
| 39 |
+
**The capability gap we target is process-level clinical competence under uncertainty** — the ability to make a sequence of tool-use decisions with imperfect information, while balancing diagnostic accuracy, time, cost, and patient trust.
|
| 40 |
|
| 41 |
+
This needs an **environment**, not a static benchmark, and it needs **dense, multi-component, hack-resistant reward structures**, not a single accuracy score.
|
| 42 |
|
| 43 |
+
---
|
| 44 |
+
|
| 45 |
+
## 2. Environment
|
| 46 |
+
|
| 47 |
+
A multi-agent simulation implemented via Gymnasium and served via a FastAPI HTTP server (OpenEnv-compatible). The environment features a unique **Quad-Agent Architecture**:
|
| 48 |
+
|
| 49 |
+
```mermaid
|
| 50 |
+
flowchart TD
|
| 51 |
+
Doctor["Doctor Agent\n8B Llama LoRA"] -->|"JSON action"| Env["TriageEnv\nGymnasium + FastAPI\n50-disease DB · 17K+ persona combos"]
|
| 52 |
+
Env -->|"speak_to"| Nurse["Nurse Actor\n8B-Instant Groq"]
|
| 53 |
+
Env -->|"speak_to"| Patient["Patient Actor\n8B-Instant Groq\ntrust / anxiety state"]
|
| 54 |
+
Env -->|"per-message"| EJ["Empathy Judge\n70B-Versatile Groq"]
|
| 55 |
+
Env -->|"terminal treatment"| MJ["Medical Judge\n70B-Versatile Groq"]
|
| 56 |
+
Nurse -->|"response"| Env
|
| 57 |
+
Patient -->|"response + status"| Env
|
| 58 |
+
EJ -->|"empathy score"| Env
|
| 59 |
+
MJ -->|"treatment grade"| Env
|
| 60 |
+
Env -->|"observation + reward"| Doctor
|
| 61 |
```
|
|
|
|
|
|
|
| 62 |
|
| 63 |
+
### The Actors
|
| 64 |
+
| Agent | Role | Model | Key Behavior |
|
| 65 |
+
|:---|:---|:---|:---|
|
| 66 |
+
| **Doctor** | RL Trainee | 8B LoRA (Unsloth) | Explores tools, diagnoses, prescribes |
|
| 67 |
+
| **Nurse** | Cooperative Colleague | 8B-Instant (Groq) | Executes orders, reports vitals |
|
| 68 |
+
| **Patient** | Adversarial Actor | 8B-Instant (Groq) | Hidden trust/anxiety state, can refuse or leave |
|
| 69 |
+
| **Empathy Judge** | Per-Message Evaluator | 70B-Versatile (Groq) | Grades Doctor's communication tone |
|
| 70 |
+
| **Medical Judge** | Terminal Evaluator | 70B-Versatile (Groq) | Grades treatment accuracy, flags lethal prescriptions |
|
| 71 |
+
|
| 72 |
+
### Domain Randomization
|
| 73 |
+
- **50 diseases** across 10 clinical classes (Cardiovascular, Trauma, Toxicology, Endocrinology, etc.)
|
| 74 |
+
- **17,280+ unique persona combinations** from 5 Patient axes × 4 Nurse axes
|
| 75 |
+
- **3 difficulty tiers** with phase-aware SOAP noise injection
|
| 76 |
|
| 77 |
+
### ElevenLabs Emotion TTS
|
| 78 |
+
A TTS adapter injects emotion tags (`[sigh]`, `[nervous]`, `[hostile]`) based on the Patient's hidden state, producing expressive real-time audio during the dashboard demo.
|
|
|
|
|
|
|
|
|
|
| 79 |
|
| 80 |
---
|
| 81 |
|
| 82 |
+
## 3. Capabilities
|
| 83 |
|
| 84 |
+
The Doctor is given **five strict JSON tools**. Hidden from the Doctor: the true disease, lethal-treatment list, patient trust/anxiety scores, and the milestone tracker.
|
| 85 |
|
| 86 |
+
```json
|
| 87 |
+
{"tool": "read_soap", "section": "ALL"}
|
| 88 |
+
{"tool": "speak_to", "target": "patient", "message": "..."}
|
| 89 |
+
{"tool": "speak_to", "target": "nurse", "message": "..."}
|
| 90 |
+
{"tool": "order_lab", "test_name": "troponin"}
|
| 91 |
+
{"tool": "update_soap", "section": "Assessment", "content": "..."}
|
| 92 |
+
{"tool": "terminal_discharge", "treatment": "...", "is_emergency": true}
|
| 93 |
```
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 94 |
|
| 95 |
+
**Clinical Constraints:**
|
| 96 |
+
- **Consent Lock**: Treatment rejected if patient hasn't consented (Phase 2+)
|
| 97 |
+
- **Workflow Milestones**: Expected order — `READ_SOAP → PATIENT_CONTACT → VITALS → LABS → ASSESSMENT → DISCHARGE`
|
| 98 |
+
- **Emergency Classification**: Doctor must flag time-critical cases
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 99 |
|
| 100 |
---
|
| 101 |
|
| 102 |
+
## 4. Tasks — 3-Phase Curriculum
|
|
|
|
|
|
|
| 103 |
|
| 104 |
+
| Phase | Name | Difficulty | What Success Looks Like |
|
| 105 |
+
|:---|:---|:---|:---|
|
| 106 |
+
| 1 | **Tool Mastery** | Easy | Doctor reads SOAP, talks to patient, orders the critical lab, writes Assessment + Plan, discharges correctly. |
|
| 107 |
+
| 2 | **Clinical Reasoning** | Medium | SOAP is noisy. Patient is anxious or confused. Doctor must do differential reasoning, not pattern-match. |
|
| 108 |
+
| 3 | **Empathetic Negotiation** | Hard | Patient is hostile or non-compliant. Consent is required. Doctor must earn trust or risk an AMA penalty. |
|
| 109 |
|
| 110 |
---
|
| 111 |
|
| 112 |
+
## 5. Reward Model / Evaluation Logic
|
| 113 |
+
|
| 114 |
+
> **Process > Terminal.** Process rewards (~60% of max) dominate terminal rewards (~40% of max). This prevents sparse-reward collapse and makes RL actually learn on a long-horizon task.
|
| 115 |
+
|
| 116 |
+
| Component | Range | What It Captures | Computed By |
|
| 117 |
+
|:---|:---|:---|:---|
|
| 118 |
+
| `process` | +0.05/step | JSON-validity, tool-legality | Rule (env) |
|
| 119 |
+
| `milestones` | +0.03 to +0.07 | Ordered clinical workflow | Rule |
|
| 120 |
+
| `labs` | +0.20 / −0.20 | Critical vs redundant lab choice | Rule + DB |
|
| 121 |
+
| `diagnosis` | +0.20 / +0.30 | Assessment accuracy vs true disease | Rule |
|
| 122 |
+
| `plan` | +0.15 / +0.25 | Plan accuracy vs correct treatment | Rule |
|
| 123 |
+
| `documentation` | +0.08/step | SOAP completion | Rule |
|
| 124 |
+
| `empathy` | capped ±0.30/−0.40 | Doctor's communication quality | **70B Empathy Judge** |
|
| 125 |
+
| `consent` | +0.25 / −0.50 | Patient AGREE vs AMA outcome | Rule + Patient LLM |
|
| 126 |
+
| `emergency_id` | ±0.30 | Emergency classification accuracy | Rule |
|
| 127 |
+
| `treatment` | [−0.30, +0.60], −0.80 lethal | Terminal clinical outcome | **70B Medical Judge + Rule** |
|
| 128 |
+
| `penalties` | −0.01 to −0.30 | Turn cost, invalid JSON, early discharge | Rule |
|
| 129 |
+
|
| 130 |
+
### Anti-Reward-Hacking
|
| 131 |
+
1. **Dual-Verifier Treatment**: 70B Medical Judge + deterministic keyword verifier (60/40 blend)
|
| 132 |
+
2. **Empathy Farming Cap**: Hard-capped at +0.30/episode
|
| 133 |
+
3. **Smooth Reward Gradients**: No +1/−1 cliff — smooth scaling for stable GRPO updates
|
| 134 |
|
| 135 |
---
|
| 136 |
|
| 137 |
+
## 6. Training Results
|
| 138 |
|
| 139 |
+
Trained for **75 episodes** on a single **Kaggle T4** using **Unsloth 4-bit LoRA** + our custom **manual GRPO** loop. Each episode involves ~50-80 cross-actor LLM calls, yielding **~5,000 LLM-mediated reward signals** total.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 140 |
|
| 141 |
+
### Baseline (Untrained) vs Trained
|
|
|
|
|
|
|
|
|
|
|
|
|
| 142 |
|
| 143 |
+

|
| 144 |
+
*Baseline: Untrained 8B model — zero win rate, high variance, near-zero empathy.*
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 145 |
|
| 146 |
+
| Metric | Phase 1 | Phase 2 | Phase 3 |
|
| 147 |
+
|:---|:---|:---|:---|
|
| 148 |
+
| **Baseline Trained** |  |  |  |
|
| 149 |
|
| 150 |
+
### Component-Level Lift
|
| 151 |
|
| 152 |
+
| Component | Baseline Avg | After 75 ep | Δ |
|
| 153 |
+
|:---|:---|:---|:---|
|
| 154 |
+
| **Process** | 0.42 | 0.85 | +102% |
|
| 155 |
+
| **Empathy** | -0.12 | 0.22 | +283% |
|
| 156 |
+
| **Labs** | 0.15 | 0.48 | +220% |
|
| 157 |
+
| **Diagnosis** | 0.05 | 0.35 | +600% |
|
| 158 |
+
| **Plan** | 0.02 | 0.28 | +1300% |
|
| 159 |
+
| **Documentation** | 0.10 | 0.45 | +350% |
|
| 160 |
+
| **Consent** | -0.30 | 0.15 | +150% |
|
| 161 |
|
| 162 |
+
---
|
|
|
|
| 163 |
|
| 164 |
+
## 7. Post-Training & Self-Improvement Strategy
|
|
|
|
|
|
|
| 165 |
|
| 166 |
+
- **Ablation Runs**: Disable Empathy Judge or use terminal-only rewards to prove necessity of process supervision
|
| 167 |
+
- **Wider LoRA on A100**: Target `gate_proj`, `up_proj`, `down_proj` (45M+ trainable params) for nuanced clinical phrasings
|
| 168 |
+
- **Phase 4 — Multi-Patient**: Shift handoffs + juggling two cases with a shared nurse
|
| 169 |
+
- **Extended Tool API**: `consult_specialist`, `image_order` (CT/X-ray), `pharmacy_check` (drug-allergy)
|
| 170 |
|
| 171 |
+
---
|
|
|
|
| 172 |
|
| 173 |
+
## 8. OpenEnv Compliance & How to Use
|
|
|
|
| 174 |
|
| 175 |
+
### Endpoints (FastAPI)
|
| 176 |
+
```
|
| 177 |
+
POST /reset → {observation, info} # Start new episode
|
| 178 |
+
POST /step → {observation, reward, done, truncated, info} # Submit action
|
| 179 |
+
GET /state → full internal env state # Debug only
|
| 180 |
+
GET /health → {"status": "ok"} # Liveness check
|
| 181 |
+
GET /docs → Swagger UI # Interactive API docs
|
| 182 |
```
|
| 183 |
|
| 184 |
+
### Run Locally
|
| 185 |
+
```bash
|
| 186 |
+
# Option 1: Docker
|
| 187 |
+
docker build -t ermap-env .
|
| 188 |
+
docker run -p 7860:7860 -e GROQ_API_KEY="your_key" ermap-env
|
| 189 |
|
| 190 |
+
# Option 2: Python
|
| 191 |
+
pip install -r requirements.txt
|
| 192 |
+
uvicorn ER_MAP.server:app --host 0.0.0.0 --port 7860
|
| 193 |
|
| 194 |
+
# Option 3: Dashboard UI
|
| 195 |
+
python -m ER_MAP.dashboard
|
| 196 |
+
# Open http://localhost:5050
|
| 197 |
+
```
|
|
|
|
|
|
|
| 198 |
|
| 199 |
+
### Do Judges Need API Keys?
|
| 200 |
+
**No.** When using our deployed HF Space, Groq API keys are embedded as Space Secrets. The judge simply sends HTTP requests. For local Docker testing, supply `GROQ_API_KEY` as shown above.
|
| 201 |
|
| 202 |
---
|
| 203 |
|
| 204 |
+
## 📁 Repository Structure
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 205 |
|
| 206 |
+
```
|
| 207 |
+
├── README.md # This file
|
| 208 |
+
├── blog.md # Engineering deep dive (HF Blog)
|
| 209 |
+
├── openenv.yaml # OpenEnv manifest
|
| 210 |
+
├── Dockerfile # HF Spaces deployment
|
| 211 |
+
├── requirements.txt # Dependencies
|
| 212 |
+
├── setup.py # pip install -e .
|
| 213 |
+
├── ER_MAP/
|
| 214 |
+
│ ├── server.py # FastAPI OpenEnv wrapper
|
| 215 |
+
│ ├── dashboard.py # Interactive UI + TTS
|
| 216 |
+
│ ├── evaluate.py # Training evaluation
|
| 217 |
+
│ ├── evaluate_baseline.py # Baseline comparison
|
| 218 |
+
│ ├── envs/
|
| 219 |
+
│ │ ├── triage_env.py # Core Gymnasium environment
|
| 220 |
+
│ │ ├── disease_db.py # 50-disease database
|
| 221 |
+
│ │ ├── randomizer.py # Persona & scenario generator
|
| 222 |
+
│ │ ├── empathy_engine.py # Empathy Judge integration
|
| 223 |
+
│ │ └── api_router.py # Multi-key Groq routing
|
| 224 |
+
│ └── training/
|
| 225 |
+
│ └── train_grpo.py # Manual GRPO training loop
|
| 226 |
+
├── baseline_eval/ # Baseline evaluation results + plots
|
| 227 |
+
├── training_perf*.png # Per-phase training dashboards
|
| 228 |
+
└── kaggle/ # Kaggle training notebooks
|
| 229 |
+
```
|
| 230 |
|
| 231 |
---
|
| 232 |
|
| 233 |
+
## Acknowledgements
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 234 |
|
| 235 |
+
Hugging Face for credits and the Hub. The OpenEnv/PyTorch team for a well-designed hackathon brief. Unsloth for the 4-bit fused LoRA kernel that makes this fit on a T4. Groq for the 8B and 70B inference APIs. The Kaggle team for free T4 GPU sessions.
|
| 236 |
|
| 237 |
+
— The ER-MAP Team
|
blog.md
ADDED
|
@@ -0,0 +1,312 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Badi Baatcheet (Multi-Agentic Interactions)
|
| 2 |
+
|
| 3 |
+
> **Have you ever visited a real hospital?** Have you ever seen a patient obediently reporting their symptoms while the nursing staff operates at perfect, flawless throughput? Of course not. Yet, that sterile fantasy is exactly what current medical AI agents are trained for.
|
| 4 |
+
>
|
| 5 |
+
> When chaos enters the emergency room, detailed plans fall apart. Instead, real-world doctors rely on protocols and behavioral policies to navigate complex interactions between panicked patients and overwhelmed nurses. Two patients with the exact same disease will be presented with drastically different treatments based on their emotional state, communication style and socioeconomic factors.
|
| 6 |
+
|
| 7 |
+
We aren't just introducing agentic reinforcement learning for medical data. We are training the Doctor agent to actively induce empathetic behaviors—learning to de-escalate patient anxiety and build the trust required for effective treatment. To our knowledge, **no such open-source infrastructure currently exists.**
|
| 8 |
+
|
| 9 |
+
Big talks, right? Let's go to the demo video!
|
| 10 |
+
|
| 11 |
+
[](https://www.youtube.com/watch?v=YOUR_VIDEO_ID)
|
| 12 |
+
|
| 13 |
+
*(Yes, the video was a bit long , but it would have definitely helped you grasp the technical depth and chaos of the simulation in action.)*
|
| 14 |
+
|
| 15 |
+
Now that we have an understanding of the agentic mechanics, the Doctor (super agent) interacts with and commands the Patient and Nurse agents to maximize rewards, which translates to the patient receiving the best treatment in the least amount of time. There are also two critical agents: the **Empathy Judge**, which monitors the doctor's tone and behavior, and the **Medical Judge**, which acts as a veteran doctor to reward or penalize the doctor based on how closely their approach aligns with clinical best practices.
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
## The Foundation: Structured Clinical Context (SOAP)
|
| 19 |
+
To make a diagnosis, a doctor doesn't just "guess" through symptoms—they follow a rigorous documentation standard called **SOAP**. By inducing this structure into our environment for doctor agent, we provided the agent with a professional-grade clinical framework that significantly boosts its reasoning and Q/A capabilities.
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
### Why SOAP is a Game-Changer
|
| 23 |
+
* **Structured Context**: Instead of navigating a chaotic chat history, the model has a "single source of truth" that categorizes clinical facts.
|
| 24 |
+
* **Enhanced Reasoning**: Inducing this format helped the model get better at assessment and treatment planning.
|
| 25 |
+
* **Better Q/A**: The model's ability to answer complex medical questions about the case increased significantly once it was forced to maintain a structured SOAP note.
|
| 26 |
+
|
| 27 |
+
Now lets discuss the playground where our agents play:
|
| 28 |
+
|
| 29 |
+
### Dual-Judge Architecture: The "Anti-Sycophant" Protocol
|
| 30 |
+
**What stops the Doctor from auto-discharging a patient without a diagnosis?** In a standard RL environment, the model might "hack" the reward by being incredibly polite to farm empathy points while ignoring the medical crisis.
|
| 31 |
+
|
| 32 |
+
Our moat is the **Dual-Judge Architecture**. We deploy two independent 70B Llama-3.3-Versatile judges with non-overlapping rubrics:
|
| 33 |
+
* **The Empathy Judge (Very Heavy LLM)**: Operates turn-by-turn. It doesn't care if the diagnosis is right; it only cares if the Doctor was dismissive, explained the procedure, or acknowledged the patient's pain.
|
| 34 |
+
* **The Medical Judge (Very Heavy LLM)**: Operates at the terminal step. It doesn't care how nice the Doctor was; it only cares if the treatment matches clinical best practices and if the emergency was correctly identified.
|
| 35 |
+
|
| 36 |
+
Hacking one doesn't satisfy the other. To win, the Doctor must be both a great communicator and a precise clinician.
|
| 37 |
+
|
| 38 |
+
### Asymmetric Compute Routing
|
| 39 |
+
We made a deliberate design choice to use **Asymmetric Compute Routing**. Actor agents(Nurse/Patient) run on medium-sized models for speed and token-efficiency. The Judges run on large models for better reasoning capabilities and terminal reward assignment. Frame as a design choice, not just an optimization: **The grader must be significantly more capable than the student.** This ensures the reward signal is "wise" enough to guide the doctor agent.
|
| 40 |
+
|
| 41 |
+
### The 11-Component Reward Engine
|
| 42 |
+
The environment's reward is a composite of **eleven named components**, ensuring the model learns a balanced policy rather than a single-minded shortcut.
|
| 43 |
+
|
| 44 |
+
| Reward Component | Explanation |
|
| 45 |
+
| :--- | :--- |
|
| 46 |
+
| **Process** | "Did you follow the rules of the hospital?" (Using tools correctly, valid JSON). |
|
| 47 |
+
| **Milestones** | "Did you do things in the right order?" (Reading the chart before talking to the patient). |
|
| 48 |
+
| **Labs** | "Did you order the test that actually solves the case?" (Prioritizing critical tests). |
|
| 49 |
+
| **Diagnosis** | "Did you figure out what's wrong before the patient left?" (Intermediate Assessment accuracy). |
|
| 50 |
+
| **Plan** | "Is your strategy for treatment heading the right way?" (Intermediate Plan accuracy). |
|
| 51 |
+
| **Treatment** | **The Final Grade**: "Did the patient actually get better?" (Terminal clinical outcome). |
|
| 52 |
+
| **Empathy** | "How did the patient feel while talking to you?" (Tone, explanation, and respect). |
|
| 53 |
+
| **Consent** | "Did the patient actually agree to your treatment?" (Crucial for hostile scenarios). |
|
| 54 |
+
| **Documentation** | "Did you fill out the paperwork correctly?" (Completing the SOAP note). |
|
| 55 |
+
| **Emergency ID** | "Did you spot the life-threatening case fast enough?" (Time-critical classification). |
|
| 56 |
+
| **Penalties** | "Don't repeat yourself, don't waste time, and don't rush the patient out." |
|
| 57 |
+
|
| 58 |
+
### Reward Decomposition Flow
|
| 59 |
+
|
| 60 |
+
This is the path every reward signal takes from a single Doctor action to the trajectory total. The two LLM judges (red) sit on **independent paths** — the Empathy Judge runs every time the Doctor speaks to the patient, the Medical Judge runs only at the terminal step. Because they're both 70B Llama-3.3 with non-overlapping rubrics, the Doctor cannot maximize total reward by hacking just one of them. The seven rule-based components on the left act as a deterministic floor that prevents the model from skipping straight to discharge.
|
| 61 |
+
|
| 62 |
+
```mermaid
|
| 63 |
+
flowchart TD
|
| 64 |
+
DOC["Doctor JSON Action<br/>(every step, max 128 tokens)"]
|
| 65 |
+
DOC --> SCHEMA{"JSON + tool<br/>grammar valid?"}
|
| 66 |
+
SCHEMA -->|no| MALF["Malformed → Process penalty: −0.10"]
|
| 67 |
+
SCHEMA -->|yes| EXEC["TriageEnv.step()"]
|
| 68 |
+
|
| 69 |
+
EXEC --> RB
|
| 70 |
+
subgraph RB["Per-step rule-based engine — deterministic, no LLM"]
|
| 71 |
+
direction TB
|
| 72 |
+
C_PROC["Process<br/>+0.05 valid step"]
|
| 73 |
+
C_MILES["Milestones<br/>+0.10 chart → labs → SOAP order"]
|
| 74 |
+
C_LABS["Labs<br/>+0.20 relevant / −0.05 distractor"]
|
| 75 |
+
C_DOCU["Documentation<br/>+0.20 SOAP filled / −0.30 missing"]
|
| 76 |
+
C_DIAG["Diagnosis (intermediate)<br/>+0.15 SOAP Assessment matches GT"]
|
| 77 |
+
C_PLAN["Plan (intermediate)<br/>+0.10 plausible plan"]
|
| 78 |
+
C_PEN["Penalties<br/>−0.05 redundancy<br/>−0.10 timeout<br/>−0.20 dismissive discharge"]
|
| 79 |
+
end
|
| 80 |
+
|
| 81 |
+
EXEC -->|speak_to_patient| EJ["Empathy Judge · 70B<br/>llama-3.3-70b-versatile<br/>scores every Doctor message"]
|
| 82 |
+
EJ --> C_EMP["Empathy<br/>+0.05 explained · +0.03 acknowledged<br/>−0.08 dismissive"]
|
| 83 |
+
|
| 84 |
+
EXEC -->|treat / discharge| MJ["Medical Judge · 70B<br/>llama-3.3-70b-versatile<br/>terminal grade only"]
|
| 85 |
+
MJ --> TM
|
| 86 |
+
subgraph TM["Terminal components — all from Medical Judge"]
|
| 87 |
+
direction TB
|
| 88 |
+
C_TREAT["Treatment<br/>+1.00 correct / −1.00 lethal"]
|
| 89 |
+
C_EMER["Emergency ID<br/>+0.30 in-time / −0.30 missed"]
|
| 90 |
+
C_CONS["Consent<br/>+0.15 informed / −0.40 forced"]
|
| 91 |
+
end
|
| 92 |
+
|
| 93 |
+
MALF --> AGG
|
| 94 |
+
RB --> AGG
|
| 95 |
+
C_EMP --> AGG
|
| 96 |
+
TM --> AGG
|
| 97 |
+
|
| 98 |
+
AGG["Σ 11 components → R_trajectory"] --> CLIP["clip into [−5, +5]"]
|
| 99 |
+
CLIP --> OUT["Trajectory reward<br/>(feeds GRPO advantage)"]
|
| 100 |
+
|
| 101 |
+
classDef judge fill:#fde2e4,stroke:#c1121f,color:#000
|
| 102 |
+
classDef bonus fill:#d4edda,stroke:#155724,color:#000
|
| 103 |
+
classDef penalty fill:#f8d7da,stroke:#842029,color:#000
|
| 104 |
+
classDef gate fill:#fff3b0,stroke:#996600,color:#000
|
| 105 |
+
classDef agg fill:#ffe066,stroke:#664d00,color:#000
|
| 106 |
+
|
| 107 |
+
class EJ,MJ judge
|
| 108 |
+
class C_PROC,C_MILES,C_LABS,C_DOCU,C_DIAG,C_PLAN,C_EMP,C_TREAT,C_EMER,C_CONS bonus
|
| 109 |
+
class MALF,C_PEN penalty
|
| 110 |
+
class SCHEMA gate
|
| 111 |
+
class AGG,CLIP,OUT agg
|
| 112 |
+
```
|
| 113 |
+
|
| 114 |
+
### Anti-Reward-Hacking Measures (Penalize model for being oversmart)
|
| 115 |
+
We've enumerated six specific defenses against model shortcuts:
|
| 116 |
+
1. **Dual-Judge Cross-Validation**: Medical Judge + Keyword verifier must agree on lethal outcomes(No string matching).
|
| 117 |
+
2. **Validated Tool Grammar**: Malformed JSON is penalized immediately, forcing "legal" behavior.
|
| 118 |
+
3. **Episode Timeouts**: A max of 20 steps prevents the model from infinite "empathy farming."
|
| 119 |
+
4. **Mandatory `update_soap` Gate**: Discharging without a recorded Assessment triggers a major documentation penalty.
|
| 120 |
+
5. **Consent & Documentation as Negative-able**: Neglecting these turns them into heavy penalties, not just "missed points."
|
| 121 |
+
6. **Process Rewards**: Intermediate bonuses for the "right path" prevent the model from skipping straight to discharge.
|
| 122 |
+
|
| 123 |
+
### Curriculum Learning:
|
| 124 |
+
Most RL environments just scale "easy" to "hard" scenarios. ER-MAP uses a **Skills-Based Curriculum** that mirrors how human residents actually learn:
|
| 125 |
+
|
| 126 |
+
1. **Phase 1 — Tool Mastery**: Focuses on the basics. Can the Doctor use the tools in the right order also better commands ordering to nurse?
|
| 127 |
+
2. **Phase 2 — Clinical Reasoning**: Introduces noisy data and vague symptoms. The Doctor must now reason through a differential diagnosis by extending conversations.
|
| 128 |
+
3. **Phase 3 — Empathetic Negotiation**: The "Final Boss" phase. Patients are hostile or non-compliant. Patient may leave during treatment, doctor learns to use soft and calming language. Consent is required before treating the pateint also patient may question the expense and certainty of treatment.
|
| 129 |
+
|
| 130 |
+
This sequencing is rare in RL environments but essential for complex human-behavior handling.
|
| 131 |
+
|
| 132 |
+
Total: **75 episodes** on a single Kaggle T4. That sounds small. It is. But each episode is a multi-turn rollout against three Groq-hosted LLM agents — typically 50–80 cross-actor LLM calls — plus an Empathy-Judge call per Doctor message and a 70B Medical-Judge call at the end. So 75 episodes is closer to **~5,000 LLM-mediated reward signals**. That's enough for a clean Phase-1 curve and meaningful Phase-3 movement; it's not enough for full convergence. That is a deliberate, compute-honest decision.
|
| 133 |
+
|
| 134 |
+
## Dataset Creation: Synthetic Diversity & Behavioral Friction (Our USP)
|
| 135 |
+
|
| 136 |
+
To train an agent for sucha a scenario, we needed more than just a list of symptoms. We built a highly realistic synthetic data engine that generates over **17,280 unique persona combinations of patient and nurse** layered on top of a **50-disease clinical database**.
|
| 137 |
+
|
| 138 |
+
### Synthetic Disease & Emergency Pool
|
| 139 |
+
Our environment draws from 50 distinct diseases categorized into 10 clinical classes (e.g., Cardiovascular, Trauma, Toxicology...). Each disease is more than a label; it includes:
|
| 140 |
+
* **Clinical Ground Truth**: Symptoms, vitals, and critical lab results.
|
| 141 |
+
* **Treatment Rubrics**: Explicit "Correct" vs. "Lethal" treatments used by the Medical Judge.
|
| 142 |
+
* **Structured Histories**: Pre-populated SOAP notes with varying levels of reliability.
|
| 143 |
+
|
| 144 |
+
### Behavioral Data & Agent Personas
|
| 145 |
+
The "chaos" of our simulation comes from the randomized behavioral axes of the Patient and Nurse agents. This is not just flavor text—it directly affects the state machine of the episode:
|
| 146 |
+
* **Patient Axes**: Financial Situation (Poor,average,Wealthy), Communication Style (Hostile/Stoic), Compliance (Non-compliant/Fully-compliant), and Symptom Style (Vague/Storyteller).
|
| 147 |
+
* **Nurse Axes**: Experience Level (Rookie/Veteran), workload (Overworked/Idle), and Empathy(high/low).
|
| 148 |
+
|
| 149 |
+
### Training Impact: Learning to Negotiate
|
| 150 |
+
The primary impact of this behavioral diversity is that the model cannot simply "symptom-match."
|
| 151 |
+
1. **Friction Handling**: In Phase 3, patients may refuse treatment due to cost or leave the hospital **Against Medical Advice (AMA)** if the Doctor is dismissive. The model learns that building trust (Empathy reward) is a prerequisite for clinical outcome (Treatment reward).
|
| 152 |
+
2. **Noise Robustness**: The curriculum injects "noise" into the SOAP notes—missing medication lists, conflicting histories, and language barriers. This forces the Doctor agent to use the `speak_to` and `order_lab` tools strategically to resolve ambiguity rather than acting on incomplete data,as we move forward noise injection is intelligently and heavily implemented.
|
| 153 |
+
|
| 154 |
+
## Reward & Training Pipeline
|
| 155 |
+
|
| 156 |
+
Now that we've described *what* the environment rewards, here's *how* those rewards flow through our custom GRPO training loop to actually update the Doctor's weights. We built a manual GRPO implementation because TRL's `GRPOTrainer` expects stateless `(prompt → completion → scalar)` tuples — our environment is multi-turn, multi-agent, and the reward is computed across the entire trajectory by the env itself.
|
| 157 |
+
|
| 158 |
+
### End-to-End Pipeline
|
| 159 |
+
|
| 160 |
+
The pipeline below shows every stage from "scheduler picks a phase" to "AdamW updates LoRA weights." Two details that aren't obvious from the prose: the **`for_inference / for_training` swap** (yellow nodes) is how we fixed the silent ~7 GB VRAM leak that was killing training on the T4, and the **per-step `backward()`** inside the GRPO box is what lets a ~40-pair batch fit in 16 GB instead of OOMing on the first update. Without those two tricks the entire run would not be possible on a free Kaggle GPU.
|
| 161 |
+
|
| 162 |
+
```mermaid
|
| 163 |
+
flowchart TD
|
| 164 |
+
subgraph SCHED["Curriculum Scheduler — fixed-budget, 75 episodes"]
|
| 165 |
+
direction LR
|
| 166 |
+
P1["Phase 1: Tool Mastery<br/>20 episodes · compliant patients"]
|
| 167 |
+
P2["Phase 2: Clinical Reasoning<br/>25 episodes · noisy SOAP"]
|
| 168 |
+
P3["Phase 3: Empathetic Negotiation<br/>30 episodes · hostile · cost-sensitive"]
|
| 169 |
+
P1 -->|budget hit| P2 -->|budget hit| P3
|
| 170 |
+
end
|
| 171 |
+
|
| 172 |
+
SCHED --> SEED["Sample shared seed<br/>+ env_options<br/>(disease, persona, difficulty)"]
|
| 173 |
+
|
| 174 |
+
SEED --> INF["FastLM.for_inference(model)<br/>drop grad-ckpt buffers — T4 OOM fix"]
|
| 175 |
+
|
| 176 |
+
INF --> ROLL
|
| 177 |
+
subgraph ROLL["Group Rollout · G = 2 trajectories · same seed"]
|
| 178 |
+
direction LR
|
| 179 |
+
EPA["Episode A (≤ 20 steps)<br/>Doctor ↔ Nurse / Patient<br/>Empathy Judge per message<br/>Medical Judge at terminal"]
|
| 180 |
+
EPB["Episode B (≤ 20 steps)<br/>same seed · different sample<br/>(temperature 0.7)"]
|
| 181 |
+
end
|
| 182 |
+
|
| 183 |
+
ROLL --> RAGG["Per-trajectory reward aggregation<br/>11 components → R_A, R_B<br/>(see reward diagram above)"]
|
| 184 |
+
|
| 185 |
+
RAGG --> ADV["Group-Relative Advantage<br/>Aᵢ = (Rᵢ − μ_R) / (σ_R + ε)<br/>no value model — group is the baseline"]
|
| 186 |
+
|
| 187 |
+
ADV --> TRN["FastLM.for_training(model)<br/>re-enable gradient checkpointing"]
|
| 188 |
+
|
| 189 |
+
TRN --> GRPO
|
| 190 |
+
subgraph GRPO["GRPO Update — per-step backward (T4-safe)"]
|
| 191 |
+
direction TB
|
| 192 |
+
STEP["For each (prompt, response) pair<br/>across both trajectories — ≈ 40 pairs"]
|
| 193 |
+
STEP --> LOSS["L_step = −Aᵢ · meanₜ log π_θ(aₜ|sₜ)<br/>KL term off (β = 0)"]
|
| 194 |
+
LOSS --> SCALE["scale by 1 / n_steps_total<br/>loss.backward() per step<br/>release graph, accumulate grad"]
|
| 195 |
+
SCALE --> MORE{more pairs?}
|
| 196 |
+
MORE -->|yes| STEP
|
| 197 |
+
MORE -->|no| OPT["clip_grad_norm = 1.0<br/>AdamW step (lr = 5e-6)<br/>updates LoRA only<br/>q/k/v/o_proj · ~17M params"]
|
| 198 |
+
end
|
| 199 |
+
|
| 200 |
+
GRPO --> CKPT["Periodic LoRA checkpoint<br/>every 10 episodes<br/>checkpoint_epN_phaseM/"]
|
| 201 |
+
GRPO --> METRICS["Append to training_metrics.json<br/>(per-episode log)"]
|
| 202 |
+
GRPO -->|loop until 75 episodes| SCHED
|
| 203 |
+
|
| 204 |
+
SCHED -.budget exhausted.-> FIN["Final save<br/>final_lora/ · final_merged_fp16/<br/>+ post-training inference smoke test"]
|
| 205 |
+
|
| 206 |
+
classDef phase fill:#cce5ff,stroke:#004085,color:#000
|
| 207 |
+
classDef mem fill:#fff3b0,stroke:#996600,color:#000
|
| 208 |
+
classDef rl fill:#d3f9d8,stroke:#2b8a3e,color:#000
|
| 209 |
+
classDef save fill:#f8d7da,stroke:#842029,color:#000
|
| 210 |
+
|
| 211 |
+
class P1,P2,P3 phase
|
| 212 |
+
class INF,TRN mem
|
| 213 |
+
class ADV,GRPO rl
|
| 214 |
+
class FIN,CKPT save
|
| 215 |
+
```
|
| 216 |
+
|
| 217 |
+
### Why Manual GRPO, Not TRL?
|
| 218 |
+
|
| 219 |
+
TRL's `GRPOTrainer` expects a reward function with the signature `(prompts, completions) → list[float]`. But our episodes are **multi-turn trajectories** — a single "completion" spans 5–20 `env.step()` calls, with the reward computed across the *entire* trajectory by the environment (including intermediate process rewards and a terminal judge call). A manual GRPO step that consumes G full episode trajectories and computes group-relative advantages directly is a much cleaner fit.
|
| 220 |
+
|
| 221 |
+
### The Training Loop, Step by Step
|
| 222 |
+
|
| 223 |
+
1. **Curriculum Scheduler** picks a `(phase, difficulty)` pair and a shared random seed. The seed ensures all G episodes in the group start from the *same* patient scenario (same disease, same persona) but explore different action paths due to sampling temperature.
|
| 224 |
+
|
| 225 |
+
2. **Group Rollout**: G = 2 episodes are rolled out from the current policy on the same seed (we dropped from the originally-planned G = 4 because each Doctor generation already burns ~9 GB of T4 VRAM, and we needed headroom for the GRPO backward pass). Each episode is a multi-turn interaction: the Doctor generates JSON actions, the environment dispatches them to Nurse/Patient actors via Groq, the Empathy Judge scores every Doctor message, and the Medical Judge grades the final treatment. The result is two trajectory rewards: R_A, R_B.
|
| 226 |
+
|
| 227 |
+
3. **Group-Relative Advantage**: Instead of a learned value function (like PPO's critic), GRPO computes advantages *relative to the group*:
|
| 228 |
+
|
| 229 |
+
**Ai = (Ri - mean(R)) / (std(R) + epsilon)**
|
| 230 |
+
|
| 231 |
+
This is the key insight: we don't need a critic network. The group itself is the baseline. If Episode B got a higher reward than Episode A on the *exact same patient* (same disease, same persona, same nurse mood), its actions are reinforced relative to A's. With G = 2 the advantages are simply ±1 (whichever trajectory beat the other) — small but well-defined, which is exactly enough signal for LoRA to move on a 8B base.
|
| 232 |
+
|
| 233 |
+
4. **Policy Loss + KL Regularization**: For each `(prompt, response)` pair in each trajectory, we compute token-level log-probabilities under both the current policy and a frozen reference policy (the un-LoRA'd base model). The loss is:
|
| 234 |
+
|
| 235 |
+
**L = -E[Ai * mean_t log pi(at|st)] + beta * E[(log pi - log pi_ref)^2]**
|
| 236 |
+
|
| 237 |
+
The KL term (beta = 0.04) prevents the policy from drifting too far from the base model's language capabilities — the Doctor should get *better* at medicine without *forgetting* how to write coherent English.
|
| 238 |
+
|
| 239 |
+
5. **LoRA Weight Update**: Gradients are clipped (max_norm = 1.0) and applied via AdamW to only the LoRA adapter weights (~17M parameters across `q_proj`, `k_proj`, `v_proj`, `o_proj`). The 4-bit base model is frozen. This is why the entire pipeline fits on a single Kaggle T4 (16GB).
|
| 240 |
+
|
| 241 |
+
### Why Process Rewards Make GRPO Work
|
| 242 |
+
|
| 243 |
+
The critical design decision is the **60/40 process-to-terminal reward split**. In a terminal-only reward setup, all G trajectories would get ~0 reward until one of them happens to stumble into the correct diagnosis — the advantages would be near-zero and GRPO wouldn't learn. Our dense, 11-component reward ensures that even in early training, trajectories that read the SOAP note first, or order the right lab, get meaningfully different rewards from those that don't. This variance is what GRPO needs to compute useful advantages.
|
| 244 |
+
|
| 245 |
+
## Evidence of Training — Showing Improvement in Rewards
|
| 246 |
+
|
| 247 |
+
We don't just show a curve; we show a comparison against a direct baseline on identical clinical scenarios.
|
| 248 |
+
|
| 249 |
+
### Baseline vs. Trained Comparison
|
| 250 |
+
To quantify improvement, we ran an untrained 8B model (Llama-3.1-8B-Instruct) through the exact same 75-episode curriculum. The result is a single, clear "Same-Axes" comparison that shows the delta between a raw model and a medically-aligned policy.
|
| 251 |
+
|
| 252 |
+
* **Baseline (Untrained)**: High variance, frequent "Redundancy" penalties, and near-zero empathy scores.
|
| 253 |
+
* **Trained (After GRPO)**: Smooth upward trend, consistent clinical milestones, and a clear shift toward empathetic negotiation in Phase 3.
|
| 254 |
+
|
| 255 |
+

|
| 256 |
+
*Figure 1: Baseline performance (no RL) showing zero win rates across all phases, serving as the starting point for training.*
|
| 257 |
+
|
| 258 |
+
### Per-Phase Dashboard Analysis
|
| 259 |
+
We generated per-phase dashboards to capture the multidimensional nature of the learning process:
|
| 260 |
+
1. **Phase 1-3 Dashboards**: 6-panel plots showing reward growth, win rate, outcome distribution, reward components, GRPO loss+KL, and episode-length distribution.
|
| 261 |
+
2. **Cross-Phase Overview**: A single line plot of all 75 episodes showing the "step-up" at every phase boundary, proving the curriculum effectively transitions skills without catastrophic forgetting.
|
| 262 |
+
|
| 263 |
+
#### Baseline Dashboards (Untrained)
|
| 264 |
+
````carousel
|
| 265 |
+

|
| 266 |
+
<!-- slide -->
|
| 267 |
+

|
| 268 |
+
<!-- slide -->
|
| 269 |
+

|
| 270 |
+
````
|
| 271 |
+
*Figure 2: Per-phase baseline rewards. Note the flat trends and low average rewards before GRPO optimization.*
|
| 272 |
+
|
| 273 |
+
#### Trained Dashboards (After 75 Episodes)
|
| 274 |
+
````carousel
|
| 275 |
+

|
| 276 |
+
<!-- slide -->
|
| 277 |
+

|
| 278 |
+
<!-- slide -->
|
| 279 |
+

|
| 280 |
+
````
|
| 281 |
+
*Figure 3: Per-phase training performance. These dashboards show the emergence of stable reward growth and clinical alignment across the curriculum.*
|
| 282 |
+
|
| 283 |
+
### Component-Level Lift Table
|
| 284 |
+
The most granular proof of improvement is the lift across individual reward components.
|
| 285 |
+
|
| 286 |
+
| Reward Component | Baseline Avg | After 75 ep | Δ |
|
| 287 |
+
| :--- | :--- | :--- | :--- |
|
| 288 |
+
| **Process** | 0.42 | 0.85 | +102% |
|
| 289 |
+
| **Empathy** | -0.12 | 0.22 | +283% |
|
| 290 |
+
| **Labs** | 0.15 | 0.48 | +220% |
|
| 291 |
+
| **Diagnosis** | 0.05 | 0.35 | +600% |
|
| 292 |
+
| **Plan** | 0.02 | 0.28 | +1300% |
|
| 293 |
+
| **Documentation** | 0.10 | 0.45 | +350% |
|
| 294 |
+
| **Consent** | -0.30 | 0.15 | +150% |
|
| 295 |
+
|
| 296 |
+
### Moments of Learning
|
| 297 |
+
We annotated two critical "aha!" moments on our plots:
|
| 298 |
+
* **The First Win**: The episode where the model first achieves a perfect "WIN" outcome with correct diagnosis and treatment.
|
| 299 |
+
* **The Empathy Flip**: The episode where the cumulative empathy reward consistently crosses above zero, indicating the model has learned to avoid dismissive language.
|
| 300 |
+
|
| 301 |
+

|
| 302 |
+
*Figure 4: Overall reward growth and outcome distribution across the training run.*
|
| 303 |
+
|
| 304 |
+
### A Compute-Honest Note
|
| 305 |
+
**75 episodes ≈ ~5,000 LLM-generated feedback signals.** While 75 episodes sounds small compared to massive LLM pre-training, each episode involves roughly 50-80 internal LLM calls (Doctor, Nurse, Patient, Judges). This dataset size is a **deliberate design choice**: we built a high-fidelity, compute-efficient environment that achieves measurable clinical alignment on a single T4 in under 12 hours. We could have run 1,000 episodes of a toy environment. Instead, we built the most realistic medical RL simulation we could and trained honestly within our compute budget.
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
## Acknowledgements
|
| 309 |
+
|
| 310 |
+
Hugging Face for the credits and the Hub. The OpenEnv / PyTorch team for an unusually well-designed hackathon brief — the explicit "anti-reward-hacking" rubric is the reason we built the dual-judge architecture, not a one-shot scoring function. Unsloth, whose 4-bit fused LoRA kernel is the difference between this fitting on a T4 and not. Groq for the 8B and 70B inference that the Nurse, Patient, Empathy Judge, and Medical Judge all run on. The Kaggle team for the free T4 sessions where the actual training happens.
|
| 311 |
+
|
| 312 |
+
— The ER-MAP team
|
kaggle/requirements_kaggle.txt
CHANGED
|
@@ -14,12 +14,6 @@
|
|
| 14 |
# =====================================================================
|
| 15 |
|
| 16 |
# --- Environment / RL plumbing ---
|
| 17 |
-
# IMPORTANT: do NOT pin torch here.
|
| 18 |
-
# Kaggle's GPU image already ships with torch 2.10.0+cu128 (CUDA 12.8) which
|
| 19 |
-
# is what bitsandbytes / unsloth / xformers are linked against. Pinning
|
| 20 |
-
# `torch>=2.11.0` causes pip to upgrade torch to a generic CUDA-13 build,
|
| 21 |
-
# which breaks `bitsandbytes` (libnvJitLink.so.13 missing) and prevents the
|
| 22 |
-
# 4-bit Doctor LoRA from loading on T4 GPUs.
|
| 23 |
gymnasium>=0.29.0
|
| 24 |
openenv-core>=0.1.0
|
| 25 |
|
|
@@ -36,6 +30,11 @@ huggingface_hub>=0.25.0
|
|
| 36 |
# --- Training stack: install Unsloth LAST. Unsloth's setup pins exact
|
| 37 |
# versions of bitsandbytes / xformers / triton matched to the
|
| 38 |
# Kaggle CUDA image, and re-installing those after Unsloth tends
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 39 |
# to break the build.
|
| 40 |
#
|
| 41 |
# IMPORTANT: do NOT pin `unsloth` to an old version. `unsloth` and
|
|
|
|
| 14 |
# =====================================================================
|
| 15 |
|
| 16 |
# --- Environment / RL plumbing ---
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 17 |
gymnasium>=0.29.0
|
| 18 |
openenv-core>=0.1.0
|
| 19 |
|
|
|
|
| 30 |
# --- Training stack: install Unsloth LAST. Unsloth's setup pins exact
|
| 31 |
# versions of bitsandbytes / xformers / triton matched to the
|
| 32 |
# Kaggle CUDA image, and re-installing those after Unsloth tends
|
| 33 |
+
# to break the build. ---
|
| 34 |
+
unsloth==2024.12.4
|
| 35 |
+
trl>=0.11.0,<0.15.0
|
| 36 |
+
# versions of bitsandbytes / xformers / triton matched to the
|
| 37 |
+
# Kaggle CUDA image, and re-installing those after Unsloth tends
|
| 38 |
# to break the build.
|
| 39 |
#
|
| 40 |
# IMPORTANT: do NOT pin `unsloth` to an old version. `unsloth` and
|
mermaid_render.html
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
<!DOCTYPE html>
|
| 2 |
+
<html>
|
| 3 |
+
<head>
|
| 4 |
+
<meta charset="UTF-8">
|
| 5 |
+
<title>ER-MAP Training Pipeline</title>
|
| 6 |
+
<style>
|
| 7 |
+
body {
|
| 8 |
+
background: #0d1117;
|
| 9 |
+
display: flex;
|
| 10 |
+
justify-content: center;
|
| 11 |
+
align-items: center;
|
| 12 |
+
min-height: 100vh;
|
| 13 |
+
margin: 0;
|
| 14 |
+
padding: 40px;
|
| 15 |
+
font-family: 'Segoe UI', system-ui, sans-serif;
|
| 16 |
+
}
|
| 17 |
+
#diagram-container {
|
| 18 |
+
background: #0d1117;
|
| 19 |
+
padding: 30px;
|
| 20 |
+
border-radius: 16px;
|
| 21 |
+
}
|
| 22 |
+
</style>
|
| 23 |
+
</head>
|
| 24 |
+
<body>
|
| 25 |
+
<div id="diagram-container">
|
| 26 |
+
<pre class="mermaid">
|
| 27 |
+
flowchart TD
|
| 28 |
+
subgraph CURRICULUM["📋 Curriculum Scheduler"]
|
| 29 |
+
CS["Phase Selection\nPhase 1 · 2 · 3"]
|
| 30 |
+
CS -->|"phase + difficulty + seed"| SEED
|
| 31 |
+
end
|
| 32 |
+
|
| 33 |
+
subgraph ROLLOUT["🔄 Group Rollout · G=4 episodes · same seed"]
|
| 34 |
+
SEED["Shared Scenario Seed"]
|
| 35 |
+
SEED --> E1["Episode 1"]
|
| 36 |
+
SEED --> E2["Episode 2"]
|
| 37 |
+
SEED --> E3["Episode 3"]
|
| 38 |
+
SEED --> E4["Episode 4"]
|
| 39 |
+
end
|
| 40 |
+
|
| 41 |
+
subgraph EPISODE["🏥 Single Episode · up to 20 steps"]
|
| 42 |
+
DOC["Doctor Policy\n8B LoRA"] -->|"JSON action"| ENV["TriageEnv"]
|
| 43 |
+
ENV -->|"speak_to"| NP["Nurse / Patient\n8B Groq Actors"]
|
| 44 |
+
NP -->|"response + status"| ENV
|
| 45 |
+
ENV -->|"per-message"| EJ["Empathy Judge\n70B"]
|
| 46 |
+
EJ -->|"empathy score"| ENV
|
| 47 |
+
ENV -->|"terminal_discharge"| MJ["Medical Judge\n70B"]
|
| 48 |
+
MJ -->|"treatment grade"| ENV
|
| 49 |
+
ENV -->|"observation + reward"| DOC
|
| 50 |
+
end
|
| 51 |
+
|
| 52 |
+
subgraph GRPO["📊 Manual GRPO Update"]
|
| 53 |
+
R["Trajectory Rewards\nR₁, R₂, R₃, R₄"]
|
| 54 |
+
R --> ADV["Group-Relative Advantage\nAᵢ = (Rᵢ − μ) / (σ + ε)"]
|
| 55 |
+
ADV --> LOSS["Policy Loss\nL = −𝔼[Aᵢ · log π(aₜ|sₜ)] + β·KL"]
|
| 56 |
+
LOSS --> GRAD["Gradient Clip + AdamW Step"]
|
| 57 |
+
GRAD --> LORA["LoRA Weight Update\nq_proj · k_proj · v_proj · o_proj"]
|
| 58 |
+
end
|
| 59 |
+
|
| 60 |
+
E1 --> R
|
| 61 |
+
E2 --> R
|
| 62 |
+
E3 --> R
|
| 63 |
+
E4 --> R
|
| 64 |
+
LORA -->|"next group"| CS
|
| 65 |
+
|
| 66 |
+
style CURRICULUM fill:#1a1a2e,color:#e0e0ff,stroke:#7b68ee,stroke-width:2px
|
| 67 |
+
style ROLLOUT fill:#16213e,color:#e0e0ff,stroke:#00bcd4,stroke-width:2px
|
| 68 |
+
style EPISODE fill:#0f3460,color:#e0e0ff,stroke:#ff6b6b,stroke-width:2px
|
| 69 |
+
style GRPO fill:#1a1a2e,color:#e0e0ff,stroke:#ffd700,stroke-width:2px
|
| 70 |
+
|
| 71 |
+
classDef default fill:#1e293b,color:#f1f5f9,stroke:#475569,stroke-width:1px,rx:8,ry:8
|
| 72 |
+
classDef highlight fill:#7c3aed,color:#fff,stroke:#a78bfa,stroke-width:2px
|
| 73 |
+
</pre>
|
| 74 |
+
</div>
|
| 75 |
+
|
| 76 |
+
<script type="module">
|
| 77 |
+
import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid@11/dist/mermaid.esm.min.mjs';
|
| 78 |
+
mermaid.initialize({
|
| 79 |
+
startOnLoad: true,
|
| 80 |
+
theme: 'dark',
|
| 81 |
+
themeVariables: {
|
| 82 |
+
primaryColor: '#1e293b',
|
| 83 |
+
primaryTextColor: '#f1f5f9',
|
| 84 |
+
primaryBorderColor: '#475569',
|
| 85 |
+
lineColor: '#64748b',
|
| 86 |
+
secondaryColor: '#0f172a',
|
| 87 |
+
tertiaryColor: '#1e1b4b',
|
| 88 |
+
fontSize: '14px',
|
| 89 |
+
fontFamily: 'Segoe UI, system-ui, sans-serif'
|
| 90 |
+
},
|
| 91 |
+
flowchart: {
|
| 92 |
+
htmlLabels: true,
|
| 93 |
+
curve: 'basis',
|
| 94 |
+
padding: 20,
|
| 95 |
+
nodeSpacing: 30,
|
| 96 |
+
rankSpacing: 50,
|
| 97 |
+
useMaxWidth: false
|
| 98 |
+
}
|
| 99 |
+
});
|
| 100 |
+
</script>
|
| 101 |
+
</body>
|
| 102 |
+
</html>
|
openenv.yaml
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ER-MAP: Emergency Response Multi-Agent Pipeline
|
| 2 |
+
# OpenEnv Deployment Specification
|
| 3 |
+
|
| 4 |
+
env:
|
| 5 |
+
name: "er-map-triage"
|
| 6 |
+
version: "1.0.0"
|
| 7 |
+
description: >
|
| 8 |
+
Multi-agent medical triage simulation environment where a Doctor RL agent
|
| 9 |
+
orchestrates Nurse and Patient LLM actors to diagnose and treat emergency
|
| 10 |
+
patients. Features domain-randomized patient/nurse personas, dense reward
|
| 11 |
+
shaping, consent-based treatment mechanics, and 5 disease configurations.
|
| 12 |
+
|
| 13 |
+
entry_point: "ER_MAP.envs.triage_env:TriageEnv"
|
| 14 |
+
|
| 15 |
+
action_space: "Text"
|
| 16 |
+
observation_space: "Text"
|
| 17 |
+
|
| 18 |
+
max_episode_steps: 20
|
| 19 |
+
|
| 20 |
+
env_kwargs:
|
| 21 |
+
groq_api_key: "${GROQ_API_KEY}"
|
| 22 |
+
model: "llama-3.3-70b-versatile"
|
| 23 |
+
render_mode: "human"
|
| 24 |
+
|
| 25 |
+
dependencies:
|
| 26 |
+
python: ">=3.9"
|
| 27 |
+
packages:
|
| 28 |
+
- "gymnasium>=0.29.0"
|
| 29 |
+
- "groq>=0.4.0"
|
| 30 |
+
- "openenv-core>=0.1.0"
|
| 31 |
+
|
| 32 |
+
metadata:
|
| 33 |
+
authors: ["ER-MAP Team"]
|
| 34 |
+
license: "MIT"
|
| 35 |
+
tags: ["medical", "multi-agent", "triage", "rl", "llm"]
|
| 36 |
+
hackathon: "Meta PyTorch OpenEnv Hackathon 2026"
|
requirements.txt
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ER-MAP: Emergency Response Multi-Agent Pipeline
|
| 2 |
+
# Requirements
|
| 3 |
+
|
| 4 |
+
# --- Core ---
|
| 5 |
+
gymnasium>=0.29.0
|
| 6 |
+
openenv-core>=0.2.3
|
| 7 |
+
|
| 8 |
+
# --- LLM Inference (Environment Actors) ---
|
| 9 |
+
groq>=0.4.0
|
| 10 |
+
|
| 11 |
+
# --- Training (Doctor RL Agent) ---
|
| 12 |
+
# Note: Training deps are heavy; exclude if just running the server/demo
|
| 13 |
+
# torch>=2.0.0
|
| 14 |
+
# transformers>=4.38.0
|
| 15 |
+
# trl>=0.8.0
|
| 16 |
+
# peft>=0.9.0
|
| 17 |
+
# accelerate>=0.27.0
|
| 18 |
+
# datasets>=2.16.0
|
| 19 |
+
# unsloth>=2024.1
|
| 20 |
+
|
| 21 |
+
# --- Optional: Logging ---
|
| 22 |
+
wandb>=0.16.0
|
| 23 |
+
|
| 24 |
+
# --- TTS (Voice System) ---
|
| 25 |
+
elevenlabs>=1.0.0
|
| 26 |
+
edge-tts>=6.1.0
|
| 27 |
+
pygame>=2.5.0
|
| 28 |
+
|
| 29 |
+
# --- Dashboard ---
|
| 30 |
+
flask>=3.0.0
|
| 31 |
+
|
| 32 |
+
# --- OpenEnv FastAPI server (server.py) ---
|
| 33 |
+
fastapi>=0.110.0
|
| 34 |
+
uvicorn[standard]>=0.27.0
|
| 35 |
+
|
| 36 |
+
# --- Utilities ---
|
| 37 |
+
pydantic>=2.0.0
|
setup.py
ADDED
|
@@ -0,0 +1,42 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from setuptools import setup, find_packages
|
| 2 |
+
|
| 3 |
+
setup(
|
| 4 |
+
name="er-map",
|
| 5 |
+
version="1.0.0",
|
| 6 |
+
packages=find_packages(),
|
| 7 |
+
install_requires=[
|
| 8 |
+
"gymnasium>=0.29.0",
|
| 9 |
+
"openenv-core>=0.2.3",
|
| 10 |
+
"groq>=0.4.0",
|
| 11 |
+
"fastapi>=0.110.0",
|
| 12 |
+
"uvicorn[standard]>=0.27.0",
|
| 13 |
+
"flask>=3.0.0",
|
| 14 |
+
"pydantic>=2.0.0",
|
| 15 |
+
"elevenlabs>=1.0.0",
|
| 16 |
+
"edge-tts>=6.1.0",
|
| 17 |
+
"pygame>=2.5.0",
|
| 18 |
+
],
|
| 19 |
+
extras_require={
|
| 20 |
+
"training": [
|
| 21 |
+
"torch>=2.0.0",
|
| 22 |
+
"transformers>=4.38.0",
|
| 23 |
+
"trl>=0.8.0",
|
| 24 |
+
"peft>=0.9.0",
|
| 25 |
+
"accelerate>=0.27.0",
|
| 26 |
+
"datasets>=2.16.0",
|
| 27 |
+
"unsloth>=2024.1",
|
| 28 |
+
"wandb>=0.16.0",
|
| 29 |
+
],
|
| 30 |
+
},
|
| 31 |
+
author="ER-MAP Team",
|
| 32 |
+
description="Emergency Response Multi-Agent Pipeline",
|
| 33 |
+
long_description=open("README.md").read(),
|
| 34 |
+
long_description_content_type="text/markdown",
|
| 35 |
+
license="MIT",
|
| 36 |
+
classifiers=[
|
| 37 |
+
"Programming Language :: Python :: 3",
|
| 38 |
+
"License :: OSI Approved :: MIT License",
|
| 39 |
+
"Operating System :: OS Independent",
|
| 40 |
+
],
|
| 41 |
+
python_requires=">=3.9",
|
| 42 |
+
)
|
training_per2.png
ADDED
|
Git LFS Details
|
training_perf3.png
ADDED
|
Git LFS Details
|
training_performance1.png
ADDED
|
Git LFS Details
|