NSPG — Noisy-Space Policy Gradient for Diffusion Policies in Offline RL

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.

Motivation

A mismatch between two spaces

A diffusion policy generates actions by iteratively denoising a sequence of noisy latents into a clean sample, so optimization takes place in a noisy latent space. Reinforcement-learning objectives, however, are defined over the executed actions that interact with the environment. Prior methods either treat the diffusion model purely as an action sampler or evaluate a critic directly on noisy latents. We instead define the connection between the two spaces explicitly through $Q_{\text{noisy}}$.

where optimization happens

Noisy latent space

Policy improvement, sampling, and the denoising trajectory all live here. Latents $x^{(k)}$ share the action's dimension but are not actions — the environment never sees them.

$x^{(T)} \rightarrow \cdots \rightarrow x^{(0)}$

$Q_{\text{noisy}}$

bridges via induced actions

where value is defined

Clean action space

Rewards, transitions, and the critic $Q^{\pi}(s,a)$ are defined only over executed actions $a \in \mathcal{A}$. This is the only space with genuine MDP semantics.

$Q^{\pi}(s,a) = \mathbb{E}\!\left[\textstyle\sum_t \gamma^t r(s_t,a_t)\right]$

Method

How NSPG works

1 Background

Diffusion policies act through noisy latents

A diffusion policy does not emit an action in one pass. It starts from noise and runs a reverse denoising trajectory $x^{(T)} \rightarrow \cdots \rightarrow x^{(0)}$, where only the final sample is executed as $a=x^{(0)}$.

The key point is that an intermediate latent $x^{(k)}$ is not an action in the MDP. It is a conditioning point in the generation process, and after the remaining denoising steps it can induce several compatible clean actions.

Design principle. Value belongs to executable actions. NSPG keeps the critic in clean action space, then lifts the resulting signal back to noisy space.

latent x⁽ᵏ⁾ → compatible actions

action dim 1

action dim 2

Reading the plot. The central marker is one fixed latent $x^{(k)}$. Dashed trajectories show possible denoising outcomes, dots are sampled clean actions, and contours mark the induced clean-action distribution.

2 Core idea

Value a latent by the actions it can execute

NSPG gives a noisy latent a value only through the clean actions it can produce. Instead of asking the critic to score $x^{(k)}$ directly, we evaluate executable actions and average their returns under the distribution induced by that latent.

Definition 3.1 · noisy-space action value

$$\begin{aligned} Q_{\mathrm{noisy}}(s,x^{(k)},k) &= \mathbb{E}_{a\sim\pi_\theta(\cdot\mid s,x^{(k)},k)} \left[Q^{\pi_\theta}(s,a)\right]. \end{aligned}$$

The gradient then becomes an advantage-weighted score in noisy space: compatible actions with above-average value reinforce the latent direction, while below-average actions are suppressed.

Theorem 3.4 · noisy-space policy gradient

$$\begin{aligned} \nabla_{x^{(k)}} Q_{\mathrm{noisy}} &= \mathbb{E}_{a\sim\pi_\theta} \left[A_{\mathrm{noisy}}(s,a,x^{(k)},k)\right. \\ &\qquad\left.\cdot\nabla_{x^{(k)}}\log q_k(x^{(k)}\mid a)\right]. \end{aligned}$$

This separates NSPG from decoded guidance, as used for example in Diffusion Q-Learning-style updates. Decoding first commits to one clean estimate $\hat a(x^{(k)})$ and then optimizes $Q(s,\hat a)$; NSPG keeps the full compatible-action distribution, where multimodal behavior and high-value modes can diverge.

NSPG noisy-space gradient field — **Conceptual contrast.** NSPG preserves the compatible-action distribution; decoded guidance commits to one clean proxy before value optimization.

Decoded-action gradient field — **Conceptual contrast.** NSPG preserves the compatible-action distribution; decoded guidance commits to one clean proxy before value optimization.

3 Framework

NSAC: actor–critic in two spaces

Noisy-Space Actor–Critic alternates between a pessimistic action-space critic (a lower-confidence-bound over an ensemble) and a diffusion actor improved by the noisy-space gradient.

The improvement objective rewrites into a denoising regression with a value-guidance term, so the actor trains exactly like a diffusion model — with no backpropagation through the denoising trajectory, and the critic never evaluated on noisy latents.

Each iteration updates the critic toward a conservative Bellman target, then nudges the actor along $\nabla_{x^{(k)}}Q_{\text{noisy}}$. Value estimation stays entirely in action space; policy improvement happens in noisy latent space.

Algorithm 1 · NSAC

Require: dataset $\mathcal{D}$; actor $\pi_\theta$; critic ensemble $\{Q_{\phi_h}\}$; $\lambda$, $\eta$, $\rho$.

