Source-linked AI summary

BAGEN: Are LLM Agents Budget-Aware?

Yuxiang Lin, Zihan Wang, Mengyang Liu, Yuxuan Shan, Longju Bai, Junyao Zhang, Xing Jin, Boshan Chen, Jinyan Su, Xingyao Wang, Jiaxin Pei, Manling Li

arXiv:2606.00198v1cs.LGcs.AIcs.CL

TL;DR

The paper addresses the gap between post-hoc cost tracking and agents’ ability to estimate resource needs during execution. It evaluates progressive interval estimation through rollout replay across models and environments, finding weak alignment with task success, systematic optimism, and late failure recognition. Early stopping can save tokens on failed trajectories, while precise interval calibration remains difficult.

  • Problem

    Agent benchmarks commonly measure token consumption after execution and rarely assess whether agents can estimate budget needs as feasibility changes turn by turn.

    Source: §1 · pp. 1–2
  • Method

    The paper evaluates progressive interval estimation by replaying unconstrained rollouts and querying agents at each prefix for a remaining-budget interval or an impossible declaration.

    Source: §1 · p. 2
  • Results

    Task success correlates only weakly with interval hit rate (r≈0.35); models are systematically optimistic, recognize failure late, and early stopping saves 28–64% of tokens on failed trajectories.

    Source: §1 · pp. 2–3
  • Takeaways & Limitations

    Budget awareness is distinct from task performance and can serve as a resource-control signal, while early stopping provides a practical intervention on failed trajectories.

    Source: §1 · pp. 2–3

    The evaluation avoids online estimation because estimating during rollout consumes tokens and would conflate task-completion cost with self-assessment cost.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

While agents are increasingly spending more resources, today agent cost is mostly measured only after execution. A Budget-Aware Agent (BAGEN) should treat budget as an active control signal, rather than a passive cost metric. We first systematically define budget estimation as internal budgets (from agent computation) and external budgets (from agent actions). We then formalize budget-awareness as progressive interval estimation: at each step of a plan, an agent should predict an upper and lower bound on remaining budget, and alert when completion is unlikely. Scoring with a rollout-replay protocol, we find consistent failure patterns on four environments and five frontier agents: (1) strong agents do not necessarily have strong budget-awareness, with correlation r=0.35. (2) frontier models are consistently over-optimistic, continue spending on tasks that are unlikely to succeed, instead of alerting the user early. (3) budget-aware signal is actionable and trainable. Early stop saves 28-64% tokens on failed trajectories, and SFT+RL strengthens early stop and alert behavior. (4) precise interval calibration remains challenging, with interval coverage capping at 47% after SFT+RL. Project page: https://ragen-ai.github.io/bagen/

1. Introduction

The paper frames budget awareness as an agent’s ability to estimate remaining resources and task feasibility during execution, addressing evaluations that measure spending only after the fact. It introduces progressive interval estimation and finds that budget awareness diverges from task performance, with systematic optimism, late failure recognition, and actionable early stopping.

1. Introduction

  • The paper defines budget awareness as estimating remaining budget and whether a task is still finishable, because cost is usually tracked only after execution.It distinguishes internal budgets consumed by agent computation from external budgets committed by agent actions.Source: §1 · pp. 1–2
  • Progressive interval estimation re-queries an agent at every rollout prefix for a remaining-budget interval or an impossible declaration, scored against realized budget and outcome.The evaluation covers five frontier models and four environments spanning internal and external budgets.Source: §1 · p. 2, Figure 1 · p. 2
  • Task success correlates only weakly with interval hit rate (r≈0.35), and no model dominates all three budget-awareness sub-capabilities.On Search-R1, Opus has the highest task success rate (75.8%), while Sonnet produces better intervals (36.5% vs. 23.1%).Source: §1 · p. 2
  • All twenty model–environment pairs underestimate remaining budget more often than they overestimate it, and weaker models are more optimistic, not less.Feasibility improves with fine-tuning, but interval coverage reaches only 47% after SFT+RL.Source: §1 · p. 2
  • On failed trajectories, models predict feasibility above 70% even after 60% of the budget is consumed; alarms fire only in the final 20%.Estimates shift during execution, but not reliably toward the truth.Source: §1 · p. 2
  • Early stopping on impossible predictions saves 28–64% of tokens on failed trajectories, at a cost of 1.6–4.2 percentage points in success rate.SFT followed by RL improves estimation on Sokoban.Source: §1 · p. 3
  • The authors argue budget should serve as a control signal for resource-limited agents, not merely an after-the-fact accounting metric.Source: §1 · p. 3

