Skip to content

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

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

Tool-augmented math agent combining internal and external computation; efficient and accurate

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

AgentMath is a tool-augmented agent framework that integrates LLM reasoning with sandboxed Python code execution for mathematical problem-solving. It introduces automated trajectory synthesis, an agentic reinforcement learning paradigm with GRPO, and a 4–5× faster asynchronous training infrastructure. AgentMath-30B-A3B achieves 90.6% / 86.4% / 73.8% on AIME24 / AIME25 / HMMT25, surpassing OpenAI-o3-mini and Claude-Opus-4.0-Thinking.

Problem

Pure-text Large Reasoning Models (LRMs) suffer from computational errors, redundant self-correction loops, and token inefficiency on problems requiring precise arithmetic, symbolic manipulation, or geometric reasoning. Three compounding bottlenecks existed: (1) scarcity of high-quality tool-augmented training trajectories, (2) no RL mechanism for autonomously learning optimal tool-use strategies in an agentic, multi-turn interactive setting, and (3) no training infrastructure capable of handling ultra-long sequences (up to 96k tokens) with up to 96 tool calls per trajectory under traditional batch-synchronous RL.

Method

Stage 1 — Tool-Driven Data Synthesis (SFT data). Pure-text chain-of-thought data (from AM-Thinking, OpenThoughts, AceReason, filtered to 346k samples via n-gram/MinHash/semantic deduplication against eval sets) is transformed into tool-augmented trajectories via an injection function \(F_{\text{inject}}\):

\[\tau'_{\text{agent}} = F_{\text{inject}}(\tau_{\text{text}}), \quad \text{where } s_{\text{calc}} \Rightarrow (c,\, o_{\text{sim}})\]

A teacher model (DeepSeek-V3) injects executable code blocks into computationally intensive reasoning steps (3k-token segments). Four refinement procedures follow: format normalization, sandbox executability checks with bounded resampling, ground-truth execution output replacement (\(o_{\text{real}} = E(c)\), verified by Qwen3-32B), and AST-depth-based tool-rationality filtering. Failed trajectories are recycled to inject self-correction traces (diagnose → repair → re-execute), producing 14k correction examples. Final SFT corpus: 316k samples, avg 8.3 tool calls, avg 16.9k tokens.

Stage 2 — Agentic Reinforcement Learning. The MDP formulation interleaves <think>, <code>, and <interpreter> segments. GRPO (with dynamic sampling, asymmetric gradient clipping, token-level loss, KL removal following DAPO) optimizes only over <think> and <code> tokens; <interpreter> tokens are masked. Reward:

\[R_{\text{total}} = R_{\text{acc}} + \mathbf{1}[R_{\text{acc}}=1] \cdot R_{\text{tool}}\]

where \(R_{\text{acc}} \in \{0,1\}\) is binary correctness and \(R_{\text{tool}} = \min(R_{\text{max}},\, \alpha + \beta \cdot N_{\text{code}})\) incentivizes tool use conditioned on correctness. RL training uses 42k hard problems in three stages with progressively increasing max length (48k→72k→96k tokens) and tool-call caps (48→72→96).

Training Infrastructure. Three innovations enable scalable agentic RL: (1) Request-level asynchronous rollout scheduling — coroutine-driven decoupling of inference from agent logic eliminates head-of-line blocking; (2) Agentic Partial Rollout — each trajectory \(\tau\) is decomposed into budget-limited segments \(\tau^{(1)} \oplus \dots \oplus \tau^{(N)}\) capped by per-segment length \(L_{\text{seg}}\) and tool-call count \(T_{\text{seg}}\), yielding 2.2–2.5× speedup; (3) Prefix-Aware Weighted Load Balancing — requests routed by \(k^* = \arg\min_k W_k\) with weight \(w_j = L_j/L_{\text{base}} + w_{\text{base}}\) and LRU-cache sticky sessions for KV-cache reuse.

Key Contributions

  • Automated three-stage trajectory synthesis pipeline converting pure-text CoT into verified, executable tool-augmented agent trajectories at scale (316k samples), including self-correction capability injection.
  • Agentic RL paradigm with GRPO that dynamically interleaves NL generation and real-time code execution, enabling emergent code self-correction and autonomous tool-use strategy learning.
  • Scalable asynchronous training system (async scheduler + partial rollout + prefix-aware load balancing) achieving 4–5× end-to-end training throughput improvement for ultra-long, multi-turn RL.
  • State-of-the-art results on AIME24/25 and HMMT25, surpassing o3-mini and Claude-Opus-4.0-Thinking at 30B scale.

Results

  • AgentMath-30B-A3B: 90.6% / 86.4% / 73.8% on AIME24 / AIME25 / HMMT25 (avg@32); surpasses OpenAI-o3-mini and Claude-Opus-4.0-Thinking; competitive with OpenAI-o3, Gemini-2.5-Pro, DeepSeek-R1-671B-0528.
  • AgentMath-8B: 89.8% / 84.7% / 71.3%; outperforms tool-augmented CIR-Qwen3-NT8-8B and text-only DS0528-Qwen3-8B (86.0% / 76.3% / 61.5%).
  • AgentMath-1.7B: 59.6% / 48.1% / 40.2%; surpasses tool-augmented CoRT-1.5B (43.1% / 30.2% / 20.1%) and text-only OpenReasoning-1.5B (55.5% / 45.6% / 31.5%).
  • STILL-3-TOOL-32B (tool-augmented baseline, comparable scale): 81.7% / 64.2% / 45.4% — substantially below AgentMath-30B-A3B.
  • Infrastructure: parallel sandbox cluster reduces tool-call latency from 175 s to 1.2 s; partial rollout yields 2.2–2.5× speedup; overall system 4–5× faster end-to-end.

Limitations

  • AgentMath-235B-A22B is trained with SFT only (no RL) due to computational cost constraints; the RL gains demonstrated at smaller scales may not be reflected at this size.
  • Evaluation uses avg@32 inference runs, which may not match deployment settings where fewer samples are feasible.
  • The trajectory synthesis pipeline depends on proprietary teacher models (DeepSeek-V3, DeepSeek-R1-0528); reproducibility requires access to these or equivalent models.
  • The composite reward \(R_{\text{tool}}\) linearly incentivizes tool call count conditioned on correctness; pathological over-invocation behaviors under distribution shift are not analyzed.
  • Scope is limited to mathematical competition problems; generalization of the agentic RL paradigm to other tool-use domains is not demonstrated.

Relevance to Agentic AI / LLM Agents

AgentMath is a concrete instantiation of the agentic AI paradigm in which an LLM policy learns to orchestrate external tool calls within an MDP framework, acquiring emergent self-correction behaviors through multi-turn RL rather than static demonstrations. The agentic RL design — masked policy gradients over environment tokens, group-relative optimization, and reward shaping for tool efficiency — is a transferable blueprint for any agent that must interleave generation with external feedback. The infrastructure contributions (async rollout, partial rollout, prefix-aware load balancing) directly address the open engineering challenge of running RL at scale over long-horizon, multi-action-step agent trajectories, which is the central bottleneck for training capable tool-use agents. Together, these contributions advance the state of scalable agent training beyond the code-execution domain.