Source-linked AI summary

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

Wenxuan Ding, Nicholas Tomlin, Greg Durrett

arXiv:2602.16699v3cs.CLcs.AI

TL;DR

LLM agents must acquire information under uncertainty while paying for exploration, yet standard policies can remain insensitive to these costs. The paper formalizes this problem across synthetic, QA, and file-reading tasks and introduces Calibrate-Then-Act, which supplies inferred environmental priors for action selection. CTA changes agent behavior toward more cost-sensitive and optimal decision-making, with RL-trained CTA outperforming baseline RL.

  • Problem

    LLM agents must decide when information gathering is worth its cost, but their exploration policies are frequently static across changing uncertainty and cost profiles.

  • Method

    CTA formalizes exploration as sequential decision-making and presents agents with inferred priors about latent environment state to separate uncertainty calibration from action selection.

  • Results

    Across synthetic, QA, and file-reading settings, CTA enables more dynamic cost-sensitive decisions, and RL-trained CTA outperforms baseline RL.

  • Takeaways & Limitations

    Explicit meta-level priors can induce LLM agents to reason about Pareto-optimal behavior and navigate action costs effectively.

  • Takeaways & Limitations

    In FILEREADING, the agent is assumed to be trained on task data linking filenames to format information, and priors may be supplied by a format predictor.

Abstract

from arXiv · show

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.

1 Introduction

LLM agents must balance information-gathering costs against uncertainty, but their exploration policies are often static. Calibrate-Then-Act (CTA) makes environmental priors explicit so agents can select actions more optimally across cost and uncertainty profiles.

  • Exploration improves information but increases API costs, interaction latency, and user burden.
  • Agent policies often remain static despite settings where confidence should determine whether to explore or act directly.
  • CTA decouples uncertainty calibration from action selection by presenting priors that prompt abstract reasoning about the sequential decision problem.
  • Across QA and coding, CTA uses confidence or environmental-structure priors to support dynamic decisions, while RL further improves performance over baseline RL.
  • The paper frames exploration across three settings and targets better cost-performance tradeoffs than baseline agents.

2 Background: Agentic Exploration with Costs

The paper formalizes cost-aware exploration as sequential decision-making under partial observability, where agents infer latent environment features and choose costly exploration or commitment actions. The objective is to maximize expected discounted reward while adapting to uncertainty and action costs.

  • The environment is modeled as a partially observable Markov decision process with states, actions, observations, transitions, rewards, and cost-integrating discounts.
  • Problem-critical unobserved environment features are represented by a latent variable Z that affects agent performance.
  • At each timestep, the agent updates a posterior over Z and chooses among exploration actions or a terminating commit action.
  • Exploration costs discount successful task-completion value before commitment, with different actions carrying different costs.
  • The objective is to maximize expected discounted reward by adapting exploration and commitment to uncertainty and environmental cost constraints.

3 Tasks

The paper instantiates cost-aware exploration in Pandora’s Box, optional-retrieval QA, and file-reading coding tasks. These tasks require deciding whether information-gathering actions justify their costs before committing to an answer.

  • 3.1 Pandora’s Box: Pandora’s Box tests whether agents can reason about uncertainty-cost tradeoffs by verifying boxes or committing to a box under discounted reward.
  • 3.1 Pandora’s Box: The optimal Pandora’s Box policy commits when a box’s posterior probability exceeds γ; otherwise it verifies the highest-probability remaining box.
  • 3.2 Task QA: In QA, the agent chooses between retrieving external evidence and answering directly, trading potential accuracy gains against latency and API costs.
  • 3.2 Task QA: QA priors estimate direct-answer and post-retrieval correctness probabilities, with the oracle retrieving when pret · γ ≥ pda.
  • 3.3 Task FILEREADING: FILEREADING requires inferring latent CSV formatting attributes, then selecting unit tests, code execution, or final answering under action-specific discounts.
  • 3.3 Task FILEREADING: The file-format prior may come from conventions, past experience, or an explicitly provided format predictor.

