Stratified GRPO: Handling Structural Heterogeneity in Reinforcement Learning of LLM Search Agents

Large language model (LLM) agents increasingly rely on external tools such as search engines to solve complex, multi-step problems, yet their rollouts are structurally heterogeneous: variations in tool-call number, placement, and outcomes induce distinct behaviors and reward distributions. As a result, policy gradient methods with a single global baseline suffer from cross-stratum bias, an "apples-to-oranges'' comparison that distorts credit assignment and impedes exploration. To address this issue, we propose Stratified GRPO. Its core component, Stratified Advantage Normalization (SAN), partitions trajectories into homogeneous strata based on structural properties and computes advantages locally within each stratum, ensuring comparisons only among true peers. We show that SAN eliminates cross-stratum bias, yields conditionally unbiased unit-variance estimates within strata, and preserves the global unbiasedness and unit-variance properties of standard normalization, resulting in a more reliable learning signal. To improve robustness in finite-sample regimes, we further linearly blend SAN with the global estimator. Experiments on factual QA and deep-research agent benchmarks demonstrate that Stratified GRPO consistently outperforms GRPO by up to 12.6 points, achieving higher training rewards, improved training stability, and more effective search policies. These results establish structure-aware advantage normalization as an effective correction for RL of search agents with structurally heterogeneous rollouts.

Large language model agents often use external tools, such as search engines, to answer complex questions. However, different attempts at the same task can involve very different search behaviors, such as using different numbers of searches or retrieving information of different sources. Using the same feedback rule for all these different attempts can therefore mislead the learning process. This paper proposes Stratified GRPO, a training method that makes these comparisons fairer by grouping similar agent behaviors and comparing attempts within each group. This helps the model better identify which actions are truly useful. Experiments on factual question answering and deep research tasks show that our method improves performance, training stability, and search behavior, helping language model agents learn to use tools more effectively.