Source-linked AI summary

HiDiffTIR: Hierarchical Difficulty-Aware Policy Optimization for Multi-Turn Tool-Integrated Reasoning

Yucan Guo, Xiaohan Wang, Miao Su, Saiping Guan, Zhongni Hou, Jiajun Chai, Wei Lin, Guojun Yin, Xiaolong Jin, Jiafeng Guo, Xueqi Cheng

arXiv:2608.21863v1cs.CLcs.AI

TL;DR

TIR agents need learning signals that distinguish the difficulty and importance of trajectories and reasoning steps, because uniform credit assignment treats heterogeneous tool-use decisions alike. HiDiffTIR introduces supervision-free hierarchical difficulty-aware credit assignment at trajectory and turn levels, and experiments show consistent gains in multi-turn TIR performance and tool invocation accuracy across multiple benchmarks.

  • Problem

    Existing TIR approaches often assign uniform advantages and treat correct tool calls equally, overlooking variation in difficulty and learning value across trajectories and reasoning steps.

  • Method

    HiDiffTIR performs difficulty-aware credit assignment at trajectory and turn levels using group-level statistics from standard RL rollouts without additional supervision.

  • Results

    HiDiffTIR consistently improves multi-turn TIR performance and tool invocation accuracy across multiple tool-using benchmarks.

  • Takeaways & Limitations

    Difficulty-aware signals provide a promising direction for training tool-using LLM agents by capturing heterogeneous importance among tool-calling decisions.

  • Takeaways & Limitations

    Evaluations were restricted to Qwen3-4B and Qwen3-8B, so scalability and behavioral dynamics for larger-parameter LLMs remain insufficiently validated.

Abstract

from arXiv · show

Tool-Integrated Reasoning (TIR) is a fundamental capability for LLM agents to solve complex tasks by interacting with external tools iteratively. Reinforcement Learning (RL) has become the dominant paradigm for enabling this capability. However, existing approaches typically assign uniform trajectory-level advantages and treat all correct tool calls equally, ignoring the varying difficulty and learning value across trajectories and reasoning steps. This can lead to imprecise learning signals that do not adequately distinguish between trivial and challenging tool-use patterns. To address this limitation, we propose HiDiffTIR, a Hierarchical Difficulty-aware policy optimization framework for multi-turn TIR. HiDiffTIR performs difficulty-aware credit assignment at both trajectory and turn levels, enabling the policy to focus on more informative trajectories and harder reasoning steps. Notably, this fine-grained optimization is achieved without additional supervision, relying solely on group-level statistics derived from standard RL rollouts. Extensive experiments on three tool-using benchmarks demonstrate that HiDiffTIR consistently improves multi-turn TIR performance and tool invocation accuracy over strong RL baselines, highlighting the necessity of difficulty-aware credit assignment for effective policy optimization in tool-integrated LLM agents.

1 Introduction

Existing TIR training often uses coarse or uniform credit assignment, despite different trajectories and reasoning steps varying in difficulty and learning value. HiDiffTIR addresses this with supervision-free, hierarchical difficulty-aware optimization at trajectory and turn levels, improving multi-turn reasoning and tool-use accuracy across benchmarks.

  • Motivation: Existing TIR methods commonly assign uniform advantages across steps, obscuring the contributions of individual reasoning decisions.This is especially problematic in multi-turn settings, where errors can accumulate and only some decisions are critical to success.
  • Motivation: Finer-grained methods combine trajectory- and turn-level signals but still treat correct tool calls equivalently, leaving learning signals insensitive to step difficulty.They typically preserve trajectory rewards while directly applying or heuristically propagating turn-level signals.
  • HiDiffTIR: HiDiffTIR performs difficulty-aware credit assignment at both trajectory and turn levels, prioritizing informative trajectories and harder tool-calling decisions.It differentiates rollouts by relative difficulty and redistributes advantages across reasoning steps according to their difficulty.
  • HiDiffTIR: The framework achieves fine-grained policy optimization without additional supervision, using only group-level statistics from standard RL rollouts.Its difficulty judgments do not require external judges or rubrics.
  • Results: Extensive experiments show improved tool-use accuracy and multi-turn reasoning performance across multiple benchmarks.The reported improvements cover both performance and tool invocation accuracy.