4 Method: Calibrating Agent Environment Exploration

CTA explicitly supplies estimated priors about latent environment state, separating uncertainty calibration from action selection so agents can reason about exploration costs and information value. The framework uses confidence-based priors for QA and filename-based priors for file-reading tasks, with prompted and RL variants.

  • Decision-making under uncertainty: Agents must compare exploration costs with the expected value of information when deciding whether to continue exploring or commit.The value of information depends on current beliefs represented by prior and posterior distributions.
  • CTA framework: CTA provides an estimated prior ˆp(Z | x) to the agent, conditioning its decision loop on explicit information about the environment.The prior can be supplied at inference and/or training time.
  • CTA variants: CTA-Prompted adds prior information to the baseline prompt, while CTA-RL additionally fine-tunes the model with a discounted-reward objective.Both variants use the same prior-conditioned prompting structure, with CTA-RL adding end-to-end GRPO training.
  • Prior estimation: For QA, CTA estimates answerability using verbalized confidence because the probability of answering correctly without retrieval is not directly observable.The approach is designed to work with both black-box and glass-box model usage.
  • Prior estimation: Confidence regression reduces ECE from 0.618 to 0.029 on PopQA, addressing poor initial calibration that end-to-end RL may not systematically learn.The paper contrasts recalibration with memorizing task-specific facts during long RL training.
  • Prior estimation: For file reading, MBERT predicts filename-conditioned distributions over delimiter, quote character, and skiprows, then supplies these probabilities to the agent.This decouples uncertainty estimation from action selection during RL training or testing.

5 Experimental Setup

The experiments evaluate CTA against prompted and RL baselines on retrieval-augmented QA and filename-informed file-reading tasks under varied exploration costs. Performance is assessed through exploration behavior, task accuracy, and cost-discounted reward.

  • Datasets: QA uses 1,000 PopQA questions with Contriever retrieval and samples γ uniformly from [0.1, 0.65] to vary retrieval costs.The discount factor controls the retrieval-cost regime studied for each question.
  • Datasets: FILEREADING hides the true file format while exposing filenames, and varies unit-test and code-execution discounts across task instances.The relative code-cost settings use ρ ∈{0.5, 1.0, 2.0, 4.0}.
  • Metrics: The evaluation reports exploration statistics, final task accuracy, and reward that discounts correctness by exploration costs.QA uses retrieval rate, while FILEREADING uses unit-test and code-attempt counts.
  • Models and baselines: Qwen3-8B is the base model, and CTA-PROMPTED and CTA-RL are compared against directly prompted and end-to-end RL baselines.The prompted baselines include a thinking and a non-thinking condition.
  • QA evaluation: Figure 3 compares retrieval-before-answering versus direct-answer decisions across γ bins and reports oracle policy match rate.CTA-PROMPTED is identified as making adaptive cost-sensitive decisions with the highest oracle match rates across settings.
  • Training setup: RL training in FILEREADING duplicates instances across four relative cost values, producing a 4× larger training set than MBERT training.This exposes RL settings to diverse cost trade-offs.

6 Results

Across synthetic, QA, and FILEREADING settings, explicit priors help LLM agents adapt exploration to uncertainty and costs, whereas baselines often follow static or poorly calibrated policies.

  • Pandora’s Box: 94.0% optimal match rate: CTA-PROMPTED adapts to discount factors and prior distributions when explicit priors make environment constraints available.Average reward also increases relative to settings without explicit priors or without thinking mode.
  • Cross-task findings: Baseline LLM agents fail to form calibrated decision rules from task descriptions alone across QA and FILEREADING, even after RL.The issue is not limited to a single task setting.
  • FILEREADING: Both PROMPTED and RL collapse to static test-first behavior across FILEREADING cost regimes, producing 0% code-first traces.RL without explicit priors defaults to a suboptimal exploration policy.
  • QA: CTA-PROMPTED adapts retrieval frequency to retrieval costs and achieves the highest QA reward and oracle match rate across cost constraints.It retrieves more often at high discount factors and less often at low discount factors.
  • FILEREADING: CTA-RL improves FILEREADING discounted reward from 0.259 for RL to 0.268, a 3.5% increase, across aggregated cost settings.The result is reported on unseen test data and is attributed to conditioning training on estimated priors.
  • Cost sensitivity: CTA-RL remains on the Pareto frontier across relative cost values, while static strategies perform well only in limited cost regimes.RL shows an advantage mainly at large ρ, whereas code-first performs well mainly at small ρ.

