DiagGym

DiagGym: Virtual Clinical Environment (EHR World Model)

<div align="center"> <img src="https://raw.githubusercontent.com/MAGIC-AI4Med/DiagGym/main/assets/logo.png" width="150"/> <div align="center"></div> </div>

DiagGym is a high‑fidelity virtual clinical environment that simulates examination outcomes conditioned on evolving patient states. It enables closed‑loop training and evaluation of diagnostic agents with reinforcement learning, without real‑world risk.

Overview

In clinical practice, diagnosis is iterative: clinicians decide which examinations to order, interpret new findings, and commit to a final diagnosis when sufficient evidence is available. Most medical LLMs collapse this into a single step using static records, which prevents learning when to order tests and when to stop.

DiagGym addresses this by generating exam results conditioned on a patient profile, prior exams, and the next requested exam. Diagnostic agents can interact in a closed loop:

• Request any exam (including those not present in real EHR trajectories)
• Receive realistic, consistent results instantly
• Adapt diagnostic plans across multiple turns

This enables safe, scalable training and evaluation of end‑to‑end diagnostic strategies.

Details can be found in our paper https://arxiv.org/abs/2510.24654

Quickstart

Option A: Serve with vLLM and query via OpenAI‑compatible API

1. Start a server:

vllm serve Henrychur/DiagGym --served-model-name DiagGym

1. Inference example (Python):

import os
from typing import Optional
from openai import OpenAI

# Separator for concatenating past events
SEP = "<SEP>"
stop_tokens = [SEP, "<endoftext>"]

# API credentials
API_KEY = os.getenv("OPENAI_API_KEY", "None")
API_BASE = os.getenv("OPENAI_API_BASE", "http://localhost:8079/v1")

class DiagGym:
    def __init__(self, model_name_or_path: str, api_key: str = API_KEY, api_base: str = API_BASE) -> None:
        self.model_name_or_path = model_name_or_path
        self.client = OpenAI(api_key=api_key, base_url=api_base)

    def simulate(self, context: str, past_events_list: list, exam_name: str) -> Optional[str]:
        """Generate exam results based on patient context and past events."""
        context += "\nThe following summarizes the results from the patient's medical examination:\n"
        
        if len(past_events_list) == 0:
            input_prompt = context + "Exam name:\n" + exam_name + "\nExam results:\n"
        else:
            past_events_str = [
                f"Exam name:\n{event_name}\nExam results:\n{resp}"
                for (event_name, resp) in past_events_list
            ]
            input_prompt = context + SEP.join(past_events_str) + SEP + "Exam name:\n" + exam_name + "\nExam results:\n"
        
        response = self.client.completions.create(
            model=self.model_name_or_path,
            prompt=input_prompt,
            max_tokens=8192,
            temperature=1.0,
            stop=stop_tokens
        )
        return response.choices[0].text.strip()

# Example inputs
diaggym = DiagGym("DiagGym")
context = "Name: ___ ... Chief Complaint: right lower quadrant pain ... (details omitted)"
past_events_list = [("Anion Gap", "Value: 16; Numeric Value: 16.0; Units: mEq/L; Reference Range: 8.0 - 20.0")]
exam_name = "CHEST (PORTABLE AP)"

print(diaggym.simulate(context, past_events_list, exam_name))

Option B: Use transformers locally (no server)

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "Henrychur/DiagGym"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype="auto",
    device_map="auto"
)

SEP = "<SEP>"
context = "Name: ___ ... Chief Complaint: right lower quadrant pain ..."
past_events_list = [("Anion Gap", "Value: 16; Numeric Value: 16.0; Units: mEq/L; Reference Range: 8.0 - 20.0")]
exam_name = "CHEST (PORTABLE AP)"

if len(past_events_list) == 0:
    prompt = (
        context
        + "\nThe following summarizes the results from the patient's medical examination:\n"
        + "Exam name:\n" + exam_name + "\nExam results:\n"
    )
else:
    past_events_str = [
        f"Exam name:\n{event_name}\nExam results:\n{resp}"
        for (event_name, resp) in past_events_list
    ]
    prompt = (
        context
        + "\nThe following summarizes the results from the patient's medical examination:\n"
        + SEP.join(past_events_str) + SEP
        + "Exam name:\n" + exam_name + "\nExam results:\n"
    )

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
    **inputs,
    max_new_tokens=512,
    temperature=1.0,
    do_sample=True,
    eos_token_id=tokenizer.eos_token_id
)
text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(text)

Use Cases

• Reinforcement learning of diagnostic agents (e.g., GRPO): interactively recommend exams, receive simulated results, and decide when to finalize diagnosis.
• End‑to‑end evaluation of agentic diagnostic reasoning in a closed‑loop environment.
• Research simulation and data augmentation on examinations not present in real EHR trajectories.

For an end‑to‑end agent example, see Henrychur/DiagAgent-14B.

Evaluation Summary

On 863 MIMIC‑IV cases, DiagGym shows strong step‑level similarity and full‑chain consistency, low distribution gaps for both numeric and free‑text exams, and competitive efficiency.

Training and Data

• Data construction:

  • Source: MIMIC‑IV
  • Each case is reorganized into a patient profile (chief complaint, physical exam, history, allergies, family/social history, discharge diagnosis) and a time‑ordered set of examinations (laboratory, microbiology, radiology).
  • Filters remove cases without physical exams or with pre‑established diagnoses; exam names are standardized; only exams within 24 hours before admission are included to avoid diagnosis leakage.

• Training setup:

  • Objective: treat all exam results as plain text and train with autoregressive, token‑level NLL.
  • Model initialization: Qwen2.5‑Instruct‑7B.
  • Typical configuration: AdamW, learning rate 4e‑5 with linear warmup, max sequence length 8192; batch/accumulation tuned to hardware.
  • Hardware: NVIDIA A100 (example: 8 × 80GB).

Citation

@misc{qiu2025evolvingdiagnosticagentsvirtual,
      title={Evolving Diagnostic Agents in a Virtual Clinical Environment}, 
      author={Pengcheng Qiu and Chaoyi Wu and Junwei Liu and Qiaoyu Zheng and Yusheng Liao and Haowen Wang and Yun Yue and Qianrui Fan and Shuai Zhen and Jian Wang and Jinjie Gu and Yanfeng Wang and Ya Zhang and Weidi Xie},
      year={2025},
      eprint={2510.24654},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2510.24654}, 
}

Contact

• For implementation detailes, please refer to our GitHub: https://github.com/MAGIC-AI4Med/DiagGym
• Email: [email protected]