PaperlayerView exact revision on arXiv

Source-linked AI summary

One Policy, Any Budget: Internalizing Budget-Aware Search via Reinforcement Learning

Xiaowei Sun, Jin Li, Yili Hong, Yikun Fu, Yanghua Xiao

arXiv:2609.00813v1cs.AI

TL;DR

Existing RL-based search agents are predominantly trained under fixed or unconstrained budgets and cannot adapt when resource availability changes at deployment. AnySearch internalizes budget-aware search through a progressively removed scaffold, two-phase curriculum reinforcement learning, and adaptive reward weighting. Across seven benchmarks and three backbone models, it consistently outperforms baselines across budget levels, generalizes to unseen budgets, and achieves high tool productivity with low token consumption.

  • Problem

    Existing RL-based search agents are predominantly trained under fixed or unconstrained budgets and cannot adapt when resource availability changes at deployment.

  • Method

    AnySearch progressively removes a training scaffold through two-phase curriculum reinforcement learning and uses adaptive reward weighting to jointly optimize answer accuracy and search efficiency.

  • Results

    Across seven benchmarks and three backbone models, AnySearch consistently outperforms baselines at all tested budget levels, generalizes to unseen constraints, achieves highest tool productivity, and produces lowest total token consumption.

  • Takeaways & Limitations

    A single policy can internalize budget-aware search and autonomously handle allocation decisions across varying budgets without retraining.

  • Takeaways & Limitations

    The framework models budget as a discrete count of external search calls and evaluates with a static 2018 Wikipedia dump, leaving continuous multi-objective costs and temporal adaptation for future work.

Abstract

from arXiv · show

While reinforcement learning has enabled LLM-based search agents to invoke external tools, existing methods train under fixed budgets and cannot adapt when constraints vary at deployment. We propose AnySearch, a framework that enables a single policy to perform budget-aware search under any budget constraint through a training scaffold and curriculum reinforcement learning. In the first phase, we train the agent with explicit budget state injection and structured reasoning prompts that guide efficient allocation under linearly decaying budgets. In the second phase, the scaffold is removed and the agent learns to operate autonomously under adaptively sampled budget constraints, matching inference conditions. Both phases are optimized with a composite reward that couples answer accuracy with budget efficiency through absolute and relative signals, where an adaptive weight amplifies the efficiency signal for high-accuracy queries and attenuates it for low-accuracy ones. Extensive experiments on seven general and multi-hop QA benchmarks show that our method outperforms baselines across all budget scales, generalizes to unseen constraints beyond the training range, and achieves superior tool productivity without excessive token overhead. Our code is available at https://github.com/xwsun01/AnySearch.

1 Introduction

Existing search agents trained under fixed budgets cannot adapt when resource constraints change, motivating AnySearch’s single policy for budget-aware allocation. AnySearch progressively removes training guidance while jointly optimizing answer accuracy and search efficiency, and experiments show consistent gains across budgets.

  • Existing RL search agents trained under fixed or unconstrained budgets cannot adapt when resource availability changes at deployment.
  • AnySearch internalizes budget-aware search so inference requires only a total budget specification, with the policy autonomously handling allocation decisions.
  • The two-phase curriculum begins with explicit budget tracking and structured prompts, then removes the scaffold while adaptively sampling weak budget levels.Phase I uses linearly decaying budgets; Phase II removes the scaffold and maintains coverage through uniform smoothing.
  • Extensive experiments demonstrate consistent improvements, robust generalization to unseen budgets, and superior tool productivity.Ablations also validate scaffold internalization and the necessity of each component.
  • A composite reward couples answer accuracy with budget efficiency through absolute and relative signals, using adaptive weighting based on query difficulty.

2 Related Work

Related work spans budget-aware reasoning and reinforcement-learning search agents, but prior methods do not train search policies under dynamic tool-call budgets. This gap motivates AnySearch’s focus on internalized adaptation across constraints.

  • Budget-Aware Reasoning: Budget-aware reasoning methods train under varying token constraints but still require explicit budget signals at inference.
  • Agentic RL with Search Engines: Recent search-agent methods use reinforcement learning to interleave reasoning with external search and optimize multi-turn interaction trajectories.
  • Agentic RL with Search Engines: No prior work trains search agents under dynamic tool-call budgets.

3 Method

