Improving Search Agent with One Line of Code

Jian Li; Dongsheng Chen; Zhenhua Xu; Yizhang Jin; Jiafu Wu; Chengjie Wang; Xiaotong Yuan; Yabiao Wang

cs.LGMar 10, 2026

Improving Search Agent with One Line of Code

Authors: Jian Li, Dongsheng Chen, Zhenhua Xu, Yizhang Jin, Jiafu Wu, Chengjie Wang, Xiaotong Yuan, Yabiao Wang

Organizations: Nanjing University · Tencent YoutuLab

Abstract

Tool-based Agentic Reinforcement Learning (TARL) has emerged as a promising paradigm for training search agents to interact with external tools for a multi-turn information-seeking process autonomously. However, we identify a critical training instability that leads to catastrophic model collapse: Importance Sampling Distribution Drift(ISDD). In Group Relative Policy Optimization(GRPO), a widely adopted TARL algorithm, ISDD manifests as a precipitous decline in the importance sampling ratios, which nullifies gradient updates and triggers irreversible training failure. To address this, we propose Search Agent Policy Optimization (SAPO), which stabilizes training via a conditional token-level KL constraint. Unlike hard clipping, which ignores distributional divergence, SAPO selectively penalizes the KL divergence between the current and old policies. Crucially, this penalty is applied only to positive tokens with low probabilities where the policy has shifted excessively, thereby preventing distribution drift while preserving gradient flow. Remarkably, SAPO requires only one-line code modification to standard GRPO, ensuring immediate deployability. Extensive experiments across seven QA benchmarks demonstrate that SAPO achieves +10.6% absolute improvement (+31.5% relative) over Search-R1, yielding consistent gains across varying model scales (1.5B, 14B) and families (Qwen, LLaMA).

Figures & tables

Figure 1: Comparison of training dynamics between SAPO and GRPO regarding (a) Importance Sampling Ratio, (b) Clip Ratio, (c) Entropy, and (d) Reward.

w/o Retrieval
	Single-Hop QA				Multi-Hop QA					Avg.
Methods	NQ	TriviaQA	PopQA	Avg.	HotpotQA	2Wiki	Musique	Bamboogle	Avg.	Avg.
Direct Generation	0.106	0.288	0.108	0.167	0.149	0.244	0.020	0.024	0.109	0.134
SFT	0.249	0.292	0.104	0.215	0.186	0.248	0.044	0.112	0.148	0.176
R1-Instruct Guo et al. (2025)	0.210	0.449	0.171	0.277	0.208	0.275	0.060	0.192	0.184	0.224
R1-Base Guo et al. (2025)	0.226	0.455	0.173	0.285	0.201	0.268	0.055	0.224	0.187	0.229

Table 1: Accuracy comparison of SAPO-3B against baseline methods using Qwen2.5-3B Qwen et al. (2025) across multiple QA benchmarks. Bold indicates best results, and underline denotes second best results.

Method	Single-Hop QA			Multi-Hop QA				Avg.
Method	NQ	TriviaQA	PopQA	HotpotQA	2Wiki	Musique	Bamboogle	Avg.
Search-R1-GRPO	0.402	0.662	0.436	0.374	0.376	0.134	0.333	0.388
+ $KL$	0.425	0.625	0.443	0.428	0.371	0.175	0.342	0.398
	(+0.023)	(-0.037)	(+0.007)	(+0.054)	(-0.005)	(+0.041)	(+0.009)	(+0.010)
+ $KL_{r}$	0.439	0.660	0.456	0.457	0.369	0.169	0.369	0.417
	(+0.014)	(+0.035)	(+0.013)	(+0.029)	(-0.002)	(-0.006)	(+0.027)	(+0.019)

Table 2: Ablation results on different KL terms in SAPO . The values in parentheses denote the absolute performance change compared to the previous row.

Figure 2: (a) Hyperparameter sensitivity analysis. (b,c) Scaling trends of Qwen2.5-Instruct with SAPO across different model sizes (1.5B to 14B). We report both EM and F1 scores.

Backbone: LLaMA-3.2-3B-Base
Method	Single-Hop QA			Multi-Hop QA				Avg.
Method	NQ	TriviaQA	PopQA	HotpotQA	2Wiki	Musique	Bamboogle	Avg.
Search-R1-GRPO	0.394	0.596	0.437	0.280	0.264	0.056	0.105	0.305
Search-R1-SAPO	0.471	0.630	0.452	0.355	0.301	0.102	0.248	0.366
Backbone: LLaMA-3.2-3B-Instruct
Search-R1-GRPO	0.357	0.578	0.378	0.314	0.233	0.090	0.306	0.322

Table 3: Generalization results on LLaMA-3.2-3B (Base and Instruct). SAPO consistently outperforms the GRPO baseline across both model variants.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 3 — Figure 3: Prompt template for SAPO.

Size	Metric	Single-Hop QA			Multi-Hop QA				Avg.
Size	Metric	NQ	TriviaQA	PopQA	HotpotQA	2Wiki	Musique	Bamboogle	Avg.
1.5B	EM	0.440	0.591	0.432	0.413	0.381	0.160	0.384	0.400
1.5B	F1	0.528	0.671	0.480	0.532	0.445	0.247	0.487	0.484
3B	EM	0.469	0.629	0.457	0.456	0.448	0.203	0.432	0.442
3B	F1	0.562	0.708	0.504	0.581	0.517	0.299	0.553	0.533
7B	EM	0.473	0.658	0.480	0.471	0.439	0.231	0.480	0.462

