Organizations: School of Computer Science and Engineering, Southeast University, Nanjing, China · Shandong Jianzhu University, China · UniSA STEM, University of South Australia, Adelaide, Australia
Large language models (LLMs) rely on long-horizon tool invocation sequences for complex tasks, where each invocation can alter the task state and condition subsequent decisions. In long-horizon tool use, final-outcome rewards provide weak credit assignment over long interaction traces. Step-level rewards can offer more targeted feedback, but obtaining reliable step supervision often requires human or LLM judgment, or additional rollouts to estimate the downstream effect of an intermediate decision. In this paper, we argue that effective tool-use agents should estimate the long-horizon value of a possible next tool invocation before executing it. This objective requires comparative supervision over alternative invocations under the same context, while logged trajectories only contain the invocation that was actually taken. Therefore, we propose Comparative Inference for Tool-use Agents (CITA). CITA trains a Comparative Inference Model (CIM) from paired signals that combine observed tool behavior, scalable supervision from a Bayesian tool-graph simulator, and semantic judgments from LLM-based comparison. The resulting CIM learns to estimate how likely a possible next tool invocation is to support final task success under the current context. Across three tool-use benchmarks and multiple backbone LLMs, CITA consistently improves Tool F1 and task success. Additional analysis shows that CIM learns accurate step-level value estimates for comparative tool choices.
Figures & tables
Figure 1: Pilot analysis on long-horizon tool-use tasks. (a) Task success drops sharply as the required number of tool invocations increases. (b) Most failed trajectories remain unrecoverable even after several backtracking steps. (c) At many fork points, the model-preferred tool is not the tool with the higher empirical success rate under the same context. These patterns motivate comparative value estimation for possible next-tool choices before execution.
Figure 2: Overview of CITA. CITA trains a Comparative Inference Model (CIM) from paired signals built from real trajectories, a Bayesian tool-graph simulator, and LLM-based comparison. The trained CIM guides GRPO rollouts, provides dense comparative rewards, and supports iterative next-tool selection during inference.
Method
Toolathlon
TOUCAN
TRAJECT-Bench
Qwen2.5-7B
Llama3.1-8B
Qwen2.5-7B
Llama3.1-8B
Qwen2.5-7B
Llama3.1-8B
Tool F1 ↑
TSR ↑
Tool F1 ↑
TSR ↑
Tool F1 ↑
TSR ↑
Tool F1 ↑
TSR ↑
Tool F1 ↑
TSR ↑
Tool F1 ↑
TSR ↑
Prompting and search-based tool use
ReAct
33.00
22.73
43.79
25.41
42.74
32.68
44.00
34.22
43.64
30.96
33.71
23.58
ToolChain*
38.21
27.84
42.22
26.17
25.59
16.93
48.11
37.26
41.97
29.74
47.85
36.77
ToolLLM
21.10
11.62
35.43
20.43
26.28
17.59
33.34
27.11
50.10
31.50
54.68
32.00
Table 1: Main results on three benchmarks. We report Tool F1 (%) and Task Success Rate (TSR, %) with Qwen2.5-7B and Llama3.1-8B as backbone models. The best results are highlighted in bold.
Model
TOUCAN
Toolathlon
TRAJECT-Bench
MAE ↓
Dir. ↑
AUROC ↑
MAE ↓
Dir. ↑
AUROC ↑
MAE ↓
Dir. ↑
AUROC ↑
CIM w/ Lvalue only
0.4752
0.6923
0.9615
0.2601
0.3974
0.4198
0.3846
0.4872
0.5124
Full CIM
0.2083
0.9231
0.9620
0.2515
0.6154
0.6980
0.1768
0.7949
0.7816
Table 2: CIM value estimation quality on held-out paired examples.
Figure 3: CIM value-gap estimation on held-out paired examples. Each point compares the predicted value gap Δv^ with the observed value gap Δv . Compared with SFT, GRPO-only, and CIM trained with Lvalue only, CITA shows clearer alignment with the ground truth.
Variant
Tool F1
TSR
CITA (full)
60.47
49.70
w/o CIM guide
55.12
43.68
w/o rcmp
56.21
44.37
w/o Lgap
54.86
41.92
w/o confidence
57.08
45.72
w/o Dsandbox
51.78
38.64
Table 3: Ablation results averaged over six dataset-backbone settings.
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Benchmark
Avg. tool invocations / episode
# Unique tools
Main difficulty
Toolathlon
∼13.5
1,250
Long horizon and multi-step
TOUCAN
∼3.2
15,294
Large tool pool
TRAJECT-Bench
∼6.4
715
Medium horizon and trajectory quality
Appendix
Table 4: Statistics of the processed benchmark splits used in our experiments.
Figure 4: Training reward curves during CIM-guided GRPO. Rewards generally increase over training steps on Toolathlon, TOUCAN, and TRAJECT-Bench for both Qwen2.5-7B and Llama3.1-8B, showing that CITA provides useful optimization signals across different tool-use settings.
Dataset
Model
AUPRC ↑
CvS ↑
Pearson r↑
Spearman r↑
TOUCAN
CIM w/ Lvalue only
0.9945
0.2667
0.7652
0.8965
Full CIM
0.9945
0.9333
0.9004
0.8700
Toolathlon
CIM w/ Lvalue only
0.5401
0.4091
-0.1307
-0.1884
Full CIM
0.7560
0.5303
0.3077
0.2193
Appendix
Table 5: Additional CIM value estimation metrics. CIM w/ Lvalue only uses only the value prediction loss, while Full CIM uses the complete objective with value-gap learning, consistency regularization, and confidence calibration.
Variant
Toolathlon-Qwen
Toolathlon-Llama
TOUCAN-Qwen
TOUCAN-Llama
TRAJECT-Qwen
TRAJECT-Llama
Avg.
F1
TSR
F1
TSR
F1
TSR
F1
TSR
F1
TSR
F1
TSR
F1
TSR
CITA (full)
58.56
47.82
59.29
45.05
59.67
54.23
56.44
47.08
59.30
48.31
69.58
55.72
60.47
49.70
w/o CIM guide
53.21
39.62
54.38
38.97
55.40
48.10
52.16
40.55
54.09
41.23
61.48
53.61
55.12
43.68
w/o rcmp
54.06
40.88
55.73
39.76
56.82
48.62
53.39
41.31
55.84
42.05
61.42
53.60
56.21
44.37
w/o Lgap
52.94
37.91
53.28
36.85
55.11
47.06
51.72
38.77
53.68
39.65
62.43
51.28
54.86
41.92
w/o Lcons
56.21
43.20
57.46
41.92
58.20
50.61
55.30
44.08
56.94
45.60
63.89
54.37
58.00
46.63
Appendix
Table 6: Detailed ablation results on each dataset-backbone setting. Avg. is the arithmetic mean over all six settings.
Phase
Source weight
Sample ratio
Dobs
Dsandbox
Dcmp
Dobs
Dsandbox
Dcmp
Phase 1
0.8
1.4
0.0
0.6
1.0
0.0
Phase 2
1.0
1.0
1.0
0.8
1.0
1.0
Phase 3
1.8
0.3
1.2
1.0
0.3
1.0
Appendix
Table 7: Phase-aware mixing of the three paired data sources. Source weights affect the training weight of examples, while sample ratios affect how many examples are sampled from each source.
Benchmark
Backbone
Prec.
Rec.
Tool F1
TSR
NTAcc.
Prog.
ArgQ.
Avg. Steps
TOOLATHLON
Qwen2.5-7B
66.20
52.50
58.56±1.62
47.82±1.95
44.12
100.00
72.12
8.93
TOOLATHLON
Llama3.1-8B
65.40
54.22
59.29±1.75
45.05±2.03
49.58
100.00
67.15
8.26
TOUCAN
Qwen2.5-7B
63.80
56.04
59.67±1.88
54.23±2.22
55.00
100.00
55.08
2.50
TOUCAN
Llama3.1-8B
64.10
50.42
56.44±1.76
47.08±2.05
60.00
100.00
62.58
2.05
TRAJECT-Bench
Qwen2.5-7B
66.17
53.72
59.30±1.64
48.31±2.10
60.00
100.00
23.60
4.80
TRAJECT-Bench
Llama3.1-8B
75.30
64.67
69.58±1.42
55.72±1.84
62.50
100.00
32.64
5.00
Appendix
Table 8: Additional results and confidence intervals for CITA across benchmarks and backbones. Tool F1 and TSR are reported with estimated 95% confidence intervals.
Figure 5: Confidence calibration of CIM. We compare a CIM variant trained without the confidence loss Lconf and the full CIM on held-out context-action pairs. In the upper row, blue bars show empirical accuracy in each confidence bin, the red line shows average predicted confidence, and the dashed diagonal indicates perfect calibration. In the lower row, histograms show the distribution of predicted confidence scores. The full CIM achieves much better calibration, reducing ECE from 0.460 to 0.087.
Method
Inference strategy
Time / task (s) ↓
TSR (%) ↑
Time / success (s) ↓
SFT
Greedy decoding
3.10±0.08
23.84
13.00
ToolRL
RL policy, greedy decoding
3.16±0.09
26.60
11.88
StepTool
Step-RL policy, greedy decoding
3.22±0.10
26.59
12.11
WKM
Knowledge-guided planning
4.65±0.14
36.41
12.77
WebDreamer
Online future-state simulation
11.92±0.36
35.11
33.95
CITA
Actor candidates + CIM scoring
4.84±0.16
49.70
9.74
Appendix
Table 9: Inference cost and task success on the same 2 H100 hardware. TSR is averaged over the six dataset-backbone settings. Time per successful task is computed from average inference time and TSR.
Candidate next-tool invocation
Vϕ(s,a)
cϕ(s,a)
Sϕ(s,a)
start_quiz_attempt
0.74
0.82
0.61
list_quiz_questions
0.28
0.71
0.20
get_quiz
0.33
0.65
0.21
Appendix
Table 10: CIM scores at the critical fork point in the Canvas quiz example.
Candidate next-tool invocation
Vϕ(s,a)
cϕ(s,a)
Sϕ(s,a)
geocode
0.71
0.79
0.56
search_places
0.38
0.62
0.24
directions
0.25
0.55
0.14
Appendix
Table 11: CIM scores at the key map-planning fork point in the trip planning example.
Large language model (LLM) agents increasingly interleave natural language reasoning with external tools such as web search and code execution. These tool-use policies are often optimized via reinforcement learning (RL), which can amplify spurious correlations in the training data. In this work, we study when and why RL-trained agents learn shortcut tool-selection policies: invoking tools based on superficial prompt cues rather than genuine task requirements. We construct controlled synthetic environments combining factual question answering and mathematical reasoning tasks, and inject cues that are strongly correlated with specific tools during training but causally irrelevant to tool necessity. Across counterfactual evaluations where cues are present but the associated tools are not required, agents exhibit substantial shortcut behavior, with spurious tool invocation rates increasing by up to 39 percent. However, shortcut formation is not universal: across the conditions we test, it arises only when the agent has already learned to use the target tool reliably, suggesting that task competence, rather than dataset imbalance alone, is a key factor in shortcut learning. A swapped-cue analysis further shows that semantic alignment between cues and tools substantially amplifies this effect. To mitigate these failures, we introduce a dense, decision-level reward in which an LLM judge evaluates the necessity of each tool call. This tool-necessity reward effectively suppresses cue-driven tool use while preserving task performance, providing a practical approach to improving the robustness of LLM agent tool-use policies.
Yiwei Yang, Haoxiang Zhang, Bingbing Wen +6
University of Washington · University of California San Diego · Stanford University
Large language model (LLM) agents must generalize from tools seen during training to unseen tools at deployment. A key challenge is tool trialing, i.e., excessive trials waste the interaction budget, whereas selective trials enable exploration of unfamiliar tools. Existing outcome-based post-training leaves wasteful trials unguided, while turn-level supervision may suppress necessary exploration. We introduce ToolCompass, a post-training framework that guides tool trialing by organizing tool-call representations according to shared functions. Specifically, ToolCompass models each function class as a von Mises--Fisher distribution and jointly reduces intra-function variation across domains and increases inter-function separation. This structure transfers experience from seen tools to functionally similar unseen tools, directing exploration away from unrelated alternatives. ToolCompass requires no ground-truth call traces or unseen-tool access and incurs no inference overhead. Experiments on AppWorld and FTRL show consistent gains across GRPO, RFT, and DMPO. improves AppWorld OOD task success by up to 10.71 percentage points over vanilla post-training and performs best among competitive baselines on both benchmarks.
Junlin Fang, Chong Zhang, Do Nguyen-Thanh +3
College of Computing and Data Science, Nanyang Technological University · Zhejiang University · Australian Artificial Intelligence Institute, University of Technology Sydney +1
Tool-augmented LLM agents tend to call tools indiscriminately, even when the model can answer directly. Each unnecessary call wastes API fees and latency, yet no existing benchmark systematically studies when a tool call is actually needed. We propose When2Tool, a benchmark of 18 environments (15 single-hop, 3 multi-hop) spanning three categories of tool necessity -- computational scale, knowledge boundaries, and execution reliability -- each with controlled difficulty levels that create a clear decision boundary between tool-necessary and tool-unnecessary tasks. We evaluate two families of training-free baselines: Prompt-only (varying the prompt to discourage unnecessary calls) and Reason-then-Act (requiring the model to reason about tool necessity before acting). Both provide limited control: Prompt-only suppresses necessary calls alongside unnecessary ones, and Reason-then-Act still incurs a disproportionate accuracy cost on hard tasks. To understand why these baselines fail, we probe the models' hidden states and find that tool necessity is linearly decodable from the pre-generation representation with AUROC 0.89--0.96 across six models, substantially exceeding the model's own verbalized reasoning. This reveals that models already know when tools are needed, but fail to act on this knowledge during generation. Building on this finding, we propose Probe&Prefill, which uses a lightweight linear probe to read the hidden-state signal and prefills the model's response with a steering sentence. Across all models tested, Probe&Prefill reduces tool calls by 48% with only 1.7% accuracy loss, while the best baseline at comparable accuracy only reduces 6% of tool calls, or achieves a similar tool call reduction but incurs a 5× higher accuracy loss. Our code is available at https://github.com/Trustworthy-ML-Lab/when2tool