---
title: 'DiagGym: MDP Diagnostic Simulator'
url: https://www.emergentmind.com/topics/diaggym
type: topic
---

# DiagGym: MDP Diagnostic Simulator

DiagGym is a transformer-based diagnostics world model formulated as a Markov Decision Process (MDP) and designed to simulate multi-turn clinical diagnosis scenarios. It is built upon real electronic health record (EHR) data and provides a high-fidelity virtual environment where large language model (LLM) agents are trained to optimize both diagnostic accuracy and efficiency via reinforcement learning. DiagGym supports dynamic, outcome-driven agent training and evaluation, enabling emergent diagnostic management strategies unattainable through passive or static case summarization approaches [2510.24654].

## 1. Formalization as Conditional World Model and MDP

DiagGym models the diagnostic process as an MDP in which an agent repeatedly requests and interprets clinical examinations, then commits to a final diagnosis. The environment is formally specified as follows:

- **State space $S$:** At time $t$, $s_t = (\mathcal{B}, E_t)$ for simulation and $s_t = (\mathcal{I}, E_t)$ for the agent. Here, $\mathcal{B}$ is a comprehensive patient background profile including chief complaint, history, and the held-out diagnosis, while $\mathcal{I}$ omits the diagnosis for agent inputs. $E_t$ denotes the ordered list of $(a_i, e_i)$ pairs, each corresponding to an examination $a_i$ and observed result $e_i$.
- **Action space $A$:** $A = \{a_1, ..., a_N\} \cup \{\text{final diagnosis action}\}$, where $a_i$ can be a specific clinical test or the terminal action to commit a diagnosis.
- **Transition dynamics:** On requesting an exam, the outcome is sampled from the learned conditional generator $\Phi_{\text{env}}(\cdot| a_{t+1}, E_t, \mathcal{B})$. The next state is updated to include this exam-outcome pair.
- **Reward function:** $R(s_t, a_t) = \lambda_1 r_{\text{diag}} + \lambda_2 r_{\text{exam}} + \lambda_3 r_{\text{turn}}$ with
  - $r_{\text{diag}}$ rewarding correct diagnoses,
  - $r_{\text{exam}}$ using $F_1$ alignment with reference exam sequences,
  - $r_{\text{turn}}$ penalizing dialogues exceeding $T_{\max}$.
Typical coefficients: $\lambda_1=1$, $\lambda_2=0.5$, $\lambda_3=1$, $\gamma=0.99$.

Key formulas:
- Joint transition:
  $$
  P(s_{t+1}, o_{t+1} \mid s_t, a_{t+1}) = \delta[s_{t+1} = s_t \cup \{(a_{t+1}, o_{t+1})\}] \cdot \Phi_{\text{env}}(o_{t+1} \mid a_{t+1}, E_t, \mathcal{B})
  $$
- World-model training:
  $$
  \mathcal{L}_{\mathrm{sim}} = -\sum_{t=0}^{T-1} \log\,\Phi_{\mathrm{env}}\left(\hat{e}_{t+1}\mid a_{t+1}, E_t, \mathcal{B}\right)
  $$
- RL agent training:
  $$
  \max_{\theta} \mathbb{E}_{\tau \sim \pi_\theta}\left[\sum_{t=0}^{T-1}\gamma^t R(s_t, a_t)\right]
  $$

## 2. Data Sources and Preprocessing

DiagGym leverages EHR sequences from the MIMIC-IV database, encompassing tens of thousands of ICU patients.

- **Profile extraction:** Discharge summaries are systematically parsed into fields: chief complaint, present/past history, social/family history, physical exam, and discharge diagnosis.
- **Filtering heuristics:** Cases missing physical exam data or where the diagnosis appears in the prior histories are excluded (using LLM checks for verbatim overlaps).
- **Examination chain reconstruction:** Free-text physical exams are converted to structured form via LLM prompting, followed by aggregation of lab, microbiology, and radiology results obtained within one day prior to admission.
- **Chronological sorting and de-duplication:** All interventions are time-ordered, retaining only the first pre-admission instance per examination.
- **Text encoding:** Patient profiles and exam-result tuples are tokenized as plain text; numerical lab values are retained as text tokens.
- **Data splits:** 
  - Training set: 114,239 patients, 4,897 diseases, $\sim$29 exams per patient.
  - World-model test set: 863 disease-balanced cases.
  - DiagBench: 750 physician-validated cases; 99 cases further annotated with 973 stepwise rubrics.

## 3. Model Architecture and Implementation

DiagGym employs a transformer architecture, initialized from Qwen2.5-Instruct-7B, as the core world model:

- **Conditional next-token generator:** $\Phi_{\text{env}}$ is parameterized as an autoregressive model predicting the tokenized sequence for each exam result, conditioned on patient profile, exam history, and current action.
- **Attention mechanism:** The network applies self-attention over the concatenated prompt of the patient summary, previous $(a_i, e_i)$ pairs, and the examination request, without task-specific architectural modifications.
- **Vocabulary and representation:** Clinical text and numerical lab values are represented within a single vocabulary, supporting flexible simulation across clinical modalities.
- **No additional modules:** The world model does not incorporate specialized clinical reasoning heuristics outside the transformer backbone.

## 4. Training Procedures and Hyperparameters

Training leverages large-scale hardware and follows a deterministic protocol:

- **World model ($\Phi_{\text{env}}$):**
  - 8 $\times$ A100 80GB GPUs, DeepSpeed ZeRO Stage 2.
  - Learning rate: $4 \times 10^{-5}$, maximum input length: 8,192 tokens, 15 epochs, batch size ensuring 8K tokens per context.
- **Agent initialization ($\pi_\theta$ cold-start):**
  - Supervised learning on 1,000 curated trajectories; learning rate $1 \times 10^{-5}$, 3 epochs.
- **Reinforcement learning for $\pi_\theta$:**
  - Algorithm: GRPO (Generative RL).
  - Max dialogue turns $T_{\max} = 12$; rollouts per update: 5; batch size: 512 parallel rollouts.
  - Learning rate: $1 \times 10^{-6}$, discount $\gamma = 0.99$.
  - 4 nodes $\times$ 8 A100 GPUs; DiagGym served via vLLM.
  - Judge: Qwen2.5-72B.
  - Training converges empirically after $\sim$200 steps.

## 5. Integration with Diagnostic Agent Training Loop

Within each simulation episode, the following loop is executed:

1. The agent receives the initial state $s_0 = (\mathcal{I}, \emptyset)$.
2. It samples an action $a_t \sim \pi_\theta(\cdot | s_t)$: either an exam request or terminal diagnosis.
3. For exam actions, the environment returns an outcome $o_t \sim \Phi_{\text{env}}(\cdot)$, updating $s_{t+1}$; for final diagnosis, the episode ends.
4. Rewards $r_{\text{diag}}$, $r_{\text{exam}}$ (using $F_1$ overlap with reference tests), and $r_{\text{turn}}$ are computed as per the reward function.
5. The policy network $\pi_\theta$ is updated via GRPO using the cumulative discounted return.

This setup enables optimization of policies over complex, realistic diagnostic sequences rather than static, single-step labels.

## 6. Evaluation, Benchmarks, and Exemplary Results

DiagGym's fidelity and usefulness are quantitatively and qualitatively established:

- **World-model fidelity:** On 863 held-out cases, DiagGym achieves over 96% binary clinical consistency, a step-wise LLM-judged similarity of 3.565/5, and a 1-Wasserstein distance of 0.128 on lab test distributions—significantly surpassing baseline LLM simulators (e.g., DeepSeek-v3 yields 1.336).
- **Agentic trajectories:** Agents trained in DiagGym, exemplified by DiagAgent, accurately reproduce reference timelines (e.g., ordering CBC $\rightarrow$ identifying neutrophilia $\rightarrow$ ordering CT $\rightarrow$ diagnosing appendicitis), closely matching physician workflows.
- **Benchmark performance:** On DiagBench,
  - DiagAgent achieves 9.34% higher single-turn diagnostic accuracy and a 44.03% gain in exam recommendation hit ratio compared to state-of-the-art LLMs;
  - In multi-turn settings, a 15.12% accuracy improvement and a 23.09% boost in exam recommendation $F_1$ score;
  - Rubric-based evaluation yields a 7.1% higher weighted rubric score than the next best model.
- **Procedural quality:** In process-level benchmarks, the agent fulfills 95% of high-weight physician rubrics, demonstrating alignment with expert procedural strategies while minimizing unnecessary testing.

These results substantiate the claim that interactive exploration within a high-fidelity virtual clinical environment yields superior and more clinically plausible diagnostic management compared to purely instruction-tuned LLMs [2510.24654].

## 7. Significance and Role in Diagnostic AI Research

DiagGym establishes a methodological foundation for training and evaluating diagnostic agents in a clinically realistic, dynamic setting. By leveraging a learned world-model grounded in authentic EHR data, it enables agent learning that integrates patient-specific context, temporal exam selection, and adaptive reasoning. The framework advances the field by:

- Providing a virtual test bed with high clinical consistency and fine-grained process annotation.
- Enabling explicit multi-turn reinforcement learning for diagnostic management.
- Demonstrating quantifiable improvements over both static fine-tuning and prompt-engineered baselines—both in process and outcome metrics.
- Supporting rigorous benchmarking via physician-validated cases and rubric-based process evaluation.

This suggests that DiagGym may serve as a central component in the development of next-generation diagnostic LLMs capable of dynamic, high-quality, and efficient clinical reasoning [2510.24654].

Source: https://www.emergentmind.com/topics/diaggym