Calibrate-Then-Act: Cost-Aware Exploration in LLM Agents

Wenxuan Ding; Nicholas Tomlin; Greg Durrett

cs.CLFeb 18, 2026

Calibrate-Then-Act: Cost-Aware Exploration in LLM Agents

Authors: Wenxuan Ding, Nicholas Tomlin, Greg Durrett

Organizations: New York University

Abstract

LLM agents are deployed in environments where they must interact to acquire information. In these scenarios, the agent must reason about inherent cost-uncertainty tradeoffs in how to act, such as when to stop exploring and commit to an answer. For instance, on a programming task, an agent might run the code it generates, or it might generate tests for that code snippet; the cost of writing and running a test is nonzero, but typically lower than the cost of running buggy code. In this work, we show that we can induce LLM agents to explicitly reason about balancing these cost-uncertainty tradeoffs, then act more optimally in their environments. We formalize multiple tasks, including retrieval-augmented QA and a file reading coding task, as sequential decision-making problems under uncertainty. Each problem has latent environment state that impacts the agent's performance. We introduce a framework called Calibrate-Then-Act (CTA), where we pass the agent an inferred prior about this environment state to enable it to act more optimally. This information qualitatively changes agent behavior, and adds environment sensitivity to the agent which is not learned via standard RL training. Our results on a synthetic task, QA, and file reading show that making cost-benefit tradeoffs explicit with CTA helps agents discover more optimal decision-making strategies.

Figures & tables

Figure 1: Given the same task, a coding agent may either verify assumptions via intermediate checks carefully (right) or attempt a direct solution as soon as possible (left). The optimal choice depends on uncertainty and specific cost constraints. Calibrate-Then-Act (CTA) materializes this information for better decision-making.

Setting	$z^{*}\in\mathcal{Z}$	$p$	$\mathcal{A}$	$\mathcal{O}$	$\theta$	$R$
Unified (§ 2 )	Latent env. state	Prior over relevant latent variables	Actions	Observations	Cost	Reward
Pandora (§ 3.1 )	$z^{*}\in\{z_{1},\ldots,z_{K}\}$ Prize-containing box	$\{p_{k}=\pi(z^{*}=z_{k})\}$	Verify( $k$ ), Commit( $k$ )	Whether $z_{k}$ has prize	$\gamma$	$\gamma^{T}\mathbb{I}(z_{k}=z^{*})$
QA (§ 3.2 )	$(p_{\text{da}},p_{\text{ret}})$ Answer probabilities with and without retrieval	$p_{\text{da}}=\delta(z=k_{\text{da}}(\mathbf{x}))$ , $p_{\text{ret}}=\delta(z=k_{\text{ret}})$	Retrieve, Answer( $a$ )	Retrieved context	$\gamma$	$\gamma^{T}\mathbb{I}(a=a^{*})$
Code (§ 3.3 )	$\mathbf{z}^{}=(z_{d}^{},z_{q}^{},z_{s}^{})$ $\in\mathcal{Z}_{d}\times\mathcal{Z}_{q}\times\mathcal{Z}_{s}$	$p(\mathbf{z}\mid n)$	UNIT_TEST ( $f$ ), CODE ( $d$ , $q$ , $s$ ), ANSWER( $a$ )	Format value, stdout + stderr	$(d_{u},d_{c})$	$d_{u}^{U}d_{c}^{C}\mathbb{I}(a=a^{*})$

Table 1: Unified formalization of cost-aware decision problems and their instantiations across tasks. We characterize each by latent variables

z^{*}

, a prior belief

p

over

\mathbf{z}

, an action space

\mathcal{A}

consisting of exploration and commit actions, observations

\mathcal{O}

revealed through exploration, costs

\theta

associated with the exploration actions, and a final reward

R

that discounts task success by incurred costs.

Figure 2: Standard agentic decision loop (left) and proposed method CTA (right). In CTA, we learn a prior estimator from training data and condition the agent on estimated $\hat{p}$ at inference and/or training time, inducing more optimal decision making through explicit reasoning over prior probabilities.