2 Related Work

Related work has shifted from supervised tool-use training toward reinforcement learning and difficulty-aware reasoning methods. However, existing difficulty-aware approaches do not jointly model process-level feedback and difficulty for tool-integrated reasoning.

  • TIR training: Earlier TIR research primarily used supervised fine-tuning with solution annotations that interleave reasoning steps and tool invocations.Recent work increasingly shifts toward reinforcement learning to improve policy learning through interaction.
  • HIDIFFTIR: The HIDIFFTIR framework is presented as hierarchical difficulty-aware credit assignment for tool-integrated reasoning.Its design targets the gap between general difficulty-aware RL and process-sensitive TIR optimization.
  • Difficulty-aware RL: Reinforcement learning methods for general reasoning emphasize challenging samples using outcome-level difficulty signals.These approaches are largely designed for general reasoning tasks rather than tool-integrated reasoning.
  • Research gap: Existing difficulty-aware reasoning methods remain insufficient for jointly modeling process-level feedback and difficulty in TIR.The related-work discussion identifies this as the gap addressed by hierarchical difficulty-aware credit assignment.

3 HIDIFFTIR

HIDIFFTIR formulates multi-turn TIR as hierarchical policy optimization that assigns difficulty-aware credit at trajectory and turn levels. It derives fine-grained signals from tool-call correctness and group-level statistics, then combines them for GRPO optimization.

  • 3.1 Problem Formulation: A TIR agent generates multi-turn reasoning and tool actions, receives tool feedback, and produces a final answer through interleaved reasoning and tool use.
  • 3.2 Overall Framework: HIDIFFTIR distinguishes the relative learning importance of trajectories and reasoning turns instead of treating them uniformly.Trajectory-level difficulty considers constituent tool calls, while turn-level credit considers tool-selection and tool-using difficulty.
  • 3.3 Trajectory-Level Difficulty-Aware Credit Assignment: Tool-call similarity produces turn-level raw rewards by matching predicted and ground-truth calls, with unmatched calls contributing zero.Similarity measures agreement on tool names, parameter names, and parameter values; matching uses the Hungarian algorithm.
  • 3.3 Trajectory-Level Difficulty-Aware Credit Assignment: Group-level tool-selection success rates estimate the difficulty of ground-truth calls, and turn-level statistics emphasize the most challenging call within each turn.A positive reward indicates correct tool selection, so the success rate measures how frequently each call is correctly invoked across trajectories.
  • 3.4 Turn-Level Difficulty-Aware Credit Assignment: Difficulty-aware reweighting emphasizes harder tool-use decisions, down-weights consistently easy turns, and suppresses incorrect turns.For negative-advantage trajectories, HIDIFFTIR uses uniform weights to maintain stable penalization of unsuccessful rollouts.
  • 3.3 Trajectory-Level Difficulty-Aware Credit Assignment: Normalized reweighted turn rewards are aggregated into trajectory scores, converted to group-relative advantages, and redistributed across turns before GRPO optimization.The resulting objective emphasizes relatively harder successful trajectories while maintaining stable within-group comparisons; α balances trajectory- and turn-level signals.

4 Experiments

