AgentMath: Empowering Mathematical Reasoning for Large Language Models via Tool-Augmented Agent

AI summary

Overview

Research area: Large language model reasoning, tool-augmented agents, and reinforcement learning for mathematical problem solving.

Technical level: Advanced — assumes familiarity with reinforcement learning (policy optimization, reward shaping), chain-of-thought reasoning, and distributed training infrastructure.

Scope: The paper proposes an end-to-end framework (data synthesis, agentic RL, and distributed training infrastructure) that lets LLMs interleave natural-language reasoning with real code execution to solve competition-level math problems.

What This Paper Is About

Even the strongest reasoning models like o3 and DeepSeek-R1 still make arithmetic and symbolic errors because they do math entirely in text, forcing them into long correction loops that waste tokens and reduce accuracy. AgentMath solves this by letting the model write and execute code mid-reasoning, then use the interpreter's output to continue. The paper's real contribution is making this work at scale: generating high-quality training data automatically, training the tool-use behavior with reinforcement learning, and building the infrastructure to make that training tractable.

Key Contributions

1. Automated tool-augmented data synthesis pipeline. A three-stage process converts existing pure-text chain-of-thought data into structured agent trajectories containing executable code and real interpreter feedback — injecting code at computation-heavy steps, then refining for format consistency, code executability, and feedback alignment.

2. Agentic reinforcement learning paradigm. An RL training loop where natural-language generation and code execution dynamically interleave, with environment-generated tokens masked from the policy gradient. Models learn tool-invocation strategies autonomously and develop emergent code self-correction.

3. Scalable async training infrastructure. Three systems techniques — request-level asynchronous rollout scheduling, agentic partial rollout, and prefix-aware weighted load balancing — yield 4–5× end-to-end training throughput speedup and make RL feasible on 96k-token sequences with up to 96 tool calls.

4. State-of-the-art results across model scales. AgentMath sets new records on AIME24, AIME25, and HMMT25 for open models at 1.7B, 8B, 30B-A3B, and 235B-A22B scales, with the 30B MoE (3B active parameters) beating OpenAI-o3-mini and Claude-Opus-4.0-Thinking.

Main Findings

• AgentMath-30B-A3B hits 90.6% / 86.4% / 73.8% on AIME24, AIME25, and HMMT25 respectively, surpassing o3-mini and Claude-Opus-4.0-Thinking and approaching o3, Gemini-2.5-Pro, and DeepSeek-R1-671B-0528 — despite activating only 3B parameters.

• Tool augmentation beats pure-text reasoning consistently. At 8B scale, AgentMath-8B reaches 89.8% / 84.7% / 71.3% versus 86.0% / 76.3% / 61.5% for a text-only comparable model. At 1.7B, it scores 59.6% / 48.1% / 40.2% versus 55.5% / 45.6% / 31.5% for the strongest text-only 1.5B baseline.

• Large efficiency gains from code execution. In a controlled SFT comparison on 20k samples, tool-augmented data beat text-only data (60.5% vs. 57.1% on AIME24). During RL, AgentMath reached 76.2% on AIME24 in ~400 steps versus ~1600 steps for the text baseline — a 4× training efficiency improvement, while also shortening inference sequences by roughly 4k tokens (~14%).

• Multi-stage RL progressively unlocks capability. Extending context from 48k to 72k to 96k tokens and tool-call limits from 48 to 96 raised AIME24 accuracy from 78.4% to 89.8% and AIME25 from 72.2% to 84.7%. Average tool calls per problem rose from 27 to 31, and the fraction of problems using code rose from 70% to 95%.

• Data refinement stages each add measurable accuracy. Unrefined synthetic data yielded only 35.3% on AIME24; format correction (+12.1), executability verification (+5.4), feedback alignment (+3.5), self-correction injection (+2.3), and selective feedback masking (+1.9) accumulated to 60.5%.

• Strong data scaling law. Growing the tool-augmented dataset from 2k to 300k samples lifted AIME24 accuracy from 27.2% to 78.4% and AIME25 from 21.1% to 72.2%.

• Infrastructure is the hidden bottleneck. Tool-call latency dropped from 175s to 1.2s via a distributed sandbox cluster, and partial rollout alone gave a 2.2–2.5× speedup.

Methodology in Plain English

The approach has three layers.

Data. The team started with ~346k math problems and their text-only solutions from existing sources, filtered to avoid benchmark contamination, and kept only hard problems. Using a strong teacher model (DeepSeek-V3), they replaced heavy computational steps in each solution with executable Python code, wrapped in <code> tags, and inserted execution results wrapped in <interpreter> tags. Every code block was then actually run in a sandbox; failures were retried or reverted to text. A judge model checked that the reported outputs matched reality. They also generated ~14k examples where the model sees a failed program and learns to diagnose, repair, and re-execute — teaching recovery behavior.

Training. The model first undergoes supervised fine-tuning on ~316k of these trajectories to learn the basic code-calling format. Then RL (using GRPO) trains it further: the model generates reasoning and code in a generate–pause–execute–resume loop, and receives reward based on answer correctness plus a small bonus for efficient tool usage. Crucially, tokens produced by the interpreter are masked out of the gradient so the model only learns from its own decisions. Problems where all sampled attempts succeed or all fail are filtered out, since they carry no learning signal.

