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 Df
Generator f(x)
reweighting factor gf(⋅)
Forward KL Divergence
xlogx
exp(x)
χ2 -Divergence
x2−1
x
α -Divergence ( α>1 )
α(α−1)xα−x
xα−11
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 ℓ .
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±1543
5263±47
4400±1196
3732±388
3387±302
88±25
107.4
SAC
10938±68
5131±97
4945±186
4463±151
2509±276
61±5
100.0
Diffusion RL Baselines
DIPO
10136±2525
5184±132
977±20
3809±1112
1191±770
46±3
70.3
QSM
10444±250
5103±73
3614±360
3635±621
3200±255
86±7
103.0
Table 2: Performance of 1M steps on Gym-MuJoCo environments. We report the average returns of the last 10 checkpoints ± 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( ℓ ) with different truncation lower bound ℓ . Setting ℓ=0 corresponds to GeMPO-Linear .
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( ℓ ). The inclusion of negative samples consistently provides the highest gains up to 14.3% .
Figure 6: Ablation studies on KL tolerance and diffusion steps.
Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.
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.
Diffusion models have recently emerged as expressive policy representations for online reinforcement learning (RL). However, their iterative generative processes introduce substantial training and inference overhead. To overcome this limitation, we propose to represent policies using MeanFlow models, a class of few-step flow-based generative models, to improve training and inference efficiency over diffusion-based RL approaches. To promote exploration, we optimize MeanFlow policies under the maximum entropy RL framework via soft policy iteration, and address two key challenges specific to MeanFlow policies: action likelihood evaluation and soft policy improvement. Experiments on MuJoCo, DeepMind Control Suite and HumanoidBench benchmarks demonstrate that our method, Mean Flow Policy Optimization (MFPO), achieves performance comparable to or exceeding current diffusion-based baselines while considerably reducing training and inference time. Our code is available at https://github.com/dongxiaoyi-xyz/MFPO.
Xiaoyi Dong, Xi Sheryl Zhang, Jian Cheng
C2DL, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · AiRiA +1
Reinforcement Learning (RL) has emerged as a central paradigm for advancing Large Language Models (LLMs), where both pre-training and RL post-training stages are grounded in the same log-likelihood formulation. In contrast, recent RL approaches for diffusion models, most notably Denoising Diffusion Policy Optimization (DDPO), optimize an objective different from the pretraining objectives--score/flow matching loss. In this work, we establish a novel theoretical analysis: DDPO is an implicit form of score/flow matching with noisy targets, which increases variance and slows convergence. Building on this analysis, we introduce Advantage Weighted Matching (AWM), a policy-gradient method for diffusion. It uses the score/flow-matching loss and reweights each sample by its advantage. In effect, AWM raises the influence of high-reward samples and suppresses low-reward ones while keeping the modeling objective identical to pretraining. This simple yet effective design yields substantial benefits: on the GenEval, OCR, and PickScore benchmarks, AWM delivers up to a 34× speedup over Flow-GRPO (which builds on DDPO), when applied to Stable Diffusion 3.5 Medium and FLUX, without compromising generation quality. Code is available at https://github.com/scxue/advantage_weighted_matching