2. Two Failures of Single-Point Budget Estimation

A pilot study finds that single-point budget estimates are systematically optimistic and differ between task-start and later-turn judgments. These failures motivate intervals and repeated estimates as execution progresses.

2. Two Failures of Single-Point Budget Estimation

  • The pilot elicited total-token estimates at task start and after replaying logged prefixes from five frontier models on two internal-budget environments.Source: §2 · p. 3
  • Across five models on Sokoban and Search-R1, first-turn predictions underestimate realized budget more often than they overestimate it.The bias tracks model confidence rather than task difficulty; weaker models are more optimistic.Source: §2 · p. 3
  • On Sokoban, Gemini’s feasibility macro-F1 improves by +21.9 points from first-turn to all-turn evaluation; on Search-R1, Qwen3-235B moves the opposite way by 9.3 points.Which way first-turn judgments differ from all-turn judgments depends on the model and task.Source: §2 · p. 3
  • Intervals can express uncertainty or infeasibility, while queries at every turn capture refinements from partial progress that a task-start estimate misses.Source: §2 · p. 3

3. Agent Budget Awareness

BAGEN frames budget awareness as estimating remaining resources throughout execution and judging whether completion is still feasible. It separates internal compute costs from external action costs, evaluates the capability through replayed trajectory prefixes, and measures three sub-capabilities.

3.1. Budget Modalities: Internal and External

  • Internal budgets measure the model’s generated compute, while external budgets measure costs committed through actions and may couple dimensions such as money, time, and inventory.Internal-budget tasks test whether the model can predict the tokens needed for remaining reasoning and actions; external budgets can include constraints whose trade-offs interact.Source: §3.1 · pp. 3–4

3.2. Progressive Interval Estimation and Rollout-Replay

  • At each turn, progressive interval estimation uses the observed trajectory prefix to estimate remaining budget or declare completion impossible; rollout-replay scores predictions against actual remaining cost and outcome.The protocol first records an unconstrained rollout, then replays each non-terminal prefix for estimation, avoiding the extra token cost that interleaved estimation would add to task execution.Source: §3.2 · p. 4

3.3. Sub-Capability Metrics

  • BAGEN separates feasibility prediction, early failure detection, and remaining-budget interval estimation into distinct sub-capabilities.Interval quality is evaluated on successful trajectories only when feasibility is correctly predicted; coverage checks whether actual remaining cost falls within the interval, while tightness penalizes wide intervals.Source: §3.3 · pp. 4–5

3.4. Experimental Setup

  • The evaluation covers internal-budget tasks Sokoban, Search-R1, and SWE-bench, plus Warehouse, an external-budget supply-chain task with coupled cost, time, and occupancy constraints.Warehouse spans 24 weeks and uses challenge-conditioned feasibility probes balanced between reachable and unreachable cases; the study evaluates five frontier models and trains Qwen2.5-7B-Instruct.Source: §3.4 · pp. 5–6

4. Budget Awareness Decouples from Task Performance

Budget-estimation quality is distinct from task-completion performance across the evaluated environments. Agent estimates can outperform linear extrapolation for some Warehouse budget dimensions, while no model leads across all budget-awareness sub-capabilities.

