Expected Return Causes Outcome-Level Mode Collapse in Reinforcement Learning and How to Fix It with Inverse Probability Scaling

Abhijeet Sinha; Sundari Elango; Dianbo Liu

cs.LGJan 29, 2026

Expected Return Causes Outcome-Level Mode Collapse in Reinforcement Learning and How to Fix It with Inverse Probability Scaling

Authors: Abhijeet Sinha, Sundari Elango, Dianbo Liu

Organizations: National University of Singapore, Singapore · IIT Madras, India

Abstract

Many reinforcement learning (RL) problems admit multiple terminal solutions of comparable quality, where the goal is not to identify a single optimum but to represent a diverse set of high-quality outcomes. Nevertheless, policies trained by standard expected return maximization routinely collapse onto a small subset of outcomes, a phenomenon commonly attributed to insufficient exploration or weak regularization. We show that this explanation is incomplete: outcome level mode collapse is a structural consequence of the expected-return objective itself. Under idealized learning dynamics, the log-probability ratio between any two outcomes evolves linearly in their reward difference, implying exponential ratio divergence and inevitable collapse independent of the exploration strategy, entropy regularization, or optimization algorithm. We identify the source of this pathology as the probability multiplier inside the expectation and propose a minimal correction: inverse probability scaling, which removes outcome-frequency amplification from the learning signal, fundamentally changes the learning dynamics, and provably yields reward-proportional terminal distributions, preventing collapse in multimodal settings. We instantiate this principle in Group Relative Policy Optimization (GRPO) as a drop-in modification, IPS-GRPO, requiring no auxiliary models or architectural changes. Across different reasoning and molecular generation tasks, IPS-GRPO consistently reduces outcome-level mode collapse while matching or exceeding baseline performance, suggesting that correcting the objective rather than adding exploration heuristics is key to reliable multimodal policy optimization.

Figures & tables

Figure 1: Learned terminal-state distributions on the 2D hyper-grid. (a) Target reward-induced distribution $p(x)\propto R(x)$ . (b) GRPO and (c) FlowRL collapse onto a single high-reward mode despite multiple symmetric optima. (d) IPS-GRPO successfully recovers all modes and closely matches the target distribution. Reported values denote $\ell_{1}$ distance between the learned and target distributions.

Method	$\boldsymbol{n=2}$	$\boldsymbol{n=4}$
GRPO	$2.8\times 10^{-2}$	$5.2\times 10^{-4}$
FlowRL	$2.8\times 10^{-2}$	$5.2\times 10^{-4}$
IPS-GRPO	$\mathbf{3.0\times 10^{-3}}$	$\mathbf{3.7\times 10^{-4}}$

Table 1:

\ell_{1}

distance between the learned terminal-state distribution and the target distribution on the hyper-grid domain. Lower is better.

Panel of Figure 2 — Figure 2: Average number of distinct admissible modes recovered as a function of the number of samples drawn from the trained policy. Across all three HypoSpace domains, IPS-GRPO steadily discovers new modes with increased sampling, while GRPO and FlowRL saturate early, indicating outcome-level mode collapse.

Method	Causal Inference	3D Reconstruction	Boolean/DNA
GRPO	$16.02\%\pm 0.35\%$	$31.61\%\pm 0.01\%$	$9.50\%\pm 0.00\%$
FlowRL	$12.82\%\pm 0.05\%$	$26.85\%\pm 0.37\%$	$10.74\%\pm 0.00\%$
IPS-GRPO	$\mathbf{43.91\%\pm 1.39\%}$	$\mathbf{90.00\%\pm 1.24\%}$	$\mathbf{60.74\%\pm 0.78\%}$

Table 2: Final recovery rate (%) on HypoSpace tasks. Average over 5 random seeds. Higher is better.

SYNTH
Threshold	Algorithm	Yield $\uparrow$	OB100 $\downarrow$
0.80	REINVENT	$5695\pm 173$	$605\pm 41$
0.80	IPS-GRPO	$\mathbf{7371\pm 212}$	$\mathbf{554\pm 34}$
0.85	REINVENT	$418\pm 78$	$2934\pm 68$
0.85	IPS-GRPO	$\mathbf{1580\pm 58}$	$\mathbf{1178\pm 43}$
ALL-AMIDE

Table 3: Molecular optimization results. Yield denotes the number of unique molecules exceeding the reward threshold; OB100 denotes oracle calls required to find 100 such molecules. The oracle budget for SYNTH is 10,000 and for ALL-AMIDE is 5,000. Average over 3 random seeds

