Research map / Multi-agent systems and games / LLM agents

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

ICML 2026 · Multi-agent systems and games › LLM agents · approach: Agent data synthesis and behavior studies · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Stratified GRPO, which normalizes advantages within strata of structurally similar search-agent trajectories (e.g., by search-call count) and blends them with global GRPO advantages, improving multi-hop QA performance over GRPO.

Problem

Policy gradient methods for LLM search agents use a single global baseline across structurally heterogeneous trajectories (varying search-call number, placement, outcomes), which yields cross-stratum bias and hinders exploration of multi-step search strategies.

Approach

Stratified Advantage Normalization (SAN) partitions a prompt's rollouts into homogeneous strata by structural properties and computes mean/std-normalized advantages within each stratum; a linear blend with the global GRPO advantage stabilizes small strata.

Key results

• Stratified GRPO reaches 41.6 average EM on Qwen2.5-3B-Base vs 33.3 for GRPO (+8.3), with multi-hop gains up to 14.5 points over strongest baseline.
• On Qwen2.5-3B-Instruct, Stratified GRPO averages 40.5 EM vs 29.2 for GRPO, where GRPO training collapses.
• Ablation: SAN alone reaches 36.7 average on Base and 39.9 on Instruct, with blending adding further gains.

Surprising findings

• Standard GRPO on the instruct model suffers training collapse, while stratified advantages keep rewards stable and monotonically increasing.

Limitations

• Theoretical guarantees are stated in the large-sample/idealized limit (epsilon=0); finite-sample behavior needs blending.
• Stratification relies on a predefined partition function (e.g., search count), so choice of structure is manual.
• Experiments limited to Qwen2.5-3B models and seven QA benchmarks.

Experimental setup

Datasets: Natural Questions HotpotQA 2018 Wikipedia dump

Benchmarks: NQ TriviaQA PopQA HotpotQA 2WikiMultiHopQA MuSiQue Bamboogle

Metrics: exact match

Baselines: Direct Generation SFT RAG Search-o1 IRCoT Search-R1 R1 ReSearch GRPO

Builds on: GRPO Search-R1 RLOO PPO

Closest papers by meaning

• TIC-GRPO: Provable and Efficient Optimization for Reinforcement Learning from Human Feedback (NeurIPS 2026 · similarity 0.89)
• Solving the Granularity Mismatch: Hierarchical Preference Learning for Long-Horizon LLM Agents (ICLR 2026 · similarity 0.89)
• Group-in-Group Policy Optimization for LLM Agent Training (NeurIPS 2025 · similarity 0.88)
• Stronger-MAS: Multi-Agent Reinforcement Learning for Collaborative LLMs (ICLR 2026 · similarity 0.88)
• StraTA: Incentivizing Agentic Reinforcement Learning with Strategic Trajectory Abstraction (NeurIPS 2026 · similarity 0.88)
• Group-Graph Policy Optimization for Long-Horizon Agentic Reinforcement Learning (NeurIPS 2026 · similarity 0.88)