PaperlayerView exact revision on arXiv

Source-linked AI summary

Healthcare AI GYM for Medical Agents

Minbyul Jeong

arXiv:2605.02943v1cs.LGcs.AI

TL;DR

Medical AI lacks a unified environment for training agents on broad, multi-turn clinical reasoning with specialized tools. This paper introduces HEALTHCARE AI GYM and TT-OPD, which achieves the best performance on 10 of 18 benchmarks with average +3.9 pp improvement over the non-RL baseline.

  • Problem

    Medical AI training lacks a unified environment spanning broad clinical domains, specialized tools, and multi-turn reasoning needed for agentic clinical execution.

  • Method

    The paper builds HEALTHCARE AI GYM and proposes TT-OPD, which uses outcome-aware teacher regularization to provide dense guidance at every conversation turn.

  • Results

    TT-OPD achieves the best performance on 10 of 18 benchmarks, with an average +3.9 pp improvement over the non-RL baseline.

  • Takeaways & Limitations

    Multi-turn agentic RL can exhibit response explosion and collapse into verbose single-turn monologues, while TT-OPD maintains controlled lengths and sustained tool use.

  • Takeaways & Limitations

    The evaluation assumes single-turn generation produces zero accuracy on all benchmarks, limiting interpretation of comparisons involving single-turn evaluation.

Abstract

from arXiv · show

Clinical reasoning demands multi-step interactions -- gathering patient history, ordering tests, interpreting results, and making safe treatment decisions -- yet a unified training environment provides the breadth of clinical domains and specialized tools to train generalizable medical AI agents through reinforcement learning remains elusive. We present a comprehensive empirical study of multi-turn agentic RL for medical AI, built on \gym{}, a gymnasium-compatible environment spanning 10 clinical domains with 3.6K+ tasks, 135 domain-specific tools, and a knowledge base of 828K medical passages. Our analysis reveals that agentic multi-turn structure degrades into verbose single-turn monologues, characterized by monotonic length explosion and a simultaneous erosion of tool-use frequency. We characterize how this collapse, alongside distillation instability, stems from the misalignment of sparse terminal rewards with sequential clinical trajectories. We find that vanilla GRPO achieves strong final accuracy on some benchmarks but suffers from training instability, evidenced by significant oscillations in response length and prolonged convergence periods. To improve training efficiency and stability, we propose Turn-level Truncated On-Policy Distillation (TT-OPD), a self-distillation framework where a gradient-free EMA teacher leverages outcome-privileged information to provide dense, outcome-aware KL regularization at every conversation turn. TT-OPD achieves the best performance on 10 of 18 benchmarks with an average +3.9~pp improvement over the non-RL baseline with faster early convergence, controlled response length, and sustained multi-turn tool use.

1 Introduction

The paper introduces HEALTHCARE AI GYM to train medical agents for multi-step clinical reasoning and studies the resulting challenges in multi-turn reinforcement learning. It proposes TT-OPD, which improves benchmark performance while stabilizing training and preserving agentic behavior.

  • Motivation: Frontier medical models remain largely limited to passive, single-turn benchmarks, motivating agentic reinforcement learning for multi-step clinical decision-making.The proposed transition targets navigation of high-stakes uncertainty rather than static question answering.
  • Related work: Existing medical-agent environments cover only fragments of clinical reasoning and generally lack integrated tool use or explicit RL policy optimization.The cited systems variously emphasize diagnostic dialogue, multi-agent workflows, or code-centric tools.
  • Training pathologies: Multi-turn RL produces response explosion, with outputs growing monotonically as models use token-level coverage to compensate for absent intermediate feedback.The passage characterizes this as a pathology absent from single-turn settings.
  • Method and results: +3.9 pp is TT-OPD’s average improvement over the non-RL baseline, with the method best on 10 of 18 benchmarks; its EMA teacher and outcome-conditioned hints provide dense turn-level regularization.Vanilla GRPO achieves +9.4 pp training accuracy but exhibits the reported training instabilities.
  • Environment: 10 clinical domains, 3.6K+ tasks, 135 domain-specific tools, and 828K medical passages comprise HEALTHCARE AI GYM, alongside a safety-aware 5D reward function.The environment is Gymnasium-compatible and designed for broad medical-agent training.

2 Related Work

