Green-White MIMIC Diagnostic Replay

This project is a diagnostic-agent experiment built on the local 10% MIMIC subset in data/mimic-iv/ and data/mimic-iv-notes/.

The main idea is simple:

CSV data -> DuckDB -> safe cohort -> frozen episodes -> live agent experiment/UI

Most of the code before the UI is not the experiment itself. It is preprocessing that turns large MIMIC CSV files into safe, fixed admission-level cases. The real Green-White agent experiment starts after outputs/episodes/mimic_greenwhite_episodes.jsonl exists.

Live diagnosis requires an OpenAI-compatible API key. There is no heuristic White Agent and no heuristic fallback.

What The Project Demonstrates

The project demonstrates a two-agent diagnostic loop:

• A Green Agent acts as the clinical chart environment.
• A White Agent acts as the diagnostic reasoning policy.

The White Agent is intentionally not allowed to browse the raw tables directly. It can only ask the Green Agent for evidence through approved tools. The Green Agent returns chart-grounded observations and hides diagnosis labels until the White Agent submits a final diagnosis.

This makes the demo more like a controlled diagnostic replay than a normal chatbot over a database.

Why Two Agents

The Green Agent is deterministic. It owns:

• the frozen admission episode;
• the hidden ICD-derived labels;
• the safe evidence tools;
• turn and tool-cost budgets;
• source provenance;
• scoring;
• event logs.

The Green Agent is the only component allowed to know the ground truth before final scoring. It enforces the rule that the White Agent cannot see ICD codes, diagnosis titles, or discharge diagnosis sections before submit_diagnosis.

The White Agent is model-driven. It owns:

• the running case summary;
• the current diagnosis list;
• the reasoning summary;
• the next evidence request;
• the final diagnosis submission.

The White Agent is backed by an OpenAI-compatible chat-completions API. You can use OpenAI directly, or an OpenAI-compatible local/server provider such as vLLM by setting OPENAI_BASE_URL.

Before Experiments Vs During Experiments

There are two phases.

Phase 1: Build the experiment substrate

This is the preprocessing pipeline:

MIMIC CSVs
  -> DuckDB database
  -> readable joined tables
  -> usable admission cohort
  -> sanitized note sections
  -> frozen episode JSONL

This phase answers questions like:

• Which admissions have discharge notes?
• Which admissions have strict ICD-10 hidden diagnosis labels for scoring?
• Which admissions have radiology, microbiology, procedures, or OMR evidence?
• Which note sections are safe to expose?
• Which sections might leak the final diagnosis and must be blocked?

Phase 2: Run the agent experiment

This starts from frozen episodes:

episode -> Green initial observation -> White action -> Green tool result -> ... -> final diagnosis -> score

The model is only involved in Phase 2. The preprocessing pipeline does not call the OpenAI API.

Why The Pipeline Exists

MIMIC data is table-based and not directly shaped like an agent case. A model should not receive a raw discharge note with the discharge diagnosis inside it, and it should not see diagnoses_icd.csv before final scoring.

The pipeline creates a safe case file where each admission has:

• safe patient/admission context;
• sanitized note sections;
• available radiology reports;
• microbiology records;
• OMR measurements;
• procedure evidence;
• hidden ICD-10 labels used only for final scoring.

That final case object is called a MimicEpisode.

The One-Command Workflow

Use this as the normal preprocessing command:

source .venv/bin/activate
python -m src.data.run_pipeline --config configs/paths.yaml --limit 50 --max-turns 20 --max-cost 30

This builds the normal 50-episode demo pool:

outputs/mimic_subset.duckdb
outputs/schema_report.md
outputs/cohort_report.md
outputs/episodes/mimic_greenwhite_episodes.jsonl

--max-turns is the maximum number of White Agent actions before the Green Agent stops the case. --max-cost is the maximum cumulative tool cost. These values are saved as defaults inside newly built episodes.

Episode generation uses a strict ICD-10 cohort: every diagnosis label for an admission must be ICD-10, and at least one diagnosis must be a scored ICD-10 disease/injury code. Symptom, status, and external-cause chapters such as R, Z, and V-Y are excluded from scoring labels. ICD-9 admissions are not used as fallback cases.

