Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AWPO: Enhancing Tool-Use of Large Language Models through Explicit Integration of Reasoning Rewards

Zihan Lin, Xiaohan Wang, Hexiong Yang, Jiajun Chai, Jie Cao, Guojun Yin, Wei Lin, Ran He

TL;DR

AWPO tackles integrating explicit reasoning rewards into tool-use RL for LLMs, addressing instability when mixing reasoning and outcome rewards. It introduces variance-aware gating, difficulty-aware weighting, and adaptive clipping, together with an LLM-as-a-Judge to provide fine-grained reasoning scores and increase signal variance. The authors derive a theoretical upper bound on expected policy improvement and show how their design improves the signal-to-noise ratio. Empirically, AWPO achieves state-of-the-art multi-turn tool-use on BFCL and API-Bank across Qwen3 backbones while preserving MMLU-Pro performance; a 4B AWPO model even outperforms Grok-4 on multi-turn accuracy by 16 percentage points. Collectively, AWPO offers a robust recipe for integrating fine-grained reasoning feedback into tool-use RL for LLMs.

Abstract

While reinforcement learning (RL) shows promise in training tool-use large language models (LLMs) using verifiable outcome rewards, existing methods largely overlook the potential of explicit reasoning rewards to bolster reasoning and tool utilization. Furthermore, natively combining reasoning and outcome rewards may yield suboptimal performance or conflict with the primary optimization objective. To address this, we propose advantage-weighted policy optimization (AWPO) -- a principled RL framework that effectively integrates explicit reasoning rewards to enhance tool-use capability. AWPO incorporates variance-aware gating and difficulty-aware weighting to adaptively modulate advantages from reasoning signals based on group-relative statistics, alongside a tailored clipping mechanism for stable optimization. Extensive experiments demonstrate that AWPO achieves state-of-the-art performance across standard tool-use benchmarks, significantly outperforming strong baselines and leading closed-source models in challenging multi-turn scenarios. Notably, with exceptional parameter efficiency, our 4B model surpasses Grok-4 by 16.0 percent in multi-turn accuracy while preserving generalization capability on the out-of-distribution MMLU-Pro benchmark.

AWPO: Enhancing Tool-Use of Large Language Models through Explicit Integration of Reasoning Rewards

TL;DR

AWPO tackles integrating explicit reasoning rewards into tool-use RL for LLMs, addressing instability when mixing reasoning and outcome rewards. It introduces variance-aware gating, difficulty-aware weighting, and adaptive clipping, together with an LLM-as-a-Judge to provide fine-grained reasoning scores and increase signal variance. The authors derive a theoretical upper bound on expected policy improvement and show how their design improves the signal-to-noise ratio. Empirically, AWPO achieves state-of-the-art multi-turn tool-use on BFCL and API-Bank across Qwen3 backbones while preserving MMLU-Pro performance; a 4B AWPO model even outperforms Grok-4 on multi-turn accuracy by 16 percentage points. Collectively, AWPO offers a robust recipe for integrating fine-grained reasoning feedback into tool-use RL for LLMs.

Abstract

While reinforcement learning (RL) shows promise in training tool-use large language models (LLMs) using verifiable outcome rewards, existing methods largely overlook the potential of explicit reasoning rewards to bolster reasoning and tool utilization. Furthermore, natively combining reasoning and outcome rewards may yield suboptimal performance or conflict with the primary optimization objective. To address this, we propose advantage-weighted policy optimization (AWPO) -- a principled RL framework that effectively integrates explicit reasoning rewards to enhance tool-use capability. AWPO incorporates variance-aware gating and difficulty-aware weighting to adaptively modulate advantages from reasoning signals based on group-relative statistics, alongside a tailored clipping mechanism for stable optimization. Extensive experiments demonstrate that AWPO achieves state-of-the-art performance across standard tool-use benchmarks, significantly outperforming strong baselines and leading closed-source models in challenging multi-turn scenarios. Notably, with exceptional parameter efficiency, our 4B model surpasses Grok-4 by 16.0 percent in multi-turn accuracy while preserving generalization capability on the out-of-distribution MMLU-Pro benchmark.
Paper Structure (40 sections, 5 theorems, 81 equations, 2 figures, 5 tables, 1 algorithm)

This paper contains 40 sections, 5 theorems, 81 equations, 2 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Group Relative Policy Optimization for LLMs
  4. Post-Training for Tool-Use
  5. LLM-as-a-Judge for Reward Design
  6. Method
  7. Theoretical Analysis
  8. Advantage Weighted Policy Optimization
  9. Reward Design.
  10. Algorithm Design
  11. Experiments
  12. Experiment Setup
  13. Training Details
  14. Benchmarks
  15. Baselines
  16. ...and 25 more sections

Key Result

Lemma 3.2

To analyze the optimization landscape, we assume the objective $J:\mathbb{R}^d\to\mathbb{R}$ admits a smooth approximation satisfying the $L$-smoothness condition bottou2018optimization, i.e., for all $\theta,\theta^+\in\mathbb{R}^d$, Consider the update where $\alpha>0$ is the learning rate and $\widehat{g}$ is an unbiased gradient estimator, $\mathbb{E}[\widehat{g}] = \nabla J(\theta)$. Then t

Figures (2)

  • Figure 1: Overview of the AWPO framework. AWPO incorporates variance-aware gating, difficulty-aware weighting, and dynamic clipping mechanisms to effectively achieve reward integration.
  • Figure 2: Overview of the structured system prompt designed for the Tool-use Consistency Evaluator. The evaluation protocol is divided into four weighted dimensions: Reasoning Path (35%), Tool Selection (30%), Parameter Setting (25%), and Execution Strategy (10%). To ensure rigorous assessment, the prompt incorporates hard constraints based on ground truth tool trajectories and utilizes a discrete six-tier scoring rubric (ranging from Tier VI to Tier I) to map qualitative judgments to fixed numerical values (0.00--1.00), thereby reducing variance in the LLM judge's output.

Theorems & Definitions (11)

  • Definition 3.1: Objective Function
  • Lemma 3.2: Expected improvement upper bound
  • Definition 3.3: Fisher-Normalized Correlation
  • Lemma 3.4: Upper Bound for the Gradient Norm
  • Lemma 3.5: Variance Bound for the Stochastic Gradient
  • Theorem 3.6: Upper-bound signal--variance decomposition
  • Theorem 3.7: Sufficient condition for expanded optimization potential
  • proof
  • proof
  • proof
  • ...and 1 more