Rewards for Budget-Aware Multi-Turn Agents That Adapt to Arbitrary Test-Time Constraints

Tom A. Lamb; Adam Golinski; Eleonora Gualdoni; Sinead Williamson; Michael Kirchhof

cs.LGOct 20, 2026

Rewards for Budget-Aware Multi-Turn Agents That Adapt to Arbitrary Test-Time Constraints

Authors: Tom A. Lamb, Adam Golinski, Eleonora Gualdoni, Sinead Williamson, Michael Kirchhof

Organizations: Apple · University of Oxford

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.

Figures & tables

Figure 1: Cost-conditioned rewards. [Left] Return

G

against episode cost

C

for successful (

R{=}1

) and failed (

R{=}0

) trajectories across different cost-conditional rewards. Shading marks

C>B

. Bands illustrate

\Lambda_{b}\in[0,1]

for Rate and Dual . [Right] Reward properties: Budget Cutoff stops rollouts before exceeding

B

, Budget Jump denotes a return discontinuity at

B

, In-Budget Cost penalises resource use before

B

, Overspend Size makes the return depend on the extent of overspending, Failure Return

<0

means failed trajectories can receive negative return, and Learned Price denotes a separate learned violation price for each training budget.

Figure 2: Budget-augmented POMDP. The policy receives the task $x$ , total budget $B_{0}$ , and remaining budget $b_{t}$ alongside the environment observation $o_{t}$ . Each action’s cost is deducted from $b_{t}$ , and the updated budget $b_{t+1}$ is appended in natural language to the next observation ( Section 3.2 ).

Figure 3: In-domain budget utility and compliance. Budget utility

\mathcal{U}(b)

and natural compliance

\mathrm{Comp}(b)

against stated budget

b

for Qwen3.5-9B and 4B. Rows show AppWorld test-normal under step, token, and API-call budgets, followed by BabyAI under turn budgets. Each policy is trained for one environment and resource type. Vertical dashed lines mark training budgets, with other points unseen during training. Bands show

\pm 1

standard deviation of the binomial rate estimates.

Figure 4: Spend by task difficulty and overspend magnitude. AppWorld test-normal under step, token, and API-call budgets. (a) Mean spend by task difficulty for Qwen3.5-9B. (b) Mean relative overspend

\mathbb{E}_{\pi}[\bar{C}-1\mid\bar{C}>1]

for Qwen3.5-9B and 4B, shown as a percentage of the budget. Horizontal axes show the stated budget

B

. Overspend axes are capped at

100\%

	$B=10$	$B=15$	$B=32$
Base	$1.65\pm 0.56$	$1.45\pm 0.53$	$1.46\pm 0.40$
GRPO-Std	$1.02\pm 0.05$	$1.03\pm 0.16$	$1.05\pm 0.17$
Gate	$\mathbf{1.93\pm 0.66}$	$1.73\pm 0.50$	$1.53\pm 0.39$
Sigmoid-Gate	$1.87\pm 0.69$	$1.63\pm 0.45$	$1.53\pm 0.36$
Dual	$1.92\pm 0.62$	$\mathbf{1.76\pm 0.48}$	$\mathbf{1.61\pm 0.38}$

Table 1: Mean API calls per step by budget. Qwen3.5-9B, restricted to episodes completed within budget on AppWorld test-normal with step budgeting. Bold: highest per budget.

Method	Win (%) $\uparrow$	Near-best (%) $\uparrow$	Regret (pp) $\downarrow$
Linear-Mag	3.1	6.5	14.01
Trunc	6.6	20.5	7.56
Rate	3.1	26.0	7.46
Gate-Neg	16.1	26.0	6.00
Exp-Mag	7.8	22.9	4.32
Sigmoid	14.0	39.2	4.31

Table 2: Utility across budgets. Win and near-best (

\leq

2pp) rates and mean regret against the per-budget best method, averaged equally over the eight AppWorld and BabyAI experiments ( Section 12.1 ).

Figure 5: Raw task success against spend. AppWorld test-normal pass@1 without a budget cutoff against mean spend for Gate , Gate-Neg , Sigmoid-Gate , and Dual . Each point corresponds to a stated budget. Policies retain budget feedback, but successes count regardless of compliance. Bands show

\pm 1

standard deviation of the binomial rate estimates.

