Healthcare AI GYM for Medical Agents

2026-05-01

AI Summary

This work addresses the lack of a unified reinforcement learning environment for medical AI agents that supports multi-turn interaction and generalization training. To this end, the authors develop a Gymnasium-compatible medical training platform encompassing ten clinical domains, over 3,600 tasks, 135 specialized tools, and 828,000 medical knowledge snippets. They further propose Turn-level Truncated Online Policy Distillation (TT-OPD), which employs a gradient-free EMA teacher model to provide outcome-aware KL regularization at each dialogue turn, thereby mitigating sparse rewards and training instability. Experimental results demonstrate that TT-OPD achieves state-of-the-art performance on 10 out of 18 benchmarks, with an average improvement of 3.9 percentage points, significantly accelerates convergence, and enables controllable response length and sustained multi-turn tool usage.

Abstract

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.

Problem

• medical AI agents
• clinical reasoning
• reinforcement learning
• multi-turn interaction
• training stability

Innovation

• Healthcare AI GYM
• multi-turn agentic RL
• tool-use stability
• on-policy distillation
• outcome-aware KL regularization