Density-aware reward scaling improves diversity in reinforcement learning outcomes
Density-Aware Reward Scaling: A Reinforcement Learning Framework for Reward-Proportional Sampling
Machine LearningArtificial Intelligence
Summary
Reinforcement learning often focuses too much on a few top-scoring solutions, ignoring other good but different options. The authors introduce Density-Aware Reward Scaling (DARS), a method that adjusts rewards based on how common an outcome is, encouraging a wider range of good results. This method works without extra models or changing the core learning process. Testing shows DARS helps systems find more diverse successful solutions in tasks like navigation and generating images from text.
What this means in practice
- •For machine learning engineers: Increase the variety of high-quality solutions in reinforcement learning tasks to avoid mode collapse in continuous outcome settings.
- •For creative ai developers: Improve diversity and semantic relevance of images generated from text prompts while maintaining reward-based guidance in generative models.
Authors
Ria Shekhawat, Abhijeet Sinha, Yingtao Zhu, Dianbo Liu
Abstract
Reinforcement learning objectives that maximize expected return can concentrate probability mass on a small subset of high-reward outcomes, limiting diversity when multiple distinct solutions are desirable. We introduce Density-Aware Reward Scaling (DARS), a simple reward-processing mechanism that promotes reward-proportional sampling in continuous outcome spaces while retaining standard policy-gradient training. DARS estimates the local density of outcomes within each on-policy rollout group and inversely scales rewards by this density. The method requires no learned density model, auxiliary objective, critic, or modification to the underlying policy-gradient objective. In a controlled multimodal navigation task, DARS-GRPO avoids the single-mode collapse exhibited by GRPO and closely recovers the target mode allocation. Applied to text-to-image prompt adaptation, DARS-GRPO consistently improves image and prompt diversity across datasets, Policy LMs, and reward settings while maintaining good semantic alignment. These results show that density-aware reward scaling provides a lightweight way to improve coverage of diverse high-reward outcomes within existing reinforcement-learning pipelines.