7 Related Work

Related work studies incomplete-information decision making and cost-aware tool use, but the joint treatment of uncertainty priors and cost constraints remains underexplored.

  • Decision making under incomplete information: Prior work addresses incomplete information through clarification, environment queries, and collaboration in partially observed or ambiguous settings.These approaches span underspecified queries, ambiguity, and partially observed environments.
  • Agents in cost-aware deployment: Cost-aware agent research examines tool planning and reducing unnecessary retrieval or tool use while preserving answer quality.Examples include abstention, selective search, efficiency-oriented policies, and steerable clarification policies.
  • Research gap: The joint treatment of uncertainty priors and cost constraints is identified as an underexplored aspect of efficient exploration.Both factors together determine Pareto-optimal behavior in the paper’s framing.

8 Conclusion

The paper shows that presenting priors over unobserved environment features can induce LLMs to reason about Pareto-optimal behavior and navigate action costs.

  • Conclusion: Providing priors over unobserved environment features enables LLMs to balance uncertainty-cost tradeoffs and navigate action costs effectively.The paper presents this as a way to shape agent policies through meta-level information.
  • Conclusion: Meta-level information about capabilities may help induce agents to think and act more optimally.This conclusion is stated within the paper’s studied environment-interaction settings.

A Limitations

The evaluation uses controlled QA and FILEREADING settings that are simpler than open-ended deployments. CTA also depends on useful prior estimates and does not guarantee optimal behavior in arbitrary environments.

  • Scope: Controlled QA and FILEREADING settings isolate cost-aware exploration but omit richer action spaces, longer horizons, and noisier observations found in open-ended deployments.Extending CTA to more complex environments remains future work.
  • Prior quality: CTA is robust to moderate prior-estimation error, but performance can degrade when priors are highly miscalibrated or confidently wrong.The approach demonstrates prior estimation from model confidence and historical task data, while broader-domain estimators remain future work.
  • Reasoning requirements: CTA does not guarantee optimal behavior in arbitrary environments because the model must correctly interpret priors and reason about their implications for action selection.More complex scenarios may require posterior updates, multi-step value of information, and reasoning over dependencies among latent variables.

C Sensitivity to Prior Quality

The FILEREADING sensitivity analysis degrades CTA-RL’s prior quality systematically and finds gradual performance loss rather than a sharp failure transition. Confidently wrong priors are substantially more harmful than uncertain ones, while CTA-RL remains statistically competitive across cost regimes.

  • Prior-quality sensitivity: As prior quality decreases, both reward and task accuracy degrade gradually rather than through a sharp failure transition.Even the BERT prior achieves only 67% accuracy, yet approximate priors still yield meaningful gains.
  • Prior-quality sensitivity: Adversarial priors produce the largest performance drop because being confidently wrong is substantially more harmful than being uncertain.The adversarial prior inflates the incorrect option to 95% confidence.
  • Statistical comparison: CTA-RL achieves significant gains over different baselines in several cost regimes and is never statistically worse than any baseline.The tests use paired per-task rewards and 10,000 bootstrap resamples across cost settings ρ ∈ {0.5, 1.0, 2.0, 4.0}.
  • Statistical comparison: CTA-RL significantly outperforms the no-prior test-first baseline at ρ ∈ {0.5, 1.0, 2.0} and is statistically indistinguishable at ρ = 4.0.Against the prior-aware code-first baseline, it significantly outperforms at ρ = 4.0 and is indistinguishable at lower ρ.