Prior medical-agent environments cover complementary parts of clinical reasoning but differ in tool support, workflow modeling, policy optimization, and clinical grounding. Related RL and multi-turn methods provide relevant foundations, while highlighting challenges in importance weighting and cross-turn credit assignment.

  • Medical AI Agents: Medical-agent environments address distinct fragments: diagnostic dialogue, multi-agent workflows, Gymnasium code tools, or professional-tool integration.AgentClinic lacks tool use and RL training; Agent Hospital lacks policy optimization; MedAgentGym uses code-centric rather than clinically grounded tools.
  • RL for LLMs and On-Policy Distillation: Policy-gradient RL underpins LLM alignment, while DPO bypasses reward models and GRPO, DAPO, and Dr. GRPO modify reward or clipping strategies.GRPO uses group-relative rewards; DAPO adds dynamic sampling and asymmetric clipping; Dr. GRPO removes length-normalization bias.
  • RL for LLMs and On-Policy Distillation: Online single-iteration GRPO has πθ/πold ≡1.0, limiting the relevance of DAPO’s clipping mechanism.The passage identifies this importance-ratio property specifically for online single-iteration GRPO.
  • Multi-Turn Agent Optimization: Multi-turn RL requires credit assignment across turns; process reward models offer step-level feedback but assume linear reasoning chains.Self-RAG and Self-BioRAG use adaptive retrieval and self-reflection, with Self-BioRAG targeting biomedical reasoning.

3 HEALTHCARE AI GYM: Environment Design

HEALTHCARE AI GYM is a Gymnasium-compatible environment for training medical agents across 10 clinical domains, unifying clinical execution with reinforcement-learning integration. It combines clinically grounded tools with a multidimensional reward framework emphasizing diagnostic precision, procedural safety, and coherence.

  • Environment scope and interface: 10 clinical domains are exposed through a unified Gymnasium API, including step(action)/render(), for integration with modern reinforcement-learning pipelines.The environment is designed to bridge static medical knowledge retrieval and agentic clinical execution.
  • Clinically grounded tools: 135 domain-specific tools are consolidated into 25 user-facing categories spanning evidence retrieval, clinical assessment, intervention actions, and reasoning scaffolds.The tools include BM25-based knowledge-base querying and 22 validated scoring instruments.
  • Reward design: A 5D Reward Function formalizes clinical priorities into a single optimization objective beyond binary accuracy.The dimensions are accuracy, procedure, safety, format, and coherence.
  • Reward design: wacc=0.25, wproc=0.20, wsafe=0.20, wfmt=0.10, and wcoh=0.10 make diagnostic precision and procedural safety the primary drivers of policy updates.An optional assertion dimension uses wassert=0.15 when rubric annotations are available.

4 Turn-Level Truncated On-Policy Distillation

TT-OPD regularizes medical-agent policies at every conversation turn using an EMA teacher and outcome-conditioned privileged hints. The method addresses sparse terminal-reward credit assignment while supporting stable, multi-turn behavior and controlling response-length growth.

  • Problem formulation: Clinical-agent episodes are modeled as POMDP trajectories whose sparse terminal reward is computed only after submission or horizon termination.Each action may contain natural-language reasoning or structured tool calls, and the episode ends at successful submit_answer() or horizon T.
  • Credit assignment: Sparse terminal rewards create severe credit-assignment problems because structured tool-call actions are difficult to annotate and tool execution dynamically changes observations.These factors make step quality context-dependent and complicate adapting process reward models to agentic environments.
  • GRPO baseline: In online single-iteration GRPO, πθ = πold makes the importance ratio identically 1.0, rendering multi-iteration clipping ineffective.GRPO replaces the learned value function with group-relative advantages computed across G rollouts per prompt.
  • TT-OPD objective: TT-OPD computes turn-level KL regularization across the trajectory, truncating contributions from turns exceeding context limit Lmax.Its EMA teacher tracks the student without explicit teacher gradient updates, providing structural regularization for multi-turn behavior.
  • Outcome-aware regularization: Outcome-conditioned hints modify the teacher distribution while remaining hidden from the student, stabilizing correct behaviors and penalizing procedural errors through outcome-aware KL gradients.A cosine length-controlled reward additionally discourages monotonic response-length growth as responses approach Lmax.

5 Experiments

