Stratified Advantage Normalization (SAN)

• Stratified Advantage Normalization (SAN) is a technique that partitions trajectories into homogeneous groups, enabling unbiased, locally normalized advantage estimation.
• By computing statistics within each stratum, SAN ensures that policy gradient updates avoid cross-stratum bias and yield stable, scale-consistent learning signals.
• Empirical evaluations show that SAN enhances training reward and convergence stability, underscoring its practical impact in heterogeneous RL environments.

Stratified Advantage Normalization (SAN) is a technique developed to address the statistical and optimization challenges arising from structural heterogeneity in reinforcement learning (RL), particularly in settings such as LLM search agents where agent-generated trajectories vary dramatically in their structure, reward distributions, and operational complexity. SAN ensures that credit assignment and normalization for policy gradient updates are performed within homogeneous strata of trajectories, thereby eliminating systematic bias that arises when using global baselines over fundamentally incomparable samples. The method has been mathematically analyzed and empirically validated as central to the Stratified GRPO algorithm, establishing stratification as a principled solution for RL in structurally heterogeneous environments (Zhu et al., 7 Oct 2025).

1. Motivation and Definition

Stratified Advantage Normalization is designed to solve the problem of cross-stratum bias—deterministic offsets resulting from direct comparison of heterogeneous trajectories in policy optimization. Standard advantage normalization methods, which compute baseline and scaling statistics globally across all trajectories, inadvertently perform "apples-to-oranges" comparisons when the population of trajectories is structurally diverse (e.g., differing in search count, branching factor, or action outcomes).

SAN partitions a batch of trajectories $\mathcal{B}$ into disjoint strata $\{\mathcal{B}_k\}$, each defined by a shared structural property (e.g., the same number of search engine calls). Within each stratum, advantages are normalized according to the local empirical mean and standard deviation. This yields:

$$A_{\text{SAN}}(\tau) = \frac{R(\tau) - \hat{\mu}_k(x)}{\hat{\sigma}_k(x) + \epsilon}$$

where $A_{\text{SAN}}(\tau)$ is the normalized advantage for trajectory $\tau$ in stratum $k$, $R(\tau)$ is its reward, $\hat{\mu}_k(x)$ and $\hat{\sigma}_k(x)$ are the mean and standard deviation for group $k$ (possibly further conditioned on external variables such as prompt $x$), and $\epsilon$ is a small constant for numerical stability.

2. Mechanism and Statistical Properties

The SAN mechanism involves three key steps: (a) stratum assignment of each trajectory based on a discrete structural property, (b) computation of local stratum statistics, and (c) normalization of each sample's advantage within its stratum.

For trajectory $\tau \in \mathcal{B}_k$ (where $k$ is the stratum index), the SAN-normalized advantage is directly centered and scaled within its group.

Stratum Index ($k$)	Empirical Mean ($\hat{\mu}_k(x)$)	Empirical Std. Dev. ($\hat{\sigma}_k(x)$)
1	mean of $R(\tau)$ in $\mathcal{B}_1$	std of $R(\tau)$ in $\mathcal{B}_1$
2	mean of $R(\tau)$ in $\mathcal{B}_2$	std of $R(\tau)$ in $\mathcal{B}_2$
...	...	...

Mathematically, SAN guarantees:

$$\mathbb{E}[A_{\text{SAN}} | k] = 0 \qquad \text{and} \qquad \mathrm{Var}[A_{\text{SAN}} | k] = 1$$

for any stratum $k$, ensuring unbiasedness and unit variance locally. These conditional properties are not assured by global normalization, which yields:

$$\mathbb{E}[A_{\text{GN}}|k] = \frac{\mu_k - \mu}{\sigma}, \qquad \mathrm{Var}[A_{\text{GN}}|k] = \frac{\sigma_k^2}{\sigma^2}$$

where $(\mu, \sigma)$ are global mean and standard deviation, and $(\mu_k, \sigma_k)$ are local statistics for stratum $k$.

3. Elimination of Cross-Stratum Bias

A key contribution of SAN is the provable elimination of deterministic stratum offsets in advantage calculation. Considering the global advantage estimator $A_G(\tau) = R(\tau) - R_{\text{global}}$, it can be decomposed as:

