Uddiii commited on
Commit
7a90355
·
1 Parent(s): a3804d9

Submission-ready: README, blog, training pipeline, baseline evidence, OpenEnv compliance

Browse files
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="flex-1 flex flex-col">
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
- <span>▶</span> Start Case
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
- __all__ = ["TriageEnv", "generate_ground_truth", "construct_prompts", "AgentRouter"]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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>=0.1.0
 
 
 
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
- # ER-MAP: Emergency Room Multi-Agent Protocol
2
-
3
- > **A multi-agent RL environment for training medical triage AI with curriculum learning, empathy-aware rewards, and realistic patient simulation.**
4
-
5
- Built for the [Meta × PyTorch OpenEnv Hackathon](https://pytorch.org/blog/openenv/).
6
-
 
 
7
  ---
8
 
9
- ## Overview
10
 
11
- ER-MAP simulates a realistic Emergency Department where a **Doctor agent** (the RL policy) must diagnose and treat patients by orchestrating two auxiliary LLM agents (**Nurse** and **Patient**) through structured clinical tools. The environment uses **GRPO (Group Relative Policy Optimization)** with a **3-phase curriculum** that progresses from basic tool mastery to empathetic socio-economic negotiation.
12
 
13
- ### Key Innovation
14
 
15
- Unlike traditional medical QA benchmarks, ER-MAP tests *process-level clinical competence*:
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
- ## Architecture
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
- ## Disease Database
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
- ## 3-Phase Curriculum Learning
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
- ## Empathy Engine (Intent-Based)
107
 
108
- The empathy system uses a **causal chain** instead of keyword matching:
109
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
110
  ```
111
- Doctor message → classify_intent() → PatientState.update() → consent_decision() → reward
112
- ```
113
 
114
- ### Intent Classification (Heuristic, No LLM Call)
115
- - **Empathetic**: "I understand", "you're safe", "that must be scary" → trust ↑, anxiety ↓
116
- - **Explanatory**: "let me explain", "this test will", "because we need" → trust ↑
117
- - **Dismissive**: "just calm down", "that's not important", "hurry up" trust ↓↓, anxiety ↑↑
118
- - **Acknowledgment**: "tell me more", "when did this start" trust (mild)
 
 
 
 
 
 
 
 
119
 
120
- ### Patient Trust/Anxiety Model
121
- - Trust (0-100): Starts based on persona. Modified by Doctor behavior.
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
- ## Milestone Tracker
129
 
130
- Tracks clinical workflow compliance:
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
- ## Reward Architecture
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
- ## SOAP Noise Injection
160
-
161
- Phase-dependent noise applied to patient history:
162
 
163
- | Phase | Noise Type | Examples |
164
- |-------|-----------|----------|
165
- | 1 | None | Clean data, all fields accurate |
166
- | 2 | Clinical | Missing allergies, vague ROS, inconsistent PMH |
167
- | 3 | Behavioral | "Patient homeless, med history unknown", "Language barrier", "Anxious about billing" |
168
 
169
  ---
170
 
171
- ## Project Structure
172
-
173
- ```
174
- ER_MAP/
175
- ├── envs/
176
- │ ├── triage_env.py # Gymnasium environment (core)
177
- │ ├── randomizer.py # Ground truth + persona generation
178
- │ ├── disease_db.py # 50-disease database (10 classes)
179
- │ ├── empathy_engine.py # Intent classifier + trust model + milestones
180
- │ └── api_router.py # LLM API routing (Groq)
181
- ├── training/
182
- │ ├── train_grpo.py # GRPO training with curriculum scheduler
183
- │ └── train_ppo.py # Legacy PPO script (deprecated)
184
- ├── tts_engine.py # ElevenLabs TTS with speech markers
185
- ├── autoplay.py # Demo episode runner
186
- ├── evaluate.py # Evaluation harness
187
- └── dashboard.py # Metrics visualization
188
- ```
 
 
 
 
189
 
190
  ---
191
 
192
- ## Quick Start
193
 
194
- ### Installation
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
- ### Run a Demo Episode
202
- ```bash
203
- export GROQ_API_KEY="your_key"
204
- python -m ER_MAP.autoplay
205
- ```
206
 
207
- ### Train with GRPO + Curriculum
208
- ```bash
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
- #### Train on Kaggle Free Tier (recommended $0)
 
 
221
 
222
- The repo ships with a complete Kaggle workflow under `kaggle/`:
223
 
224
- - `kaggle/train_ermap_grpo_kaggle.ipynb` the notebook (clone repo install deps run GRPO → push checkpoints to HF Hub)
225
- - `kaggle/kaggle_helpers.py` — secret loader, HF push/pull, env-summary printer
226
- - `kaggle/requirements_kaggle.txt` Kaggle-image-aware dependency list (Unsloth pinned last)
227
- - `kaggle/KAGGLE.md` full setup guide with hardware feasibility table and gotchas
 
 
 
 
 
228
 
229
- Tested target: **single Tesla T4 16 GB**, Llama-3.1-8B 4-bit + LoRA(r=16), 120 episodes, ~6-8 h per session.
230
- See `kaggle/KAGGLE.md` for the full step-by-step.
231
 
232
- ### Environment API
233
- ```python
234
- from ER_MAP.envs.triage_env import TriageEnv
235
 
236
- env = TriageEnv(groq_api_key="your_key")
 
 
 
237
 
238
- # Phase 2 with medium difficulty
239
- obs, info = env.reset(options={"phase": 2, "difficulty": "medium"})
240
 
241
- action = '{"tool": "read_soap"}'
242
- obs, reward, done, truncated, info = env.step(action)
243
 
244
- # info now includes:
245
- # info["patient_state"] = {"trust": 55.0, "anxiety": 40.0, ...}
246
- # info["milestones"] = {"achieved": {"READ_SOAP": True, ...}, "completion": 0.17}
 
 
 
 
247
  ```
248
 
249
- ---
 
 
 
 
250
 
251
- ## Training Budget Estimate (HuggingFace $200 Credits)
 
 
252
 
253
- | Phase | Episodes | Est. Time (A100) | Est. Cost |
254
- |-------|----------|------------------|-----------|
255
- | Phase 1 | 40–60 | ~2–3 hours | ~$15 |
256
- | Phase 2 | 60–80 | ~3–4 hours | ~$25 |
257
- | Phase 3 | 80–120 | ~4–6 hours | ~$35 |
258
- | **Total** | **200** | **~10–13 hours** | **~$75** |
259
 
260
- Recommended model: **Qwen3-4B** (via Unsloth 4-bit) — best performance/cost ratio.
 
261
 
262
  ---
263
 
264
- ## TTS Engine (Presentation Only)
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
- Speech markers (breathing, sighs, pauses) are injected into the TTS pipeline but **never fed back to the RL agents** — strict separation maintained.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
274
 
275
  ---
276
 
277
- ## OpenEnv Compliance
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
- ## License
293
 
294
- MIT License. Built for the Meta × PyTorch OpenEnv Hackathon 2026.
 
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
+ ![OpenEnv](https://img.shields.io/badge/OpenEnv-Compatible-green) ![License](https://img.shields.io/badge/License-MIT-blue) ![Python](https://img.shields.io/badge/Python-3.9+-yellow)
 
 
 
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
+ ![Baseline Phase Comparison](baseline_eval/baseline_phases_comparison.png)
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** | ![Phase 1](training_perf3.png) | ![Phase 2](training_per2.png) | ![Phase 3](training_performance1.png) |
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
+ [![Watch the Demo Video](https://img.youtube.com/vi/YOUR_VIDEO_ID/maxresdefault.jpg)](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
+ ![Baseline Phase Comparison](baseline_eval/baseline_phases_comparison.png)
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
+ ![Phase 1 Baseline](baseline_eval/baseline_phase1_rewards.png)
266
+ <!-- slide -->
267
+ ![Phase 2 Baseline](baseline_eval/baseline_phase2_rewards.png)
268
+ <!-- slide -->
269
+ ![Phase 3 Baseline](baseline_eval/baseline_phase3_rewards.png)
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
+ ![Phase 1 Trained](training_perf3.png)
276
+ <!-- slide -->
277
+ ![Phase 2 Trained](training_per2.png)
278
+ <!-- slide -->
279
+ ![Phase 3 Trained](training_performance1.png)
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
+ ![Overall Reward Growth](ER_MAP/reward_curve.png)
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

  • SHA256: 5e963846b49c7c14dc17df5ceb0bd56ee5934f3dbdc996621f12de74b8b7936d
  • Pointer size: 130 Bytes
  • Size of remote file: 68.2 kB
training_perf3.png ADDED

Git LFS Details

  • SHA256: 843d406e5108377d99605655abcd503908aff391bd71e8c96f37b83dddff10d4
  • Pointer size: 130 Bytes
  • Size of remote file: 64.5 kB
training_performance1.png ADDED

Git LFS Details

  • SHA256: d41a8e9fe0b379c642d010daf3785b7ef468e8ec704de3eabc3de6479bad93e3
  • Pointer size: 130 Bytes
  • Size of remote file: 63.2 kB