Experiments evaluate RL-trained medical agents with the same multi-turn AgentRunner and domain tools used during training, revealing broad TT-OPD strengths alongside systematic agentic overhead on knowledge-recall tasks. TT-OPD also offers more stable training dynamics than vanilla GRPO, despite GRPO reaching a slightly higher peak accuracy.

  • Evaluation protocol: Single-turn evaluation yields zero accuracy because it truncates the learned search → assess → submit pipeline before submit_answer.The experiments therefore use the multi-turn AgentRunner and domain tools used during training.
  • Benchmark results: TT-OPD achieves the best performance on 12 of 18 benchmarks across multiple-choice QA, visual QA, EHR, and long-form QA.Its strongest reported results include MedQA 87.1%, MedMCQA 66.2%, MIMIC-III 62.7%, and eICU 57.1%.
  • Benchmark results: 83.8% text-based MMLU-Med accuracy falls to 60.6% with Base+AR and 65.5% with TT-OPD, showing systematic multi-turn overhead on knowledge-recall tasks.Tool calls can introduce errors on knowledge-recall benchmarks, while retrieval-augmented reasoning trades off parametric precision.
  • Training dynamics: 62.0% is GRPO’s peak accuracy at step 55, compared with 61.1% for TT-OPD at step 60, but GRPO oscillates between 7.7K–10.8K response tokens.TT-OPD’s cosine reward controls response-length instability while preserving sustained multi-turn tool use.
  • Training dynamics: 61.1% validation accuracy at step 60 is +8.5 pp over the 52.6% base model, while TT-OPD averages 59.5% (±1.4 pp) over steps 40–60.Both methods converge non-monotonically, but TT-OPD’s main advantage is training stability rather than raw peak accuracy.

6 Analysis

The ablation isolates distinct failure modes in multi-turn distillation: periodic resets cause KL collapse, unconditioned EMA yields generic regularization, and outcome hints without length control cause response explosion. Full TT-OPD combines EMA, outcome-aware hints, and cosine length control to achieve sustained convergence with stable turns, confirming multi-turn collapse as agentic-specific.

  • Ablation failure modes: Periodic teacher resets cause catastrophic KL collapse, with KL dropping 2.637 →0.343 at step 10 and accuracy declining 56.9% →49.3%.Copying student weights every T steps destroys the distillation gradient that guides the student.
  • Ablation failure modes: 53.8% accuracy at step 40 marks non-monotonic convergence after EMA replaces periodic resets, eliminating KL collapse but leaving turns eroding 7.82 →Without outcome-aware conditioning, the EMA teacher provides only a generic regularization signal.
  • Ablation failure modes: 91.7% clipping by step 40 follows outcome hints without length control, as responses grow monotonically toward Lmax.Positive hints reinforce detailed reasoning, causing response explosion.
  • Full TT-OPD: 61.1% accuracy with stable turns 7.0–7.4 over 60 steps shows full TT-OPD resolving response explosion through cosine length-controlled reward.EMA prevents KL collapse, outcome hints provide outcome-aware guidance, and cosine reward prevents response explosion.
  • Implication: Multi-turn collapse is agentic-specific, unlike single-turn OPD settings where response lengths are naturally bounded and turn structure is not a concern.The analysis extends prior instability patterns to the multi-turn agentic setting.

7 Discussion and Conclusion … A.4 Knowledge Base: 828K Passages

The paper introduces TT-OPD as a stable approach to multi-turn medical-agent RL and documents HEALTHCARE AI GYM’s modular environment, clinical tasks, tools, and knowledge infrastructure. Across 18 benchmarks, TT-OPD achieves broad competence while the appendix details the Gymnasium interface and construction supporting 10 clinical domains.

  • 7 Discussion and Conclusion: Process-level reward models, hierarchical sub-goal conditioning, and additional turn-level signals are proposed to improve credit assignment and training in long clinical episodes.Suggested sub-goals include achieving a correct diagnosis before treatment.
  • 7 Discussion and Conclusion: TT-OPD achieves the best performance on 10 of 18 benchmarks across multiple-choice, visual, EHR, and long-form medical QA.Reported scores include MedQA 87.1%, MedMCQA 66.2%, PathVQA 45.3%, Quilt-VQA 30.7%, MIMIC-III 62.7%, eICU 57.1%, LiveQA 62.5%, and MedicationQA 60.9%.
  • 7 Discussion and Conclusion: Vanilla GRPO gains +9.4 pp in training accuracy, peaking at 62.0% at step 55, but oscillates between 7.7–10.8K tokens in response length.It leads on KQA-Golden 65.3% and KQA-Silver 64.9%, while response explosion, multi-turn collapse, and distillation instability characterize multi-turn agentic RL.
  • A Healthcare AI GYM: Detailed Construction: HEALTHCARE AI GYM comprises ∼30K lines of code across 10 clinical domains in a modular construction detailed by the appendix.The appendix provides comprehensive design, implementation, and construction details.
  • A.1 Gymnasium Interface: The environment implements the Gymnasium API with text-based conversation observations and actions, tool execution through step(), and termination at submit_answer() or max_turns.The scalar reward is computed from 5 dimensions, while observations can contain up to 100000 characters and actions up to 10000 characters.
  • A.2 Domain Design: The task system defines structured patient scenarios, expected tool interactions, natural-language assertions, and reward bases, with 3,631 instantiated tasks and 135 registered tools.Of the tasks, 2,657 are used for RL training and 307 for validation; six multi-phase clinical pathways support cross-domain journeys.
  • A.3 Tool System Implementation: Tools use decorator-registered READ, WRITE, THINK, and GENERIC types, with automatic function schemas and error-returning execution for invalid actions.Domain tools span knowledge search, clinical assessment, patient data access, clinical actions, reasoning, and documentation across 9 domains.
  • A.4 Knowledge Base: 828K Passages: The knowledge base indexes 828,473 passages in an SQLite FTS5 database using Porter tokenization and BM25 ranking, with all search tools sharing one thread-safe backend.Sources include 581K MedCPT passages, 122K biomedical QA pairs, 83K generator passages, and 52K MedInstruct passages.

