DiagAgent-8B

DiagAgent-8B: RL-Optimized Diagnostic Agent

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

DiagAgent‑8B is a reinforcement learning‑optimized large language model for interactive, multi‑turn diagnostic reasoning. Unlike one‑shot medical LLMs, it can:

• recommend the most informative examinations,
• update the working diagnosis as new evidence arrives,
• decide when to stop and commit to a final diagnosis.

DiagAgent‑8B is trained end‑to‑end inside the DiagGym virtual clinical environment with multi‑turn RL (GRPO), enabling safe, closed‑loop learning without real‑world risk.

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

Quickstart

Run locally with transformers

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "Henrychur/DiagAgent-8B"
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",
    attn_implementation="flash_attention_2"
)

def chat(messages, max_new_tokens=1024, temperature=0.0):
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        do_sample=(temperature > 0),
        eos_token_id=tokenizer.eos_token_id
    )
    gen_ids = outputs[0][inputs.input_ids.shape[-1]:]
    return tokenizer.decode(gen_ids, skip_special_tokens=True).strip()

SYSTEM_PROMPT = (
    "You are a medical AI assistant. Analyze patient information, suggest relevant tests, "
    "and provide a final diagnosis when sufficient information is available.\n\n"
    "RESPONSE FORMAT:\n"
    "If more information is needed:\n"
    "```\n"
    "Current diagnosis: <your current best diagnosis>\n"
    "Based on the patient's initial presentation, the following investigation(s) should be performed: <one additional test>\n"
    "Reason: <reason for the test>\n"
    "```\n"
    "If sufficient information exists for diagnosis:\n"
    "```\n"
    "The available information is sufficient to make a diagnosis.\n"
    "Diagnosis: <final diagnosis>\n"
    "Reason: <brief justification>\n"
    "```"
)

initial_inquiry = (
    "- Patient Information: ___ y/o F\n"
    "- Chief Complaint: Early satiety, weight loss, abdominal pain\n"
    "- HPI: 1-month weight loss (10 lbs), early satiety, fatigue; prior emesis; reduced intake; denies fever/chills.\n"
    "- PMH: Asthma, hyperlipidemia, HTN, osteoarthritis, polymyalgia rheumatica, CAD (NSTEMI), osteoporosis, H. pylori, s/p TAH/USO.\n"
    "- Allergy: Lisinopril."
)

messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": initial_inquiry}
]

print(chat(messages))

Tips:

• Use device_map="auto" to spread weights; for constrained VRAM consider 4‑bit loading (bitsandbytes).
• Keep temperature low (0–0.3) for more deterministic decision‑making.

Serve with vLLM and use the OpenAI‑compatible Chat API

vllm serve Henrychur/DiagAgent-8B --served-model-name DiagAgent-8B

from openai import OpenAI
client = OpenAI(base_url="http://localhost:8000/v1", api_key="NONE")

messages = [
    {"role": "system", "content": "You are a medical AI assistant... (same format as above)"},
    {"role": "user", "content": "Initial inquiry text..."}
]

resp = client.chat.completions.create(
    model="DiagAgent-8B",
    messages=messages,
    temperature=0.0,
    max_tokens=1024
)
print(resp.choices[0].message.content)

Evaluation Results

The following tables are taken directly from the project evaluation. For evaluation details and scripts, see the paper and the GitHub repository.

Single‑Turn Evaluation

End‑to‑End Evaluation (Automatic Metrics)

End‑to‑End Evaluation (Rubric-based Metrics)

Training Details

DiagAgent‑8B is optimized with multi‑turn RL (GRPO) inside DiagGym.

• Trajectory construction:
  • Initial inquiry (structured patient history without final diagnosis)
  • Iterative steps: preliminary diagnosis → recommended exam + rationale → exam result
  • Final diagnosis focused on a single primary condition
• Reward design:
  • Diagnosis Accuracy (exact match)
  • Examination Recommendation F1 (overlap with reference EHR exams)
  • Turn Penalty (discourages >12 interaction turns)

For implementation and scripts, see DiagAgent/train/rl/ in the GitHub.

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

• Email: [email protected]
• GitHub: https://github.com/MAGIC-AI4Med/DiagGym