AnySearch formulates search as a resource-constrained sequential decision problem and internalizes budget-aware allocation through scaffolded, two-phase reinforcement learning. Its curriculum removes explicit guidance, samples weak budget levels adaptively, and uses rewards that jointly promote answer accuracy and efficient search.

  • Task Formulation: The agent must answer accurately while using at most B external search calls, with search consuming budget and reasoning still possible after exhaustion.The policy observes accumulated context and remaining budget, choosing among reasoning, search, and answer actions.
  • Training Scaffold: The training scaffold injects remaining, used, and total budget states and prompts separate information sufficiency from budget-conditioned search strategy.These signals make budget availability observable and guide decisions about whether another search is necessary and worthwhile.
  • Two-Phase Curriculum: Phase I linearly decays budgets from Bmax to 1 with the scaffold active, while Phase II removes it and samples budgets adaptively from weaker-performing levels.Phase I initializes per-level accuracy statistics; Phase II receives only the total budget at episode start and maintains coverage through minimum sampling exposure.
  • Two-Phase Curriculum: The adaptive sampling distribution emphasizes budget levels furthest behind the best performance while retaining minimum probability for every level.A sliding window tracks recent accuracy or reward statistics at each budget level, with smoothing and minimum-probability controls.
  • Reward Design: The composite reward combines accuracy, format validity, length, and tool-efficiency signals, with absolute and relative rewards encouraging fewer searches and efficient correct trajectories.The absolute signal rewards correct answers for saved budget, while the relative signal compares against the most efficient correct trajectory within the sampling group.
  • Reward Design: Adaptive efficiency weighting strengthens tool-efficiency pressure for high-accuracy queries and attenuates it for low-accuracy queries so accuracy remains primary.The efficiency weight controls the proportion of tool reward in the total reward based on group accuracy.

4 Experiments

Across seven QA benchmarks and three backbone models, AnySearch consistently improves accuracy across budget scales while generalizing to unseen constraints and reducing search and token costs. Ablations support the role of scaffold removal, adaptive sampling, two-phase training, and reward design.

  • Main results: AnySearch achieves an average EM of 0.431 with Qwen2.5-7B-Instruct, surpassing StepSearch at 0.403 across the evaluated benchmarks.BATS and Search-o1 achieve 0.227 and 0.266, respectively.
  • Budget generalization: AnySearch maintains dominant performance across varying budgets and generalizes to unseen budgets beyond the maximum training budget of 5.On Bamboogle, accuracy rises from about 0.40 at budget 5 to about 0.42 at budget 8, while Search-R1 plateaus around 0.37.
  • Efficiency: At budget 3 on Bamboogle, AnySearch scores 0.38 versus Search-R1 at 0.34 and reaches about 0.35 accuracy with one fewer search.These results place AnySearch beyond the baselines on the cost-performance frontier.
  • Efficiency: AnySearch achieves the highest Tool Productivity and lowest total token consumption across evaluated backbones and datasets.The paper attributes lower token use to fewer redundant retrievals, while the two-phase curriculum has comparable single-run training cost to fixed-budget baselines.
  • Ablations: The adaptive tool reward balances accuracy and efficiency by reducing efficiency pressure on low-accuracy queries and emphasizing it on high-accuracy queries.Removing either absolute or relative efficiency signals degrades the combined outcome, while removing the tool reward causes the sharpest decline.
  • Ablations: Ablations show that scaffold removal, adaptive budget sampling with λ=0.6, and the full two-phase curriculum produce the most stable performance across budgets.Training with the scaffold throughout creates a train-inference gap, while uniform or unsmoothed sampling and either single phase underperform.
  • Training behavior: Removing the format reward causes reward deterioration, shorter responses, and reduced policy entropy, whereas the full reward maintains stable learning trajectories.Qualitative examples also show adaptive search depth and selective stopping under different budgets.

5 Conclusion

AnySearch internalizes budget-aware search in one policy through progressive scaffold removal, two-phase curriculum reinforcement learning, adaptive sampling, and composite rewards. Across seven QA benchmarks and three backbones, it improves performance across budgets, generalizes to unseen budgets, and improves tool and token efficiency.

  • Conclusion: AnySearch internalizes budget-aware search through a progressively removed training scaffold and two-phase curriculum reinforcement learning with adaptive budget sampling.A composite reward jointly optimizes answer accuracy and search efficiency.
  • Conclusion: Across seven general and multi-hop QA benchmarks and three backbone models, AnySearch improves over baselines at all budget levels while generalizing to unseen budgets.It also achieves the highest tool productivity and lowest token consumption.

Limitations

AnySearch’s limitations concern its discrete cost model, dependence on the backbone and retrieval corpus, and reliance on a static 2018 Wikipedia environment.

  • The framework models budget as a discrete count of external search calls, whereas deployment costs also include latency, monetary expense, and system load.The authors identify continuous, multi-objective cost modeling as a natural extension.
  • AnySearch cannot recover answers absent from both the backbone model’s parametric knowledge and the retrieval corpus.Its efficiency improves search decisions but does not overcome missing answer information.
  • Training and evaluation use a static 2018 Wikipedia dump, so temporal adaptation is not tested.Extension to a continuously updated open-web setting is left for future work.

Ethics Statement