4. Budget Awareness Decouples from Task Performance

  • Across all four environments and three sub-capabilities, task-completion strength does not reliably predict budget-awareness quality.Across 20 model-environment pairs, task success rate correlates weakly with feasibility prediction (r ≈0.35).Source: §4 · p. 6, §4.1 · p. 6

4.1. Task Success Rate and Budget Estimation is Decoupled

  • On Search-R1, Opus has the highest task success rate (75.8%), but Sonnet has a higher interval hit rate (36.5% vs. 23.1% for Opus); Warehouse interval hit rates range from Gemini’s 50.2% to Sonnet and Qwen’s 17.3%.On SWE-bench, Opus leads task success, Gemini feasibility prediction, and GPT-5.2 interval estimation; Warehouse has no separate rollout-success outcome.Source: §4.1 · p. 6
  • On Warehouse, agent midpoint estimates reduce occupancy error in most model-progress bins, especially early, while linear extrapolation often remains competitive for cumulative cost.Figure 3 encodes positive cells where the agent has lower absolute error than the extrapolation baseline.Source: §4.1 · p. 7, Figure 3 · p. 7

4.2. No Model Dominates All Sub-Capabilities

  • No model leads all three sub-capabilities: Gemini has the highest Sokoban binary F1 (61.9%) but the lowest interval hit rate (8.8%), while Warehouse’s highest Fail-F1 is Gemini’s 62.8%.Interval hit rate correlates more strongly with midpoint bias (r ≈−0.67) and width adequacy (r ≈0.62) than with feasibility F1 (r ≈0.35).Source: §4.2 · p. 7

5. Why Does Budget Estimation Fail?

Budget-estimation failures mix calibration and reasoning limits: optimism and late failure recognition persist, while estimates update but do not reliably improve. Training helps some predictions, but precise interval estimates and cross-task transfer remain difficult.

5.1. Binary Feasibility Is Calibration; Interval Estimation Is Reasoning (Hypothesis 1)

  • RL without an SFT warm-start collapses into always-impossible predictions or invalid formats; SFT provides a format prior that sparse-reward RL does not recover alone.Figure 4 indicates RL improves estimation only from a suitable SFT initialization.Source: Figure 4 · p. 8, §5.1 · p. 9

5. Why Does Budget Estimation Fail?

  • Across five frontier models, optimism and late recognition are supported; interval estimation is reasoning-limited, binary feasibility is calibration-limited, and estimates update but not uniformly toward truth.The four tests support Hypotheses 2 and 4, reject Hypothesis 3, and distinguish the interval and feasibility findings for Hypothesis 1.Source: §5 · p. 8, §5.3 · p. 10

5.1. Binary Feasibility Is Calibration; Interval Estimation Is Reasoning (Hypothesis 1)

  • SFT raises Qwen-7B feasibility accuracy from 25.5% to roughly 90%, but interval coverage reaches only 47% after SFT+RL, separating calibration from reasoning limits.For interval estimation, SFT coverage varies from 26% to 53% with target width; SFT+RL has 28% median midpoint relative error.Source: §5.1 · p. 8, Figure 4 · p. 8
  • Training improves reward in every in-task setting, but cross-task evaluation retains only 17–36% of in-task reward, suggesting budget estimation is largely task-specific.The strongest token-budget interval quality is on Search-R1; Sokoban and SWE-bench remain harder, with feasibility learned more readily than calibrated intervals.Source: Table 2 · p. 9, §5.1 · p. 9

5.2. Optimistic Bias Is Universal Across Models and Tasks (Hypothesis 2)

  • Across all 20 model-environment pairs, optimistic misses outnumber conservative misses at every rollout-progress bin; weaker models are more optimistic, not less.Sonnet and Opus are closest to calibrated, but still skew low; averaging over more turns does not eliminate the bias.Source: §5.2 · p. 9