A.5 5D Reward Implementation … D.1 Log-Probability Baseline (Text-Only)

The paper combines a composable five-dimensional reward system and behavioral policies with TT-OPD, a turn-level EMA-teacher distillation method for stable multi-turn tool-augmented learning. Domain toolkits support standardized interactions, while text-only log-probability results show that GRPO preserves parametric knowledge.

  • A.5 5D Reward Implementation: The reward system composes accuracy, process quality, safety, format, and coherence functions into a scalar suitable for GRPO training.The implementation is approximately 400 lines, and grpo_reward_fn() computes the weighted reward dimensions for GRPOTrainer.
  • A.5 5D Reward Implementation: Accuracy uses exact match for MCQs, ROUGE-1 plus BLEU-1 overlap for open-ended answers, or BiomedBERT BERTScore with soft fallback.Exact-match rewards are 1.0 for correct answers and 0.0 otherwise.
  • A.5 5D Reward Implementation: Process quality weights coverage at 60%, diversity at 20%, and thoroughness at 20%, with rubric scoring weighted 70% when provided.Coverage measures expected tool calls with matching arguments; rubric checks required elements, required tools, and forbidden elements.
  • A.5 5D Reward Implementation: Safety detects 50+ violation patterns across five severity levels, capping total reward at 0.1 for critical violations and applying a −0.3 penalty for severe violations.Patterns map to AMA principles including nonmaleficence, beneficence, and autonomy.
  • A.6 Behavioral Policies: Domain policies inject safety-first, evidence-based behavioral guidelines, tool-use rules, and restrictions such as reviewing patient data before diagnosis.Policies also require allergy checks before prescribing and specialist transfer when cases exceed scope.
  • B TT-OPD Algorithm: TT-OPD filters for mixed-outcome prompts, computes group-relative advantages, and combines GRPO with turn-level KL distillation from an outcome-informed EMA teacher.Privileged teacher context is removed from outputs, while the teacher is updated periodically using an exponential moving average.
  • B TT-OPD Algorithm: TT-OPD stabilizes multi-turn optimization by anchoring the student to a smoothly updated teacher while preserving intermediate reasoning and tool-use behavior.Its design also regulates response length and mitigates policy-collapse and KL-instability failure modes.
  • C Domain Tool Inventory: Each domain toolkit combines domain-specific tools with shared PubMed, evidence-retrieval, and medical-wiki tools using OpenAI-compatible function calling.Tools return JSON-serializable outputs; think() captures reasoning, and submit_answer() triggers reward evaluation.

D.2 Multi-Turn Agentic Evaluation

Across clinical evaluation domains, TT-OPD generally delivers the strongest results, especially for visual QA, EHR action execution, and several long-form tasks. Multi-turn agentic evaluation nevertheless introduces substantial overhead on some visual benchmarks, while GRPO remains strongest on selected knowledge-intensive long-form tasks.

  • Multiple-Choice QA: 87.1%: TT-OPD outperforms GRPO (85.5%), Base+AR (78.8%), and the no-tool base model (70.7%) on MedQA.This represents a +16.4 pp improvement over the no-tool base model.
  • Visual QA: 5 of 6: TT-OPD achieves the best or near-best result across visual QA benchmarks, including 45.3% on PathVQA.On PathVQA, TT-OPD exceeds base text (40.5%) and GRPO (41.5%), while agentic evaluation trails text-based evaluation on SLAKE and PMC-VQA.
  • EHR Reasoning: 62.7% and 57.1%: TT-OPD leads action-based EHR scoring on MIMIC-III and eICU, respectively.These scores exceed Base+AR (62.1%, 55.9%) and GRPO (61.1%, 55.5%) on the corresponding datasets.
  • Long-Form QA: 3 of 5: TT-OPD leads MedLFQA benchmarks, while GRPO leads KQA-Golden (65.3%) and KQA-Silver (64.9%).TT-OPD leads LiveQA (62.5%), MedicationQA (60.9%), and HealthSearchQA (45.3%); LiveQA improves from 53.2% with Base text to 62.5%.