Figure 3: Model retrieval decision across retrieval discount factors

\gamma

on the QA task. Red bars show the percentage of examples where the model retrieves before answering; green bars show direct answers. The black line reports the oracle policy match rate. CTA-Prompted makes adaptive cost-sensitive decisions and achieves highest oracle match rates across settings.

Method	Optimal Match Rate (%)	Avg. Reward
Oracle policy	100.0	0.649
Prompted-NonThink	11.0	0.441
Prompted	23.0	0.476
CTA-Prompted-NonThink	20.0	0.436
CTA-Prompted	94.0	0.625

Table 2: Performance comparison on the Pandora’s Box task. A thinking model given access to priors can solve the problem near-optimally.

Figure 4: Action distribution for agents on FileReading across relative cost parameters

\rho

. Each stacked bar shows the proportion of decision traces corresponding to different action patterns, with the reward

R

labeled above. Annotated percentages indicate the fraction of tasks where the agent attempts code execution before any unit tests.

Table 7

Figure 5: Average reward under varying costs for our approaches. Static strategies (test-first or code-first) achieve high reward only in limited regimes, whereas CTA-RL with estimated priors is consistently on the Pareto frontier

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Prior	Prior Acc	Avg Reward	Task Acc	# Code	$R_{\rho=0.5}$	$R_{\rho=1.0}$	$R_{\rho=2.0}$	$R_{\rho=4.0}$
BERT	0.67	0.269	99.1%	1.46	0.382	0.315	0.230	0.147
Noisy	0.52	0.251	96.3%	1.71	0.365	0.295	0.206	0.136
Uniform	0.46	0.251	96.8%	1.53	0.349	0.300	0.219	0.134
Adversarial	0.33	0.203	94.4%	2.26	0.299	0.240	0.173	0.099

Appendix

Table 5: Sensitivity to prior quality in the FileReading task. Performance degrades smoothly as priors become less accurate, without a sharp failure transition.

$\rho$	vs. Test first (B1)	vs. Code first (B2)	vs. RL
$0.5$	$+0.0845^{**}$	$+0.0149^{=}$	$+0.0224^{**}$
$1.0$	$+0.0508^{**}$	$+0.0059^{=}$	$+0.0117^{*}$
$2.0$	$+0.0169^{**}$	$+0.0153^{=}$	$+0.0030^{=}$
$4.0$	$-0.0021^{=}$	$+0.0287^{**}$	$-0.0010^{=}$

Appendix

Table 6: Paired bootstrap tests for CTA-RL on FileReading .

\Delta

denotes the mean per-task reward gain of CTA-RL over each baseline. Tests use

n=300

paired tasks per

\rho

and

B=10{,}000

bootstrap resamples.

Panel of Figure 6 — Figure 6: Example interaction trace on a 3-bag Pandora’s Box instance with priors (0.04, 0.68, 0.28) and discount factor $\gamma=0.2$ with thinking mode disabled . In this setting, the model explores all bags before committing and follows a suboptimal verification order, rather than prioritizing the highest-probability option.

Panel of Figure 7 — Figure 7: Example interaction trace on a 3-bag Pandora’s Box instance with priors (0.04, 0.68, 0.28) and discount factor $\gamma=0.2$ , where the model is not given access to the prior probabilities. In this setting, the model implicitly treats the bags as equally likely and follows a suboptimal strategy that deviates from the optimal policy.

Panel of Figure 8 — Figure 8: Example model reasoning trace on a 3-bag Pandora’s Box instance with priors (0.04, 0.68, 0.28) and discount factor $\gamma=0.2$ . The model explicitly compares the expected value of immediate guessing versus verification and then chooses to guess $B$ immediately, which is the optimal strategy in this case. Key reasoning steps, including the explicit comparison between action value and exploration cost, are highlighted in yellow. This example illustrates that when priors are provided explicitly, the model can reason about uncertainty and exploration cost to select an optimal strategy.

Panel of Figure 9 — Figure 9: Prompt template for Pandora’s Box setting.