Experiments evaluate HIDIFFTIR across multiple benchmarks, backbone sizes, ablations, efficiency measures, and tool-use accuracy. The method consistently improves overall performance, with gains varying by model size, tool-environment complexity, and credit-assignment configuration.

  • Experimental Setting: Experiments evaluate HIDIFFTIR on FTRL, ToolHop, and BFCL using Qwen3-4B and Qwen3-8B backbones against several TIR baselines.The evaluation includes in-domain and out-of-domain benchmarks, with RL and non-RL comparison methods.
  • Main Results: HIDIFFTIR achieves the best overall performance across benchmarks and outperforms strong baselines including FTRL and MatchTIR.The reported gains appear in both in-domain and out-of-domain settings for both backbone sizes.
  • Main Results: HIDIFFTIR improves over corresponding baselines with both Qwen3-4B and Qwen3-8B, while gains on in-domain FTRL are more pronounced for Qwen3-4B.The results are described as robust across backbone sizes and scalable to the larger model.
  • Main Results: HIDIFFTIR brings larger improvements on FTRL and ToolHop, whose environments require reasoning over thousands of candidate tools, than on BFCL.The reported pattern associates stronger gains with datasets containing more tools, while BFCL improvements are comparatively modest.
  • Main Results: Performance gains remain modest on Missing Functions and Missing Parameters subsets, which require recognizing insufficient tools or user-provided information.The reported limitation concerns missing-information scenarios and possible overemphasis on successful tool execution.
  • Ablation Study: Removing either trajectory-level or turn-level difficulty-aware credit assignment degrades performance, with trajectory-level removal causing a larger overall drop on FTRL.The ablation supports contributions from both hierarchical components.
  • Ablation Study: At the final training step, HIDIFFTIR reduces average response length from 2,085 to 1,893 tokens and generation latency by 12.8% versus the w/o turn-level variant.The authors attribute this efficiency to identifying critical reasoning steps and suppressing redundant tool calls.
  • Tool-Use Accuracy Analysis: All HIDIFFTIR variants outperform MatchTIR (KM) in tool-use accuracy during training, while α = 0.7 achieves the highest average FTRL test score.Lower α values reach higher peak training accuracy, but the best test performance occurs at α = 0.7.

5 Conclusions

HIDIFFTIR performs difficulty-aware credit assignment at both trajectory and turn levels, redistributing advantages across reasoning steps without additional supervision. Experiments show consistent gains in multi-turn TIR performance and tool invocation accuracy.

  • HIDIFFTIR models relative trajectory and turn difficulty to provide fine-grained learning signals for heterogeneous tool-calling decisions.The framework redistributes advantages across reasoning steps at both levels.

Limitations

HIDIFFTIR’s evaluation was limited to moderately sized Qwen3 models, and its scalability to larger models remains unvalidated. Its rollout-derived difficulty estimates also assume foundational instruction-following and may become sparse when initial tool use completely fails.

  • Evaluations used only Qwen3-4B and Qwen3-8B, leaving scalability to larger models insufficiently validated.The algorithm is described as model-agnostic, but larger-parameter behavioral dynamics remain untested.
  • The framework assumes foundational instruction-following capability to derive difficulty signals from online group-level rollout statistics.When the initial policy produces no valid tool invocations, reward sparsity may require supervised warm-up.
  • Difficulty estimation captures operational rollout bottlenecks rather than absolute static semantic complexity.

A.1 Datasets and Evaluation Metrics

The appendix describes benchmark datasets and metrics, compares HIDIFFTIR with several RL baselines, and details tool-call similarity computation. It also references training prompts and group-size comparisons.

  • Datasets and Evaluation Metrics: Dataset statistics are reported in Table 5, while benchmark-specific metrics are used for evaluation.The benchmarks include ToolHop and BFCL categories such as Base Multi-Turn, Missing Functions, Missing Parameters, and Long-Context Multi-Turn.
  • Datasets and Evaluation Metrics: FTRL evaluates tool-invocation precision and task-execution completeness, combining them into a total reward.Solve-F1 is described as the harmonic mean of Solve-P and Solve-R.
  • Baselines: HIDIFFTIR is compared with Vanilla, GRPO, ToolRL, FTRL, and MatchTIR using hard Kuhn-Munkres or soft Optimal Transport alignment variants.GRPO and several baselines broadcast a single trajectory-level advantage, whereas MatchTIR derives turn-level advantages.
  • Training Materials: The appendix provides the policy training prompt template and reports a performance comparison across different group sizes.
  • Tool Call Similarity Computation: Tool-call similarity uses a matching matrix whose components assess tool names, parameter names, and parameter values.Parameter-name similarity uses Jaccard similarity when tool names match, and the components are combined into a normalized score.

