Noisy-Space Policy Gradient for Diffusion Policies in Offline Reinforcement Learning

Mahmoud Selim; Cristina Cipriani; Karl H. Johansson

cs.LGSep 7, 2026

Noisy-Space Policy Gradient for Diffusion Policies in Offline Reinforcement Learning

Authors: Mahmoud Selim, Cristina Cipriani, Karl H. Johansson

Organizations: TRATON CV AB, Stockholm, Sweden · KTH Royal Institute of Technology, Stockholm, Sweden

Abstract

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 $Q_{\mathrm{noisy}}$ , 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,\hat{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{\scriptstyle\pm 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{\scriptstyle\pm 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{\scriptstyle\pm 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{\scriptstyle\pm 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{\scriptstyle\pm 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{\scriptstyle\pm 1.5}$

Table 1: Average normalized scores of NSAC over four runs compared with baseline methods (mean

\pm

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 $\pm$ 12	83 $\pm$ 6	55 $\pm$ 8	49 $\pm$ 7	81 $\pm$ 12	83.7 $\pm$ 10	88 $\pm$ 7
visual-cube-double-play-singletask-task1-v0	34 $\pm$ 23	4 $\pm$ 4	6 $\pm$ 2	8 $\pm$ 6	21 $\pm$ 11	34 $\pm$ 12	36 $\pm$ 14
visual-puzzle-3x3-play-singletask-task1-v0	7 $\pm$ 15	88 $\pm$ 4	7 $\pm$ 2	100 $\pm$ 0	94 $\pm$ 1	98.2 $\pm$ 2	100 $\pm$ 0
visual-puzzle-4x4-play-singletask-task1-v0	0 $\pm$ 0	26 $\pm$ 6	0 $\pm$ 0	8 $\pm$ 15	33 $\pm$ 6	40 $\pm$ 5.7	47.5 $\pm$ 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

\pm

standard deviation). Best results, as well as results within 5% of the best performance for each task, are highlighted .

Panel of Figure 2 — 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.

Panel of Figure 3 — 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 $\beta_{k}$	Variance preserving (VP)
Critic ensemble size $H$	$10$
Batch size	$256$
Actor learning rate	$3\times 10^{-4}$
Critic learning rate	$3\times 10^{-4}$

Appendix

Table 3: Shared hyperparameters used across all tasks and experiments.

Domain	Task	$\boldsymbol{\lambda}$	$\boldsymbol{\eta}$	$\boldsymbol{\rho}$
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 $\nabla_{x^{(k)}}Q_{\mathrm{noisy}}(x^{(k)},k)$ , the gradient of the expected clean-action value under the diffusion-induced conditional. Bottom: the decoded baseline follows $\nabla_{x^{(k)}}Q(\hat{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.

Panel of Figure 5 — 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.

Panel of Figure 6 — Figure 6: Ablation on the pessimism coefficient $\rho$ across three representative environments. Curves report normalized return as a function of training steps for different values of $\rho$ .

Panel of Figure 7 — 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 ( $n_{\mathrm{samples}}=2$ )	3.2	$0.90\times$
NSAC ( $n_{\mathrm{samples}}=4$ )	3.4	$0.98\times$
NSAC ( $n_{\mathrm{samples}}=8$ )	3.9	$1.07\times$
NSAC ( $n_{\mathrm{samples}}=16$ )	5.8	$1.40\times$
DAC ( Fang et al., 2024 )	3.6	$1.00\times$

Appendix

Table 5: Computational cost of NSAC for different numbers of Monte Carlo samples, compared to DAC. Wall-clock time is normalized by DAC.

Diffusion Policy

Diffusion policies generate multimodal action sequences by iteratively denoising noise, enabling imitation, planning, and control in robotics.

See all 202 papers

Diffusion Policy through Conditional Proximal Policy Optimization
Mar 5, 2026Ben Liu, Shunpeng Yang, Hua ChenDiffusion Policy RL for Diffusion Models
Southern University of Science and Technology, Shenzhen, China · Limx Dynamics · Hong Kong University of Science and Technology, Hongkong, China +1
Sample-Efficient Diffusion-based Reinforcement Learning with Critic Guidance
May 28, 2026Shutong Ding, Zejia Zhong, Zhongyi Wang +4Diffusion Guidance Robotic RL
1ShanghaiTech University.
Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View
Aug 14, 2026Yixian Xu, Yuanrui Zhang, Shengjie Luo +2Diffusion Model Alignment Policy Gradient Methods
State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · ByteDance Seed

Mar 5, 2026J/K move · Enter open · S save

Diffusion Policy through Conditional Proximal Policy Optimization

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

Diffusion Policy RL for Diffusion Models Proximal Policy Optimization

Open paper AskPDF arXiv

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.

Abstract

Diffusion Policy

Explore similar work

Diffusion Policy through Conditional Proximal Policy Optimization

Sample-Efficient Diffusion-based Reinforcement Learning with Critic Guidance

Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View