One Policy, Any Budget: Internalizing Budget-Aware Search via Reinforcement Learning
Xiaowei Sun, Jin Li, Yili Hong, Yikun Fu, Yanghua Xiao
cs.AI
2026-09-01
AnySearch puts budget-aware search in one policy via a two-phase curriculum. Qwen2.5-7B-Instruct hits 0.431 mean EM vs 0.403 for StepSearch, including unseen budgets.
RL search agents such as Search-R1 usually learn when and what to search under a fixed budget. At deployment the budget moves and the policy does not: tight budgets over-search, loose budgets sit idle. BATS tracks budget state each turn, but still needs an external tracker at inference. Budget awareness never enters the policy.
AnySearch, from Fudan, Southeast University, and Shanghai Jiao Tong, wants a single policy that is told only a total budget B and then spends it on its own.
Search is a sequential decision with a hard cap. Each step may think, search, or answer. Each search costs one budget unit. After the budget hits zero the agent can still think, but cannot search.
Training uses GRPO in two phases.
Phase I turns on a scaffold. Every turn starts with <budget> remaining=R; used=U; total=T </budget>, and the <think> block must judge both information sufficiency and whether another search is worth the remaining budget. The budget anneals linearly from Bmax=5 to 1 with equal steps per level, so the agent first learns how to search, then when to stop.
Phase II removes the scaffold and states only the total budget B at episode start, matching inference. Budgets are sampled from a sliding-window accuracy tracker: weak levels get more mass, uniform smoothing keeps strong levels from being forgotten. Default λ=0.6.
The reward ties accuracy to search efficiency. The tool term is an absolute factor times a relative factor: the absolute term pays for a correct answer in proportion to budget saved; the relative term compares against the cheapest correct trajectory in the same group. An adaptive weight γq scales with group accuracy, amplifying efficiency on easy queries and muting it on hard ones so correctness comes first. Format and length terms sit alongside. On Qwen3-4B, dropping the format reward produces shorter replies and entropy collapse around step 175.
Three backbones, seven QA sets (NQ, TriviaQA, PopQA, HotpotQA, 2Wiki, MuSiQue, Bamboogle), inference budget B=5.
| Backbone | AnySearch mean EM | Strongest RL baseline |
| Qwen2.5-7B-Instruct | 0.431 | StepSearch 0.403 |
| Llama-3.1-8B-Instruct | 0.448 | ZeroSearch 0.424 |
| Qwen3-4B | 0.407 | StepSearch 0.383 |
Prompt-only BATS and Search-o1 lag far behind (0.227 and 0.266 on Qwen2.5-7B-Instruct). A budget cap without trained allocation is not enough.
Sweeping budgets from 1 to 8, AnySearch's accuracy rises monotonically and holds on B=6–8, which never appeared in training. On Bamboogle with Qwen2.5-7B-Instruct it sits near 0.40 at B=5 and 0.42 at B=8; Search-R1 plateaus near 0.37 at high budgets. At B=3 AnySearch is about 0.38 versus Search-R1 at 0.34; Search-R1 needs B=4 to reach about 0.35.
Tool productivity (correct answers per search call) is highest as well. On HotpotQA with Qwen2.5-7B at B=6, AnySearch TP is 0.354 versus 0.276 for StepSearch and 0.212 for Search-R1. Retrieved documents dominate token count, so fewer searches cut inference cost.
Ablations say the scaffold should be on in Phase I and off in Phase II and at inference. Keeping it through training and stripping it only at inference hurts low-budget accuracy; the agent leans on explicit budget tags. Never using a scaffold leaves early exploration without an allocation pattern. Putting the scaffold back at inference does not help a trained policy.
The efficiency reward is correctness-gated. A counterfactual on 2Wiki with Qwen3-4B at budget 5 finds that cases where AnySearch stops early, errs, and the no-efficiency control is right, cover 0.33% of the test set. The 3.32% of extra correct answers more than offset that.
Product search budgets already move: a latency-critical path may allow one or two calls, deep research many more. Fixed-budget training means one model per band. AnySearch spends about the same single-run cost as a fixed-budget baseline (500 steps on one 8×H800 node) and returns one policy that covers the trained range and extrapolates a little past it.
For people training search agents, the portable pieces are: teach budget as state, then remove the scaffold, and gate efficiency on correctness so hard queries do not stop early to save calls.
Budget is a discrete search count. Real cost mixes latency, money, and load; a continuous multi-objective cost is not studied.
Allocation quality is capped by the backbone and the corpus. If the answer is in neither parameters nor the index, searching less will not recover it. Training and evaluation use the static 2018 Wikipedia dump, matching Search-R1, with no open-web or temporal-drift test.
Bmax=5 is a narrow training range. Extrapolation to 6–8 looks stable; larger budgets and a true zero-search extreme are not developed.