Agents learn to manage resources within user-set budgets during tasks
Rewards for Budget-Aware Multi-Turn Agents That Adapt to Arbitrary Test-Time Constraints
Machine LearningArtificial Intelligence
Summary
Many language-based AI helpers need to finish tasks without running out of resources like time or steps. The authors studied if these helpers can learn to wisely use their given budget instead of relying on external controls. They trained different models with special rewards that encouraged staying within budget and found these models adjusted their efforts better to task difficulty and limits. This approach helped the agents follow the rules naturally and succeed more often, even on new challenges and budgets they hadn't seen before.
What this means in practice
- •For conversational ai developers: Create AI assistants that better manage their computing resources to complete tasks reliably under changing budget limits.
- •For robotics software teams: Develop control policies for robots that adapt their action budgets dynamically to task difficulty and operational constraints.
Authors
Tom A. Lamb, Adam Golinski, Eleonora Gualdoni, Sinead Williamson, Michael Kirchhof
Abstract
Large language model agents are often given an inference budget and asked to maximise success within it. A multi-turn agent has to allocate this budget across many environment interactions. An external harness can enforce the budget, but it remains unclear whether the agent itself can learn when and how much to spend. We ask whether agents can learn this allocation across dynamically specified budgets and which rewards best teach this behaviour. We adapt existing length penalties to budget constraints spanning multiple resource types, introduce new reward formulations, and compare eleven objectives in a shared training pipeline. We train Qwen3.5 models at 4B and 9B on AppWorld with budgets over tokens, API calls, and environment steps, and on the BabyAI gridworld with turn budgets. We train a separate policy for each environment and resource type, with each policy trained across budget values. Training with rewards that penalise overspending substantially improves within-budget success and natural compliance over untrained models and over training on task success alone, with agents learning to adapt their resource use to stated budgets and task difficulty. These gains extend to unseen budgets and, zero-shot, to a shifted task distribution and a new benchmark. Reward objectives trade off success within budget and natural compliance. A sigmoid penalty for overspending most often achieves the highest within-budget success at an individual budget, while a fixed negative reward for overspending generally yields the strongest natural compliance. Combining penalties for overspend frequency and magnitude gives the smallest average gap to the best method at each budget. Overall, we show that a single agent can learn to allocate resources across turns and adapt its behaviour to user-specified budgets using budget feedback without external budget-allocation controllers.