Infrastructure. Because trajectories average 24k tokens and ~27 tool calls, standard batch-synchronous RL stalls badly. The team decoupled inference from agent logic, ran code in a distributed sandbox cluster, scheduled each trajectory as an independent asynchronous request, and broke long trajectories into budget-limited segments so no single sample hogs resources. A prefix-aware load balancer routes segments back to the same inference engine (via LRU caching) to reuse KV-cache and avoid recomputation.

Why This Matters

Impact on research. This work shows that tool use is not just a prompting trick but something that can be learned through RL — and that the marginal gains from tool augmentation compound with scale. The infrastructure contributions (partial rollout, async scheduling, prefix-aware balancing) address a real and under-discussed bottleneck: most agentic RL papers hand-wave the systems problem that makes long-horizon, tool-heavy training feasible. The emergent self-correction behavior also suggests agentic RL yields capabilities that SFT alone does not.

Real-world applications:

• Education technology — math tutors and problem-solving assistants that show verifiable, step-by-step work rather than confidently wrong arithmetic.
• Scientific and engineering computation — agents that can chain symbolic manipulation, equation solving, and numerical analysis while self-correcting when a step fails.
• Financial and quantitative analysis — workflows where a model must produce auditable calculations with reproducible intermediate results, not just plausible-sounding numbers.
• Enterprise automation and data analysis — general tool-using agents that reliably invoke APIs, query databases, and recover from execution errors.

Industry relevance. The 30B-A3B result is significant commercially: a sparse MoE with 3B active parameters matching much larger dense models means lower serving cost for frontier-level math reasoning. The 4–5× training speedup directly reduces the compute bill for anyone training agentic RL systems. And the fully automated data synthesis pipeline means teams don't need expensive human annotation to build tool-use training sets — the teacher-model-plus-verification loop is reproducible.

Future Directions

1. Extending beyond mathematics. The framework is domain-agnostic in principle. Applying it to physics, chemistry, formal verification, data science, and general software engineering — where correctness is also checkable — is the obvious next step and would test whether the reward design generalizes.

2. RL training at the largest scale. AgentMath-235B-A22B was trained with SFT only due to compute constraints, yet still reached 93.4% / 90.8% / 81.7%. Running the full agentic RL pipeline at that scale is untested and could close the remaining gap to OpenAI's tool-augmented o3 and o4-mini (which score 95.2–99.5% on AIME24/25).

3. Multi-tool and multi-environment agents. The current agent talks only to a code interpreter. Combining code execution with search, symbolic solvers, theorem provers, and diagram tools — and learning when to use which — is an open problem the current reward design does not address.

4. Reducing teacher-model dependence and improving robustness. The synthesis pipeline leans on DeepSeek-V3 as both generator and judge. Whether smaller or self-hosted models can sustain the same data quality, and how the learned tool-use policies degrade under distribution shift or adversarial problem formulations, remain open questions.

5. Pushing the length and call budgets further. The paper notes that truncation rates were held below 10% by expanding budgets. Whether 128k+ contexts and 200+ tool calls yield further gains — and whether the partial rollout segmentation scales gracefully — is unresolved.

Target Audience

This paper is most valuable to applied ML researchers and engineers working on agentic LLM systems, particularly those building RL training pipelines for tool-using models. It is also directly relevant to infrastructure and systems engineers who need to make long-horizon, multi-turn agent RL feasible on real hardware. Competition math and formal reasoning researchers will care about the benchmark results and the data synthesis pipeline. Product and strategy leads in AI will find the 30B-A3B cost/performance result compelling. Readers without background in RL policy optimization or distributed training will find the method sections dense, though the introduction, results tables, and cognition analysis are accessible to a broader technical audience.

Authors’ abstract

Large Reasoning Models (LRMs) like o3 and DeepSeek-R1 have achieved remarkable progress in reasoning tasks with long cot. However, they remain computationally inefficient and struggle with accuracy when solving problems requiring complex mathematical operations. In this work, we present AgentMath, an agent framework that seamlessly integrates language models' reasoning capabilities with code interpreters' computational precision to efficiently tackle complex mathematical problems. Our approach introduces three key innovations: (1) An automated method that converts natural language chain-of-thought into structured tool-augmented trajectories, generating high-quality supervised fine-tuning (SFT) data to alleviate data scarcity; (2) A novel agentic reinforcement learning (RL) paradigm that dynamically interleaves natural language generation with real-time code execution. This enables models to autonomously learn optimal tool-use strategies through multi-round interactive feedback, while fostering emergent capabilities in code refinement and error correction; (3) An efficient training system incorporating innovative techniques, including request-level asynchronous rollout scheduling, agentic partial rollout, and prefix-aware weighted load balancing, achieving 4-5x speedup and making efficient RL training feasible on ultra-long sequences with scenarios with massive tool invocation. The evaluations show that AgentMath achieves state-of-the-art performance on challenging mathematical competition benchmarks including AIME24, AIME25, and HMMT25. Specifically, AgentMath-30B-A3B attains 90.6%, 86.4%, and 73.8% accuracy respectively, surpassing OpenAI-o3-mini and Claude-Opus-4.0-Thinking while remaining competitive with OpenAI-o3, Gemini-2.5-Pro, and DeepSeek-R1-671B-0528.These results validate the effectiveness of our approach and pave the way for building scalable mathematical reasoning agents.