Appendix

Table 5: Scaling trends of Qwen2.5-Instruct with SAPO across different model sizes (1.5B to 14B). We report both EM and F1 scores. The best results are marked in bold.

Figure 4: Evolution of EM accuracy over training steps for SAPO across seven benchmarks.

Question: Which film has the director who died earlier, Deuces Wild or Cavalcade Of The West?

Ground Truth: [”Cavalcade of the West”, ”Cavalcade Of The West”]

Model Output: Qwen2.5-3B-Instruct-SAPO

To determine which film has the director who died earlier between ”Deuces Wild” and ”Cavalcade Of The West”, I need to gather information about the directors of these films and their death dates. Let’s start by searching for this information.

<search> Director of Deuces Wild </search>

<documents> Doc 1 (Title: ”Deuces Wild”): Deuces Wild is a 2002 American crime drama film directed by Scott Kalvert …

Doc 3 (Title: ”Scott Kalvert”): Scott Kalvert (August 15, 1964 – March 5, 2014 ) was an American film director, known mainly for his 1995 film ”The Basketball Diaries”… and 2002’s ”Deuces Wild”. </documents>

Appendix

Table 6: Case Studies illustrating SAPO trajectories from the 2wikimultihopqa benchmark. Key evidence is extracted from the original searched documents.

Question: Who is the president of the newly declared independent country that is part of the Commission of Truth and Friendship with the country that Tony Gunawan is from?

Ground Truth: [”Francisco Guterres”]

Model Output: Qwen2.5-14B-Instruct-SAPO

<think> To answer this question, I need to identify the country that Tony Gunawan is from, find out if it has a Commission of Truth and Friendship with any newly declared independent country, and then find out who the president of that newly declared independent country is. </think>

<search> Tony Gunawan country of origin </search>

<documents> Doc 1 (Title: ”Tony Gunawan”): Tony Gunawan (, born 9 April 1975) is a Chinese descended, Indonesian born badminton player… Gunawan is a former Olympic gold medalist and world champion for Indonesia…

Doc 2 (Title: ”Tony Gunawan”): …Tony played for Indonesia from 1992–2001…

Appendix

Table 7: Case Study illustrating SAPO from the Musique benchmark. Key evidence is extracted from the original searched documents.

Explore similar work

Beyond Outcome Rewards: Step-Level Self-Distilled Policy Optimization for Deep Search Agents
Aug 13, 2026Haoze Wu, Chuqiao Kuang, Tianyi Zhuang +1Group Relative Policy Optimization Credit Assignment in RL
The Hong Kong University of Science and Technology · Huawei Technologies Ltd.
Enhancing LLM-based Search Agents via Contribution Weighted Group Relative Policy Optimization
Apr 15, 2026Junzhe Wang, Zhiheng Xi, Yajie Yang +4LLM Agents RL for LLM Agents
Fudan University · Shanghai Artificial Intelligence Laboratory · Shanghai Key Laboratory of Intelligent Information Processing
Search-E1: Self-Distillation Drives Self-Evolution in Search-Augmented Reasoning
May 21, 2026Zihan Liang, Yufei Ma, Ben Chen +4Tool-Augmented Language Model Agents On-Policy Self-Distillation

Aug 13, 2026J/K move · Enter open · S save

Beyond Outcome Rewards: Step-Level Self-Distilled Policy Optimization for Deep Search Agents

Haoze Wu, Chuqiao Kuang, Tianyi Zhuang, Xiaoguang Li

The Hong Kong University of Science and Technology · Huawei Technologies Ltd.

Group Relative Policy Optimization Credit Assignment in RL Step-Level Credit Assignment in RL On-Policy Distillation Sparse-Reward RL Policy Learning+1

Open paper AskPDF arXiv

Deep search agents operate over trajectories spanning dozens of steps, yet standard reinforcement learning provides only a single outcome reward per trajectory, which is far too sparse for effective credit assignment. On-policy self-distillation (OPSD) addresses this by using the model's own logits as dense token-level teachers, but extending it to search agents introduces a fundamental tension: the teacher, having access to privileged information such as the correct answer, produces a distribution that differs systematically from the student's exploration-based reasoning, and naive distillation causes the student to inherit this information asymmetry rather than learn better search strategies. We resolve this tension through two contributions. First, we construct Evidence Anchors, which are concise, step-level evidence snippets extracted from the web, as privileged information that captures key reasoning steps without revealing the entire answer path. Second, we propose Step-Level Self-Distilled Policy Optimization (SSPO), which converts teacher-student disagreement into step-level advantage weights within GRPO, applied exclusively to incorrect trajectories. This design decouples what to update from how much to update: the outcome reward determines the direction of policy change, while the teacher modulates its magnitude at each step. Correct trajectories are left untouched, preserving their diversity. On Qwen3-8B, SSPO consistently outperforms GRPO across BrowseComp, GAIA, and FRAMES, surpassing or matching GRPO trained with twice as many gradient steps while adding only about 5 percent overhead per step from a single additional forward pass.

Abstract

Explore similar work

Beyond Outcome Rewards: Step-Level Self-Distilled Policy Optimization for Deep Search Agents

Enhancing LLM-based Search Agents via Contribution Weighted Group Relative Policy Optimization

Search-E1: Self-Distillation Drives Self-Evolution in Search-Augmented Reasoning