GeMPO: Generalized Measure Matching for Online Diffusion Reinforcement Learning

Haitong Ma; Chenxiao Gao; Tianyi Chen; Na Li; Bo Dai

cs.LGMar 10, 2026

GeMPO: Generalized Measure Matching for Online Diffusion Reinforcement Learning

Authors: Haitong Ma, Chenxiao Gao, Tianyi Chen, Na Li, Bo Dai

Organizations: Harvard University · Georgia Institute of Technology

Abstract

A commonly used family of RL algorithms for diffusion policies conducts softmax reweighting over samples from the behavior policy, which often induces an overgreedy policy and fails to utilize feedback from negative samples. In this work, we introduce GeMPO, a simple and unified framework that generalizes reweighting scheme in diffusion RL from softmax to general monotonic functions. GeMPO revisits diffusion RL via a measure matching perspective: First, we construct a virtual target policy measure via solving a regularized policy optimization objective; Second, we minimize the divergence between the current policy and this target measure through reweighted flow matching. This formulation offers two key advantages: i) It extends weight design beyond traditional exponential reweighting, allowing it to be tailored to diverse reward landscapes; and ii) by relaxing the non-negativity constraint on the target measure, our framework provides a principled justification for negative reweighting. We provide interpretations of how negative reweighting actively repels the policy from suboptimal actions and thus facilitates exploration. Extensive empirical evaluations demonstrate that GeMPO achieves competitive or superior performance by leveraging these flexible weighting schemes, and we provide practical guidelines for selecting reweighting methods in practice.

Figures & tables

Divergence $D_{f}$	Generator $f(x)$	reweighting factor $g_{f}(\cdot)$
Forward KL Divergence	$x\log x$	$\exp(x)$
$\chi^{2}$ -Divergence	$x^{2}-1$	$x$
$\alpha$ -Divergence ( $\alpha>1$ )	$\frac{x^{\alpha}-x}{\alpha(\alpha-1)}$	$x^{\frac{1}{\alpha-1}}$

Table 1: Examples of

f

-divergence, their generator functions, and resulting reweighting factors.

Figure 1: Bandit problem to answer Q2 . Left: Reward functions and initial policy. Middle: Regret curves comparing different reweighting with and without negative reweighting functions over 200 epochs. Right: Regret curves with different lower bound

\ell

Figure 2: Bandit problem exploring Q3 . Left: Two reward functions and initial policy. Right: Regret curves comparing different reweighting schemes over 100 epochs. Linear excels in sharp rewards and Square excels at broad rewards, both outperform Exp .

Method		HalfCheetah	Humanoid	Ant	Walker2d	Hopper	Swimmer	Norm. Avg.(%)
Classic Model-Free RL
	TD3	$9820\pm 1543$	$5263\pm 47$	$4400\pm 1196$	$3732\pm 388$	$\mathbf{3387}\pm 302$	$88\pm 25$	$107.4$
	SAC	$\underline{10938}\pm\ \ \ \ 68$	$5131\pm 97$	$\underline{4945}\pm\ \ 186$	$4463\pm 151$	$2509\pm 276$	$61\pm 5$	$100.0$
Diffusion RL Baselines
	DIPO	$10\,136\pm 2525$	$5184\pm 132$	$977\pm 20$	$3809\pm 1112$	$1191\pm 770$	$46\pm 3$	$70.3$
	QSM	$10\,444\pm 250$	$5103\pm 73$	$3614\pm 360$	$3635\pm 621$	$\underline{3200}\pm\ \ 255$	$86\pm 7$	$103.0$

Table 2: Performance of 1M steps on Gym-MuJoCo environments. We report the average returns of the last 10 checkpoints

\pm

standard error over 5 random seeds, and normalized average score normalized by SAC score. Bold numbers indicate best mean returns and underlined numbers indicate mean returns within standard error of the best runs.

Figure 3: Performance with different reweighting and reward functions on the Cheetah Run (left) and Walker-Run (right) task. Boxes and error bars show mean returns and standard error over 5 seeds at the last checkpoint.

