Train agents to respect token and tool-call budgets — Overspend penalties — What does it mean?
The news. On October 20, 2026, researchers at Apple and the University of Oxford posted Rewards for Budget-Aware Multi-Turn Agents That Adapt to Arbitrary Test-Time Constraints (arXiv 2610.28409). They train Qwen3.5 models at 4B and 9B parameters on AppWorld, with budgets on tokens, API calls and environment steps, and on BabyAI with turn budgets, and compare eleven reward functions. Read the paper →
Picture climbing into a taxi and saying, "I can pay at most 40 for this ride." The driver now makes a choice at every junction: the short road, or the scenic one. An agent faces the same choice. Every turn costs something, whether it is generated tokens, an API call or one more environment step, and the budget is the fare cap the user names before the run starts.
The usual fix is a meter outside the driver: a harness that counts spend and cuts the ride off at the cap. The paper asks whether the driver can learn to read the meter instead. The setup is plain. The total budget and the remaining budget are written into the text of every observation, so at each turn the policy sees how much is left and decides how much to spend. Because the budget is just input text, a budget never seen in training needs no change to the policy or to inference. What differs between the eleven methods is only the reward: the rollouts, prompts, budget grids and seed are shared, and each method turns the binary task reward into a budget-shaped return.
The rewards differ in what happens to the driver's pay when the meter passes the cap. Read the last column as a score for how far each one falls short of the best method at each budget, averaged over the paper's eight AppWorld and BabyAI experiments:
| Reward | Taxi version | What it does past the budget | Mean regret in utility (lower is better) |
|---|---|---|---|
| Trunc | ride stops at the cap | rollout is cut at the budget; reward is task success inside it | 7.56 pp |
| Gate | no pay for any ride over the cap | zero reward if the budget is exceeded, even when the task was solved | 3.33 pp |
| Gate-Neg | a fixed fine for any ride over the cap | a fixed negative reward whenever the budget is exceeded; generally the strongest natural compliance | 6.00 pp |
| Sigmoid-Gate | a fine that grows smoothly the farther past the cap | sigmoid penalty subtracted only after the budget; most wins at an individual budget (20.1%) | 3.82 pp |
| Rate | a fine that grows with how often drivers go over | learned price per training budget on the overspend rate | 7.46 pp |
| Dual | both fines at once | learned step price on overspend frequency plus a linear penalty on overspend size | 2.92 pp |
| Linear-Mag | pay shrinks linearly past the cap | no step at the budget; reward falls linearly beyond it | 14.01 pp |
Two choices stand out. Gate is the bluntest reward and it still beats the truncation baseline: Trunc ends the rollout at the budget, so the tokens past the cap never reach training, while Gate lets the overspending rollout finish and scores it zero, which, whenever another rollout in the same group scored a success, gives it a negative advantage on every action it took, including those beyond the cap. And the fixed fine of Gate-Neg buys compliance at a price: it generally gives the strongest natural compliance, yet its regret is nearly double Gate's, and the authors suggest that the strong penalty may make the policy hesitant near the boundary and cost it some in-budget success.
Dual is closest to the per-budget best on average (2.92 pp), while Sigmoid-Gate wins at the most individual budgets. The paper reads Rate's weaker showing as evidence that a price on the overspend event works better when it is combined with information about the overspend size. Rewards that press on cost before the cap (Linear-Cost, Sigmoid) did not beat Gate, so early pressure to finish does not necessarily improve in-budget success.
Walk one training group through three rewards (illustrative numbers). The budget is 10 steps and three rollouts attempt the same task: A solves it in 8 steps, B solves it in 14, and C fails after 9. Take each rollout's advantage as its reward minus the group mean.
- Task reward only: rewards are 1, 1, 0, the mean is 2/3, so A and B both get +1/3. The overspending run is reinforced exactly as much as the compliant one.
- Gate: rewards are 1, 0, 0, the mean is 1/3, so A gets +2/3 while B and C get −1/3 each. B now pays for its extra 4 steps, even though it solved the task.
- Gate-Neg with a −1 fine: rewards are 1, −1, 0, the mean is 0, so A gets +1, B gets −1 and C gets 0. Failing inside the budget now beats solving outside it, which is why this reward favors compliance and can make the agent cautious near the cap.
Trunc would have stopped B at step 10 with no success, so the policy would never see what steps 11 to 14 did.
Trained agents adapt their spending to the stated budget and to task difficulty. Before training, the 9B base model complied on fewer than half of the episodes across the token-budget range, and the task-reward-only baseline used about one API call per turn whatever the budget. The budget-shaped policies (Gate, Sigmoid-Gate and Dual) batched more calls per turn, and batched more as the step budget tightened. Step and token spend rose with the stated allowance and with difficulty, while API-call spend followed difficulty more than the allowance. The gains held at budgets between and beyond the six training budgets, on a shifted set of AppWorld tasks, and zero-shot on a new benchmark, ASTRA.
The limits are plain too. The paper trains a separate policy for each environment and resource type, the models are small (4B and 9B), and no method reached perfect compliance at every budget. Raw task success also generally rose with spend, so a budget is a real trade against accuracy, not a free saving. This is the training-side partner of the reasoning budget idea: instead of a rule outside the model that decides when to think longer, the model learns the rule.
Goes deeper in: Agent Engineering → Cost & Latency Engineering → The Cost Profile of an Agent
Related explainers
- TRIAGE — role-typed credit assignment — another reward-design fix for agent RL, changing who gets credit instead of who gets fined.
- Tev-Gym — progress-potential reward shaping — shaping a long-horizon agent's reward with progress instead of the final outcome only.
- Token Budgets — affine-typed budget ownership — the opposite route: enforce a shared cap in the type system instead of training the agent to respect it.