5.3. Estimates Update Over Turns, but Not Always Toward the Truth (Hypothesis 3)

  • Later estimates differ from first-turn predictions but are not consistently better: Gemini’s Sokoban macro-F1 rises 21.9 points, while Qwen’s SWE-bench macro-F1 falls 12.5 points.Figure 6 scatters on both sides of the equality line, so later turns cannot be assumed to refine predictions toward truth.Source: Figure 6 · p. 10, §5.3 · p. 10

5.4. Failure Is Recognized Too Late to Act On (Hypothesis 4)

  • On failed trajectories, models predict feasibility above 70% even after 60% of the budget is spent; predictions drop sharply only in the final 20% of execution.Figure 7 shows the late-recognition pattern across environments.Source: Figure 7 · p. 10, §5.4 · p. 11

6. The Signal Is Actionable and Trainable

The section tests whether budget estimates can guide early stopping and examines how training choices shape interval quality. Acting on impossible predictions saves tokens with a small success-rate cost, while interval width trades coverage against precision.

6.1. Early Stopping Saves Tokens at Low Risk

  • The policy terminates a trajectory when the estimator predicts impossible; false aborts trade success for compute, while false continues miss opportunities to save compute.The signal is already present in model predictions; the policy acts on it.Source: §6.1 · p. 11
  • Early stopping saves 28% to 64% of tokens on failed trajectories while reducing overall success by only 1.6 to 4.2 percentage points.Savings are near zero on Search-R1 but largest on Warehouse and SWE-bench; GPT-5.2 saves 64% with a 6.6% false-abort rate, while Opus saves 28% with 2.2% false aborts.Source: §6.1 · p. 11

6.2. Training Dynamics: Width, Epoch, and Initialization

  • 100-token fixed intervals achieve 47% coverage with 49% midpoint relative error; 1,000-token intervals exceed 90% coverage but reach 170% error, while narrow intervals yield 26% coverage and 45% error.Wider targets buy coverage at the cost of precision; narrower targets do the reverse.Source: §6.2 · pp. 11–12

7. Related Work

Related work spans resource-constrained decision-making, budget-aware LLM agents, efficiency methods, self-monitoring, and long-horizon agent failures. This paper positions budget awareness as mid-trajectory self-estimation, distinct from static-data interval calibration.

7. Related Work

8. Conclusion and Limitations

BAGEN frames remaining-budget estimation as a distinct agent capability and evaluates it across five frontier models and four environments. The results show structured failures and actionable early stopping, while interval estimation remains an open challenge.

8. Conclusion and Limitations

  • Across four environments, BAGEN separates budget awareness from task completion and identifies optimistic bias, late failure recognition, and calibration-versus-reasoning limits.The paper reports early stopping and SFT-then-RL as actionable training directions, while identifying interval-estimation improvement as a central open problem.Source: §8 · p. 12
  • The study evaluates progressive interval estimation across five frontier models and four internal- and external-budget environments using rollout replay.The training analysis is limited to Sokoban; future work includes additional model scales and environments, feedback from estimator predictions into actor decisions, and fungible budgets across dimensions.Source: §8 · p. 12
  • A simple early-stop policy makes budget estimates actionable, and SFT-then-RL strengthens early-stop and alert behavior on Sokoban.The interval-estimation gap remains the central open problem.Source: §8 · p. 12

A. Extended Related Work

The related-work discussion distinguishes estimating remaining budget online from managing known budgets, static interval prediction, compute-control, self-correction, and progress monitoring. The appendices describe four evaluation environments, replay-based estimation, pilot probes, Warehouse feasibility construction, and model deployment metadata.

A. Extended Related Work

  • Prior work largely models known budgets, static prediction intervals, compute control, correctness monitoring, or progress; BAGEN instead estimates remaining resources online with uncertainty and infeasibility warnings.The paper positions this capability at the intersection of dimensions highlighted in Figure 9.Source: A. Extended Related Work · pp. 18–20

B.1. Environments and Tasks

  • The four benchmarks cover token budgets in SearchR1, Sokoban, and SWE-bench, plus Warehouse’s joint financial and operational constraints.Warehouse tests whether an agent can estimate feasibility across external budgets rather than internal reasoning cost.Source: §B.1 · p. 21