The study reports software, model, and dataset license checks, uses existing public resources, and involves neither human annotation nor primary data collection.

  • The authors verified licenses for all software, models, and datasets used in the study.
  • The work uses existing public QA datasets and a publicly available Wikipedia corpus without human annotation or primary data collection.
  • The authors report no privacy concerns or personally identifiable information in the training or evaluation data.

A.2 Full Efficiency Results

Across multi-hop QA benchmarks, backbone models, and budget levels, AnySearch consistently achieves the highest accuracy and tool productivity.

  • AnySearch achieves the highest accuracy across all four multi-hop QA benchmarks, three backbone models, and budget levels B = 4, 5, 6.The comparison extends the main-text efficiency analysis across all reported configurations.
  • AnySearch achieves the highest tool productivity across all four multi-hop QA benchmarks, three backbone models, and budget levels B = 4, 5, 6.The authors describe this as evidence that its efficiency advantage generalizes beyond one dataset or model.
  • The efficiency advantage generalizes beyond a single dataset or model across the reported multi-hop QA configurations.

A.3 Token Consumption Analysis

Token analysis measures total output and retrieved tokens at B = 5 across seven benchmarks. AnySearch records the lowest token consumption on every dataset, consistent with reduced redundant search and retrieval overhead.

  • Figure 7 measures average total token count across seven benchmarks on Llama-3.1-8B-Instruct with B = 5.Total tokens include both output tokens and retrieved tokens.
  • AnySearch achieves the lowest token consumption on all seven datasets.The reported measure includes output and retrieved tokens.
  • Reduced redundant search calls lower the associated retrieval overhead.

B.1 Training Dynamics

Training remains stable across Qwen2.5-7B-Instruct and Llama-3.1-8B-Instruct, while the evaluation suite spans general and multi-hop question-answering settings.

  • Training dynamics: Training reward increases steadily, response length remains reasonable, and policy entropy decreases gradually without sudden collapse across both model families.These dynamics are consistent with the Qwen3-4B results reported in the main text.
  • Benchmarks: The evaluation suite includes Natural Questions, TriviaQA, and PopQA for open-domain, compositional, and long-tail factual question answering.Natural Questions uses Wikipedia-annotated answers, TriviaQA contains independently collected evidence, and PopQA derives questions from Wikidata triples.
  • Benchmarks: HotpotQA, 2WikiMultiHopQA, MuSiQue, and Bamboogle test multi-hop reasoning across noisy, compositional, and multi-step evidence settings.The benchmarks cover two-hop through four-hop reasoning, evidence aggregation, and questions designed to resist keyword lookup.

B.3 Baselines Implementation

The implementation compares prompt-based, retrieval-based, and reinforcement-learning baselines under search budgets, while AnySearch removes its training scaffold before inference. Case studies illustrate selective allocation and lower token overhead.

  • Baselines: The baselines include BATS, Search-o1, Search-R1, ZeroSearch, and StepSearch, representing prompt-based, RAG-based, and RL-based approaches.All methods receive a search budget limiting the maximum number of search actions.
  • Cost analysis: Retrieval tokens account for 51–68% of total token consumption across all datasets, identifying returned search documents as the primary inference-cost contributor.This motivates defining the budget over external search calls and reducing unnecessary searches.
  • Prompt implementation: The Phase I scaffold structures budget observation, reasoning, and action, while the Phase II and inference prompt specifies only the total budget and basic interaction format.The scaffold is removed after training, and inference retains the reasoning, search, observation, and answering interaction cycle.
  • Prompt implementation: At each reasoning step, the explicit budget state reports remaining, used, and total searches, requiring the agent to assess information sufficiency and search value.The search phase costs 1 unit per search, and the agent answers once further searching is unnecessary.
  • Case studies: With budget 3, AnySearch identifies the person and relevant religious information in two searches, then stops with one search remaining instead of verifying redundantly.The case attributes this selective stopping and query prioritization to internalized budget-aware allocation.

D.1 Consolidating Evidence for Internalization

Three complementary observations support the claim that AnySearch internalizes budget state and conditions search decisions on it without external tracking at inference.

  • Scaffold redundancy at inference: Reintroducing the full scaffold at inference produces virtually identical accuracy to standard inference without it.The result indicates that the policy has absorbed the budget-tracking capability previously supplied by the scaffold.
  • Active budget conservation: AnySearch uses substantially fewer searches per correct answer than baselines under the same budget ceiling.The authors interpret this tool-productivity result as active budget conservation rather than indiscriminate exhaustive search.
  • Early stopping without external signals: AnySearch stops early when confident, using only 2 of 3 searches in Case 4 and 4 of 6 in Case 2.These behaviors require implicit awareness of remaining resources without external state injection.
  • Conclusion: Together, the observations suggest that the trained policy maintains an internal representation of budget state and conditions its search decisions accordingly.The conclusion concerns autonomous inference-time allocation rather than reliance on external budget tracking.
Loading 2609.00813v1…