# An AI agent for treatment reasoning over a biomedical tool universe

arXiv: https://arxiv.org/abs/2606.28692 · HTML: https://mlsanity.com/paper/2606.28692/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Clinical and biomedical applications, one of 17 directions in LLM reasoning.
- Builds on 21 of its 64 references that are in the map.
- First posted on arXiv 27 Jun 2026.

## Placed in

- LLM reasoning / [Clinical and biomedical applications](https://mlsanity.com/lr/clinical-and-biomedical-applications/)
- LLM agents and tool use / [Medical and scientific domain agents](https://mlsanity.com/agents/medical-and-scientific-domain-agents/)

## Builds on (most cited first, 21 of 21)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Self-Instruct: Aligning Language Models with Self-Generated Instructions](https://mlsanity.com/paper/2212.10560/) (2022)
- [Retrieval-Augmented Generation for Large Language Models: A Survey](https://mlsanity.com/paper/2312.10997/) (2023)
- [Large Language Models Encode Clinical Knowledge](https://mlsanity.com/paper/2212.13138/) (2022)
- [Chatbot Arena: An Open Platform for Evaluating LLMs by Human Preference](https://mlsanity.com/paper/2403.04132/) (2024)
- [OpenAI GPT-5 System Card](https://mlsanity.com/paper/2601.03267/) (2025)
- [STaR: Bootstrapping Reasoning With Reasoning](https://mlsanity.com/paper/2203.14465/) (2022)
- [ToolRL: Reward is All Tool Learning Needs](https://mlsanity.com/paper/2504.13958/) (2025)
- [MEDITRON-70B: Scaling Medical Pretraining for Large Language Models](https://mlsanity.com/paper/2311.16079/) (2023)
- [ToolACE: Winning the Points of LLM Function Calling](https://mlsanity.com/paper/2409.00920/) (2024)
- [Time Travel in LLMs: Tracing Data Contamination in Large Language Models](https://mlsanity.com/paper/2308.08493/) (2023)
- [Empowering Biomedical Discovery with AI Agents](https://mlsanity.com/paper/2404.02831/) (2024)
- [Towards Accurate Differential Diagnosis with Large Language Models](https://mlsanity.com/paper/2312.00164/) (2023)
- [Geometric-Mean Policy Optimization](https://mlsanity.com/paper/2507.20673/) (2025)
- [A Toolbox for Surfacing Health Equity Harms and Biases in Large Language Models](https://mlsanity.com/paper/2403.12025/) (2024)
- [KGARevion: An AI Agent for Knowledge-Intensive Biomedical QA](https://mlsanity.com/paper/2410.04660/) (2024)

## Abstract (arXiv)

Treatment reasoning underpins every therapeutic decision, integrating disease context, comorbidities, medications, contraindications, and evolving biomedical knowledge to select an appropriate therapy. It is inherently iterative: candidates are weighed against many constraints, revised as evidence emerges, and grounded in verifiable sources. Here we introduce ATHENA-R1, an AI agent for treatment reasoning across all FDA approved drugs since 1939, trained by reinforcement learning over a universe of 212 biomedical tools. At each step it identifies missing information, selects and runs relevant tools, and incorporates the evidence. To train it without human-annotated traces, we build a two-level self-learning framework: multi-agent systems construct the tools, tasks, and reasoning trajectories for supervised fine-tuning, then reinforcement learning with scientific feedback rewards reasoning quality (evidence gathering, grounded tool use, logical non-redundancy). Across five benchmarks of 3,168 drug reasoning tasks and 456 patient treatment cases, ATHENA-R1 outperforms language models and tool-use systems, reaching 94.7% accuracy on open-ended drug reasoning and 82.9% on treatment reasoning, 17.8 and 10.7 points above GPT-5. In blinded evaluations by experts from 28 rare disease organizations, it is preferred over reference models on all criteria, and physicians rated it favorably on complex hospitalized cardiovascular and infectious-disease cases. Adverse-event hypotheses it generated, tested in electronic health records from 5.4 million patients, reached adjusted odds ratios of 1.48-1.84, with no elevation among negative controls. Because it requires knowing what evidence to seek before concluding, treatment reasoning has long been hard for AI; we show it can be reframed as a learnable process of iterative evidence gathering that reinforcement learning can train AI to perform.