RAGEN-2: Mitigating Reasoning Collapse in RL

• RAGEN-2 is a framework for diagnosing and mitigating template collapse in multi-turn agentic RL using an information-theoretic approach.
• It combines within-input diversity and cross-input mutual information metrics to detect input-agnostic reasoning that standard entropy measures miss.
• The method employs SNR-Aware Filtering to select high-signal prompts, yielding significant improvements across diverse RL benchmarks.

RAGEN-2 is a framework for diagnosing and mitigating reasoning collapse—specifically, template collapse—in reinforcement learning (RL) training of multi-turn LLM agents. Unlike previous diagnostics that rely solely on entropy to measure the diversity of agent reasoning, RAGEN-2 introduces an information-theoretic approach that combines within-input diversity and cross-input distinguishability. The framework defines and addresses a failure mode where agents preserve local diversity while producing input-agnostic reasoning chains, leading to degraded task performance that standard metrics cannot detect (Wang et al., 7 Apr 2026).

1. Problem Setting and Definitions

Multi-turn agentic RL involves agents interacting over multiple turns, observing context $x_t$ (system prompt, observation history, prior reasoning tokens $z_{1:t-1}$, and actions $a_{1:t-1}$), generating a reasoning chain $z_t$ and action $a_t$, and receiving scalar reward $r_t$. Trajectories $\tau = (x, z, a, r)$ are optimized under a regularized policy-gradient objective, such as PPO or GRPO:

$$L(\theta) = \mathbb{E}_{x, \tau}[A(\tau, x)] - \lambda_{KL} D_{KL}(\pi_\theta \parallel \pi_{ref}) + \lambda_H H(\pi_\theta)$$

Reasoning collapse refers to any reduction in input sensitivity of the agent’s chain-of-thought or plan. RAGEN-2 introduces the more specific phenomenon of template collapse, defined as the agent maintaining high within-input entropy $H(Z|X)$ while having low mutual information $I(X; Z)$—generating diverse yet input-agnostic reasoning across different prompts.

2. Information-Theoretic Decomposition and Metrics

RAGEN-2 applies Shannon’s decomposition: total output diversity $z_{1:t-1}$0 splits into within-input entropy $z_{1:t-1}$1 and cross-input mutual information $z_{1:t-1}$2:

$z_{1:t-1}$3

While $z_{1:t-1}$4 (reasoning entropy) is commonly used as the diversity metric, it is blind to losses in input-specific adaptation. $z_{1:t-1}$5 quantifies how much reasoning varies in response to different inputs but is intractable to compute for high-dimensional token sequences.

To address this, RAGEN-2 deploys minibatch-based mutual information proxies:

Proxy	Brief Description
Retrieval-Acc	Fraction of test samples where paired reasoning z is best matched to its own prompt x via log-likelihood.
MI-Est	Average length-normalized log-likelihood difference between matched and marginal assignments.
MI–ZScore–EMA	Rolling z-normalization of MI-Est; increases robustness to scaling effects.

Empirically, Retrieval-Acc and MI–ZScore–EMA exhibit strong positive correlation with final task performance (Spearman ≈ +0.39), while entropy-based metrics show negligible or negative correlation (–0.11 to –0.14). Thus, MI proxies are more reliable indicators of actual reasoning quality (Wang et al., 7 Apr 2026).

3. Mechanistic Explanation: Signal-to-Noise Imbalance

Template collapse is linked to gradient-level imbalances between signal (task-driven reasoning) and regularization. The total gradient per prompt decomposes as:

$z_{1:t-1}$6

where $z_{1:t-1}$7 drives task learning and $z_{1:t-1}$8 comprises regularization gradients. By Cauchy–Schwarz, $z_{1:t-1}$9, so as within-prompt reward variance $a_{1:t-1}$0, $a_{1:t-1}$1 vanishes and regularization dominates.

The local signal-to-noise ratio is defined as:

$a_{1:t-1}$2

Low-SNR prompts are pushed toward input-agnostic behavior, suppressing $a_{1:t-1}$3 even if $a_{1:t-1}$4 is high.

4. SNR-Aware Filtering Algorithm