Initialize actor $\theta$ and critics $\{\phi_h\}$.
while not converged do
1. Sample minibatch $(s,a,r,s') \sim \mathcal{D}$.
2. Sample step $k$, noise $\epsilon$; form $x^{(k)}$ via $q(x^{(k)}\mid a)$.
3. Sample $a' \sim \pi_\theta(\cdot \mid s')$.
4. Update critic with the LCB Bellman objective.
5. Compute $\nabla_{x^{(k)}}Q_{\text{noisy}}$ (NSPG).
6. Update actor with the value-guided denoising loss.
end while
return $\pi_\theta$.

Value enters the actor only through $Q_{\text{noisy}}$ and its gradient.

Contributions

This work makes the following contributions to value-based optimization of diffusion policies in offline reinforcement learning:

We identify a mismatch in value-based optimization for diffusion policies, arising between noisy latent-space policy representations and action-space value functions.
We introduce a noisy-space action-value function $Q_{\text{noisy}}$ that maps diffusion latents to expected action-space returns, and derive a noisy-space policy gradient (NSPG) that enables value-based learning without evaluating critics on noisy variables or backpropagating through the denoising process.
We develop Noisy-Space Actor–Critic (NSAC), an actor–critic algorithm that performs policy improvement in noisy latent space while maintaining value estimation in action space.
We evaluate NSAC on state-based D4RL benchmarks and vision-based OGBench tasks, and report performance relative to prior methods across locomotion, navigation, and manipulation settings.

Results

Empirical evaluation

We evaluate NSAC on state-based D4RL benchmarks (locomotion, AntMaze, Adroit) and vision-based OGBench tasks, comparing against representative offline RL baselines spanning policy-regularization, conservative value estimation, and diffusion- and flow-based methods. DAC is included as the closest score-based diffusion actor–critic baseline. All results use four random seeds and a fixed budget of two million gradient steps.

Evaluation environments

Representative benchmark environments from the evaluation suite: dexterous manipulation and vision-based OGBench tasks.

environment visuals

AdroitPen-Cloned

OGBenchPuzzle 3×3

OGBenchCube double

aggregate comparison

Aggregate performance

Switch between benchmark groups to compare aggregate scores across methods. This view summarizes the headline results; the per-task tables and diagnostic plots are kept below for a more granular readout.

D4RL state-based aggregate normalized scores.

State-based control (D4RL)

The D4RL suite spans well-behaved locomotion, sparse-reward AntMaze navigation with many suboptimal trajectories, and dexterous Adroit manipulation with limited data coverage. These settings stress the coupling between value estimation and policy optimization differently — from smooth continuous control to long-horizon, multimodal decision-making.

Benchmark (total)	CQL	IQL	DQL	IDQL-A	DAC	NSAC
Locomotion (9 tasks)	698.5	692.6	791.2	739.2	836.4	861.7
AntMaze (4 tasks)	145.6	228.3	258.2	300.4	275.4	337.2
Adroit (2 tasks)	62.4	108.8	130.1	—	145.2	168.5
State-based total	906.5	1029.7	1179.5	—	1257.0	1367.4

Average normalized scores on D4RL (mean over four runs). Gains are most pronounced on AntMaze, where value guidance through executable actions matters most. Full per-task numbers with standard deviations are in the paper.

Additional plots and ablations Training curves · guidance strength · Monte-Carlo samples · pessimism · critic ensemble

The ablations isolate the key operational levers in NSAC. Guidance strength is placed here with the other diagnostics so the experimental story stays consolidated rather than fragmented across the page.

HalfCheetah Medium learning curve — HalfCheetah medium

Hopper Medium learning curve — HalfCheetah medium

Guidance strength ablation on Walker2d — Guidance · Walker2d

Guidance strength ablation on AntMaze — Guidance · Walker2d

Monte Carlo sample ablation on Walker2d — MC samples · Walker2d

Monte Carlo sample ablation on AntMaze — MC samples · Walker2d

Pessimism coefficient ablation on Walker2d — Pessimism · Walker2d

Pessimism coefficient ablation on AntMaze — Pessimism · Walker2d

Critic ensemble ablation on AntMaze — Ensemble · AntMaze

Critic ensemble ablation on Pen — Ensemble · AntMaze

Vision-based control (OGBench)

OGBench requires learning policies directly from high-dimensional visual observations, where representation learning, multimodal action distributions, and value-based improvement are jointly challenging. We follow recent flow-policy evaluations, using a targeted baseline set of strong methods for high-dimensional visual offline RL.

OGBench visual task	IQL	ReBRAC	FQL	DAC	NSAC
visual-cube-single	70	83	81	83.7	88
visual-cube-double	34	4	21	34	36
visual-puzzle-3x3	7	88	94	98.2	100
visual-puzzle-4x4	0	26	33	40	47.5
OGBench total	111	201	229	255.9	271.5

Success rates (%) on vision-based OGBench. NSAC is best or near-best on every task and highest in aggregate, extending the noisy-space formulation beyond state-based control.

Noisy-Space Policy Gradient for diffusion policies in offline RL