A.5 Training Details

Training used reproducibility-oriented infrastructure and fixed generation and optimization settings. The setup employed eight H20 GPUs, 16 rollouts per prompt, and five optimization epochs.

  • Reproducibility: Complete hyperparameters and infrastructure configurations are provided to facilitate reproducibility.
  • Infrastructure and Hardware: Training used one node with 8 NVIDIA H20 GPUs, FSDP for model-weight distribution, and vLLM for rollout generation.GPU memory utilization for vLLM was set to 0.7.
  • Data and Generation Settings: Rollouts sampled 16 responses per prompt with temperature 1.0 and a maximum of 10 turns.Maximum model and response lengths were 32,768 and 8,192 tokens, respectively.
  • Optimization and Loss Formulation: Optimization ran for 5 epochs with a global batch size of 256 prompts, mini-batches of 32, and actor learning rate 10^-6.The entropy and KL penalty coefficients were both 0.001, with no critic warmup.

B.1 Group Size Analysis

Group size affects HiDiffTIR’s empirical difficulty estimates and produces different in-domain and out-of-domain performance trends. G = 16 achieves the strongest overall balance across both settings.

  • Setup: Group size G controls within-group reward normalization and the robustness of trajectory- and turn-level difficulty estimation.The experiments evaluate G = 4, 8, and 16 on Qwen3-4B with α = 0.7.
  • In-Domain Results: The ID average score rises from 33.03 at G = 4 to 39.48 at G = 16 as group size increases.Larger groups provide more statistically robust group-level estimates during each training iteration.
  • Out-of-Domain Results: The OOD average score is non-monotonic, decreasing from 40.19 at G = 4 to 39.77 at G = 8 before reaching 41.79 at G = 16.The framework’s targeted optimization of seen tool-use patterns can produce behavioral variation on entirely unseen tools.
  • Overall Finding: G = 16 provides the best balance, producing the highest performance on both in-domain tasks and overall OOD generalization.This setting combines the strongest reported ID and OOD outcomes despite the non-monotonic intermediate OOD trend.

B.2 Case Study

The case study examines a complex five-hop query requiring a chain of implicit entities and compares HiDiffTIR with an untrained backbone and MatchTIR. HiDiffTIR avoids parameter hallucination and redundant trial-and-error while chaining the required tool calls precisely.

  • Case Study Setup: The case study evaluates a complex 5-hop query requiring identification of a local handicraft through five implicit entities.It compares HiDiffTIR with the untrained Qwen3-4B backbone and MatchTIR (KM).
  • Parameter Hallucination: Both baseline models hallucinate unnecessary geographic constraints when initially calling the famous_peak_identifier tool.Qwen3-4B and MatchTIR append “Alps” or “Himalayas” instead of passing empty arguments; MatchTIR corrects itself only after two wasted turns.
  • Feedback Adaptation: The untrained Qwen3-4B resists self-correction after consecutive empty tool responses and attributes the problem to the tool itself.The trajectory records this overconfidence in Turns 6 and 9.
  • Reasoning Efficiency: Baseline trajectories contain redundant exploration through repeated over-specification, whereas HiDiffTIR performs precise five-call chaining without hallucinations or trial-and-error.Qwen3-4B ultimately fails within the turn limit, while MatchTIR reaches the correct answer but remains inefficient early in the trajectory.
Loading 2608.21863v1…