Figure 4: Training curves of GeMPO-Neg(

\ell

) with different truncation lower bound

\ell

. Setting

\ell=0

corresponds to GeMPO-Linear .

Figure 5: Reward functions.

Baselines Method Pred.

\uparrow

Improv. Pretrained

0.17\pm 04

– CG

3.30\pm 01

– TDS

4.64\pm 21

– CFG

5.04\pm 06

– DRAKES

6.44\pm 04

-1.2% DRAKES-KL

5.61\pm 07

-13.9% WD1

3.78\pm 12

-42.0% RL-D 2 (best baseline) 6.52(03) –

GeMPO Variants Method Pred.

\uparrow

Improv. Linear 6.54(01) +0.3% Square 6.77(03) +3.8% Neg(-0.05) 6.75(05) +3.5% Neg(-0.1) 6.95(07) +6.6% Neg(-0.3) 7.45(03) +14.3% Neg(-0.5) 5.62(113) -13.8%

Table 3: Main Results on Pred-Activity. Comparison of GeMPO variants against baselines. Standard GeMPO variants are highlighted in orange, while variants with negative sample awareness are highlighted in blue and named as Neg(

\ell

). The inclusion of negative samples consistently provides the highest gains up to

14.3\%

Figure 6: Ablation studies on KL tolerance and diffusion steps.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Name	Value
Critic learning rate	3e-4
Policy learning rate	1e-4
Diffusion steps	20
Diffusion noise schedules	Cosine
Policy network hidden layers	2
Policy network hidden neurons	256

Appendix

Table 4: Hyperparameters

Reinforcement Learning

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,717 papers

Sample-Efficient Diffusion-based Reinforcement Learning with Critic Guidance
May 28, 2026Shutong Ding, Zejia Zhong, Zhongyi Wang +4Diffusion Guidance Robotic RL
1ShanghaiTech University.
Mean Flow Policy Optimization
Apr 16, 2026Xiaoyi Dong, Xi Sheryl Zhang, Jian ChengMeanFlow Mean Flow
C2DL, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · AiRiA +1
Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models
Sep 29, 2025Shuchen Xue, Chongjian Ge, Shilong Zhang +2Score Matching Diffusion Model Alignment
UCAS · Adobe Research · HKU +1

May 28, 2026J/K move · Enter open · S save

Sample-Efficient Diffusion-based Reinforcement Learning with Critic Guidance

Shutong Ding, Zejia Zhong, Zhongyi Wang, Ke Hu+3

1ShanghaiTech University.

Diffusion Guidance Robotic RL Diffusion Policy RL for Diffusion Models Sample-Efficient RL

Open paper AskPDF arXiv

Recent advances in reinforcement learning (RL) have achieved great successes by leveraging the multimodality and exploration capability of diffusion policies. Among these approaches, one representative branch focuses on the sampling-based policy optimization. This design enables better exploration capability of the diffusion model, particularly at the beginning of training, but suffer from low exploitation in Q-value information, resulting in a slow policy convergence. Another branch pays attention to gradient-based policy optimization, which sufficiently exploits the gradient of the Q function yet tends to collapse into a unimodal policy with low diversity. To address this issue, we propose CGPO, Critic-Guided diffusion Policy Optimization, which effectively balances exploration and exploitation with the training-free guidance technique integrated into the denoising process of diffusion policy. Concretely, CGPO steers action generation toward high-value regions defined by the critic network and uses the guided actions as regression objectives. In this manner, CGPO reduces the time required to obtain high-quality actions and improves final performance with better balance between the exploration-exploitation tradeoff. We validate the effectiveness of CGPO on 5 MuJoCo locomotion tasks, and CGPO achieves state-of-the-art performance compared with existing diffusion-based RL methods. Notably, CGPO is the first success to incorporate diffusion policy into real-world RL, with its superior performance on Franka robot arm grasping tasks. Our official page is released at https://dingsht.tech/cgpo-webpage.

Abstract

Reinforcement Learning

Explore similar work

Sample-Efficient Diffusion-based Reinforcement Learning with Critic Guidance

Mean Flow Policy Optimization

Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models