Density-Aware Reward Scaling: Making Reinforcement Learning Sampling More Balanced

When pursuing high returns, reinforcement learning tends to concentrate probability mass on a few high-reward outcomes, leading to a loss of diversity. The recently proposed Density-Aware Reward Scaling (DARS) attempts to make sampled results closer to the reward-proportional distribution without altering the policy gradient framework.

Core Problem: The Conflict Between Reward Maximization and Diversity

The standard reinforcement learning objective maximizes expected return, which makes the policy tend to repeatedly output known high-reward samples. When multiple distinct but equally excellent solutions exist, this mechanism suppresses exploration and causes mode collapse. The original paper points out that DARS is designed precisely to solve the problem of reward-proportional sampling in continuous outcome spaces—it retains standard policy gradient training and only processes the reward signal.

Mode collapse is a classic failure phenomenon in generative models and reinforcement learning: when the target distribution itself has multiple peaks (multimodal), the optimization process covers only one or a few of them while completely discarding the other equally high-probability regions. In reinforcement learning, it usually stems from the policy's greedy maximization of expected return—once an action or output is confirmed to yield high reward, the policy gradient keeps amplifying its probability, causing the sampling distribution to rapidly contract into a spike.

In contrast, reward-proportional sampling requires that the probability of sampling an outcome x be proportional to its reward r(x), i.e., p(x) ∝ r(x). This objective naturally preserves the relative weights of all high-reward solutions rather than picking out only the largest one. What DARS aims to solve is precisely how to achieve this sampling distribution in continuous outcome spaces (where outputs are vectors or images rather than discrete actions) without disrupting the standard policy gradient training process.

How DARS Works: Reverse-Scaling Rewards Using Local Density

The key step of DARS is to estimate the local density of outcomes within each on-policy rollout group, then reverse-scale the rewards according to density. Rewards in high-density regions (i.e., common outcomes) are suppressed, while rewards in low-density regions (i.e., rare outcomes) are relatively boosted. This way, the sampling process naturally converges toward the reward-proportional distribution.

The method's design is very lightweight: it requires no learned density model, no auxiliary objective, no critic, and does not modify the underlying policy gradient objective. The original paper emphasizes that DARS is merely a reward processing mechanism that can be seamlessly embedded into existing reinforcement learning pipelines.

On-policy rollout group refers to a batch of outcomes sampled from the current policy for the same prompt or state during each round of training (e.g., multiple images generated from the same prompt). Since they all come from the current policy, this group of samples approximately follows the current output distribution, and can therefore serve as local empirical samples of that distribution for density estimation without additional training of a density model.

Local density estimation is typically accomplished using distances or kernel functions between samples: the denser the similar samples near an outcome, the more the current policy tends to generate it, and the higher its density value. DARS uses this density to reverse-scale the reward (e.g., r/d or r·f(1/d)), which is equivalent to penalizing common outcomes and compensating rare ones. This approach based on within-group relative density is a non-parametric estimation—low in cost but with accuracy affected by group size: the smaller the group, the higher the variance of density estimation, which is a practical limitation to note in high-dimensional tasks.

Experimental Validation: From Multimodal Navigation to Text-to-Image

In controlled multimodal navigation tasks, DARS-GRPO avoids the single-mode collapse exhibited by GRPO and reasonably recovers the target mode allocation. On text-to-image prompt adaptation tasks, DARS-GRPO consistently improves the diversity of images and prompts across different datasets, policy language models, and reward settings while maintaining good semantic alignment.

These results demonstrate that density-aware reward scaling provides a lightweight means for existing reinforcement learning pipelines to improve coverage of diverse high-reward outcomes.

GRPO (Group Relative Policy Optimization) is a critic-free policy gradient variant: it samples a group of outputs for the same prompt, normalizes rewards using the mean and standard deviation of within-group rewards to obtain relative advantages, and then updates the policy accordingly. Because it relies entirely on within-group comparisons, GRPO is insensitive to the absolute scale of rewards, simple to implement, and widely used for preference alignment in large language models and diffusion models.

For this very reason, GRPO's within-group normalization only cares about relative ranking, making it easy to push probability mass toward the best-performing output type within the group, exacerbating mode collapse. As a reward processing layer, DARS can be directly stacked on top of GRPO (i.e., DARS-GRPO): first reverse-scale rewards by density, then pass them to GRPO for within-group normalization, thereby introducing a diversity constraint without modifying its objective function.

Significance and Limitations

The value of DARS lies in improving diversity solely through reward scaling without introducing additional models or complex objectives. However, the original paper does not discuss the computational overhead and density estimation accuracy of this method in high-dimensional complex tasks, which may be points of concern for practical applications.

One-line Summary

DARS lightly improves reinforcement learning sampling diversity by reverse-scaling rewards using within-group density.

Paragraph Summary

This article introduces Density-Aware Reward Scaling (DARS): standard reinforcement learning that maximizes expected return concentrates probability mass on a few high-reward outcomes, causing mode collapse, whereas DARS aims to make the sampling distribution approximate the reward-proportional distribution. Its approach is to estimate the local density of outcomes within each on-policy rollout group and reverse-scale rewards by density—common outcomes are suppressed and rare outcomes are boosted, so sampling naturally converges toward p(x) ∝ r(x). This method requires no density model, auxiliary objective, or critic, and does not modify the policy gradient objective, allowing seamless integration into existing pipelines. In experiments, DARS-GRPO avoids GRPO's single-mode collapse on multimodal navigation tasks and recovers the target mode allocation, and consistently improves image and prompt diversity while maintaining semantic alignment in text-to-image prompt adaptation. The article also notes that its computational overhead and density estimation accuracy in high-dimensional tasks remain to be discussed.