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, trained via a three-stage automated data synthesis pipeline and a novel agentic reinforcement learning paradigm. The system also introduces an asynchronous distributed training infrastructure that achieves 4–5× speedup for ultra-long, multi-turn RL rollouts. AgentMath-30B-A3B reaches 90.6%/86.4%/73.8% on AIME24/AIME25/HMMT25, surpassing OpenAI-o3-mini and Claude-Opus-4.0-Thinking.

Problem

Pure-text LRMs (o3, DeepSeek-R1) make frequent computational errors on problems requiring symbolic manipulation or large-number arithmetic, and existing tool-augmented approaches face three bottlenecks: (1) scarcity of high-quality tool-use training trajectories, (2) no RL paradigm for learning optimal tool-invocation strategies interactively, and (3) batch-synchronous RL infrastructure that cannot handle ultra-long sequences (up to 96k tokens, 96 tool calls) with heterogeneous latency from concurrent code execution.

Method

Data synthesis (SFT): A three-stage pipeline transforms pure-text CoT data into verified agent trajectories. Stage 1 uses DeepSeek-V3 as a teacher to inject <code>/<interpreter> blocks into computational reasoning steps via a code-injection function F_inject. Stage 2 applies four quality filters: format normalization, sandbox executability verification (with bounded resampling), ground-truth execution output replacement (verified by Qwen3-32B), and AST-depth heuristics for tool-usage rationality. Stage 3 injects self-correction traces by prompting the teacher to diagnose/repair failed code executions, producing 14k correction trajectories added to a 302k main set (316k total, avg 8.3 tool calls, 16.9k tokens/sample).

Agentic RL: Uses GRPO with a generate–pause–execute–resume loop, masking <interpreter> tokens from policy gradient updates. Reward combines binary answer correctness (R_acc) with a tool-usage efficiency bonus (R_tool = min(R_max, α + β·N_code)) conditioned on correctness. Multi-stage RL progressively expands context from 48k→72k→96k tokens and tool-call caps from 48→72→96, triggered when truncation rate exceeds 10%.

Training infrastructure: Three systems optimizations: (1) request-level asynchronous rollout scheduling (coroutine-driven, decoupled inference from tool I/O), reducing tool-call latency from 175s to 1.2s; (2) Agentic Partial Rollout that decomposes trajectories into budget-limited segments (max length L_seg, max invocations T_seg) to eliminate head-of-line blocking (2.2–2.5× speedup); (3) prefix-aware weighted load balancing with LRU sticky sessions for KV-cache reuse across partial rollout segments.

Key Contributions

  • Automated code-injection pipeline converting pure-text CoT into verified, executable agent trajectories at scale (316k samples), including self-correction trace injection.
  • Agentic RL paradigm with interleaved generation/execution, loss masking for environment tokens, adaptive batch construction with backfilling, and composite accuracy+tool-efficiency reward.
  • Asynchronous distributed training system (async scheduler + partial rollout + prefix-aware load balancing) achieving 4–5× end-to-end throughput improvement for ultra-long agentic RL.
  • Multi-stage RL curriculum that adaptively expands sequence length and tool-call budget based on truncation rate triggers.

Results

  • AgentMath-30B-A3B (3B active params): 90.6% AIME24, 86.4% AIME25, 73.8% HMMT25 — surpasses OpenAI-o3-mini, 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%, outperforming text-only DS-0528-Qwen3-8B (86.0% / 76.3% / 61.5%) and tool-augmented CIR-Qwen3-NT8-8B.
  • 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%).
  • AgentMath-235B-A22B-SFT: 93.4% / 90.8% / 81.7%, exceeding DeepSeek-R1-671B-0528 (91.4% / 87.5% / 77.0%) with SFT only.
  • Tool-augmented RL reaches 76.2% AIME24 in ~400 steps vs. ~1600 steps for text-based RL to reach an inferior 68.7% — 4× efficiency gain.
  • Data scaling from 2k→300k samples: AIME24 improves from 27.2% to 78.4%.
  • Each refinement stage contributes: unrefined data yields 35.3% AIME24; full pipeline reaches 60.5%.
  • Multi-stage RL training: AIME24 improves from 78.4% to 89.8% (+11.4%) and AIME25 from 72.2% to 84.7% (+12.5%); code utilization rises from 70% to 95%.

Limitations

  • AgentMath-235B-A22B was trained with SFT only due to computational cost; RL gains at that scale are unexplored.
  • Reward design for tool-usage efficiency (linear in N_code) may not generalize beyond mathematical computation to broader agentic tasks.
  • The data synthesis pipeline depends on a powerful teacher model (DeepSeek-V3) and a judge model (Qwen3-32B), limiting accessibility for resource-constrained researchers.
  • Evaluation restricted to three competition math benchmarks (AIME24/25, HMMT25); generalization to other reasoning domains is not assessed.
  • The asynchronous training infrastructure requires a 128-node sandbox cluster and verl/VLLM stack, limiting reproducibility.

Relevance to Harnesses / Meta-Harnesses

AgentMath is a concrete instantiation of a training harness for tool-augmented agents: its three-stage data synthesis pipeline, agentic RL loop with sandboxed execution, and asynchronous rollout scheduler collectively form an end-to-end harness that orchestrates model, tools, and infrastructure. The partial rollout mechanism and prefix-aware scheduler are directly analogous to meta-harness concerns — managing heterogeneous, multi-turn agent trajectories at scale without global synchronization barriers. The design pattern of decomposing long agentic rollouts into budget-bounded segments with dynamic curriculum expansion is a transferable template for any harness supporting long-horizon, tool-calling agents.