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 integrating LRMs with symbolic execution harness [ICLR 2026]

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 tightly couples LLM chain-of-thought reasoning with sandboxed code execution to solve competition-level mathematics. It introduces a three-stage automated trajectory synthesis pipeline for SFT data, an agentic RL paradigm using interleaved generate–pause–execute–resume rollouts, and a scalable async training infrastructure achieving 4–5× throughput over synchronous baselines. AgentMath-30B-A3B reaches 90.6% / 86.4% / 73.8% on AIME24/25/HMMT25, surpassing OpenAI-o3-mini and Claude-Opus-4.0-Thinking.

Problem

Pure-text LRMs like o3 and DeepSeek-R1 incur frequent computational errors on problems requiring precise arithmetic, symbolic manipulation, or geometric reasoning, leading to redundant corrections and inflated token costs. Three compounding bottlenecks remain unaddressed: (1) high-quality tool-augmented training trajectories are scarce—existing methods either rely on brittle prompt engineering or costly manual annotation; (2) agentic RL for joint natural-language/code-execution reasoning is unexplored; (3) competition-level problems produce ultra-long trajectories (≥96k tokens, ≥96 tool calls) that cause severe long-tail effects in conventional batch-synchronous RL frameworks.

Method

Stage 1 — SFT Data Synthesis (three sub-stages): 1. Agentic Trajectory Generation via Code Injection: Start from 392k pure-text CoT samples (AM-Thinking, OpenThoughts) filtered via 4-gram/MinHash overlap against evaluation sets. A teacher model (DeepSeek-V3) applies injection function \(F_{\text{inject}}\) to each segment: complex computational steps \(s_{\text{calc}}\) are replaced with \((c, o_{\text{sim}})\) pairs (executable code + simulated output). Resulting in initial agentic trajectories \(\tau'_{\text{agent}}\). 2. Multi-Faceted Quality Refinement: (i) regex + teacher regeneration for <code>/<interpreter> format compliance; (ii) sandbox execution of each code block with bounded resampling on failure; (iii) replacement of simulated outputs \(o_{\text{sim}}\) with ground-truth \(o_{\text{real}} = \mathcal{E}(c)\), verified by Qwen3-32B as judge; (iv) AST-depth/line-count heuristics to eliminate unnecessary tool calls. 302k samples survive. 3. Self-Correction Capability Injection: 30k failed-execution trajectories are reprocessed; teacher generates diagnose→repair→re-execute→continue traces, yielding 14k valid self-correction traces injected into the final 316k-sample SFT dataset (avg 8.3 tool calls, 16.9k tokens/sample).

Stage 2 — Agentic RL: - Optimization: GRPO (no critic) with dynamic sampling, asymmetric gradient clipping, token-level loss, and KL-removal following DAPO. - Rollout protocol: generate–pause–execute–resume loop; <think> and <code> tokens receive gradient; <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 efficient tool use. - Three-stage curriculum: max response grows from 48k → 72k → 96k tokens; tool-call limits 48 → 72 → 96; partial-rollout segments 2 → 3 → 4.

Infrastructure: - Request-Level Async Scheduler: coroutine-driven; inference engine and agents decouple via async channels so GPU serves other requests while a trajectory awaits tool execution. Tool-call latency: 175 s → 1.2 s. - Agentic Partial Rollout: trajectory \(\tau = \tau^{(1)} \oplus \cdots \oplus \tau^{(N)}\); each segment bounded by \((L_{\text{seg}}, T_{\text{seg}})\), preventing any trajectory from monopolizing the batch. 2.2–2.5× speedup. - Prefix-Aware Weighted Load Balancing: weight \(w_j = L_j/L_{\text{base}} + w_{\text{base}}\); new requests route to \(\arg\min_k W_k\); LRU sticky sessions maximise KV-cache reuse for consecutive segments.

Key Contributions

  • Automated three-stage SFT pipeline converting pure-text CoT into verified, executable tool-augmented trajectories at scale (316k samples).
  • Agentic RL framework with interleaved NL/code execution and emergent self-correction, validated to improve with increasing tool invocations.
  • Async training infrastructure (decoupled sandbox cluster + partial rollout + prefix-aware LB) achieving 4–5× end-to-end throughput on ultra-long, multi-turn agent sequences.
  • SOTA results across AIME24/25 and HMMT25 at 30B scale, outperforming closed-source o3-mini and Claude-Opus-4.0-Thinking.

Results

  • AgentMath-30B-A3B: 90.6% AIME24 / 86.4% AIME25 / 73.8% HMMT25 (avg@32); surpasses OpenAI-o3-mini and Claude-Opus-4.0-Thinking; competitive with Gemini-2.5-Pro and DeepSeek-R1-671B-0528.
  • AgentMath-8B: 89.8% / 84.7% / 71.3%; outperforms tool-augmented CIR-Qwen3-NT8-8B and text-only DS-0528-Qwen3-8B (86.0% / 76.3% / 61.5%).
  • AgentMath-1.7B: 59.6% / 48.1% / 40.2%; exceeds tool-augmented CoRT-1.5B (43.1% / 30.2% / 20.1%) and text-only OpenReasoning-1.5B (55.5% / 45.6% / 31.5%).
  • Infrastructure: 4–5× training throughput gain vs. synchronous baselines; sandbox parallelization cuts tool-call latency from 175 s to 1.2 s; partial rollout alone yields 2.2–2.5× speedup.

Limitations

  • AgentMath-235B-A22B trained with SFT only (no RL) due to computational budget constraints, leaving the largest model under-optimized.
  • Data synthesis depends on powerful proprietary/large teacher models (DeepSeek-V3, DeepSeek-R1-0528, Qwen3-32B), creating a bootstrap dependency that may be hard to replicate without equivalent resources.
  • Self-correction injection uses only 14k traces; broader coverage of failure modes is unverified.
  • Evaluation limited to three competition-math benchmarks; generalization to other tool-use or reasoning domains not demonstrated.
  • The paper text is truncated, so additional ablations and limitations in the appendix are not fully accessible.

Relevance to Harnesses / Meta-Harnesses

AgentMath's three-stage SFT synthesis pipeline is itself a multi-model meta-harness: it orchestrates a teacher LLM (DeepSeek-V3), a sandboxed code executor, and a judge/verifier model (Qwen3-32B) in a pipelined loop to transform unstructured reasoning traces into structured, verified training data—exactly the pattern of a data-generation harness wrapping heterogeneous tools. The agentic RL training infrastructure (async scheduler + partial rollout + prefix-aware LB) represents a harness architecture for managing multi-turn, multi-tool agent rollouts at scale, decoupling inference, environment, and optimization components. The generate–pause–execute–resume rollout loop is a canonical harness micro-pattern for tool-augmented agents, and the paper's explicit engineering of throughput, scheduling, and KV-cache management under long-tail workloads is directly instructive for anyone building production-grade reasoning harnesses.