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)∝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 ℓ1 distance between the learned and target distributions.
Method
n=2
n=4
GRPO
2.8×10−2
5.2×10−4
FlowRL
2.8×10−2
5.2×10−4
IPS-GRPO
3.0×10−3
3.7×10−4
Table 1: ℓ1 distance between the learned terminal-state distribution and the target distribution on the hyper-grid domain. Lower is better.
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%±0.35%
31.61%±0.01%
9.50%±0.00%
FlowRL
12.82%±0.05%
26.85%±0.37%
10.74%±0.00%
IPS-GRPO
43.91%±1.39%
90.00%±1.24%
60.74%±0.78%
Table 2: Final recovery rate (%) on HypoSpace tasks. Average over 5 random seeds. Higher is better.
Threshold
Algorithm
Yield ↑
OB100 ↓
SYNTH
0.80
REINVENT
5695±173
605±41
0.80
IPS-GRPO
7371±212
554±34
0.85
REINVENT
418±78
2934±68
0.85
IPS-GRPO
1580±58
1178±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
ϵ=0.01
ϵ=0.1
ϵ=0.2
4
17.63±0.33
17.63±0.33
17.63±0.33
8
27.50±0.67
27.50±0.67
41.99±0.80
16
41.02±0.60
33.05±0.62
43.91±1.39
32
32.13±0.42
35.27±1.01
25.62±1.65
64
36.52±0.32
39.12±1.33
42.19±0.14
Table 4: Recovery rate (%) on HypoSpace causal inference for different group sizes G and probability clipping thresholds ϵ . Average over 5 random seeds. Higher recovery rate is better.
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.
Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.
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.
Armaan Sandhu, Abhijeet Sinha, Dianbo Liu
University of Massachusetts Amherst · National University of Singapore
On-policy reinforcement learning methods like GRPO suffer from mode collapse: they exhibit reduced solution diversity, concentrating probability mass on a single solution once discovered and ceasing exploration of alternative strategies. We show this stems from reverse KL minimization's mode-seeking behavior, which reinforces the first high-reward trajectory found rather than maintaining a distribution over multiple diverse solutions. We propose DMPO (Distribution-Matching Policy Optimization), which prevents mode collapse through principled approximation of forward KL minimization. DMPO constructs a group level target distribution over sampled trajectories proportional to their rewards, then aligns the policy distribution to this target. This provides mode-covering behavior without requiring sampling from the intractable global target distribution, enabling sustained exploration throughout training. We validate DMPO on NP-hard combinatorial optimization, where exponentially many feasible solutions exist but only a few approach optimality, an ideal testbed for evaluating exploration. DMPO achieves 43.9% Quality Ratio on text-based NP-Bench (vs. GRPO's 40.1%) and 43.1% on vision-based NP-Bench (vs. 38.4%), demonstrating 9% and 12% relative improvements respectively. These gains generalize to mathematical reasoning (+2.0%) and out-of-domain tasks (+2.3%), showing that diversity-preserving training enhances general reasoning capabilities across modalities. Our work establishes distribution matching as a practical, principled approach to preventing mode collapse in on-policy RL, with consistent quality improvements demonstrating sustained exploration across diverse reasoning tasks.
Xiaozhe Li, Yang Li, Xinyu Fang +10
1Tongji University, China · 6Independent · 2Shanghai AI Laboratory, China +3
A language model (LM) can usually answer the same question in more than one way, but reinforcement learning with verifiable rewards (RLVR) is indifferent to which correct answer a model produces. A solution will earn the same reward whether it is the thousandth copy of a familiar answer or one the model has never produced before. Yet, there is potential value in having the model retain multiple correct solutions as it is trained. For instance, multiple modes may give users a choice and provide problem-solving strategies that improve overall model performance. Here, we introduce ModeBench, a benchmark of multi-solution tasks in which the verifier returns both correctness and mode discovered. We then use ModeBench to measure how solution diversity changes under RLVR post-training. We find that RLVR post-training concentrates probability onto fewer correct modes even as accuracy holds or improves, and moreover, that frontier models are already highly concentrated. We then introduce our solution, Re:Max, which stores one verified example per discovered mode in a replay buffer and trains on those stored modes uniformly. A solution found once is, therefore, practiced as often as one found repeatedly. Across three model scales, two RL objectives, and harder task constructions, replay improves both how often a policy succeeds and how many different ways it can succeed.
Liv G. d'Aliberti, Marwa Abdulhai, Sofiia Druchyna +2