Skip to content

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

🕒 Published (v1): 2025-12-23 19:57 UTC · Source: Arxiv · Venue: ICLR 2026 · link

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 code execution for mathematical problem-solving. It introduces automated tool-use trajectory synthesis, an agentic RL paradigm with interleaved code generation and execution, and a scalable asynchronous training infrastructure achieving 4–5× throughput. The result is a 30B MoE model that surpasses OpenAI-o3-mini and Claude-Opus-4.0-Thinking on AIME24/25 and HMMT25.

Problem

Pure text-based LRMs (e.g., DeepSeek-R1, o3) make frequent arithmetic/symbolic errors and waste tokens on redundant corrections. Tool-augmented approaches face three bottlenecks: (1) scarcity of high-quality tool-use trajectory data; (2) no established agentic RL paradigm for learning optimal tool-invocation strategies; (3) standard synchronous batch-RL frameworks are computationally infeasible for ultra-long sequences (~96k tokens, ~96 tool calls per trajectory).

Method

Stage 1 — SFT via Tool-Augmented Trajectory Synthesis: Pure-text CoT chains (from AM-Thinking, OpenThoughts) are partitioned into 3k-token segments; a teacher model (DeepSeek-V3) injects executable code blocks (<code>) and simulates outputs (<interpreter>) for computationally intensive steps. A three-stage refinement pipeline follows: (i) format normalization, (ii) sandbox execution verification with bounded resampling, (iii) teacher-simulated outputs replaced by ground-truth execution results and validated by Qwen3-32B. Trajectories excluded due to execution failures are repurposed: the teacher generates error-diagnosis → code-repair → re-execution self-correction traces, producing 14k self-correction samples. Final SFT dataset: 316k trajectories averaging 8.3 tool calls and 16.9k tokens.

Stage 2 — Agentic RL with GRPO: The policy generates interleaved <think>/<code> tokens; execution pauses generation, runs the sandbox, and resumes with <interpreter> feedback. <interpreter> tokens are masked from the policy gradient — only agent-generated tokens receive advantage signals. Reward = binary correctness (math-verify) + tool-efficiency bonus scaled by invocation count (capped at R_max=1), applied only on correct trajectories. Multi-stage training progressively expands max sequence length (48k→72k→96k) and tool-call budget (48→72→96) triggered when truncation rate exceeds 10%.

Scalable Infrastructure: (i) Request-level asynchronous rollout scheduling decouples model inference from sandboxed execution via coroutines, eliminating head-of-line blocking and reducing tool-call latency from 175s to 1.2s. (ii) Agentic Partial Rollout decomposes each trajectory into budget-limited segments (max length L_seg, max tool calls T_seg), achieving 2.2–2.5× speedup. (iii) Prefix-Aware Weighted Load Balancing routes partial rollout requests by prefix length to minimize KV-cache thrashing, with LRU sticky sessions for cache reuse.

Key Contributions

  • Automated code-injection pipeline that converts pure-text CoT into verified, executable tool-augmented trajectories at scale (316k samples), with multi-stage quality refinement and self-correction trace injection.
  • Agentic RL paradigm (GRPO-based) with interleaved generation–execute–resume loops, loss masking on environment tokens, and composite accuracy + tool-efficiency reward.
  • Scalable training infrastructure: async request scheduler + partial rollout + prefix-aware load balancing → 4–5× end-to-end RL throughput on 96k-token, 96-tool-call sequences.
  • SOTA results on AIME24/25 and HMMT25 at multiple scales (1.7B to 235B), with 30B-A3B (3B active) surpassing OpenAI-o3-mini and Claude-Opus-4.0-Thinking.

Results

  • AgentMath-30B-A3B (3B active params): 90.6% / 86.4% / 73.8% on AIME24 / AIME25 / HMMT25 — outperforms OpenAI-o3-mini, Claude-Opus-4.0-Thinking; competitive with OpenAI-o3 and Gemini-2.5-Pro.
  • AgentMath-235B-A22B-SFT (SFT only): 93.4% / 90.8% / 81.7% — surpasses DS-671B-0528 (91.4% / 87.5% / 77.0%), on par with Qwen3-235B-A22B-Thinking-2507.
  • AgentMath-8B: 89.8% / 84.7% / 71.3% vs. text-only DS-0528-Qwen3-8B (86.0% / 76.3% / 61.5%).
  • AgentMath-1.7B: 59.6% / 48.1% / 40.2% vs. CoRT-1.5B (43.1% / 30.2% / 20.1%) and OpenReasoning-1.5B (55.5% / 45.6% / 31.5%).
  • RL efficiency: AgentMath-RL reaches 76.2% (AIME24) in ~400 steps vs. Text-Based-RL at 68.7% in ~1600 steps — 4× faster convergence.
  • Data refinement ablation: unrefined synthesis 35.3% → full pipeline 60.5% on AIME24 (20k data).
  • Inference efficiency: tool-augmented sequences ~4k tokens shorter (~14% reduction) vs. text-only.

Limitations

  • AgentMath-235B-A22B is SFT-only (no RL) due to computational constraints; RL gains at >32B scale are unquantified.
  • Evaluation restricted to math competition benchmarks (AIME24/25, HMMT25); generalization to other domains (science, code reasoning) is not assessed.
  • Self-correction injection relies on a teacher model and covers only a small fraction (14k/316k) of failure cases; robustness under out-of-distribution errors is untested.
  • Partial rollout and asynchronous scheduling add significant infrastructure complexity, raising the barrier to reproduction.
  • Tool-call efficiency reward hyperparameters (α, β, R_max) require tuning and their sensitivity is only partially addressed in appendices.

Relevance to Agentic AI / LLM Agents

AgentMath is a direct instantiation of the tool-using agent paradigm — the model functions as an agent that autonomously decides when and how to invoke a code interpreter within a multi-turn reasoning loop, learning this policy via RL from environment feedback rather than from human demonstrations. The automated trajectory synthesis pipeline addresses the critical data bottleneck for training tool-use agents, while the agentic RL framework demonstrates that emergent capabilities (code self-correction, adaptive tool invocation) arise from environment interaction rather than supervised imitation. The scalable async infrastructure tackles a core open challenge in agentic RL: efficient training on long-horizon, tool-heavy trajectories with heterogeneous latency profiles — directly relevant to anyone building multi-step agents that call external services.