DiagGym: MDP Diagnostic Simulator

• DiagGym is a transformer-based diagnostic simulator modeled as an MDP that enables realistic multi-turn clinical diagnosis scenarios.
• It leverages real electronic health record data from sources like MIMIC-IV to construct detailed patient profiles and sequential examination chains.
• DiagGym integrates reinforcement learning to train LLM agents, resulting in improved diagnostic accuracy and efficient exam recommendations.

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 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 (Qiu et al., 28 Oct 2025).

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 $t$0: $t$1, where $t$2 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 $t$3. The next state is updated to include this exam-outcome pair.
• Reward function: $t$4 with
  • $t$5 rewarding correct diagnoses,
  • $t$6 using $t$7 alignment with reference exam sequences,
  • $t$8 penalizing dialogues exceeding $t$9.
  • Typical coefficients: $s_t = (\mathcal{B}, E_t)$0, $s_t = (\mathcal{B}, E_t)$1, $s_t = (\mathcal{B}, E_t)$2, $s_t = (\mathcal{B}, E_t)$3.

Key formulas:

• Joint transition:

$s_t = (\mathcal{B}, E_t)$4

• World-model training:

$s_t = (\mathcal{B}, E_t)$5

• RL agent training:

$s_t = (\mathcal{B}, E_t)$6

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, $s_t = (\mathcal{B}, E_t)$729 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: $s_t = (\mathcal{B}, E_t)$8 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 $s_t = (\mathcal{B}, E_t)$9 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 ($s_t = (\mathcal{I}, E_t)$0):
  • 8 $s_t = (\mathcal{I}, E_t)$1 A100 80GB GPUs, DeepSpeed ZeRO Stage 2.
  • Learning rate: $s_t = (\mathcal{I}, E_t)$2, maximum input length: 8,192 tokens, 15 epochs, batch size ensuring 8K tokens per context.
• Agent initialization ($s_t = (\mathcal{I}, E_t)$3 cold-start):
  • Supervised learning on 1,000 curated trajectories; learning rate $s_t = (\mathcal{I}, E_t)$4, 3 epochs.
• Reinforcement learning for $s_t = (\mathcal{I}, E_t)$5:
  • Algorithm: GRPO (Generative RL).
  • Max dialogue turns $s_t = (\mathcal{I}, E_t)$6; rollouts per update: 5; batch size: 512 parallel rollouts.
  • Learning rate: $s_t = (\mathcal{I}, E_t)$7, discount $s_t = (\mathcal{I}, E_t)$8.
  • 4 nodes $s_t = (\mathcal{I}, E_t)$9 8 A100 GPUs; DiagGym served via vLLM.
  • Judge: Qwen2.5-72B.
  • Training converges empirically after $\mathcal{B}$0200 steps.

5. Integration with Diagnostic Agent Training Loop

Within each simulation episode, the following loop is executed:

1. The agent receives the initial state $\mathcal{B}$1.
2. It samples an action $\mathcal{B}$2: either an exam request or terminal diagnosis.
3. For exam actions, the environment returns an outcome $\mathcal{B}$3, updating $\mathcal{B}$4; for final diagnosis, the episode ends.
4. Rewards $\mathcal{B}$5, $\mathcal{B}$6 (using $\mathcal{B}$7 overlap with reference tests), and $\mathcal{B}$8 are computed as per the reward function.
5. The policy network $\mathcal{B}$9 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 $\mathcal{I}$0 identifying neutrophilia $\mathcal{I}$1 ordering CT $\mathcal{I}$2 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 $\mathcal{I}$3 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 (Qiu et al., 28 Oct 2025).

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 (Qiu et al., 28 Oct 2025).