Understanding Reasoning Collapse in LLM Agent Reinforcement Learning

SAI replication review · Referee report

Summary

RAGEN-2 identifies a failure mode in multi-turn LLM-agent RL that the authors call template collapse: reasoning that remains diverse within each input yet becomes input-agnostic across inputs. The conceptual move is to split reasoning entropy via $H(Z)=H(Z\mid X)+I(X;Z)$ and argue that the community has been monitoring the wrong half. To measure the neglected input-dependence term, the authors introduce a lightweight in-batch cross-scoring proxy (retrieval-accuracy and MI-ZScore-EMA variants) that reuses training rollouts. The mechanism is then explained through a signal-to-noise decomposition of policy gradients: within-input reward variance controls the task-gradient signal (via a Cauchy–Schwarz bound), while KL and entropy terms deliver an input-agnostic contractive push whose relative weight grows as RV shrinks. The mitigation — SNR-Aware Filtering — keeps a top fraction of prompts by within-input reward variance before each update, and is reported to improve peak validation success over PPO/DAPO/GRPO/Dr.GRPO baselines across Sokoban, FrozenLake, MetaMathQA, and Countdown, with further evidence on WebShop, SearchQA, DeepCoder, and multimodal setups. The framing is genuinely useful: separating within-input diversity from cross-input dependence is a clean lens on a real problem, and the RV-based filter is easy to add to any policy-gradient loop. The main conceptual weakness is that the paper repeatedly gestures at a formal chain (bound-on-gradient $\Rightarrow$ dominance of regularizers $\Rightarrow$ $I(X;Z)\!\to\!0$ ) that the mathematics does not deliver; Appendix M in fact shows KL-closeness preserves MI, complicating the story of KL as an MI-eroding force. The empirical results are reported without seed variance and with a peak-of-training paired metric that inflates gains on already-saturated tasks, while the primary MI diagnostic is scored by the policy being trained. A concurrent execution attempt could not confirm any training-time claim: the shipped stack failed at worker init on a NumPy/Numba pin, so beyond a definitional constant nothing was verified from execution. Together with several notational and internal-consistency issues, this leaves the mechanistic story understated and the empirical claim overstated relative to the reported evidence.

Strengths

Conceptual contribution. The Shannon decomposition of reasoning entropy into within-input diversity $H(Z\mid X)$ and input dependence $I(X;Z)$ , and the observation that standard entropy monitoring proxies only the former, is a genuinely clarifying reframing of the "reasoning collapse" literature. The four-regime picture in Figure 1 is a useful pedagogical anchor.
Cheap, on-line diagnostic. The in-batch cross-scoring proxy is elegant: it reuses trajectories already in the rollout buffer, requires no external scorer, and produces both a discrete (retrieval-acc) and a continuous (MI-ZScore-EMA) reading. It is one of the more practical process-monitoring diagnostics proposed for agent RL.
Actionable, minimal-overhead intervention. SNR-Aware Filtering is a single-knob addition to any policy-gradient trainer, reuses rewards already computed by the environment, and (per the released rollout_filter.py) is implemented as a drop-in filter rather than a bespoke algorithm — a low-friction contribution that can be tested in existing pipelines.
Breadth of empirical setups. Table 3 spans four RL algorithms, four model scales, two model families, and text/vision modalities. Even discounting the multimodal deltas, the directional pattern of gains is broad and diverse enough to argue against a narrow overfitting to Sokoban.
Honest boundary condition. The FrozenLake stochasticity sweep (Figure 11) explicitly documents where the method stops working (80–100% transition noise), which is unusual and valuable: it makes the mechanism's causal prediction (RV degrades as a signal proxy under high reward noise) empirically falsifiable, and it reproduces the predicted attenuation.
Bucketed gradient evidence. The RV-quantile sweep of task-gradient vs. regularizer-gradient norms (Figure 3) is a nice piece of direct evidence for the SNR mechanism, distinct from the outcome-level ablations that dominate the rest of the paper.

Weaknesses