Figure 6: Distribution shift. AppWorld-trained policies evaluated without retraining on AppWorld test-challenge (top) and ASTRA (bottom) under step budgets. Curves show budget utility

\mathcal{U}(b)

and natural compliance

\mathrm{Comp}(b)

against stated budget

b

. Bands show

\pm 1

standard deviation of the binomial rate estimates.

Task and environment ( Sections 3.1 and 3.2 )
$\mathcal{S},\tilde{\mathcal{S}},\mathcal{A},\Omega,\tilde{\Omega}$	Task and budget-augmented state spaces, action space, and observation spaces
$\mathcal{T},\tilde{\mathcal{T}},\mathcal{O},\tilde{\mathcal{O}}$	Task and budget-augmented transition and observation kernels
$\mathcal{M},\tilde{\mathcal{M}}$	Task POMDP and budget-augmented POMDP
$\gamma,\mu_{0}$	Discount factor and initial state distribution
$x,p_{\mathcal{X}}$	Task and task distribution over $\mathcal{X}$
$h_{t},\mathcal{H}$	Observation history at step $t$ , and set of all histories

Table 3: Notation. Symbols used throughout the main paper and the appendix.

Unit	$\kappa(a_{t})$	Interpretation of $B_{0}$	Domains
Steps	$1$	Model turns available to the agent	aw , as , ba
Tool calls	# API calls in $a_{t}$	API calls available to the agent	aw
Tokens	$\|a_{t}\|$ in tokens	Generated tokens, including reasoning and actions	aw

Table 4: Budgeted quantities. Quantities defined by the cost map

\kappa

. The domain abbreviations denote AppWorld, ASTRA, and BabyAI.

Unit	Model	Training grid $\mathcal{B}$	Evaluation grid $\mathcal{B}_{\mathrm{eval}}$
Steps	Qwen3.5-9B	{10, 15, 18, 22, 25, 29}	{4, 8, 10, 12, 15, 18, 20, 22, 25, 27, 29, 32, 36}
	Qwen3.5-4B	{10, 15, 17, 20, 24, 26}	{4, 8, 9, 10, 11, 12, 15, 17, 18, 20, 21, 22, 24, 26, 29, 31, 32, 34, 36}
Tool calls	Qwen3.5-9B	{10, 15, 18, 22, 26, 29}	{6, 10, 12, 15, 18, 20, 22, 24, 26, 29, 32, 36}
	Qwen3.5-4B	{13, 15, 17, 19, 22, 25}	{6, 10, 11, 12, 13, 14, 15, 17, 19, 20, 21, 22, 24, 25, 26, 27, 30, 31, 35, 36}
Tokens	Qwen3.5-9B	{896, 1408, 1856, 2432, 3136, 3904}	{896, 1152, 1408, 1600, 1856, 2112, 2432, 2752, 3136, 3520, 3904}
	Qwen3.5-4B	{1452, 2073, 2489, 3066, 4264, 5186}	{1452, 1728, 2073, 2240, 2489, 2752, 3066, 3584, 4264, 4672, 5186}

Table 5: Budget grids. Training and evaluation budget grids by cost unit and model. All methods use the same grids.

Hyperparameter	Value
Model and optimisation
Model families and scales	Qwen3.5 4B and 9B (released post-trained checkpoints)
Fine-tuning	full-parameter, no LoRA or adapters
Optimiser	Dr. GRPO, $\gamma=1$
Training updates	$120$ (Qwen3.5-9B), $80$ (Qwen3.5-4B)
Learning rate	flat, no schedule, $3\times 10^{-6}$

Table 6: Training hyperparameters. Budget grids are given by model and cost unit in Table 5 . Settings that differ between model scales or reward formulations are stated explicitly. All other settings are shared across methods and domains.

Figure 7: Seed robustness and longer training. Budget utility

\mathcal{U}(b)

on AppWorld test-normal under step budgets for Qwen3.5-9B Sigmoid-Gate and Linear-Mag . Two additional seeds are evaluated after

120

updates (left) and

200

updates (right), with the main

120

-update seed shown in both panels. Dotted lines mark training budgets. Bands show

\pm 1

standard error over tasks.

Figure 8: Resource use on BabyAI. (a) Mean turns used and (b) mean percentage overspend among rollouts that exceed the budget, for Qwen3.5-9B and 4B. Horizontal axes show the stated turn budget