Panel of Figure 10 — Figure 10: Prompt templates for QA .

Panel of Figure 11 — Figure 11: System prompt for FileReading .

Panel of Figure 12 — Figure 12: Continuation of the system prompt for FileReading .

Panel of Figure 13 — Figure 13: Instruction prompt template specifying the CSV task, reward parameters, and constraints provided to the agent.

Panel of Figure 14 — Figure 14: Instruction prompt template with estimated CSV format likelihoods, enabling the agent to use probabilistic defaults when trading off unit tests, code execution, and early commitment.

Panel of Figure 15 — Figure 15: Example reasoning trace of an RL-trained model without explicit prior conditioning in the CSV exploration task. Despite operating under the same high relative code cost setting ( $\rho=4.0$ ), the model defaults to verification-first behavior based on surface cues (e.g., file extension) and does not explicitly reason about uncertainty or cost trade-offs, illustrating a lack of adaptive decision-making compared to the CTA-RL model.

Panel of Figure 16 — Figure 16: Example reasoning trace of the CTA-RL model on the FileReading task ( $\rho=4.0$ ), illustrating cost-aware trade-offs between unit tests and code execution under a high relative code cost setting, while jointly reasoning about format uncertainty.

LLM Agents

Large language model agents use language models to plan, use tools, and act across multi-step tasks and environments.

See all 681 papers

ReDAct: Uncertainty-Aware Deferral for LLM Agents
Apr 8, 2026Dzianis Piatrashyn, Nikita Kotelevskii, Kirill Grishchenkov +7LLM Agents LLM Agent Reliability
Mohamed bin Zayed University of Artificial Intelligence · Ocapital · National University of Science and Technology (NUST) MISIS +3
Utility-Guided Agent Orchestration for Efficient LLM Tool Use
Mar 20, 2026Boyan Liu, Gongming Zhao, Hongli XuLLM Agent Orchestration LLM Tool Use
University of Science and Technology of China, Hefei, Anhui, China · Suzhou Institute for Advanced Research, University of Science and Technology of China
Scores Are Not Decisions: Cost-Aware Stopping for Tool Acquisition in LLM Agents
Jul 29, 2026Yicheng Feng, Yan Zhang, Yan Cheng +1Decision-Focused Learning Tool-Using Agents
School of Advanced Manufacturing and Robotics, Peking University · Desautels Faculty of Management, McGill University · Business School, Shanghai University of Finance and Economics +1

Apr 8, 2026J/K move · Enter open · S save

ReDAct: Uncertainty-Aware Deferral for LLM Agents

Dzianis Piatrashyn, Nikita Kotelevskii, Kirill Grishchenkov, Nikita Glazkov+6

Mohamed bin Zayed University of Artificial Intelligence · Ocapital · National University of Science and Technology (NUST) MISIS+3

LLM Agents LLM Agent Reliability LLM Routing LLM Inference Efficiency

Open paper AskPDF arXiv

Recently, LLM-based agents have become increasingly popular across many applications, including complex sequential decision-making problems. However, they inherit the tendency of LLMs to hallucinate, leading to incorrect decisions. In sequential settings, even a single mistake can irreversibly degrade the trajectory, making hallucinations an even bigger problem. Although larger LLMs hallucinate less, they incur a significantly higher per-token cost. In this paper, we address this tradeoff by proposing ReDAct (Reason-Defer-Act). In ReDAct, an agent is equipped with two LLMs: a small, cheap model used by default, and a large, more reliable but expensive model. When the predictive uncertainty of the small model exceeds a calibrated threshold, the decision is deferred to the large model. We evaluate our approach in text-based embodied environments such as ALFWorld and MiniGrid and show that deferring only about 15% of decisions to the large model can match the quality of using it exclusively, while significantly reducing inference costs.

Abstract

LLM Agents

Explore similar work

ReDAct: Uncertainty-Aware Deferral for LLM Agents

Utility-Guided Agent Orchestration for Efficient LLM Tool Use

Scores Are Not Decisions: Cost-Aware Stopping for Tool Acquisition in LLM Agents