Density-Aware Reward Scaling: A Reinforcement Learning Framework for Reward-Proportional Sampling

Ria Shekhawat; Abhijeet Sinha; Yingtao Zhu; Dianbo Liu

cs.LGNov 3, 2026

Density-Aware Reward Scaling: A Reinforcement Learning Framework for Reward-Proportional Sampling

Authors: Ria Shekhawat, Abhijeet Sinha, Yingtao Zhu, Dianbo Liu

Organizations: National University of Singapore

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.

Figures & tables

Density and transport
Metric	GFlowNet	GRPO	DARS-GRPO
Mean raw reward $\uparrow$	3.44	6.96	4.59
JSD $\downarrow$	0.020	0.649	0.041
Hist. TV $\downarrow$	0.149	0.975	0.219
SW 2 $\downarrow$	0.036	0.242	0.015
Energy $\downarrow$	0.007	0.281	0.002

Table 1: Navigation Task Results (avg over 5 seeds)

Panel of Figure 1 — Figure 1: Terminal distributions on the navigation task.

Figure 2: Overview of DARS-GRPO for text-to-image prompt adaptation. A trainable language-model policy samples $G$ rewrites of an input prompt, and a frozen text-to-image model generates one image per rewrite. Each image receives a task reward $r_{i}$ and is mapped by a frozen image encoder into a representation space, where a Gaussian kernel computes its within-group density $\rho_{i}$ . DARS forms the scaled reward $\tilde{r}_{i}=r_{i}/\rho_{i}$ , reducing the contribution of high-reward images in dense representation-space regions and increasing the relative contribution of high-reward images in sparsely represented regions. GRPO then uses these scaled rewards to update only the prompt-rewriting policy.

Figure 3: Comparison of prompt rewrites and their corresponding generated images for “A snowman standing in the yard” using GRPO, PAG, and DARS-GRPO. GRPO produces nearly identical rewrites and visually similar nighttime scenes, whereas PAG introduces moderate variation and DARS-GRPO generates more diverse rewrites, settings, and visual compositions.

Combined reward
		Image Vendi $\uparrow$		CLIP-align $\uparrow$		HPSv2 $\uparrow$		PickScore $\uparrow$		Prompt Vendi $\uparrow$		Prompt diversity $\uparrow$		DINO pairwise distance $\uparrow$
Policy LM	Method	COCO	Lexica	COCO	Lexica	COCO	Lexica	COCO	Lexica	COCO	Lexica	COCO	Lexica	COCO	Lexica
–	SD v1.4	8.60	9.52	0.27	0.27	0.26	0.25	21.53	20.91	–	–	–	–	40.58	41.84
–	CADS	12.93	14.08	0.26	0.26	0.24	0.23	20.86	20.28	–	–	–	–	46.22	46.96
Qwen3-0.6B	GRPO	8.76	9.81	0.24	0.24	0.23	0.23	20.92	20.69	1.01	1.04	0.00	0.01	40.90	42.10
Qwen3-0.6B	DARS-GRPO	13.04	14.45	0.24	0.24	0.23	0.23	20.80	20.37	1.43	1.67	0.05	0.06	46.92	47.89

Table 2: Prompt Adaptation for Text-to-Image Generation results on MS-COCO and Lexica (avg over 5 seeds) .

Width	Image Vendi $\uparrow$	CLIP-align $\uparrow$	Prompt Diversity $\uparrow$	Prompt Vendi $\uparrow$
10	9.99	0.26	0.12	1.99
15	14.12	0.25	0.21	3.34
20	18.01	0.24	0.17	2.74
25	19.77	0.23	0.22	3.47
30	12.14	0.25	0.11	2.09

Table 3: Kernel-width ablation results (avg over 5 seeds)

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Policy-induced movement fields for GRPO (left) and DARS-GRPO (right) on the four-mode navigation task. Streamlines show the estimated mean one-step displacement, color indicates its magnitude, and gray contours show the normalized reward landscape.

Panel of Figure 5 — Figure 5: Policy LM comparison for “A snowman standing in the yard” (top) and “A bicycle leaning against a brick wall” (bottom). Each embedded grid compares GRPO with DARS for the indicated language-model.

Panel of Figure 6 — Figure 6: Policy LM comparison for “A wooden bench under a tree” (top) and “A cat sleeping on a windowsill” (bottom). DARS expands the visual contexts while preserving the object and scene requested by each prompt.

Panel of Figure 7 — Figure 7: Policy LM comparison for “A scarecrow in the middle of a farm” (top) and “A castle on top of a hill” (bottom). The DARS samples show greater variation in character appearance, architecture, framing, and illumination than the corresponding GRPO samples.

Panel of Figure 8 — Figure 8: Policy LM comparison for “A lantern hanging outside a cottage” (top) and “A fox standing near a forest path” (bottom). Across both policy LMs, DARS broadens the within-prompt visual distribution while maintaining the prompt’s central subject.

Panel of Figure 9 — Figure 9: Reward comparison for “A snowman standing in the yard” (top) and “A bicycle leaning against a brick wall” (bottom), using Qwen3-0.6B.

Panel of Figure 10 — Figure 10: Reward comparison for “A wooden bench under a tree” (top) and “A cat sleeping on a windowsill” (bottom).

Panel of Figure 11 — Figure 11: Reward comparison for “A scarecrow in the middle of a farm” (top) and “A castle on top of a hill” (bottom).

