Empowering LLM Tool Invocation with Tool-call Reward Model

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

Why this paper was selected

Reward model for RL-based tool invocation in LLM agents

Ask a follow-up

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

💬 Ask ChatGPT✦ Ask Claude

TL;DR

TRM is a process reward model specialized for evaluating individual tool invocations in multi-turn LLM trajectories, addressing the gradient conflict that arises when outcome-only RL rewards penalize correct intermediate tool calls due to incorrect final answers. Integrating TRM into PPO and GRPO via turn-level credit assignment and turn-level advantage estimation consistently outperforms outcome-only RL baselines across search-based QA and code-based math tasks.

Problem

RL-trained tool-using LLMs typically receive only outcome rewards (final answer correctness), which assigns uniform credit across all intermediate tool calls regardless of their individual quality. This causes gradient conflict: a trajectory with correct intermediate tool calls but a wrong final answer is penalized, discouraging effective tool invocation learning and producing unstable tool-use behavior.

Method

Tool-call Reward Model (TRM) is a per-call process reward model trained on trajectories labeled with two binary scores per tool call: necessity (\(s_i^{ne}\) — does the call contribute substantive progress?) and quality (\(s_i^q\) — is the tool invoked with correct parameters?). The final binary score is \(s_i = s_i^{ne} \cdot s_i^q\). Training data is distilled from DeepSeek-R1 rollouts on HotpotQA/NQ (10K labeled trajectories); the TRM backbone (Qwen2.5 series) has its LM head replaced by a binary classification head trained with BCE loss on the hidden state of each tool call's last output token.

Integration with RL uses turn-level credit assignment: TRM score \(\tilde{s}^i\) is assigned to the last token of tool call \(a_i\), while outcome correctness is assigned to the final answer token (Eq. 5–6). For PPO, token-level advantages are computed via GAE over these turn-level rewards, weighted by hyperparameter \(\alpha\). For GRPO, a turn-level advantage estimation normalizes TRM rewards for each turn independently across a group of trajectories (rather than pooling all tool-call rewards together), preventing reward hacking where the model minimizes tool calls to avoid penalties (Eq. 9–10).

Key Contributions

• Proposes TRM — first systematic process reward model specifically designed for tool invocations in multi-turn agentic trajectories, with a workflow for data distillation and model training
• Establishes scaling laws for TRM: 3B parameters + 10K samples achieves robust performance; larger models (7B) overfit at this data scale
• Derives turn-level credit assignment and turn-level advantage estimation for both PPO and GRPO to prevent reward hacking
• Demonstrates cross-task generalization: models trained in search-tool settings improve on code-based math problems when TRM is used

Results

Search-based QA (7 datasets, avg accuracy): - Qwen2.5-3B: Search-R1-GRPO-TRM 43.49 vs. Search-R1-GRPO 42.33 vs. RAG 27.84 - Qwen2.5-7B: Search-R1-GRPO-TRM 48.62 vs. Search-R1-GRPO 46.90 vs. RAG 30.59 - PPO: Search-R1-PPO-TRM 34.93 vs. Search-R1-PPO 32.75 (3B); 39.95 vs. 37.94 (7B)

Code-based math (5 datasets, avg accuracy): - Qwen2.5-Math-1.5B: ToRL-GRPO-TRM 45.42 vs. ToRL-GRPO 43.18 vs. R1-GRPO 36.08 - Qwen2.5-Math-7B: ToRL-PPO-TRM 53.35 vs. ToRL-PPO 51.76; ToRL-GRPO-TRM 53.70 vs. ToRL-GRPO 52.19

TRM exploration (best-of-n selection): TRM-3B outperforms majority voting baseline on both in-domain (HotpotQA) and out-of-domain (2WikiMultiHopQA), though a gap to pass@n upper bound remains.

Ablations: Turn-level GRPO advantage estimation (42.47) > group-level (41.18); TRM outperforms StepSearch and AgentPRM process supervision baselines; necessity+quality together > either alone.

Limitations

• TRM is evaluated only on search and Python code tools; generalization to diverse tool types (APIs, databases, multimodal tools) is not demonstrated
• The training signal for TRM itself still depends on a frontier LLM (DeepSeek-R1) for annotation, introducing cost and potential label noise from the annotator model
• TRM scores entire-trajectory pass-rates as a proxy for utility; the binary \(s_i = s_i^{ne} \cdot s_i^q\) label may discard nuanced partial-utility cases
• The hyperparameter \(\alpha\) (TRM weight) requires task-specific tuning (0.05 for PPO, 0.01 for GRPO), adding a tuning burden
• Gap between TRM-guided best-of-n and pass@n upper bound indicates room for further TRM accuracy improvement

Relevance to Agentic AI / LLM Agents

TRM directly addresses a fundamental training challenge in agentic systems: how to assign credit to intermediate actions in multi-step tool-using trajectories rather than relying solely on terminal outcomes. This is the tool-use analog of the process reward model line in mathematical reasoning, and the proposed turn-level advantage estimation offers a concrete algorithmic solution to reward hacking in GRPO-trained agents. The cross-task generalization result — where TRM training on search tools transfers to code tool usage — suggests TRM may encode general tool-use competence rather than task-specific shortcuts, a property valuable for building versatile tool-using agents. This work complements search-agent training frameworks like Search-R1 and ToRL by providing the missing dense reward signal layer.