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

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/

Offline reinforcement learning trains decision-making systems from existing data, without allowing new trial-and-error interaction. Diffusion policies are promising for this setting because they can represent complex, multimodal behaviors, but they generate actions through several noisy intermediate steps, while rewards are only defined for the final executed actions. We address this mismatch by defining a noisy-space value function: each noisy intermediate variable is assigned the reward expected value of the clean actions it can produce. From this definition, we derive a noisy-space policy gradient, NSPG, which tells the diffusion policy how to update its noisy intermediate steps using only value estimates of executable actions. This leads to NSAC, an actor-critic algorithm that improves diffusion policies without evaluating critics on noisy variables or backpropagating through the full denoising process. Across control, navigation, manipulation, and vision-based benchmarks, NSAC improves over strong prior methods. These results suggest that diffusion policies can be trained more reliably when their noisy internal representations are linked explicitly to the action-space rewards used by reinforcement learning.