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.LG, cs.AI
2026-10-20
Apple trains Qwen3.5 agents with 11 overspend-penalty rewards; within-budget success and natural compliance beat task-only RL and transfer zero-shot to a new benchmark.
Frontier agent products already let users set the effort level: how many tokens to generate, how many tool calls to make, how many steps to take. A multi-turn agent has to spread one budget across dozens of interactions, deciding at every turn how much to spend given what is left. Today that allocation is handled by external machinery at inference time, whether planners, harnesses, or routers; the task-performing policy never learns it.
Prompting alone does not work. Untrained Qwen3.5 models see the total and remaining budget in every observation and still fail: at a step budget of 36 on AppWorld, the 9B base model reaches 37% within-budget success and the 4B only 19%, and the 9B base complies with token budgets on fewer than half of episodes. The paper asks two questions: can a single policy learn this allocation across dynamically specified budgets, and which rewards teach it best.
The budget is written into the environment, and reward shaping teaches the policy to spend it.
One design choice matters: except for Trunc, training rollouts may run past the budget. An overspending trajectory that earns zero still gets a negative advantage against its group mean, so its post-budget actions keep providing signal. Trunc deletes that segment and loses to Gate on every pooled metric.
Aggregated over eight experiments (AppWorld step, token, and API-call budgets plus BabyAI turn budgets), using three yardsticks: win rate at an individual budget, near-best rate (within 2pp of the best), and mean regret (average gap to the per-budget best, in percentage points).
| Method | Metric | Result |
| Sigmoid-Gate | win rate | 20.1%, highest |
| Gate | near-best rate | 39.8%, mean regret 3.33pp |
| Dual | mean regret | 2.92pp, lowest |
| Gate-Neg | natural compliance | generally strongest, at 6.00pp regret |
| Linear-Mag | mean regret | 14.01pp, worst |
Behaviour changes too. Under step budgets, GRPO-Std makes a flat 1.02–1.05 API calls per turn, while Gate, Sigmoid-Gate, and Dual average 1.87–1.93 calls per turn at a budget of 10 and drop to 1.53–1.61 at 32: the tighter the budget, the more the policy batches calls to save turns. Spend also tracks task difficulty, most visibly on hard tasks.
Generalisation holds up three ways: gains persist at interpolated and extrapolated budgets outside the six training values; on AppWorld test-challenge (unseen apps, harder tasks) the 9B utility curve keeps rising past the largest training budget of 29; and policies evaluated zero-shot on ASTRA, a personal-assistant benchmark, keep their budget responsiveness without retraining.
Budget control moves from an inference-time harness into the policy itself, dropping a layer of allocation machinery, and one policy covers a continuous range of budget values. The selection rules are ready to copy: Sigmoid-Gate for maximum within-budget success at a given budget, Gate-Neg when the agent must stop on its own without an external cutoff, Dual for robustness across budgets. Plain Gate, which only rewards in-budget success, is surprisingly competitive and nearly free to implement. The Trunc comparison adds a practical lesson: keep overspending rollouts in the training data as negative samples.
The components are old, state augmentation plus reward shaping. The value is the controlled comparison under one framework and a reusable reward selection table.
Incremental, but solid.
The authors state that no objective dominates across criteria, and that not every reward beats GRPO-Std in every setting, so the wrong choice can waste a training run. Each environment and resource type gets its own policy; a single policy handling several resources, or trained across environments, is untested, and zero-shot evaluation does not establish that those setups would work. Models stop at 4B and 9B. All experiments run in research-only simulated environments, with an explicit disclaimer that nothing reflects an Apple product system. Per-budget improvement margins over baselines mostly live in figures rather than quoted numbers, so anyone following up has to read the curves.