Improving Search Agent with One Line of Code

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 \textbf{S}earch \textbf{A}gent \textbf{P}olicy \textbf{O}ptimization (\textbf{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 \textbf{+10.6\% absolute improvement} (+31.5\% relative) over Search-R1, yielding consistent gains across varying model scales (1.5B, 14B) and families (Qwen, LLaMA).

• The paper introduces SAPO, a conditional token-level KL constraint added in one line to prevent gradient collapse in GRPO-based TARL.
• SAPO maintains stable importance sampling ratios and enhances gradient flow, yielding a +10.6 point improvement in QA benchmarks.
• The approach is model-agnostic, simple to implement, and scales robustly across different architectures for tool-based RL.

Search Agent Policy Optimization for Stabilizing Tool-Based Agentic Reinforcement Learning

Motivation and Problem Formulation

Tool-based Agentic Reinforcement Learning (TARL) frameworks empower LLMs to autonomously interact with external tools for information-seeking tasks. Group Relative Policy Optimization (GRPO), a predominant algorithm for TARL, exhibits severe training instability, especially as agent trajectories lengthen. The instability arises most notably as Importance Sampling Distribution Drift (ISDD), observed as vanishing importance sampling ratios which ultimately suppress gradient flow and precipitate catastrophic model collapse.

The paper identifies two core issues in GRPO: numerical inaccuracies from group-based advantage assignment and misleading policy updates where intermediate steps may diverge from final answer correctness. Specifically, ISDD originates when the policy shifts outside the trust region: the importance sampling ratio decays exponentially across multi-step agent actions, rendering gradient-based learning ineffective even for high-reward samples.

Figure 1: Dynamics of training between SAPO and GRPO highlighting IS ratio collapse, clipping spikes, entropy changes, and reward degradation.

SAPO: Conditional Token-Level KL Constrained Policy Optimization

To address ISDD, the paper proposes Search Agent Policy Optimization (SAPO), which augments GRPO with a token-level, conditional Kullback-Leibler (KL) divergence penalty. Rather than uniform hard clipping (as in PPO_clip), SAPO selectively penalizes excessive policy shifts only for positive tokens with low predicted probability under the current policy and positive advantage values.

Formally, SAPO functions as a soft trust region mechanism by applying a penalty term:

• Only to tokens where the importance sampling ratio $r_t < \tau$ and advantage $A_t > 0$
• The penalty scales logarithmically with $\log r_t$, permitting gradual exploration rather than abrupt suppression
• This approach preserves informative gradient flow for positive actions, enabling adaptive corrective steps and preventing irreversible collapse

The core algorithmic modification is minimal: SAPO can be implemented with a single line added to the standard GRPO objective, enhancing practical deployability with negligible overhead.

Empirical Evaluation

Extensive experimentation across seven QA benchmarks—spanning both single-hop (NQ, TriviaQA, PopQA) and multi-hop (HotpotQA, 2Wiki, Musique, Bamboogle) datasets—demonstrates SAPO's robust performance. Using the Qwen2.5 and LLaMA model families and identical agent and retrieval architectures, SAPO produces consistent gains:

• SAPO achieves an average exact match (EM) accuracy of 0.442, outperforming competitive baselines such as Search-R1, AutoRefine, and CriticSearch.
• SAPO yields a +10.6 point absolute (+31.5% relative) improvement over Search-R1, with pronounced gains on multi-hop tasks (+14.7 points, +24.9% relative).
• SAPO's improvements are consistent across parameter scales (1.5B—14B) and model families (Qwen2.5, LLaMA), confirming scalability and model-agnostic generalization.

  Figure 2: (a) Hyperparameter sensitivity; (b,c) scaling curves for EM and F1 scores as Qwen2.5-Instruct model size increases.

Ablation studies isolate the contribution of the conditional KL penalty: introducing unconditional KL, ratio-conditioned KL, and full advantage- and ratio-conditioned SAPO demonstrate that the selective constraint maximally balances stability and exploration, leading to superior aggregate accuracy.

Training Dynamics and Stability Analysis

SAPO's effect is further validated via quantitative analysis of training dynamics:

• SAPO maintains stable importance sampling ratios throughout training, preventing collapse seen in GRPO.
• Policy entropy remains well-regulated, indicating that informative exploration persists without excessive drift.
• Reward trajectories remain monotonically increasing with SAPO, unlike late-stage degradation in GRPO.

These findings underscore SAPO's effectiveness in mitigating catastrophic collapse and preserving exploration during iterative agentic learning.

Robustness, Generalization, and Implementation Simplicity

SAPO's robustness spans:

• Generalization to alternate architectures, re-aligning both base and instruction-aligned pre-trained models
• Favorable scaling behavior: performance increases monotonically with model size, confirming alignment with scaling laws for neural LLMs
• Minimal hyperparameter sensitivity: IS threshold $t=1.0$ is optimal and stable across tasks

  Figure 3: SAPO EM accuracy evolution across seven benchmarks and four backbone models, illustrating stable convergence.

SAPO is distinguished by its implementation simplicity: a single line of code suffices to add the KL penalty (see Appendix), making it readily integrable into any GRPO-based RL pipeline.

Case Studies and Qualitative Behavior

SAPO-trained agents demonstrate coherent and efficient reasoning over multi-turn search trajectories, autonomously refining queries and synthesizing retrieved evidence to produce correct, concise answers across diverse information-seeking tasks. The qualitative results underscore the agent’s capacity for compositional reasoning and autonomous query refinement.

Implications and Future Directions

Theoretical implications:

• SAPO enforces a token-level trust region, preventing accumulation of distributional drift across lengthy agentic trajectories
• The conditional KL penalty enables soft, context-sensitive policy correction, blending stability and adaptive exploration

Practical implications:

• Immediate applicability for stabilizing RL in complex agent-environment interactions (e.g., tool-augmented LLMs, web search agents)
• Suitable for wide deployment across diverse backbone models and task settings

Future work may involve:

• Extending SAPO to fine-grained credit assignment and reward propagation in more sophisticated multi-agent or hierarchical settings
• Explorations of variant penalty functions and adaptive thresholds for further stability enhancements
• Application to broader ranges of agentic RL tasks requiring dynamic tool invocation, memory management, and real-time knowledge integration

Conclusion

Search Agent Policy Optimization (SAPO) introduces a token-level, conditional KL constraint to stabilize GRPO-based post-training for search agents in tool-based agentic reinforcement learning. SAPO delivers strong empirical improvements (+10.6 points absolute) over previous methods, robustly prevents ISDD-induced collapse, and maintains efficient gradient flow throughout training. Its model-agnostic applicability, implementation simplicity, and theoretical grounding in soft trust region conditioning suggest SAPO as a foundational recipe for stable, scalable agentic RL in information-seeking environments (2603.10069).