E Qualitative Trace Analysis of Pandora’s Box Problem

Qualitative Pandora’s Box traces show that explicit priors and enabled thinking change exploration from indiscriminate verification to cost-sensitive commitment. Without priors or explicit value-cost reasoning, the model follows suboptimal strategies.

  • CTA-PROMPTED-NONTHINK: With thinking mode disabled, the model verifies all options before committing instead of comparing information value with exploration cost.This behavior occurs regardless of the prior distribution and leads to unnecessary exploration.
  • PROMPTED: Without access to prior probabilities, the model implicitly treats bags as equally likely and follows a suboptimal verification strategy.The trace deviates from the optimal policy.
  • CTA-PROMPTED: With explicit priors and thinking enabled, the model compares immediate commitment with discounted verification and selects behavior aligned with the oracle policy.In the example, it chooses to guess B immediately, which is optimal for the stated priors and discount factor.

G Dataset Construction Details for Task: Coding with Selective Testing

The FILEREADING benchmark generates CSV question-answering tasks whose filenames provide priors over latent formatting configurations. The oracle policy selects the most probable remaining option and decides whether verification is worth its discount cost.

  • Dataset construction: Each task combines a structured filename, generated CSV file, and query, with filename tokens inducing a prior over possible parsing configurations.The dataset uses indicative tokens such as _eu, _tab, _sas, and _cn alongside irrelevant strings.
  • Dataset construction: The latent configuration z = (z_d, z_q, z_s) specifies the delimiter, quote character, and number of header rows to skip.The correct answer is obtainable only when the file is parsed with the fully correct configuration.
  • Dataset construction: The oracle filename-to-format model maps each filename to a factorized prior distribution over formatting attributes, then samples the true configuration to generate CSV content and queries.Each of 16 filename feature configurations induces a corresponding prior, and the generated dataset contains 2,000 task instances.
  • Dataset construction: The dataset contains 2,000 instances split into 1,400 training, 300 validation, and 300 test examples.Incorrect delimiters, quote characters, or skipped rows cause malformed outputs, parsing failures, or inaccessible answers.
  • Oracle strategy: The oracle policy considers the maximum-posterior box, then chooses between immediate commitment and verification using their expected values under discounting.If verification succeeds it yields reward 1; otherwise the problem recurses on the remaining set, with base case V({i}) = 1.

I GRPO Training Details

The paper supplies task-specific prompts for Pandora’s Box, QA, and FILEREADING, including CSV-format priors and cost parameters. Example traces contrast RL behavior without explicit priors with CTA-RL’s joint reasoning about uncertainty and action costs.

  • Reasoning traces: Under ρ = 4.0, the RL model without explicit priors defaults to unit tests and does not explicitly reason about relative action costs.Its trace is characterized as verification-first behavior rather than adaptive cost-aware decision-making.
  • Prompted task setup: The prompts define sequential choices between information-gathering actions and final commitment for Pandora’s Box, QA, and CSV file-reading tasks.The task interfaces include VERIFY versus GUESS, RETRIEVE versus ANSWER, and unit tests or code execution before committing to a CSV answer.
  • Prompted task setup: The FILEREADING prompts provide estimated CSV-format likelihoods so agents can use probabilistic defaults when weighing unit tests, code execution, and early commitment.The stated uncertainty concerns the exact CSV formatting, which may be handled through assumptions or verification.
  • Reasoning traces: Under the same high relative code-cost setting, CTA-RL reasons jointly about CSV-format uncertainty and the relative costs of unit tests and code execution.This behavior determines whether verification is worth performing before a code attempt.
  • Review coverage: The paper states that its abstract and introduction claims are supported by the formalization, method, and experimental results, while limitations are discussed in Appendix A.The review material also states that the paper provides a full algorithm and proof for the discounted Pandora’s Box policy.
Loading 2602.16699v3…