Bound-to-dynamics leap. The paper repeatedly writes " $I(X;Z)\to 0$ " as the consequence of a static Cauchy–Schwarz bound on $\|g_{\text{task}}\|$ . The bound only controls one update's task-gradient magnitude; the limiting behavior of $I(X;Z)$ requires an additional dynamical argument (accumulation over steps + that $g_{\text{reg}}$ specifically decreases MI) that is asserted, not derived. Appendix M in fact shows that KL-closeness to the base preserves MI, complicating the story of KL as an MI-eroding force.
Self-referential MI proxy. The scorer $q$ is $\pi_\theta$ itself. This makes the proxy a moving target: as policy sharpness changes, matched–marginal gaps shift for reasons unrelated to $I(X;Z)$ , and the Shannon identity relied on in §2.2 does not hold for scorer-defined proxies (Appendix G warns of this, but the main text does not). A frozen-scorer sanity check would substantially strengthen the diagnostic.
Teacher-forced-likelihood floor. Under teacher forcing, tokens in $Z_{i,k}$ were generated with the true $X_i$ as prefix, so token-level agreement in system/format spans and near-verbatim carryover of state descriptors will make $\log p_\theta(Z_{i,k}\mid X_i)>\log p_\theta(Z_{i,k}\mid X_j)$ even under strong template collapse. The proxy therefore has a positive floor unrelated to genuine input dependence, and small nonzero MI values do not straightforwardly rule out collapse.
Definition-of- $X$ / retrieval-indexing mismatch. $X$ is defined as the full turn- $t$ context including the model's own prior reasoning $z_{1:t-1}$ , but the retrieval proxy treats $X$ as a discrete per-prompt label with chance floor $1/P$ . In stochastic environments and for trajectory-uniform sampling, two rollouts from the "same prompt" at turn $t{>}1$ have different $X$ . Which $X$ carries the label — the initial prompt, the turn- $t$ context, or a rollout id — is not stated.
Top- $p$ definition ambiguity. The prose calls $\rho$ a keep rate and describes it as a top fraction of prompts, but the formal rule is a variance-mass threshold. These are meaningfully different for skewed RV distributions and produce very different kept counts. This mismatch is load-bearing: at $\rho{=}0.9$ , a fraction-of-prompts reading predicts a $\sim 10\%$ step-time reduction while the reported reduction is 26–41% — only consistent with the variance-mass reading.
Missing seed variance. No table reports standard errors or seed counts. Several "consistent improvement" claims rest on deltas within likely seed noise: Dr. GRPO average $+0.8$ , Countdown $+0.0/+0.3$ , MetaMathQA $+0.6$ , and near-zero cells on saturated tasks. Absent multi-seed evidence, the claim of consistent gains is not statistically supportable from the tables shown.
Peak-of-training as headline metric. Table 3's paired "peak validation success rate" is a biased estimator of learning-algorithm quality; it favors noisier curves and interacts with early-stopping. This is especially concerning on tasks near ceiling (Countdown 97.9 $\to$ 97.9, MetaMathQA 92.6 $\to$ 93.2) where peak-of-training noise likely dominates the observed deltas.
Overstated "twice as reliably" framing. The reported Spearman values ( $+0.39$ vs. $-0.11$ / $-0.14$ ) support the correct-sign / larger-magnitude claim; "significantly more strongly" and "twice as reliably" import statistical language not backed by a test, and "twice" does not match the roughly 3× magnitude ratio.
Non-monotone quartile results reported as monotonic. Table 8 shows Task Perf $21.1\to 19.5\to 10.7\to\mathbf{11.0}$ (Q4 exceeds Q3) and Entropy $2.02\to 1.53\to 1.41\to 1.87$ (non-monotone); only MI Proxy is monotone. The Q3 range $[0.0,0.2]$ and Q4 range $[0.0,0.1]$ overlap at zero, so the two lowest quartiles are essentially the same near-zero-variance bin. The "establishes the full causal chain" framing therefore overreaches on the mechanism this ablation actually pins down.
Failure to acknowledge counter-cases the tables show. The trajectory-level filter in Table 9 crashes MI Proxy from 0.83 to 0.20 while still improving task perf; the Reward Sum filter in Table 5 improves MI by $+0.26$ yet still collapses. Both cases are direct challenges to the paper's implicit narrative that MI preservation and collapse-avoidance are tightly coupled, and both are omitted from the discussion.
"Signal-enhancement rather than noise-control" contradicts Appendix K. The MSE result in Theorem K.1 is precisely a noise-control argument; Table 1 also lists filtering as the mitigation for $g_{\text{task-noise}}$ . The exclusive framing in §5.1 misstates the paper's own theory.
Multi-turn framing weakened by single-turn headline tasks. Two of the four columns of Table 3 (Countdown, MetaMathQA) are 1-turn, 1-action tasks. The thesis is that template collapse is uniquely acute in multi-turn RL, but the headline evidence sits largely on single-turn tasks. The truly multi-turn tasks (Sokoban, FrozenLake, WebShop, SearchQA) should carry the argument.
KL/entropy attribution is bundled. Appendix M shows KL-closeness preserves MI, so KL toward an instruction-tuned base should not naturally erode $I(X;Z)$ ; the MI-eroding force must be entropy regularization (Appendix N shows how $\Delta_{\text{in}}>\Delta_{\text{marg}}$ decreases MI). The main-text bundling of KL and entropy as jointly "reward-agnostic" regularizers is misleading, and quantitative attribution would sharpen the mechanism.
Length decline as a "signature" of template collapse is under-argued. Figure 6's monotone shortening has a benign alternative reading — the model becoming more efficient after early exploration. Without a demonstrated covariance between length and MI (or between length and eventual task-performance drop), the length trace does not discriminate templated boilerplate from concise correct reasoning.
Std/Mean(RV) diagnostic overstated. Three data points, no threshold, no held-out setting — this is an exploratory observation, not a "cheap diagnostic which can be done before training."
Reward-variance gaming/noise risk understated. Definition F.4 acknowledges high RV can arise from evaluation noise, and Figure 11 shows the filter's advantage evaporates at high stochasticity. So RV picks up both signal and noise, and the filter preferentially selects the latter exactly when informative signal is scarcest. Appendix T handles this in one sentence; for a metric central to the training loop this deserves main-text treatment.
Notational and structural glitches. Chance level illustrated at $P{=}64$ while experiments use $P{=}8$ ; theorem citation to "G.1" when the actual theorem is I.2; Remark J.2 misplaced (its content is Appendix M's); duplicate "Core Author Contributions" appendix; a reference to Table 6 when Table 4 is what appears (Table 6 is a later duplicate); "top- $P$ filter with $P\to 1.0$ " collides with the prompt-batch-size symbol; the surrogate PPO/GRPO objective is written without its ratio/log-prob structure; the missing display equations for Retrieval-Acc and MI-ZScore-EMA in §2.3 make the primary diagnostics under-specified in the main text; Figure 6's "eight environments" does not match the seven enumerated in Table 2 and Appendix C.1; Theorem I.4 is truncated and its "linear in the number of steps" summary is ambiguous ( $\|\theta_T-\theta_0\|^2\propto T$ , RMS $\propto\sqrt{T}$ ).

Reproducibility & code

Headline experiments could not be executed. An attempt to run the shipped RAGEN stack (per §4.1's "veRL/HybridFlow … RAGEN defaults") never completed a single training step: worker init died with a NumPy 2.5 vs. Numba (requires $\le 2.2$ ) ImportError raised through vLLM's ngram speculative-decoding path. The upstream requirement is undocumented in the README, so a reproducer following the paper's setup instructions cannot resolve it without external knowledge. Additional friction: a Ray AF_UNIX temp-path length limit that required a +ray_kwargs.ray_init._temp_dir override, and a heavy from-source sglang[all] / flashinfer build path only tractable via prebuilt wheels. These are packaging/documentation gaps rather than errors in the algorithm, but they block reproduction of every Table 3 / Table 5 / Table 8 / Table 9 / Table 10 value and every Figure 5/7/8/9/10 trajectory.
Only a definitional constant was verified from code. The one verifiable item was that ragen/trainer/collapse_metrics.py sets chance_level = 1.0 / N_prompts, recovering $1/64=1.5625\%$ — a definitional match, not a training-time reproduction. All 19 substantive claims that were attempted (template-collapse dynamics, filter deltas across algorithms/scales/modalities, Spearman correlations, quartile ablation, trajectory-vs-prompt comparison, stochasticity sweep, adaptive kept-ratio dynamics, MI-vs-format diagnostic) are recorded as not attempted.
Diagnostic and filter mechanism are implemented. The MI proxy (retrieval accuracy, MI estimate, MI-ZScore-EMA, conditional entropy) and the SNR-aware filter are directly implemented in ragen/trainer/collapse_metrics.py and ragen/trainer/rollout_filter.py, with base scripts (run_main_table_diff_algo.sh, run_top_p_sweep.sh, run_kl_sweep.sh, run_entropy_sweep.sh) that exercise the intended configurations. The code path for the headline claim (template collapse dynamics under a no-filter baseline) is instrumented and would be runnable given a working environment.
No shipped logs, curves, or aggregates. No pre-computed training logs, wandb dumps, or CSVs are packaged; no plotting or aggregation script produces the reported Spearman values (Figure 9, Figure 8), the six-panel MI-proxy view (Figure 10), or the composed figures (Figure 5, 6, 7). Every empirical value in the paper — including the $+0.39/-0.11/-0.14$ correlations and every Table-3 delta — requires multi-day multi-GPU re-runs plus a reproducer-written reduction. Combined with the lack of reported seed variance, exact numeric agreement is uncertain.
Trajectory-level baseline of Table 9 has no shipped implementation. rollout_filter.py implements only whole-group / prompt-level selection. The "top-8 and bottom-8 trajectories per prompt" baseline described in §S would require a new per-trajectory filter path that is not present. Given that this row is the paper's central isolation of prompt-vs-trajectory SNR effects (and shows a striking MI collapse), its irreproducibility is a moderate gap.
Per-quartile training runs (Table 8) require a bespoke driver. split_into_buckets supports RV bucketing for gradient analysis within a single run, but the shipped scripts do not train four separate policies each restricted to one quartile. A researcher would need to add a bucketed-training driver, and the resulting $P{=}8, G{=}16$ with 25% kept configuration (only two prompt groups per step) is non-default.
Multimodal rows depend on external code. The Qwen2.5-VL-3B (T/V) rows in Table 3 — which carry the largest headline deltas (+53.5 / +59.5 on FrozenLake, +29.8 / +35.8 on average) — sit outside the shipped RAGEN configs and require the VAGEN codebase. Similarly, five of the eight environments cited in Figure 6 (spatial, logic-puzzle, visual, math agents) map to external repositories rather than first-party envs.
Sokoban-14B configuration is not obvious in the repo. Table 10 uses a Sokoban 14B setting for the Std/Mean(RV) diagnostic; the released model configs enumerate 0.5B/1.5B/3B/7B (README) with no packaged 14B setup.
Machine-specific paths. run_filtering_final.sh writes to hard-coded /mnt/permanent/... locations that must be overridden per user before running.

Recommended Changes

Essential

Unblock and document the environment. Pin numpy<2.2 (or the specific incompatible Numba constraint), document the Ray temp-dir workaround, and either package prebuilt vLLM/FlashInfer wheels or ship a Docker image, so the training loop can be started following only the paper's stated setup. (See Headline experiments could not be executed.)
Resolve the top- $p$ definition. Pick one definition of $\rho$ (fraction-of-prompts vs. fraction-of-variance-mass), use it consistently in text, formulas, figure labels, and compute-overhead numbers, and update the "26–41% step-time reduction at $\rho=0.9$ " either by re-reporting under the fixed-fraction reading or by re-deriving under the variance-mass reading. (See Top- $p$ definition ambiguity.)
Report per-seed variance for all headline tables. Add mean $\pm$ std over at least 3 seeds (ideally 5) for Tables 3, 5, 8, 9, 10, and mark cells whose delta is not statistically distinguishable from zero. (See Missing seed variance.)
Replace or supplement "peak validation success" with a less selection-biased metric. Report end-of-training performance, best-window-average, or learning-curve AUC alongside peak. (See Peak-of-training as headline metric.)
Downgrade the "full causal chain" and " $I(X;Z)\!\to\!0$ " language. Rephrase the Cauchy–Schwarz consequence as a magnitude bound rather than a limit, and either restrict the quartile-ablation conclusion to a correlational statement or add a controlled intervention that varies RV while holding prompt identity fixed. (See Bound-to-dynamics leap and Non-monotone quartile results.)
Rewrite the Table 8 monotonicity claim to match the numbers, and merge the near-zero-RV quartiles. Report Task Perf and Entropy as non-monotone; note that Q3 and Q4 are degenerate near-zero splits with overlapping ranges. (See Non-monotone quartile results.)
Add explicit treatment of Reward-Sum and Trajectory-level cases where MI and collapse dissociate. The Reward Sum row (MI $+0.26$ yet collapsed) and the Trajectory-level row (MI $0.83\to 0.20$ while task perf rises) directly challenge the paper's implicit MI-as-collapse-predictor story; discuss them rather than folding them into a summary sentence. (See Failure to acknowledge counter-cases and Trajectory-level baseline of Table 9.)
Fix the $P{=}64$ vs. $P{=}8$ chance-level example. Either use $1/P=12.5\%$ at $P{=}8$ , or state the retrieval-batch size $P_\text{eval}$ used and how prompts are aggregated. Clarify which chance floor applies to every reported MI Proxy value. (See Notational and structural glitches.)
Reframe "signal-enhancement rather than noise-control." State that SNR-Aware Filtering does both, consistent with Theorem K.1 and Table 1. (See "Signal-enhancement rather than noise-control" contradicts Appendix K.)
Add display equations for Retrieval-Acc and MI-ZScore-EMA in §2.3. The primary diagnostics that carry the paper's thesis must be defined in the main text. (See Notational and structural glitches.)
Release the trajectory-level filter, per-quartile-training driver, and aggregation/plot scripts. These three additions would move the paper from "code inspected, runs required" to genuinely reproducible for its central ablations. Also fix hard-coded paths in run_filtering_final.sh. (See Trajectory-level baseline of Table 9, Per-quartile training runs, No shipped logs.)

Suggested

Add a frozen-scorer sanity check for the MI proxy to isolate genuine input dependence from policy-sharpness drift and teacher-forced local-coherence effects. (See Self-referential MI proxy and Teacher-forced-likelihood floor.)
Clarify what $X$ is when computing the proxy (initial prompt vs. turn- $t$ context vs. rollout id) and reconcile with the definition in §2.1. (See Definition-of- $X$ / retrieval-indexing mismatch.)
Narrow the mechanism claim to entropy regularization (with quantitative KL vs. entropy attribution) since Appendix M shows KL preserves MI-closeness and the default KL coefficient is small. (See KL/entropy attribution is bundled.)
Reframe "twice as reliably as entropy" as a sign-and-magnitude directional claim with proper statistical language, and remove or defend the word "significantly." (See Overstated "twice as reliably" framing.)
Elevate the reward-variance gaming/noise risk into the main text with concrete examples of noise-inflated RV in high-stochasticity settings. (See Reward-variance gaming/noise risk understated.)
Rebalance the empirical presentation toward genuinely multi-turn tasks (Sokoban, FrozenLake, WebShop, SearchQA) rather than 1-turn Countdown/MetaMathQA in the headline table, or explicitly recharacterize template collapse as a general RL phenomenon. (See Multi-turn framing weakened by single-turn headline tasks.)
Plot length against MI or against later success within each environment, or moderate the claim that length decline is a behavioral signature of template collapse rather than of efficiency gain. (See Length decline as a "signature".)
Either strengthen the Std/Mean(RV) diagnostic with more settings and a stated threshold, or reframe it as exploratory. (See Std/Mean(RV) diagnostic overstated.)
Fix bookkeeping. Environment count (seven vs. eight), duplicate "Core Author Contributions" appendix (P and U), Table 6 reference vs. Table 4, misplaced Remark J.2, "Theorem G.1" citation, "top- $P$ filter with $P\to 1$ " notation, the truncated Theorem I.4, the missing surrogate-ratio structure in the PPO/GRPO objective, and the Table 9 arithmetic (16 vs. 8 trajectories per prompt in the trajectory-level baseline). Individually small but collectively erode confidence. (See Notational and structural glitches.)
Package a Sokoban-14B config or note that the Table 10 14B row was produced with a bespoke setup, so readers can reproduce the Std/Mean(RV) row. (See Sokoban-14B configuration.)