Evolving Diagnostic Agents in a Virtual Clinical Environment

TL;DR

The paper addresses the limitation of static, instruction-tuned LLMs in clinical diagnosis by introducing a dynamic, interactive framework. It builds DiagGym, a diagnostics world model trained on EHR data to generate exam results conditioned on evolving patient state, and DiagAgent, an LLM-based agent trained end-to-end with reinforcement learning to optimize diagnostic trajectory management, examination choices, and final diagnoses. A new benchmark, DiagBench, provides 750 physician-validated cases with trajectory references and 99 rubric-annotated cases to evaluate procedural quality. Across extensive single-turn and end-to-end tests, DiagAgent consistently outperforms 10 state-of-the-art LLMs and two agentic systems, demonstrating meaningful gains in diagnostic accuracy and examination-hit ratios, and highlighting the value of interactive policy learning in realistic clinical environments. The work also establishes a scalable, in-silico platform for developing and evaluating diagnostic agents before real-world deployment, with potential to advance robust, dynamic clinical decision support.

Abstract

In this paper, we present a framework for training large language models (LLMs) as diagnostic agents with reinforcement learning, enabling them to manage multi-turn diagnostic processes, adaptively select examinations, and commit to final diagnoses. Unlike instruction-tuned models trained on static case summaries, our method acquires diagnostic strategies through interactive exploration and outcome-based feedback. Our contributions are fourfold: (i) We present DiagGym, a diagnostics world model trained with electronic health records that emits examination outcomes conditioned on patient history and recommended examination, serving as a virtual clinical environment for realistic diagnosis training and evaluation; (ii) We train DiagAgent via end-to-end, multi-turn reinforcement learning to learn diagnostic policies that optimize both information yield and diagnostic accuracy; (iii) We introduce DiagBench, a diagnostic benchmark comprising 750 cases with physician-validated examination recommendations and 99 cases annotated with 973 physician-written rubrics on diagnosis process; (iv) we demonstrate superior performance across diverse diagnostic settings. DiagAgent significantly outperforms 10 state-of-the-art LLMs, including DeepSeek-v3 and GPT-4o, as well as two prompt-engineered agents. In single-turn settings, DiagAgent achieves 9.34% higher diagnostic accuracy and 44.03% improvement in examination recommendation hit ratio. In end-to-end settings, it delivers 15.12% increase in diagnostic accuracy and 23.09% boost in examination recommendation F1 score. In rubric-based evaluation, it surpasses the next-best model, Claude-sonnet-4, by 7.1% in weighted rubric score. These findings indicate that learning policies in interactive clinical environments confers dynamic and clinically meaningful diagnostic management abilities unattainable through passive training alone.

Figures (10)

• Figure 1: Overview of our method.a illustrates the overview of our method, we establish a virtual clinical environment, DiagGym, that can simulate examination results in real time. Then, within it, we train a diagnostic agent capable of managing multi-turn diagnostic trajectories in a long-term manner, recommendations, recommending diverse examinations until sufficient evidence is gathered for a final diagnosis. b presents the diagnostics world model we constructed based on EHRs, representing a virtual clinical environment. It receives the patient’s basic profile and performed past examinations, and the next examination query as condition input, then simulates the related results as feedback. c depicts the DiagAgent end-to-end multi-turn RL training, where the agent interacts with the virtual environment, rolls out different possible diagnostic trajectories, self-explores suitable examination recommendation chains, and iteratively evolves its decision-making policy through end-to-end reinforcement rewards.
• Figure 1: Example Case Study from DiagGym: Comparison of Predicted and Ground Truth Examinations. This case presents a single patient profile and final diagnosis, illustrating the step-wise evaluation setting. The core table compares the predicted examination results generated by DiagGym with the ground truth results in sequential order. The rightmost column analyzes key differences and discusses their clinical relevance in the diagnostic process.
• Figure 2: Overview of simulator evaluation settings. a Instance‑wise metrics: GPT‑4o assesses the quality of generated examination results on an individual patient case level. b Examination‑wise metrics: fidelity and diversity are evaluated by comparing the statistical distributions of generated examination results against those from real cases.
• Figure 2: Interactive Diagnostic Case Study with DiagAgent: Model Trajectory and Reference Timeline. This figure illustrates a multi-turn interaction between the DiagAgent model and a simulator. The table details the agent's step-wise reasoning, differential diagnosis, and subsequent actions (e.g., ordering lab tests). The bottom Referenced Multi-Turn Trajectory provides a ground-truth clinical timeline for comparison, demonstrating the established diagnostic process leading to the final diagnosis.
• Figure 3: Overview of single-turn evaluation settings and results. a shows the single-turn evaluation setting for examination recommendation measured with the hit ratio. b the single-turn evaluation setting for final diagnosis measured with the accuracy. c compares our DiagAgent variants against 10 leading LLMs and 2 agentic systems on examination recommendation and diagnostic decision-making in the single-turn setting.