Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

TL;DR

The paper reveals that the common objective of maximizing expected return inherently drives outcome-level mode collapse in multimodal RL, due to a frequency-based amplification in the learning signal. It introduces inverse probability scaling (IPS) to decouple reward attribution from outcome frequency, yielding a stationary distribution where outcomes are proportional to their rewards, i.e., p^* = r/Σr. Implementing IPS within a GRPO framework (IPS-GRPO) provides a drop-in method that preserves diversity without extra models, while KL regularization serves as interpolation rather than collapse prevention. Across Hyper-Grid, HypoSpace, and molecular generation tasks, IPS-GRPO consistently reduces outcome-level mode collapse and often improves reward-related metrics compared to strong baselines. The work highlights objective design as critical for reliable multimodal policy optimization and points to future extensions to actor–critic methods and dense-reward settings.

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.

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

TL;DR

The paper reveals that the common objective of maximizing expected return inherently drives outcome-level mode collapse in multimodal RL, due to a frequency-based amplification in the learning signal. It introduces inverse probability scaling (IPS) to decouple reward attribution from outcome frequency, yielding a stationary distribution where outcomes are proportional to their rewards, i.e., p^* = r/Σr. Implementing IPS within a GRPO framework (IPS-GRPO) provides a drop-in method that preserves diversity without extra models, while KL regularization serves as interpolation rather than collapse prevention. Across Hyper-Grid, HypoSpace, and molecular generation tasks, IPS-GRPO consistently reduces outcome-level mode collapse and often improves reward-related metrics compared to strong baselines. The work highlights objective design as critical for reliable multimodal policy optimization and points to future extensions to actor–critic methods and dense-reward settings.

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.
Paper Structure (35 sections, 3 theorems, 40 equations, 5 figures, 5 tables, 1 algorithm)

This paper contains 35 sections, 3 theorems, 40 equations, 5 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Mode Collapse in RL and Loss of Diversity
  4. Flow-based methods
  5. Collapse in LLM post-training
  6. Why Outcome-Level Mode Collapse is Inevitable under Expected Return
  7. Interpretation: a simple positive feedback loop
  8. Our Method: Inverse Probability Scaling
  9. The scaled objective
  10. Stationary Solution: reward-proportional outcomes
  11. With KL regularization: interpolation instead of collapse prevention
  12. From principle to algorithm: IPS-GRPO
  13. Experiments
  14. Hyper-Grid Task
  15. HypoSpace Benchmark
  16. ...and 20 more sections

Key Result

Theorem 3.1

Under gradient flow on $J(z)=\sum_o p_o(z)\,r(o)$ with a softmax parameterization, define the (on-policy) advantage of outcome $o$ at training time $t$ as $a_t(o) = r(o) -\mathbb{E}_{o'\sim p_t}[r(o')]$ Then for any outcomes $i,j$ with $p_i(t),p_j(t)>0$,

Figures (5)

  • 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.
  • 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.
  • Figure 3: Study of IPS-GRPO behaviour against GRPO
  • Figure 4: voxels generated by GRPO
  • Figure 5: voxels generated by IPS-GRPO

Theorems & Definitions (5)

  • Theorem 3.1
  • Theorem 4.1
  • Corollary 4.2: Reward-proportional stationary distribution
  • proof
  • proof