Diffusion policies offer a powerful and expressive parameterization for continuous control. Yet, their integration with reinforcement learning remains conceptually and algorithmically challenging. In this work, we address this gap by introducing a noisy-space action-value (Q-)function that assigns values to diffusion latents through the distribution of executed actions induced by the denoising process. We show that this construction admits a precise semantic interpretation and derive a noisy-space policy gradient (NSPG) that optimizes noisy latents using only clean action-space value estimates. Building on this result, we formulate a KL-regularized policy improvement over noisy latents and show that the resulting objective admits a diffusion-compatible regression form, avoiding backpropagation through the denoising process. Empirical results on state-based D4RL benchmarks and vision-based OGBench tasks demonstrate that the proposed noisy-space objective provides a principled and effective basis for training diffusion policies in offline reinforcement learning. Project webpage: https://mahmoud-selim.github.io/NSPG/
Figures & tables
Figure 1: Toy illustration of noisy-space gradients. We consider a two-dimensional multimodal action distribution with a high-density low-reward region and a lower-density high-reward region. NSPG follows the gradient of Qnoisy , the expected clean-action value under the diffusion-induced distribution, which induces guidance toward the higher-reward region. In contrast, the decoded gradient differentiates Q(s,a^(x(k))) after mapping each noisy latent to a single clean action. Arrows show the corresponding noisy-space gradient fields, and particle traces illustrate gradient-ascent rollouts across denoising steps.
Dataset
BC
Onestep-RL
CQL
IQL
IVR
XQL
Diffuser
DTQL
AlignIQL
SfBC
DQL
IDQL-A
FQL 1
DAC
NSAC (ours)
HalfCheetah-medium-v2
42.6
48.4
44.0
47.4
48.3
48.3
44.2
57.9
46.0
45.9
51.1
51.0
60.1
59.1
64.9±0.6
Hopper-medium-v2
52.9
59.6
58.5
66.3
75.5
74.2
58.5
99.6
56.1
57.1
90.5
65.4
74.5
101.2
103.7±0.9
Walker2d-medium-v2
75.6
81.8
72.5
78.3
84.2
84.2
79.7
89.4
78.5
77.9
87.0
82.5
72.7
96.8
97.6±0.8
HalfCheetah-medium-replay-v2
36.3
38.1
45.5
44.2
44.8
45.2
42.2
50.9
41.1
37.1
47.8
45.9
51.1
55.0
59.4±0.6
Hopper-medium-replay-v2
18.1
97.5
95.0
94.7
99.7
100.7
96.8
100.0
74.8
86.2
101.3
92.1
85.4
103.1
104.5±0.7
Walker2d-medium-replay-v2
26.0
49.5
77.2
73.9
81.2
82.2
61.2
88.5
76.5
65.1
95.5
85.1
82.1
96.8
98±1.5
Table 1: Average normalized scores of NSAC over four runs compared with baseline methods (mean ± standard deviation). Best results, as well as results within 5% of the best performance for each task, are highlighted .
Task
IQL
ReBRAC
FBRAC
IFQL
FQL
DAC
NSAC (ours)
visual-cube-single-play-singletask-task1-v0
70 ± 12
83 ± 6
55 ± 8
49 ± 7
81 ± 12
83.7 ± 10
88 ± 7
visual-cube-double-play-singletask-task1-v0
34 ± 23
4 ± 4
6 ± 2
8 ± 6
21 ± 11
34 ± 12
36 ± 14
visual-puzzle-3x3-play-singletask-task1-v0
7 ± 15
88 ± 4
7 ± 2
100 ± 0
94 ± 1
98.2 ± 2
100 ± 0
visual-puzzle-4x4-play-singletask-task1-v0
0 ± 0
26 ± 6
0 ± 0
8 ± 15
33 ± 6
40 ± 5.7
47.5 ± 4
OGBench total
111
201
68
165
229
255.9
271.5
Table 2: Offline RL results on OGBench visual environments, reported as success rates (mean ± standard deviation). Best results, as well as results within 5% of the best performance for each task, are highlighted .
Figure 2: Effect of guidance strength on learning dynamics across different tasks. NSAC is robust across a range of guidance strengths, but overly aggressive guidance can degrade performance.
Figure 3: Effect of the number of Monte Carlo samples used to approximate the noisy-space expectation on learning dynamics. NSAC remains stable with few samples, indicating that the noisy-space expectation can be estimated efficiently.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Hyperparameter
Value
Diffusion steps T
5
Noise schedule βk
Variance preserving (VP)
Critic ensemble size H
10
Batch size
256
Actor learning rate
3×10−4
Critic learning rate
3×10−4
Appendix
Table 3: Shared hyperparameters used across all tasks and experiments.
Domain
Task
λ
η
ρ
Locomotion
HalfCheetah-medium-v2
50
1
0
Hopper-medium-v2
30
1
1
Walker2d-medium-v2
3
1
1
HalfCheetah-medium-replay-v2
100
1
0
Hopper-medium-replay-v2
30
1
1
Walker2d-medium-replay-v2
10
1
1
Appendix
Table 4: Experiment-specific hyperparameters varied across tasks. OGBench task names are shortened for readability.
Figure 4: Toy illustration of noisy-space gradients across denoising steps. We consider a two-dimensional multimodal action distribution with a high-density low-reward region and a lower-density high-reward region. Denoising proceeds from the noisiest step, shown as 1, toward the clean sample, corresponding to step 0. Top: NSPG follows ∇x(k)Qnoisy(x(k),k) , the gradient of the expected clean-action value under the diffusion-induced conditional. Bottom: the decoded baseline follows ∇x(k)Q(a^(x(k))) after mapping each noisy latent to a single clean action. Columns show snapshots at different denoising steps; arrows denote gradient fields and traces show illustrative gradient-ascent rollouts.
Figure 5: Training curves for NSAC across all 15 evaluation environments. Each subplot corresponds to one environment and reports normalized return as a function of training steps. Shaded regions indicate standard deviation across seeds.
Figure 6: Ablation on the pessimism coefficient ρ across three representative environments. Curves report normalized return as a function of training steps for different values of ρ .
Figure 7: Ablation on critic ensemble size across two representative environments. Curves report normalized return for different ensemble sizes.
Method
Time / Update (ms)
Wall-clock Relative to DAC
NSAC ( nsamples=2 )
3.2
0.90×
NSAC ( nsamples=4 )
3.4
0.98×
NSAC ( nsamples=8 )
3.9
1.07×
NSAC ( nsamples=16 )
5.8
1.40×
DAC ( Fang et al., 2024 )
3.6
1.00×
Appendix
Table 5: Computational cost of NSAC for different numbers of Monte Carlo samples, compared to DAC. Wall-clock time is normalized by DAC.
Reinforcement learning (RL) has been extensively employed in a wide range of decision-making problems, such as games and robotics. Recently, diffusion policies have shown strong potential in modeling multi-modal behaviors, enabling more diverse and flexible action generation compared to the conventional Gaussian policy. Despite various attempts to combine RL with diffusion, a key challenge is the difficulty of computing action log-likelihood under the diffusion model. This greatly hinders the direct application of diffusion policies in on-policy reinforcement learning. Most existing methods calculate or approximate the log-likelihood through the entire denoising process in the diffusion model, which can be memory- and computationally inefficient. To overcome this challenge, we propose a novel and efficient method to train a diffusion policy in an on-policy setting that requires only evaluating a simple Gaussian probability. This is achieved by aligning the policy iteration with the diffusion process, which is a distinct paradigm compared to previous work. Moreover, our formulation can naturally handle entropy regularization, which is often difficult to incorporate into diffusion policies. Experiments demonstrate that the proposed method produces multimodal policy behaviors and achieves superior performance on a variety of benchmark tasks in both IsaacLab and MuJoCo Playground.
Ben Liu, Shunpeng Yang, Hua Chen
Southern University of Science and Technology, Shenzhen, China · Limx Dynamics · Hong Kong University of Science and Technology, Hongkong, China +1
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.
Reinforcement learning (RL) post-training provides a direct way to align diffusion models with human preferences and task-specific rewards. However, current RL algorithms for diffusion models remain fragmented: reverse-trajectory methods rely on discretized likelihood ratios, whereas forward-matching methods train on reward-labeled noising versions of the rollout samples. This paper shows that these seemingly different losses arise from a single path-space principle. Starting from the regularized diffusion-RL objective, we use importance sampling between sampling SDEs to obtain an explicit policy-gradient estimator on trajectory space. The estimator contains the stochastic Itô integral underlying Flow-GRPO-type updates; we derive an equivalent variance-reduced value-gradient form that recovers the forward-matching structure of AWM and DiffusionNFT. This identifies the empirical gap between these method families as a variance-reduction effect rather than a difference in RL principle. The derivation yields a unified design space organized by value-gradient estimation, weight functions, and sampling choices. Within this space, we propose a multi-sample KDE value-gradient estimator that reuses rollout groups, together with scale-bounded weight families that retain stable existing recipes while excluding singular ones. Experiments on SD3.5-M and Qwen-Image models validate the variance-reduction explanation and show that the resulting recipe improves over prior diffusion-RL baselines.
Yixian Xu, Yuanrui Zhang, Shengjie Luo +2
State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · ByteDance Seed