Group size $G$	$\boldsymbol{\epsilon=0.01}$	$\boldsymbol{\epsilon=0.1}$	$\boldsymbol{\epsilon=0.2}$
4	$17.63\pm 0.33$	$17.63\pm 0.33$	$17.63\pm 0.33$
8	$27.50\pm 0.67$	$27.50\pm 0.67$	$41.99\pm 0.80$
16	$41.02\pm 0.60$	$33.05\pm 0.62$	$\mathbf{43.91\pm 1.39}$
32	$32.13\pm 0.42$	$35.27\pm 1.01$	$25.62\pm 1.65$
64	$36.52\pm 0.32$	$39.12\pm 1.33$	$42.19\pm 0.14$

Table 4: Recovery rate (%) on HypoSpace causal inference for different group sizes

G

and probability clipping thresholds

\epsilon

. Average over 5 random seeds. Higher recovery rate is better.

Panel of Figure 3 — Figure 3: Study of IPS-GRPO behaviour against GRPO

Claim	Evidence in This Work
Expected-return optimization can induce outcome-level mode collapse under idealized learning dynamics.	Section 3 ; log-ratio dynamics under gradient flow ( Theorem 3.1 , Equation 2 ).
Probability-weighted updates contribute to collapse, including in settings with equal rewards.	Section 3.1 ; controlled equal-reward experiment ( Figure 3 ).
Inverse Probability Scaling mitigates outcome-frequency amplification at the objective level.	IPS objective and dynamics in Section 4 ( Theorems 4.1 and 4.2 ).
IPS-GRPO integrates IPS into standard policy-gradient pipelines with minimal modification.	Algorithm 1 ; hyper-grid recovery results ( Figures 1 and 1 ).
IPS-GRPO yields improved outcome coverage in multimodal reasoning and discovery tasks.	HypoSpace benchmarks ( Figures 2 and 2 ); molecular generation ( Table 3 ).

Table 5: Summary of the paper’s main claims and where they are supported.

Reinforcement Learning

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,717 papers

MIPS: Mitigating Mode Collapse in Reinforcement Learning for Large Outcome Spaces
Oct 14, 2026Armaan Sandhu, Abhijeet Sinha, Dianbo LiuMolecule Design Reinforcement Learning
University of Massachusetts Amherst · National University of Singapore
Beyond Mode Collapse: Distribution Matching for Diverse Reasoning
May 19, 2026Xiaozhe Li, Yang Li, Xinyu Fang +10KL-Regularized RL RL for Language Model Reasoning
1Tongji University, China · 6Independent · 2Shanghai AI Laboratory, China +3
Measuring and Mitigating Solution Mode Collapse in RLVR
Oct 8, 2026Liv G. d'Aliberti, Marwa Abdulhai, Sofiia Druchyna +2Mode Collapse RL Post-Training
Princeton University

Oct 14, 2026J/K move · Enter open · S save

MIPS: Mitigating Mode Collapse in Reinforcement Learning for Large Outcome Spaces

Armaan Sandhu, Abhijeet Sinha, Dianbo Liu

University of Massachusetts Amherst · National University of Singapore

Molecule Design Reinforcement Learning Policy Gradient Methods

Open paper AskPDF arXiv

Expected-reward maximization in reinforcement learning concentrates probability mass on a few outcomes, even in tasks that require many distinct, high-quality solutions, a failure known as mode collapse. Inverse Probability Scaling (IPS) counters this in group-based policy gradients such as GRPO by dividing each reward by the outcome's probability, estimated from its frequency within the sampled group. When the outcome space is large relative to the group, outcomes rarely repeat and these estimates become uninformative. Since an outcome can be constructed along many trajectories, its probability is a sum over all of them, and computing it directly is intractable. We propose Marginal IPS (MIPS), which performs marginal evaluation by combining the sampled trajectory with a learned model of the trajectories that lead to each outcome. On Hypergrids, where the exact sampling distribution is computable, MIPS stays near the reward-distribution target as the outcome space grows relative to the rollout group, while frequency-based IPS degrades. In molecular synthesis and phylogenetic inference, MIPS sustains diverse, high-quality sampling where GRPO and frequency-based IPS collapse, showing that MIPS serves as a reliable solution to mode collapse in RL.

Abstract

Reinforcement Learning

Explore similar work

MIPS: Mitigating Mode Collapse in Reinforcement Learning for Large Outcome Spaces

Beyond Mode Collapse: Distribution Matching for Diverse Reasoning

Measuring and Mitigating Solution Mode Collapse in RLVR