MIPS: Mitigating Mode Collapse in Reinforcement Learning for Large Outcome Spaces
Standard expected-reward RL suffers from mode collapse, concentrating probability on a few outcomes even when tasks require many distinct, high-quality solutions. MIPS tackles this by estimating marginal outcome probabilities through a learned trajectory-conditional model, avoiding frequency-based estimation that breaks down when the outcome space is much larger than the rollout group. Using the identity $P_\theta(o) = P_\theta(\tau) / P_\theta(\tau|o)$, it scales rewards by inverse estimated probabilities within a standard GRPO update. On enumerable Hypergrids, molecular synthesis, and phylogenetic inference, MIPS maintains reward-aligned diversity where group-count IPS collapses.
MIPS addresses a genuine limitation of frequency-based Inverse Probability Scaling by replacing within-group outcome counts with a learned trajectory-conditional model that estimates marginal outcome probabilities via $P_\theta(o) = P_\theta(\tau) / c_\phi(\tau|o)$. The idea is theoretically clean and the Hypergrid experiments with exact policy enumeration provide unusually rigorous validation. However, the approach inherits approximation errors from the auxiliary conditional model and is explicitly restricted to finite, acyclic, deterministic constructive processes, leaving its relevance to general RL unsettled. The empirical gains over frequency-based IPS are dramatic, particularly in large outcome spaces.
The paper's central theoretical device—estimating the marginal outcome probability via the identity $P_\theta(o) = P_\theta(\tau) / P_\theta(\tau|o)$ and a learned backward model—is elegant and well-motivated. The Hypergrid evaluation is exemplary: because exact terminal distributions are computable, the authors directly report JSD and TV divergence against the reward target across a full sweep of grid sizes $H$ and group sizes $G$, showing that IPS-GRPO degrades as $H^2$ grows while MIPS-GRPO does not. The ablation in Table 1 is decisive: stripping the conditional model and setting $c_\phi(\tau|o)=1$ causes JSD to jump from roughly 0.003–0.008 to 0.56–0.64, proving that the conditional model is essential rather than a superficial add-on.
The conditional model is only an approximation, and the paper explicitly warns that its estimates are neither unbiased nor normalized across trajectories to the same outcome. This approximation error propagates directly into the IPS reward scaling, yet there is no theoretical characterization of how conditional-model bias affects the GRPO update or the final stationary distribution. The method is also scoped to a narrow regime: finite, acyclic, deterministic constructive processes with strictly positive rewards. This excludes stochastic dynamics, cyclic graphs, and continuous outcome spaces, all of which are common in real-world RL.
In the molecular and phylogenetic experiments, the baseline collapse is so extreme—GRPO generates essentially one unique molecule, and IPS-GRPO recovers zero high-reward modes—that MIPS looks favorable by default rather than through a nuanced tradeoff. The comparison to specialized GFlowNet samplers (RGFN, PhyloGFN) is fairer, but MIPS trains an auxiliary model with extra hyperparameters, contradicting IPS's appeal as a drop-in, no-auxiliary modification.
The evidence strongly supports the scalability claim in the controlled Hypergrid setting, where exact outcome distributions let the authors directly verify that MIPS tracks the reward target as $H$ grows while frequency-based IPS and GRPO diverge. The ablation confirms the conditional model is doing meaningful work. In molecular synthesis and phylogenetics, MIPS clearly outperforms its own RL baselines, but those baselines collapse so completely that the comparison is closer to a sanity check than a competitive benchmark. Against GFlowNets, MIPS achieves broader coverage and better reward–probability correlation in phylogenetics ($r=0.967$ vs. PhyloGFN's $0.836$), yet it simultaneously yields lower mean log-likelihood ($-4372.9$ vs. $-4294.4$ for IPS-GRPO), a quality–coverage tradeoff that the paper mentions but does not deeply analyze.
Appendix C provides extensive hyperparameters for all three domains, including network architectures, learning rates, batch sizes, and optimizer settings, and Algorithm 1 spells out the coupled policy and conditional-model update exactly. The authors also report training seeds (1, 2, 3). However, the paper does not mention a code release or public repository, which would substantially ease independent verification. Reproduction is feasible from the descriptions—the Hypergrid is simple, and the molecular and phylogenetic environments reuse published benchmarks—but implementing the conditional-model factorization and its per-update fitting requires care. The main practical blocker is the absence of released code and the sensitivity of the coupled training loop to the auxiliary model's learning rate and capacity.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.