E Analytical Insights

TT-OPD’s EMA teacher creates a self-correcting, non-monotonic convergence pattern by restoring the student toward recent behavior. Reward-component imbalance dilutes some gradient signals, while EMA avoids the sawtooth KL collapse caused by periodic teacher resets.

  • EMA convergence: TT-OPD accuracy rises, dips, then recovers higher because the EMA teacher provides a built-in self-correcting mechanism.The non-monotonic pattern is described as distinct from random noise.
  • EMA convergence: EMA implicitly reduces GRPO’s effective learning rate as student–teacher distance grows, amplifying KL gradients that dampen large policy shifts.The restoring force is proportional to the student–teacher distance and shaped by the Fisher information.
  • Gradient signal dilution: ∼51:1 dilution arises because accuracy’s weighted variability is 0.103, versus 0.002 for format, causing accuracy to dominate the gradient signal.The stated reward parameters are w_acc=0.25, σ_acc=0.41, w_fmt=0.10, and σ_fmt=0.02.
  • KL boundedness: Periodic resets create sawtooth KL dynamics: divergence accumulates during learning, then drops near zero when the teacher is overwritten, destroying distillation guidance.EMA instead yields continuous KL growth under the stated bounded-shift and Lipschitz conditions.

F Training Hyperparameters

This section presents the training hyperparameters for vanilla GRPO and TT-OPD experiments conducted with Qwen3.5-9B.

  • Table 4 reports the training hyperparameters for the experiments.
  • The table includes vanilla GRPO as a baseline.
  • The table also includes TT-OPD.
  • Both methods are evaluated on Qwen3.5-9B.

G Benchmark Suite … J.3 Example 3: Transplant Medicine (MedQA, Correct)

The paper presents benchmark and safety-taxonomy resources, reports LLM usage, and illustrates TT-OPD-trained agents through multi-turn clinical trajectories spanning pharmacology, emergency medicine, and transplant medicine. These examples follow a reason → search → verify → submit pipeline and exhibit reason-first planning, graceful tool-failure handling, and verification loops.

  • I Usage of LLMs: Claude Opus 4.6 supported training experiments, manuscript concept development, grammatical checks, and citation bibliography checks.The passage attributes these uses to Claude Opus 4.6.
  • J Example Trajectories: Representative TT-OPD trajectories from training step 60 demonstrate the reason → search → verify → submit pipeline.The examples are presented to illustrate agentic reasoning patterns learned through reinforcement-learning training.
  • G Benchmark Suite: Table 5 defines the evaluation benchmark suite used in the study.The supplied passage identifies this resource but provides no benchmark contents or quantitative results.
  • H Safety Violation Taxonomy: Table 6 presents a five-level safety violation taxonomy.The supplied passage names the taxonomy but does not describe its categories.
  • J.1 Example 1: Pharmacology (MedQA, Correct): The pharmacology agent correctly selects A by reasoning that lisinopril increases bradykinin and decreases angiotensin II relative to losartan.It verifies the mechanism through evidence retrieval, medical-wiki search, and systematic evaluation of all four options; score: 1.18.
  • J.2 Example 2: Emergency Medicine (MedQA, Correct): The emergency-medicine agent correctly selects B, identifying pneumothorax from pleural injury as the likely consequence of the parasternal clavicular wound.It maps the wound to the lung apex, eliminates anatomically inconsistent alternatives, and searches for supporting medical literature.
  • J.3 Example 3: Transplant Medicine (MedQA, Correct): The transplant-medicine agent identifies chronic transplant rejection from the late timeline, recurrent jaundice, reduced hepatic blood flow, and severe enzyme elevation.It compares the answer choices against alternative diagnoses including acute rejection, fatty liver disease, and granulomatous inflammation.
  • J Example Trajectories: TT-OPD-trained agents reason before searching, recover from failed search tools, and verify retrieved evidence through think and analyze_answer_options.These patterns reduce irrelevant tool calls, prevent halting when specific backends fail, and cross-check evidence against the initial hypothesis.
Loading 2605.02943v1…