B.2. Evaluation Protocol

  • At each nonterminal logged prefix, the estimator predicts remaining budget or declares the task impossible, using only information available through that turn.Estimators receive replayed history and budget limits, but not future turns, terminal labels, or realized remaining budget.Source: §B.2 · pp. 21–22

B.3. Pilot Study Setup

  • The pilot compares task-start token estimates with later-turn estimates from the same five models on Sokoban and Search-R1, using 128 rollouts per model-task pair.It scores optimism against realized cost and compares feasibility macro-F1 from first-turn versus per-turn predictions.Source: §B.3 · p. 22

B.4. Warehouse Budget Construction

  • Warehouse feasibility probes pair logged rollouts with target-cash and three-resource constraints, balanced equally between reachable and unreachable cases.This 50/50 balance supports evaluation on both sides; it is not an estimate of real-world failure prevalence.Source: §B.4 · p. 23

B.5. Experiment Matrix

  • Table 4 presents the experiment matrix for four benchmarks, including budget modality, rollout count, horizon constraint, and predicted remaining-budget target.The supplied passage identifies the table’s scope but does not provide its cell values.Source: §B.5 · p. 23, Table 4 · p. 24

B.6. Model Identifiers and Inference Settings

  • The appendix records model identifiers and invocation settings, including provider, checkpoint or API identifier, query date, reasoning configuration, output limit, and sampling policy.“Low Thinking” denotes a constrained provider-native reasoning configuration, with backend-specific controls listed in the appendix table.Source: §B.6 · pp. 23–24

C.1. Data Source

The Warehouse environment uses real, unsmoothed retailer demand data and is designed to make agents balance money, space, and time. Its calibration exposes policy effects within an 11-step episode, while delayed financial payoffs create a long-horizon credit-assignment challenge.

C.1. Data Source

  • The demand panel uses 22 weeks of unsmoothed weekly sell-through data from an anonymized US distributor and five retailer accounts across five product families.At episode reset, retailer demand series are randomly permuted across game SKUs, and per-step demand combines two weeks with additive noise.Source: §C.1 · p. 25

C.2. Design Goal: Stressing Money, Space, and Time

  • The environment stresses money, space, and time through constrained cash flow, warehouse capacity, chargebacks, production lead times, and shipping delays.Mechanisms are anchored to at least one of these three resource axes.Source: §C.2 · p. 25

C.3. Calibration Decisions

  • Zero starting inventory makes outcomes policy-driven from the outset, while compressed production and transit times let consequences emerge within the 11-step episode.Ocean shipping is 32 ± 4 days, air is 6 ± 1 days, and production lead times are 25–45 days by SKU; the episode spans 11 steps.Source: §C.3 · pp. 25–26

C.4. Reward and Termination

  • The reward subtracts stockout and overstock penalties and a cash-below-zero penalty from operating profit, while manufacturing and financing costs affect cash rather than reward.Production costs can pay back through revenue four to six steps later, preserving a long-horizon coordination challenge.Source: §C.4 · p. 26

C.5. Theoretical Maximum Reward (Defined but Unused)

  • The appendix derives an unused cumulative-reward ceiling under idealized mean demand, demand-matched inventory, and no cash constraint.Under those assumptions, period revenue, holding cost, and operating expense are deterministic.Source: §C.5 · p. 26

D. Additional Experimental Visualizations

The additional visualizations show asymmetric budget-estimation errors, early-stopping tradeoffs, differences between task success and budget-estimation quality, and uncertainty updates as rollout evidence accumulates.

Interval-width analysis

  • Figure 12’s interval-width pattern suggests estimators partly update uncertainty as more rollout evidence becomes available.Interval width is normalized by total token budget for internal tasks and reported for time, warehouse item-weeks, and cost in Warehouse.Source: D. Additional Experimental… · p. 27