To mitigate template collapse, RAGEN-2 proposes SNR-Aware Filtering, which selects high-signal prompts per iteration for policy updates using reward variance as a proxy for SNR. The workflow is:

1. Collect $a_{1:t-1}$5 rollouts per prompt for a batch of $a_{1:t-1}$6 prompts.
2. Compute reward variance $a_{1:t-1}$7 for each prompt.
3. Rank prompts by descending $a_{1:t-1}$8.
4. Retain the smallest prefix $a_{1:t-1}$9 so that cumulative $z_t$0 mass $z_t$1 (with default $z_t$2).
5. Update using only trajectories from prompts in $z_t$3.

No additional model or rollouts are required (total rollout budget $z_t$4 remains fixed), and gradient computation time is reduced by 26–41% (Wang et al., 7 Apr 2026). Alternative selection rules (top-k, min-p, inverted) are explored; top-p filtering adapts to reward variance shifts and yields the best trade-off.

5. Empirical Evaluation Across Diverse Tasks

RAGEN-2 is empirically validated on seven RL benchmarks encompassing multi-turn and single-turn tasks, including gym-Sokoban (planning), FrozenLake (navigation), MetaMathQA and Countdown (symbolic math), SearchQA (web search), WebShop (e-commerce), and DeepCoder (code synthesis). Key findings include:

• In unfiltered baselines, Retrieval-Acc (I-proxy) declines early, yet $z_t$5 remains high, and reasoning chain length monotonically shrinks.
• SNR-Aware Filtering consistently improves both Retrieval-Acc (input-dependence) and final task success rates.

Summarized PPO/Qwen2.5-3B results (Table 1 of (Wang et al., 7 Apr 2026)):

Task	Baseline Success (%)	Filtered Success (%)	Δ (points)
Sokoban	12.9	27.3	+16.0
FrozenLake	67.0	77.9	+10.9
MetaMathQA	92.6	93.2	+0.6
Countdown	97.9	97.9	0
Average	—	—	+6.9

Similar benefits are reported with DAPO (+2.9 avg), GRPO (+3.7), Dr.GRPO (+0.8), different model sizes, instruction-tuned variants, and multimodal setups (Qwen2.5-VL text: +29.8, vision: +35.8).

Further analyses confirm:

• Only SNR-Aware Filtering reliably promotes high $z_t$6 and high success.
• MI proxies are reliable online diagnostics, with stronger positive correlations to performance than entropy.
• Causal connections between high prompt RV, higher MI, and increased success, as validated by controlled ablations.

6. Limitations and Practical Guidance

SNR-Aware Filtering assumes that reward variance within a prompt signals task-diagnostic information. In extremely sparse or highly stochastic environments (>80% transition stochasticity), the utility of RV as a proxy diminishes; practitioners should monitor $z_t$7 for reliability. Estimating RV requires grouped sampling ($z_t$8), but no additional rollouts are needed, and, in fact, computational costs may decrease.

The filtering hyperparameter ($z_t$9) requires per-task tuning; excessive filtering can hinder exploration. There is also a theoretical risk of agents inflating RV, so joint evolution of MI and RV should be monitored. All results pertain to single-agent settings; multi-agent coordination and collapse mechanisms remain unexplored.

Practically, RAGEN-2 recommends supplementing multi-turn agentic RL pipelines with an MI proxy monitor (such as MI–ZScore–EMA) and deploying prompt-level SNR-Aware Filtering. This combination enables online detection and mitigation of reasoning collapse with minimal system change and improved robustness over baseline procedures (Wang et al., 7 Apr 2026).

7. Significance for Agentic RL and Future Directions

RAGEN-2 exposes limitations of standard entropy-based diagnostics in agentic RL with LLMs, demonstrating that within-input stochasticity alone does not guarantee input-conditional reasoning. By foregrounding mutual information metrics and introducing SNR-Aware Filtering, it both advances the theoretical understanding of RL collapse modes and provides practical, lightweight remedies compatible with existing pipelines. Open avenues include adapting the approach to multi-agent scenarios and environments where reward variance is either unreliable or potentially manipulated. These developments underscore the growing importance of information-theoretic controls for robust agentic language modeling.