Technical Summary: TRUST – Uncertainty-Aligned Reinforcement Learning for Agentic Tool-Calling

Problem Statement

Large Language Model (LLM)-based agents frequently exhibit suboptimal tool-use decisions, characterized by two primary failure modes: invoking tools when they are unsupported or unnecessary, and fabricating direct responses without issuing required tool calls (hallucinated direct answers). Unlike standard textual hallucinations, these decision failures corrupt intermediate states and propagate errors across multi-step interactions, leading to execution failures, financial costs, and information leakage in real-world agentic tasks.

Existing approaches to mitigate these issues primarily rely on inference-time corrections or coarse-grained reward signals based on decision outcomes and structured checklists. However, these methods leave the uncertainty characteristics of agent decisions underexplored. The authors observe that standard decision-oriented Reinforcement Learning (RL) tends to weaken the uncertainty separation between correct and incorrect actions. Specifically, RL optimization often increases confidence in sampled high-reward actions while unintentionally collapsing the uncertainty gap, causing the model to become overconfident in incorrect decisions (e.g., unsupported tool calls). This phenomenon, termed "Wrong Decision But Low Uncertainty," impairs the model's calibration and reduces the exploration signals necessary for policy updates toward more reliable alternatives.

Methodology: TRUST Framework

To address these limitations, the authors propose TRUST (Tool-calling decision Reward with Uncertainty-Separated post-Training), a framework that integrates Uncertainty Quantification (UQ) directly into the reward design to act as a repulsive force, maintaining a clear separation between correct and incorrect decisions.

1. Turn-Level UQ-Aligned Reward

Instead of using uncertainty merely as a post-hoc diagnostic, TRUST incorporates it into the reward function to align decision correctness with model certainty.

• Uncertainty Metric: The method instantiates uncertainty using sequence perplexity (PPL).
• Certainty Margin ($m(s)$): For each decision point, the system calculates the margin between the expected perplexity of negative (incorrect) decisions and the ground-truth decision.
• Certainty Coefficient ($c(s)$): This margin is converted into a bounded certainty coefficient using a sigmoid function. A large margin indicates the model assigns lower perplexity to the correct decision than to negative ones.
• Reward Function ($R_{UQ}$): The total reward combines format validity, answer consistency, and action correctness, weighted by the certainty coefficient:
  $$R_{UQ}(z) = R_{fmt}(z) + R_{ans}(z, z^*) + c(s) \cdot R_{cls}(a, a^*)$$
  Here, $c(s) \cdot R_{cls}$ acts as a repulsive force. It pushes wrong decisions toward higher uncertainty (higher PPL) and pulls correct decisions toward lower uncertainty, thereby encouraging the policy to explore away from uncertain or unreliable decisions.

2. Unified Post-Training for Trajectories

Recognizing that agent failures often arise within multi-turn trajectories where early errors corrupt later states, TRUST extends this training paradigm from standalone decisions to full trajectories.

• Lightweight Key-Turn Annotation: Instead of exhaustively relabeling entire conversations, the framework augments trajectories with lightweight annotations for specific "key turns" where next-action decisions are critical. A labeling model selects up to two decision-critical turns per trajectory, assigning ground-truth actions from the space $\{DIRECT, TOOL, ASK, UNABLE\}$.
• Mixed Reward Optimization: The training utilizes a mixed reward function that combines the original trajectory-level checklist reward ($R_{CM2}$), which supervises task completion and tool execution quality, with the turn-level $R_{UQ}$ applied sparsely at annotated key turns. This allows for the joint optimization of overall task success and turn-level tool-calling calibration without redesigning existing performance-driven rewards.

Key Contributions

1. Uncertainty-Separated Reward Design: The paper introduces a novel reward mechanism that uses UQ as a repulsive force within the RL objective. This explicitly prevents the collapse of uncertainty separation between correct and incorrect actions, a common failure mode in standard decision-oriented RL.
2. Unified Training Paradigm: TRUST proposes a unified post-training framework that integrates turn-level decision calibration with trajectory-level task optimization. By annotating only key turns, it achieves efficient training that improves both immediate decision quality and long-horizon task performance.
3. Empirical Validation of Uncertainty Calibration: The work demonstrates that optimizing for decision correctness alone can degrade uncertainty calibration, whereas TRUST successfully aligns internal uncertainty with decision correctness, making erroneous decisions easier to identify and less likely to be reinforced.

Experimental Results

The authors evaluated TRUST on three benchmarks: When2Call (turn-level decision), BFCL-V4 (multi-turn function calling), and ToolSandbox (interactive tool use).

• Turn-Level Performance (When2Call): On the Qwen3-4B-Thinking backbone, TRUST improved the normalized accuracy by 11.47% over the base model and outperformed the best baseline (GRPO without TRUST) by 8.37%. Crucially, it reduced the False Direct Answer Rate (FDAR) to 5.07% and the overall hallucination metric significantly compared to standard GRPO.
• Trajectory-Level Performance (BFCL-V4 & ToolSandbox):
  • On BFCL-V4, TRUST achieved an overall score of 48.04% (on Qwen3-4B-Thinking), surpassing the base model by 9.43% and outperforming larger models like 30B-A3B-Instruct. Notable gains were observed in Multi-Turn (49.62%) and Irrelevance (84.84%) scenarios.
  • On ToolSandbox, TRUST achieved a state-of-the-art score of 68.28% (on Qwen3-8B-Base), outperforming the trajectory-level GRPO baseline by 7.07% and approaching the performance of the 235B parameter model.
• Uncertainty Calibration: Analysis of perplexity distributions showed that while standard RL increased the overlap between correct and wrong decision distributions (IoU rising from 34.50% to 70.21%), TRUST maintained a clear separation (IoU reduced to 35.29%), ensuring that incorrect decisions retained higher uncertainty.
• Ablation Studies: Removing the uncertainty repulsive component ($c(s)$) resulted in the most significant performance degradation, confirming its critical role in suppressing unsupported tool calls and hallucinated direct answers.

Significance and Claims

The paper claims that TRUST provides a robust solution to the "overconfident mistake" problem in agentic RL. By integrating uncertainty quantification directly into the reward function, TRUST ensures that the model's internal confidence aligns with decision correctness. This alignment not only improves immediate tool-calling accuracy but also enhances the reliability of multi-turn interactions by preventing error propagation.

The authors emphasize that TRUST achieves these gains without additional inference latency or computational overhead during deployment, as the decision-making logic is internalized into the model weights. Furthermore, the framework demonstrates that precise turn-level decision training can catalyze generalized improvements in complex, long-dependency tool-chain workflows. The work concludes that while current limitations exist regarding the reliance on perplexity-based UQ and predefined action spaces, the proposed approach offers a significant step toward more reliable and self-aware agentic systems.