B

, with dotted lines marking training budgets. Overspend axes are capped at

100\%

LLM Agent Efficiency

See all 54 papers

Oct 7, 2026

On the Clock: Towards Punctual and Productive Time-Budgeted AI Agents

We study whether small LLM agents can operate effectively under explicit wall-clock time budgets by both respecting the allocated runtime and using available time productively. We evaluate Qwen3.6-27B on five competitions from MLE-Bench Lite and Qwen3-4B on Zork I (Jericho), two agentic benchmarks where additional computational time can meaningfully improve performance. In the simplest setting, where the budget is stated only in the prompt, agents fail to translate the stated budget into controlled use of time. These failures arise from gaps in time awareness, since the harness provides no timing feedback, but also because they cannot reliably anticipate the duration of actions, and do not have a learned mapping from available time to an appropriate strategy. We investigate two complementary classes of interventions: harness-based mechanisms that expose timing information and enforce deadlines, and reinforcement learning with budget-aware rewards. Injecting timing information through the harness substantially improves budget adherence for Qwen3.6-27B without measurable loss in performance, while enforcement hooks tighten adherence further. RL with GRPO achieves near-perfect budget adherence on Zork I and generalizes to held-out budgets not seen during training, but does not improve task performance over the untrained harness on MLE-Bench. Once agents are made to respect the budget, they still fail to use additional time to improve task performance. RL-trained policies learn when to stop but often fill extra time with repeated actions, and GRPO training on multiple budgets tends to collapse toward the strategy learned for the shortest budget. Our results reveal a gap between time adherence and productive time allocation, which remains a central challenge for budget-conditioned agents.

Aaron Wang, Neelabh Madan, Vlad Sobal +5

University of Illinois at Chicago · AWS Agentic AI · New York University +1

Runtime Enforcement for AI Agents LLM Agent Harnesses RL for LLM Agents AI Agent Benchmarks

Ask PDF

May 29, 2026

BAGEN: Are LLM Agents Budget-Aware?

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/

Yuxiang Lin, Zihan Wang, Mengyang Liu +9

1Northwestern University · 2O2 Lab · 3Independent +5

LLM Agent Evaluation AI Agent Benchmarks LLM Agents LLM Agent Training

Ask PDF

May 19, 2026

ZEBRA: Zero-shot Budgeted Resource Allocation for LLM Orchestration

As autonomous agents increasingly execute end-to-end tasks under fixed monetary budgets, the pressing open question shifts from whether the budget is respected, to how to spend it effectively. Existing budget-aware methods typically control reasoning step-by-step within a single agent, or learn resource allocation policies via RL. None address how to split a budget across the composing phases of a multi-agent pipeline at inference time. We propose ZEBRA, a zero-shot framework that reduces multi-phase budget allocation to a continuous nonlinear knapsack problem: an LLM controller estimates per-phase utility curves, and a water-filling search on the Lagrange multiplier returns the per-phase split. Additive and multiplicative aggregations are unified under the same solver. On a

150

-task APPS coding benchmark, both ZEBRA variants outperform LLM-direct (budget allocation directly by an LLM) on every aggregate metric. At a budget of

α= 0.5

of the unconstrained spend, ZEBRA recovers

94.4\%

of unconstrained quality, versus

88.1\%

for LLM-direct. The advantage is statistically significant and transfers beyond coding: on a

3

-phase HotpotQA pipeline, ZEBRA beats LLM-direct by

14.3

pp, with allocations empirically robust to curve-estimation noise. On HotpotQA, ZEBRA arrives at a different budget split (near-balanced) compared to the APPS one (skewed towards a refinement phase), showing adaptation to the pipeline structure. More broadly, we show that lightweight algorithmic guidance at inference time can improve the economic behavior of autonomous multi-agent systems.

May Hamri, Inbal Talgam-Cohen

1Tel Aviv University

LLM Agent Orchestration Cost-Aware Inference LLM Agent Workflow Optimization Multi-Agent Systems Multi-Agent Resource Allocation

Ask PDF

Abstract

LLM Agent Efficiency

Explore similar work

On the Clock: Towards Punctual and Productive Time-Budgeted AI Agents

BAGEN: Are LLM Agents Budget-Aware?

ZEBRA: Zero-shot Budgeted Resource Allocation for LLM Orchestration