Estimation error direction

  • Optimistic misses dominate conservative misses across models and persist through execution, though conservative misses become more visible later.The reported direction of error means agents more often underestimate remaining budget than overestimate it.Source: D. Additional Experimental… · p. 27

Early-stopping tradeoff

  • Stopping at the first impossible prediction can save substantial budget on failed rollouts, especially in Warehouse and SWE-bench, with smaller savings in Search-R1.False-abort rates remain low for most model-benchmark pairs; Search-R1 savings are smaller because its rollouts are shorter.Source: D. Additional Experimental… · p. 27

Task success and budget estimation

  • Across benchmarks and models, higher rollout success does not necessarily imply better feasibility prediction.The paper presents budget awareness as distinct from task-solving ability and requiring direct evaluation.Source: D. Additional Experimental… · p. 27

E. Detailed Settings for Estimation RL and SFT

The section specifies how Sokoban budget probes are constructed and how estimation models are trained and evaluated. It also defines the interval targets and reward used in the SFT/RL experiments.

E. Detailed Settings for Estimation RL and SFT

  • The SFT/RL experiments use Qwen2.5-7B-Instruct on Sokoban 6x6 with one box and a 2500-token budget, using 794 SFT samples and a balanced 380-sample test set.Source: E. Detailed Settings for Estimation RL… · p. 27
  • Each probe follows a completed rollout prefix; its remaining-token label sums future assistant outputs and user inputs, while feasibility depends on rollout success and staying within budget.The held-out evaluation uses a deterministic trajectory split, and training samples are balanced between possible and impossible states.Source: E. Detailed Settings for Estimation RL… · p. 27
  • Interval SFT targets use percentage- or fixed-width bounds, while the reported RL interval experiments warm-start from interval-SFT checkpoints.The supplied settings also describe a point-estimation SFT variant, but state that reported RL interval experiments use interval checkpoints.Source: E. Detailed Settings for Estimation RL… · pp. 27–28
  • Figures 11–13 depict optimistic single-point misses, changes in interval width over task progress, and the distinction between estimation quality and task performance.Figure 12 reports normalized internal-task widths and separate normalized Warehouse widths for time, warehouse item-weeks, and cost.Source: Figure 11 · p. 28, Figure 12 · p. 28, Figure 13 · p. 29
  • GRPO uses rule-based rewards that require valid intervals covering the true remaining tokens on possible probes and an exact “impossible” answer on impossible probes.The reward favors tight covering intervals and weights possible-probe rewards 1.8 to 0.2 relative to impossible-state recognition.Source: E. Detailed Settings for Estimation RL… · p. 29

F. Additional Data for RL

The additional RL results indicate that training from the base model does not reliably produce valid interval outputs. The accompanying prompt templates specify progressive feasibility judgments and remaining-budget intervals across agent environments.

F. Additional Data for RL

  • RL runs started from the base model collapse to all-impossible answers, while a scalar-output run does not produce valid intervals; the results indicate SFT is needed before RL.Source: F. Additional Data for RL · pp. 29–30

G. Prompts for Agent Estimation and SFT/RL

  • Prompts ask agents to judge whether completion remains feasible within budget and, if so, estimate an interval for additional usage; they prioritize feasibility over interval tightness.The prompt examples provide completed-turn usage and require “impossible” when completion is not feasible.Source: G. Prompts for Agent Estimation and… · pp. 31, 33
  • The Warehouse prompt extends this format to remaining time, warehouse occupancy, and cost while requiring the rollout to meet a target cash threshold.Source: G. Prompts for Agent Estimation and… · p. 32
  • The reported Sokoban sweep is described with feasibility accuracy and recall, interval coverage, midpoint relative error, and combined reward.The notes define coverage on truly possible probes and MRE as the relative error of the interval midpoint.Source: Table 7 · p. 36, G. Prompts for Agent Estimation and… · p. 36
Loading 2606.00198v1…