If the heavy intermediate tables already exist, the wrapper reuses them. This keeps quick checks fast. To force a full rebuild:

python -m src.data.run_pipeline --config configs/paths.yaml --limit 50 --max-turns 20 --max-cost 30 --rebuild

For a quick smoke test without overwriting the 50-episode file:

python -m src.data.run_pipeline --config configs/paths.yaml --limit 5

That writes:

outputs/episodes/mimic_greenwhite_episodes_limit5.jsonl

What Each Output Means

outputs/mimic_subset.duckdb

The local DuckDB database containing the selected MIMIC CSV tables plus derived tables. It is used by the preprocessing pipeline, not directly by the White Agent.

outputs/schema_report.md

A human-readable report of loaded tables, row counts, columns, distinct patient/admission IDs, and missing required columns.

outputs/cohort_report.md

A short report describing how many admissions are usable for episodes and how many have each evidence type.

outputs/episodes/mimic_greenwhite_episodes.jsonl

The frozen episode file used by the UI and live runs. Each line is one admission-level diagnostic case.

outputs/experiments/<tag>/logs/*.jsonl

Event logs created when the White Agent and Green Agent run. Each file is one admission. These logs can be replayed in the demo without another model call.

outputs/experiments/<tag>_<timestamp>.json

Summary result files from batch experiment runs. Each file contains per-episode scores, tool sequences, submitted diagnoses, and judge verdicts for one experiment configuration.

How The 50 Admissions Are Selected

The current subset has 33,185 usable admissions before ICD-version filtering. An admission is usable if it has:

• an admission profile;
• at least one discharge note;
• at least one hidden ICD-derived diagnosis label for scoring.

The episode builder then applies the strict ICD-10 scoring filter. In the current rebuilt database, 12,036 admissions remain eligible after requiring all diagnosis labels to be ICD-10 and at least one scored ICD-10 disease/injury diagnosis.

Each usable admission gets an evidence_score:

+1 has radiology
+1 has microbiology
+1 has procedures
+1 has OMR measurements

The episode builder sorts admissions by:

evidence_score descending, then hadm_id ascending

Then it takes --limit 50.

So the current 50 are not random. They are the first 50 evidence-rich admissions. They are useful for demos because the White Agent has several evidence tools to request. The full cohort is much larger, and you can build more episodes by increasing --limit.

The current strict ICD-10 cohort distribution is:

12,036 strict ICD-10 admissions
3,160 have evidence_score 4
4,027 have evidence_score 3
3,676 have evidence_score 2
1,162 have evidence_score 1
11    have evidence_score 0

The 50-episode file is a demo/evaluation pool. It is not the final batch experiment over all admissions.

What Happens In One Agent Run

One live run uses one admission episode.

1. The Green Agent loads the episode.
2. The Green Agent shows safe initial context, such as age, gender, race, admission type, and admission location.
3. The White Agent receives the observation history and returns JSON with:
   • case_summary
   • current_diagnosis_list
   • reasoning_summary
   • action
4. The Green Agent checks whether the requested tool is allowed.
5. The Green Agent returns the requested evidence with source provenance.
6. The White Agent updates its current diagnosis list and asks for more evidence.
7. The loop repeats until the White Agent calls submit_diagnosis.
8. If a non-submit action would exhaust the hidden turn/cost budget, the Green Agent returns final_answer_required.
9. On that grace step, the White Agent must call submit_diagnosis; no more evidence tools are accepted.
10. The Green Agent reveals hidden labels and computes top-1/top-3/F1 correctness.

How Turns And Cost Are Calculated

A turn is one White Agent action processed by the Green Agent.

The initial patient context is not counted as a turn. Turn 1 starts when the White Agent chooses its first tool/action. Every action increments the turn counter, including valid tools, unavailable-tool requests, tool errors, and submit_diagnosis.

Examples:

White Agent asks ask_history       -> turn +1
White Agent asks get_radiology     -> turn +1
White Agent asks for unavailable labs -> turn +1
White Agent submits diagnosis      -> turn +1

Cost is a project-defined tool cost, not API dollars. Model/API spending is separate and is not tracked by this budget counter.

Tool costs are configured in configs/tools.yaml:

ask_history: 1
read_note_section: 1
get_radiology: 2
get_microbiology: 2
get_omr: 1
get_procedures: 1
submit_diagnosis: 0
tool_error: 1

So this sequence:

ask_history       cost 1
read_note_section cost 1
get_radiology     cost 2
get_microbiology  cost 2
submit_diagnosis  cost 0

has total tool cost:

1 + 1 + 2 + 2 + 0 = 6

max_turns limits how many White Agent actions are allowed. max_cost limits cumulative tool cost. You can set them when building episodes, when running a CLI episode, or in the demo sidebar. Runtime CLI/UI overrides do not require rebuilding episodes.

The White Agent does not see numeric budget values. If a non-submit action would otherwise end the case, the Green Agent gives one hidden final-answer grace observation: final_answer_required. A valid grace submit_diagnosis is scored normally and marked with grace_used=true; any other next action ends as budget_exhausted with no score.

Example:

python -m src.eval.run_episode \
  --episodes outputs/episodes/mimic_greenwhite_episodes.jsonl \
  --episode-index 0 \
  --max-turns 20 \
  --max-cost 30

Tools Available To The White Agent

ask_history

Returns broad safe history sections from the discharge note. The full broad history is returned only once per episode; later ask_history calls return history_already_returned with the section names and a suggestion to use read_note_section for targeted details.

read_note_section

Returns one safe note section, such as history_of_present_illness, physical_exam, past_medical_history, or hospital_course.

get_radiology

Returns admission radiology reports with note IDs and chart times when available.

get_microbiology

Returns culture/specimen/test/organism/interpretation evidence from hosp.microbiologyevents.

get_omr

Returns limited OMR measurements such as blood pressure, height, weight, or BMI when present.

get_procedures

Returns readable ICD procedure and HCPCS procedure/event evidence.

submit_diagnosis

Submits final ranked diagnoses. This is the only tool that reveals hidden labels and scoring.

final_answer_required is not a tool. It is a Green Agent observation that means no more evidence is available and the next White Agent action must be submit_diagnosis.

Unavailable evidence, including unrestricted labs, medications, ICU chartevents, high-resolution vitals, and diagnosis codes, returns a transparent tool error.

Running Batch Experiments

run_experiments.py runs the Green-White loop over many episodes in one batch and writes per-episode results, tool-use logs, and aggregate scores.

source .venv/bin/activate
export OPENAI_API_KEY="sk-..."

# Model comparison
python run_experiments.py --tag gpt4o-mini --model gpt-4o-mini --limit 50 --judge
python run_experiments.py --tag gpt4o --model gpt-4o --limit 50 --judge

# Budget ablation (same model, varying budgets)
python run_experiments.py --tag budget-tight --model gpt-4o --limit 50 --max-turns 5 --max-cost 8 --judge
python run_experiments.py --tag budget-generous --model gpt-4o --limit 50 --max-turns 20 --max-cost 30 --judge

# Memory ablation (cross-episode memory enabled)
python run_experiments.py --tag memory-gpt4o --model gpt-4o --limit 50 --memory --judge

# Direct baseline (single prompt, no tool loop)
python run_experiments.py --tag baseline-gpt4o --model gpt-4o --limit 50 --baseline --judge

Each run produces:

• outputs/experiments/<tag>/logs/hadm_<id>.jsonl — per-episode event logs.
• outputs/experiments/<tag>_<timestamp>.json — summary with scores, tool sequences, and judge verdicts.

Open-Source / Local Models

Any OpenAI-compatible provider works with --base-url:

# vLLM
python run_experiments.py --tag qwen35-27b --model Qwen/Qwen3.5-27B \
  --base-url http://localhost:8000/v1 --limit 50 --judge

# Ollama
python run_experiments.py --tag gemma-27b --model gemma3:27b \
  --base-url http://localhost:11434/v1 --limit 50 --judge

Provider configuration is in configs/agent_configs.yaml.

LLM-As-Judge Scoring

The --judge flag runs a secondary LLM evaluation after each episode. The judge receives the submitted diagnoses and the gold ICD-10 labels and returns semantic match verdicts independent of the fuzzy string scorer. This provides two complementary accuracy signals per episode.

Direct Baseline

The --baseline flag runs a single-prompt agent that receives the full safe chart at once and returns diagnoses without using the tool loop. This measures how much the iterative evidence-gathering strategy contributes beyond what a model can do in one shot.

Cross-Episode Memory

The --memory flag enables a JSONL-backed memory store. After each episode, the White Agent's case is written to memory. On subsequent episodes, the Green Agent retrieves similar past cases using keyword and diagnosis-group matching. This tests whether accumulated experience improves accuracy across a batch.

Analyzing Results

Two analysis scripts read the experiment JSON files and produce figures and tables.

analyze_results.py generates per-experiment visualizations:

python analyze_results.py

This writes figures to outputs/figures/ and a CSV summary to outputs/experiments/summary_table.csv.

analyze_compare.py generates cross-experiment comparison figures:

python analyze_compare.py

This reads all experiment JSON files (including open-model results under outputs/experiments/open_models/), maps each tag to a (model, setting) pair, and produces controlled-variable comparison plots covering model comparison, budget sensitivity, agentic vs. baseline, medical fine-tuning impact, memory ablation, and tool usage patterns.

Models Tested

The following models have been evaluated:

Proprietary: gpt-4o, gpt-4o-mini, gpt-5.4-mini, claude-sonnet

Open-source: Qwen3.5-27B, Gemma3-27B, MedGemma-27B

Each model was tested in the standard agentic setting. GPT-4o was additionally tested across budget ablation (tight/generous), a direct baseline (no tool loop), and cross-episode memory conditions.

Running The Interactive Demo

The interactive demo is a self-contained HTML page with a local API server. It supports both live runs against the White Agent and replay of saved experiment logs.

Start the server:

source .venv/bin/activate
export OPENAI_API_KEY="your-key"
python serve_demo.py

Then open:

http://localhost:8080

For vLLM or another OpenAI-compatible provider:

export OPENAI_BASE_URL="http://localhost:8000/v1"
export OPENAI_API_KEY="not-empty"
export OPENAI_MODEL="your-served-model"
python serve_demo.py

Generating A Static Demo

generate_demo.py builds a standalone HTML file from saved experiment logs. This is useful for sharing results without a running server or API key.

python generate_demo.py --tag gpt4o --tag gpt4o-mini --output demo_interactive.html

Running One Live Episode From The CLI

source .venv/bin/activate
export OPENAI_API_KEY="your-key"
python -m src.eval.run_episode \
  --episodes outputs/episodes/mimic_greenwhite_episodes.jsonl \
  --episode-index 0 \
  --max-turns 20 \
  --max-cost 30

Runtime budget overrides do not require rebuilding episodes. Use them when one case needs more turns than the saved episode default.

If OPENAI_API_KEY is not set, the live run fails clearly. There is no heuristic fallback.

Debugging Commands

The one-command pipeline is the normal path. The individual steps are still useful when you want to inspect or debug a specific stage:

python -m src.data.load_mimic_subset --config configs/paths.yaml
python -m src.data.inspect_schema --db outputs/mimic_subset.duckdb
python -m src.data.join_tables --db outputs/mimic_subset.duckdb
python -m src.data.build_cohort --db outputs/mimic_subset.duckdb
python -m src.data.extract_note_sections --db outputs/mimic_subset.duckdb
python -m src.data.sanitize_notes --db outputs/mimic_subset.duckdb
python -m src.data.build_episode_jsonl --db outputs/mimic_subset.duckdb --limit 50 --max-turns 20 --max-cost 30

Tests

Run:

source .venv/bin/activate
python -m pytest -q

The tests use mocked model responses. They do not require a real OpenAI API call.