An agent trained in a simulated clinic beats GPT-4o at ordering the right tests

In medicine, a clinical diagnosis usually takes several moves from the physician: form a plausible hypothesis from the patient's symptoms, order the tests that confirm or rule it out, and decide when to stop testing and commit to an answer. Most large language models (LLMs) do well at diagnosis on fixed cases, but they fall short on planning — on choosing and prioritizing the tests that matter across several rounds. That calls for active information gathering, the kind of control problem handled by agents trained with reinforcement learning (RL), which optimize not just the final diagnosis but the policy for choosing actions.

A virtual clinic as a world model

In a new paper, the authors build a virtual clinic for clinical diagnosis where agents choose which action to take and are rewarded not only for the correct diagnosis but for optimizing the whole sequence of actions. The key to the environment is imitating real testing behavior across different clinical scenarios. At bottom this is a world-modeling problem: the environment predicts what a given test will show for a given patient, conditioned on prior actions, which closes the loop between agent and environment.

The authors build DiagGym, a text-based clinical diagnosis environment grounded in action sequences from electronic health records (MIMIC-IV), that simulates diagnostic test results. To check that DiagGym does not collapse the space of possible results into a handful of templates, they compute a range of distributional metrics. DiagGym agrees with real trajectories in 96.91% of cases by LLM judgment and 95% by physician judgment. Its metrics land closer to real data than those of general-purpose LLMs, and it runs faster than they do. The diagnostic task is then to pick the right tests and reach the right diagnosis sooner and at lower cost.

An agent that learns from its actions

Inside the DiagGym simulator the authors train DiagAgent, which at each step chooses which test to order or when to stop and diagnose. That pushes the agent to diagnose accurately, gather relevant information efficiently, and finish with as few actions as possible. The agent is trained with reinforcement learning (RL) to optimize the entire action sequence inside an interactive environment. What it learns is strategic information gathering under uncertainty, which goes beyond plain fine-tuning on real cases.

DiagAgent performs better in two settings:

1) The single-step task, where at each step it has to propose the next test or give a diagnosis;

2) The full task, where it actively plays out the entire action sequence inside the environment.

On the single-step task the agent beats two strong agentic baselines and 10 LLMs: +9.34% in diagnostic accuracy and +44.03% in diagnostic test recommendation, including against GPT-4o and DeepSeek-v3, which lead the rest of the LLMs. On the full task it is ahead by +15.12% in diagnostic accuracy and +23.09% in F1.

What the case studies show

In a jaundice case the simulator reproduces the key findings (elevated bilirubin and the accompanying symptoms). That supports the claim that the world model is not flattening the picture into templates but keeping variety and context.

Appendicitis is the telling case. The agent narrows the hypothesis space quickly, first ordering a complete blood count and then, with justification, an abdominal CT; the findings confirm the diagnosis. The trajectory stays close to the reference one and saves unnecessary steps.

One case also shows a failure: in a ruptured ectopic pregnancy the agent gets the diagnosis right but does not initiate the urgent treatment, which sits outside the current task scope. The diagnostic reasoning is there; the actions are not.

Why it works

First, the environment itself returns plausible, varied results, so the agent learns from rich scenarios instead of the same templates over and over. Second, the training objective: what gets optimized is not a well-phrased answer but the whole path to it. RL over trajectories consistently beats pure instruction fine-tuning and raises the quality ceiling across different base model families.

The limitations are stated honestly. The training and evaluation sets were assembled with care, but RL was run only on models up to 14 billion parameters; larger ones could deliver another jump in performance. Only diagnosis was evaluated: extending the world model to treatment and prognosis may take further work. The authors welcome additional evaluation methods.

AI thus moves toward the role of a thinking medical investigator that does not simply answer questions but plans, tests hypotheses and learns the craft of clinical reasoning — a step toward genuinely intelligent medical assistants.

💾 Code