$$A_G(\tau) = A_{\text{SAN}}(\tau) + (\hat{\mu}_k(x) - R_{\text{global}})$$

where $(\hat{\mu}_k(x) - R_{\text{global}})$ is the cross-stratum bias introduced by global baseline comparison. By using $\hat{\mu}_k(x)$ for centering, SAN ensures that trajectory rewards are only compared within homogeneous groups, removing systematic bias and yielding what the paper terms a "pure and scale-stable learning signal" (Zhu et al., 7 Oct 2025).

4. Comparison with Standard Policy Gradient Normalization

Standard global policy gradient normalization (e.g., REINFORCE with baseline, normalized advantage actor-critic) treats all trajectories identically, computing normalization statistics across the entire batch. This leads to credit assignment distortions when the statistics pool over heterogeneously structured trajectories. In settings with large variations—typically observed in LLM search, tool invocation environments, or non-trivial exploration spaces—such methods increase variance and introduce training instability.

In contrast, SAN's per-stratum normalization ensures that each policy update is conditionally unbiased and uniformly scaled within each homogeneous group. The authors further show that global unbiasedness and unit variance are preserved when aggregating across all strata, matching the guarantees of standard normalization but avoiding its drawbacks under structural heterogeneity.

SAN has also been extended to include linear blending with the global estimator when strata are sparsely populated, further stabilizing updates under finite-sample constraints (Zhu et al., 7 Oct 2025).

5. Empirical Evaluation and Effect on RL Dynamics

Comprehensive experiments have been conducted on seven question answering benchmarks encompassing both single-hop and multi-hop search-enhanced QA tasks. The results demonstrate that Stratified GRPO, powered by SAN, consistently and substantially outperforms standard GRPO in training reward, stability, and policy effectiveness.

• Average training reward improves by up to 11.3 points over the baseline.
• Multi-hop benchmarks exhibit up to 14.5 point gains in relative reward.
• Training curves reveal smoother convergence, higher reward, and robust search policy learning (e.g., agents perform multi-hop search when previous approaches stagnate at one-hop).

Observed effects on learning dynamics include improved exploration, more effective credit assignment, and greater adaptation to complex task structures.

6. Practical Implications and Deployment Considerations

The adoption of SAN provides several practical benefits:

• More stable policy optimization, especially in scenarios with significant trajectory structure variance.
• Enhanced ability for search-augmented LLM agents to learn complex, multi-step reasoning strategies where performance would previously be stunted by biased, noisy learning signals.
• Reduction in the need for aggressive tuning or regularization to counteract instability, as SAN provides intrinsically better credit assignment.

For practitioners, the method is straightforward to implement, requiring only partitioning of trajectories into meaningful strata and standard computation of empirical statistics. Optimization routines remain otherwise unchanged, and SAN is compatible with existing actor-critic policy-gradient frameworks.

To mitigate finite-sample variance when some strata are small, the paper advocates blending SAN with the global estimator, serving as a robustness measure.

7. Extensions and Future Research Directions

The paper identifies several promising avenues for further exploration:

• Integrating SAN-style stratified normalization within actor-critic architectures such as PPO, particularly where value function approximation may suffer from analogous bias.
• Developing dynamic, data-driven stratum assignment mechanisms to automatically partition trajectories during training based on evolving structural properties.
• Refining blending strategies for robust estimation in low-data regimes.
• Generalizing SAN principles to other RL domains involving tool use, retrieval, planning, or reasoning with LLM agents, potentially unifying the stratified approach across various forms of agent-environment heterogeneity.

A plausible implication is that stratification may further benefit multi-agent RL and hierarchical decision processes, where trajectory variance impedes the effectiveness of global statistics-based normalization.

---

In summary, Stratified Advantage Normalization (SAN) constitutes a rigorous solution to the problem of cross-stratum bias in RL for LLM-based agents and similar structurally heterogeneous environments. By partitioning trajectories and confining normalization to homogeneous groups, SAN ensures unbiased, stable, and effective policy optimization. This method is empirically validated as critical for advanced search-augmented agents and is extensible to a broad range of RL scenarios where trajectory diversity is intrinsic to the problem setting (Zhu et al., 7 Oct 2025).