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

TL;DR

This work introduces Calibrate-Then-Act (CTA), a framework that equips LLM agents with explicit priors to reason about the cost-uncertainty tradeoffs involved in environment exploration. By separating uncertainty calibration from action selection, CTA induces more cost-efficient, Pareto-optimal exploration strategies across tasks such as Pandora's Box, knowledge QA with optional retrieval, and coding with selective testing. The approach yields notable improvements: higher alignment with oracle policies in toy settings, better-discounted rewards in QA, and enhanced generalization when combined with RL (CTA-RL) compared to end-to-end RL baselines. The findings suggest that meta-information in the form of priors—whether derived from model confidence or training data—can meaningfully shape agent policies and drive practical gains in real-world interactive tasks.

Abstract

LLMs are increasingly being used for complex problems which are not necessarily resolved in a single response, but require interacting with an environment to acquire information. In these scenarios, LLMs must reason about inherent cost-uncertainty tradeoffs in when to stop exploring and commit to an answer. For instance, on a programming task, an LLM should test a generated code snippet if it is uncertain about the correctness of that code; the cost of writing a test is nonzero, but typically lower than the cost of making a mistake. In this work, we show that we can induce LLMs to explicitly reason about balancing these cost-uncertainty tradeoffs, then perform more optimal environment exploration. We formalize multiple tasks, including information retrieval and coding, as sequential decision-making problems under uncertainty. Each problem has latent environment state that can be reasoned about via a prior which is passed to the LLM agent. We introduce a framework called Calibrate-Then-Act (CTA), where we feed the LLM this additional context to enable it to act more optimally. This improvement is preserved even under RL training of both the baseline and CTA. Our results on information-seeking QA and on a simplified coding task show that making cost-benefit tradeoffs explicit with CTA can help agents discover more optimal decision-making strategies.

Figures (16)

• 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.
• Figure 2: Standard agentic decision loop (left) and proposed method CTA with estimated priors (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's retrieval decision with respect to their confidence level $k_\text{da}$ and retrieval discount factor $\gamma$. Each dot corresponds to one question: green indicates the model directly answers, and red indicates it retrieves. The dashed line marks the oracle threshold: red region retrieves, green region directly answers. Models with calibrated priors closely align with the oracle decision rule, exhibiting more cost-aware retrieval behavior.
• Figure 4: Action pattern distribution for prompting and RL-trained agents, with and without calibrated priors, 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.
• Figure 5: Pareto frontier of average reward under varying costs. Static strategies (test-first or code-first) achieve high reward only in limited regimes, whereas CTA-RL with estimated priors consistently attains Pareto-optimal performance across cost settings.