🤖 machine learning

BAGEN: Are LLM Agents Budget-Aware?

The paper introduces BAGEN, a framework that defines budget-awareness as progressive interval estimation and demonstrates that while frontier LLM agents consistently fail to accurately predict or alert on budget overruns, their ability to make actionable early stops can be significantly improved through supervised fine-tuning and reinforcement learning.

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

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are sending a robot on a long, complex mission, like a delivery driver who has to navigate a city, pick up packages, and deliver them. You give the robot a budget: maybe 1,000 "energy points" (tokens) or a specific amount of real money.

The big problem this paper identifies is that today's smart robots (AI Agents) are terrible at knowing how much of their budget they have left while they are working. They usually only realize they ran out of money after the job is already ruined.

Here is the paper's story, broken down with simple analogies:

1. The Core Problem: The "Blind" Robot

Right now, if you ask a robot, "How much will this trip cost?" it usually guesses once at the very beginning.

The Flaw: It's like a hiker guessing how many miles they have left to walk before they even start walking. They don't account for getting lost, bad weather, or taking a wrong turn.
The Reality: The paper found that even the smartest robots are over-optimistic. They think, "I can definitely finish this with the money I have left!" even when they are actually running out of gas. They keep spending money on tasks that are already impossible to finish, just to keep trying.

2. The New Idea: The "Budget-Aware" Agent (BAGEN)

The authors propose a new type of robot that acts like a smart project manager. Instead of guessing once, this robot checks its wallet every single step of the way.

Progressive Interval Estimation: Instead of saying "I need exactly 500 more points," the robot says, "I probably need between 400 and 600 points, but if I don't find the answer soon, I might need to stop."
The "Impossible" Button: If the robot realizes it's in a hole it can't climb out of, it should hit a big red button that says "IMPOSSIBLE" and stop immediately, rather than wasting the last few dollars trying to dig deeper.

3. What They Found (The "Report Card")

The researchers tested five of the world's smartest AI models on four different types of tasks (like solving puzzles, searching the web, and managing a warehouse). Here is what they discovered:

Being Smart $\neq$ Being Budget-Conscious: Just because a robot is great at solving the puzzle doesn't mean it's good at counting its money. In fact, the correlation was weak. The best puzzle-solver was often the worst at estimating costs.
The "Optimism Bias": Every single robot tested was too hopeful. They consistently underestimated how much money they would need. The weaker the robot was at the task, the more optimistic it was about its budget.
Too Late to Save: The robots usually realized they were going to fail only after they had already spent 80% of their budget. By the time they said, "Oh no, we can't finish," it was too late to stop the bleeding.
It Can Be Taught: The good news is that this skill can be trained. By teaching the robots (using a method called SFT + RL), they learned to stop earlier.
- The Result: If you tell the robot to stop as soon as it says "Impossible," you can save 28% to 64% of the wasted money on failed attempts, with only a tiny drop in how many tasks they successfully finish.

4. The Two Types of "Budgets"

The paper looked at two kinds of spending:

Internal Budget (The Brain's Fuel): This is the computer power and "tokens" the AI uses just to think and talk.
External Budget (The Real World Cost): This is actual money, time, or physical resources (like warehouse space) the AI spends when it takes action in the real world.
The robots struggled with both, but the "over-optimism" problem was the same for both.

5. The Takeaway

The paper concludes that budget awareness is a separate skill from being smart. Currently, AI agents are like drivers who keep driving toward a destination even after their gas tank is empty, hoping they'll magically find a gas station.

The solution isn't necessarily to make the AI "smarter" at the task, but to train it to be a better accountant for itself. If we teach these agents to say "I can't do this" earlier, we can save a massive amount of resources without losing much in terms of success.

In short: We need to stop asking AI agents to just "do the job" and start teaching them to keep an eye on their wallet while they work.

Technical Summary: BAGEN: Are LLM Agents Budget-Aware?

1. Problem Statement

Current Large Language Model (LLM) agents are increasingly deployed in long-horizon, high-stakes tasks involving significant resource consumption, such as token usage (internal budget) and API calls, monetary costs, or time (external budget). However, existing evaluation paradigms treat budget primarily as a post-hoc metric, measuring cost only after execution completes.

The paper identifies a critical gap: agents lack budget awareness, defined as the ability to estimate, mid-execution, how much budget remains and whether a task is still finishable. Without this capability, agents cannot make critical control decisions, such as aborting hopeless tasks, requesting additional resources, or reallocating budgets across sub-goals. Furthermore, existing studies often rely on single-point estimation (a single prediction at the task start), which fails to capture the evolving feasibility of long-horizon tasks and does not allow for progressive refinement based on observed progress.

2. Methodology

2.1. Budget Modalities

The authors formalize two distinct budget types:

Internal Budget: Compute generated by the model's own reasoning (e.g., token count).
External Budget: Costs committed in the environment (e.g., money, time, inventory), which can be multi-dimensional and coupled (e.g., holding inventory increases revenue but consumes warehouse capacity).

2.2. Progressive Interval Estimation Protocol

To address the limitations of single-point estimation, the paper proposes Progressive Interval Estimation. Instead of a single point prediction, the agent is queried at every turn $k$ of a trajectory $\tau$ to predict:

An interval $[\hat{R}^{lo}_k, \hat{R}^{hi}_k]$ representing the remaining budget if the task is deemed feasible.
An impossible declaration if the task is deemed infeasible given the remaining budget.

This protocol introduces three key properties:

Uncertainty: Expressed via interval width.
Progressiveness: Estimates update as the agent observes partial progress.
Infeasibility Awareness: An explicit "impossible" option enables early stopping.

2.3. Rollout-Replay Evaluation

To decouple budget estimation ability from task completion ability, the authors employ a rollout-replay protocol:

Rollout Generation: The agent executes a task without budget constraints to generate a full trajectory and ground-truth costs.
Prefix Replay: The logged prefix up to turn $k$ is replayed to the agent, which is then asked to estimate the remaining budget and feasibility.
Scoring: Predictions are scored against the realized remaining cost and the actual outcome.

2.4. Sub-Capability Metrics

Budget awareness is decomposed into three measurable sub-capabilities:

Feasibility Prediction: A binary classification (feasible vs. impossible) measured by Macro-F1.
Early Failure Detection: The ability to recognize doomed trajectories early, measured by Fail-F1 (precision/recall on the "impossible" class).
Remaining Budget Estimation: The accuracy of interval predictions on successful trajectories, measured by Coverage (whether the true cost falls in the interval) and Tightness (penalizing overly wide intervals).

2.5. Experimental Setup

The study evaluates five frontier models (GPT-5.2, Claude Opus 4.7, Claude Sonnet 4.6, Gemini 3.1 Pro, Qwen3-235B) across four environments:

Internal Budget: Sokoban (planning), Search-R1 (multi-hop retrieval), SWE-bench (GitHub issue resolution).
External Budget: Warehouse (a supply-chain environment with coupled cost, time, and inventory dimensions).

3. Key Results

3.1. Decoupling of Task Performance and Budget Awareness

Weak Correlation: Task success rate correlates weakly with budget estimation quality ( $r \approx 0.35$ ). Strong task performers are not necessarily strong estimators. For example, on Search-R1, Opus achieved the highest task success (75.8%) but Sonnet produced better interval estimates (36.5% hit rate vs. 23.1%).
No Dominant Model: No single model excels across all three sub-capabilities (feasibility, early detection, interval calibration).

3.2. Systematic Failures

Optimistic Bias: All models consistently underestimate remaining budget (optimistic bias). This bias is more pronounced in weaker models and persists across all environments.
Late Failure Recognition: Agents recognize failure too late to act. On failed trajectories, models often predict "feasible" with >70% confidence even after 60% of the budget is consumed. The "alarm" typically fires only in the final 20% of execution.
Instability: First-turn estimates often disagree with later-turn estimates, and the direction of this disagreement varies by model and task.

3.3. Training and Actionability

Binary Feasibility is a Calibration Problem: Supervised Fine-Tuning (SFT) alone raised Qwen-7B's feasibility accuracy from 25.5% to ~90%, suggesting the capability was latent but required format calibration.
Interval Estimation is a Reasoning Problem: SFT improved interval coverage, but adding Reinforcement Learning (SFT+RL) was necessary to push coverage to 47%. However, RL without an SFT warm-start caused training collapse.
Actionable Signal: Implementing a simple early-stop policy (terminating when the model predicts "impossible") saved 28–64% of tokens on failed trajectories with a minimal cost to overall success rate (1.6–4.2 percentage points).

4. Contributions

Definition of Budget Awareness: Formalizes budget awareness as a distinct capability separate from task execution, defined via progressive interval estimation.
Evaluation Protocol: Introduces the rollout-replay protocol to isolate estimation capabilities from execution capabilities.
Empirical Analysis: Systematically demonstrates that frontier models suffer from systematic optimism and late failure recognition, and that these failures are distinct from task performance.
Training Insights: Shows that binary feasibility is a calibration issue solvable with SFT, while precise interval estimation is a reasoning bottleneck requiring SFT+RL, and that the signal is immediately actionable for resource saving.

5. Significance and Claims

The paper claims that budget should be treated as an active control signal rather than a passive accounting metric. The authors argue that current agents lack the metacognitive monitoring required to manage resources effectively in long-horizon tasks.

The significance lies in demonstrating that:

Budget awareness is a trainable and actionable capability, even if current models are imperfect at it.
Simple interventions (like early stopping based on model self-prediction) can yield substantial resource savings.
The gap between task performance and budget estimation suggests that future agent development must explicitly target self-monitoring and calibration rather than assuming these emerge naturally from improved reasoning or task success.

The paper concludes that while precise interval calibration remains a challenge (capping at 47% coverage after training), the ability to detect infeasibility and estimate costs is a critical, distinct capability that must be developed to enable resource-limited agents to operate effectively in real-world, high-stakes environments.