Panel of Figure 12 — Figure 12: Reward comparison for “A lantern hanging outside a cottage” (top) and “A fox standing near a forest path” (bottom).

Figure 13: Qualitative comparison for “A bicycle leaning against a brick wall.” GRPO repeatedly generates a similar nighttime composition, while DARS-GRPO varies lighting, seasonal context, and surrounding details while preserving the bicycle and wall.

Figure 14: Qualitative comparison for “A wooden bench under a tree.” DARS-GRPO explores more varied locations, tree appearances, and visual contexts than GRPO, while PAG and GFN produce more limited changes around the same central composition.

Figure 15: Qualitative comparison for “A cat sleeping on a windowsill.” All methods preserve the primary subject, while DARS-GRPO introduces broader variation in background scenery, lighting, and framing.

Parameter	Value
Policy LMs	SFT Qwen3-0.6B; SFT GPT-2
Training temperature	1.5
Evaluation temperature	1.0
max_new_tokens	64
Learning rate	$1\times 10^{-5}$
Weight decay	0.1

Appendix

Table 4: Policy optimization hyperparameters (DARS and GRPO).

Parameter	Value
Density embedding space	dinov2_image
Density kernel	Gaussian
Kernel width ( $h$ )	20.0

Appendix

Table 5: DARS-specific hyperparameters.

Component	Value
Operating System	Ubuntu 22.04.4 LTS (Linux 5.15.0-94-generic)
GPUs	$4\times$ NVIDIA A100-SXM4-40GB
Host RAM	$\sim$ 503 GiB
Python	3.13.9
PyTorch	2.12.1+cu130
Key Libraries	transformers 5.13.0, diffusers 0.39.0,

Appendix

Table 6: Hardware and software environment used for training.

Parameter	Value
Diffusion model	CompVis/stable-diffusion-v1-4
Resolution / Steps / CFG	$512\times 512$ , 20 steps, guidance scale 7.5
Scheduler	DPM-Solver Multistep
diffusion.seed_mode	random
run.seed	42
Rewrites per prompt ( $G=N$ )	64

Appendix

Table 7: Shared generation and evaluation configuration.

Setting	Implementation
Combined	Product-style aesthetic gain with CLIP alignment gating
Alignment only	CLIP alignment reward without aesthetic gain

Appendix

Table 8: Reward settings.

Split	Dataset
Training	dataset/coco/n_1024/train.jsonl
MS-COCO evaluation	dataset/coco/n_1024/test.jsonl
Lexica evaluation	dataset/lexica/n_1024/test.jsonl

Appendix

Table 9: Training and evaluation datasets.

Diversity-Aware Sampling

See all 30 papers

Make Sparse Rewards Count: Density-Aware Reward Aggregation for Multi-Reward RL
Sep 30, 2026Tong Zheng, Skylar Zhai, Zhan Cheng +7Multi-Objective Reinforcement Learning Reinforcement Learning
University of Chinese Academy of Sciences · University of Minnesota Twin Cities · University of Wisconsin–Madison +5
DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay
Mar 17, 2026Long Li, Zhijian Zhou, Tianyi Wang +7Behavioral Diversity in RL Reinforcement Learning
Griffith University, Australia · Fudan University, China · Shanghai Innovation Institute, China +3
Repulsive Population Policy Gradients for Behaviorally Diverse Reinforcement Learning
Oct 27, 2026Saeed Hedayatian, Harsh Shah, Celine Mneimneh +1Reinforcement Learning Behavioral Diversity in RL
Thomas Lord Department of Computer Science, University of Southern California, Los Angeles, USA · Independent Researcher

Sep 30, 2026J/K move · Enter open · S save

Make Sparse Rewards Count: Density-Aware Reward Aggregation for Multi-Reward RL

Tong Zheng, Skylar Zhai, Zhan Cheng, TianMing Sha+6

University of Chinese Academy of Sciences · University of Minnesota Twin Cities · University of Wisconsin–Madison+5

Multi-Objective Reinforcement Learning Reinforcement Learning RL for Language Model Reasoning Sparse-Reward RL Policy Learning

Open paper AskPDF arXiv

Multi-reward reinforcement learning trains large language models to satisfy multiple behavioral objectives simultaneously. Reward-wise normalization, as used in GDPO, preserves reward-specific relative information within rollout groups, but different objectives can still exhibit uneven learning progress. We study this behavior through advantage energy, the sum of a reward's squared advantages over a batch. Under idealized GDPO normalization, we show that this energy is proportional to active-group density: the fraction of rollout groups in which the reward provides nonzero relative advantages. This reveals a residual batch-level signal imbalance and provides a basis for calibrating reward contributions. Based on this relation, we propose Density-Aware Reward Aggregation (DARA). We derive an inverse-square-root density correction that gives greater weight to signals from less frequently active rewards. DARA computes its weights from each rollout batch, adapting to changes in reward activity throughout training without modifying the underlying policy optimization objective. Experiments on tool calling and mathematical reasoning show that DARA learns the targeted behaviors faster than GDPO, reaching high format compliance in up to 26% fewer training steps on tool calling and near-saturated length compliance in up to 65% fewer steps on mathematical reasoning, while remaining competitive in final performance. Our code is available at https://github.com/zhaihaotian/DARA.

Abstract

Diversity-Aware Sampling

Explore similar work

Make Sparse Rewards Count: Density-Aware Reward Aggregation for Multi-Reward RL

DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay

Repulsive Population Policy Gradients for Behaviorally Diverse Reinforcement Learning