🤖 machine learning

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

This paper introduces Noisy-Space Policy Gradient (NSPG), a novel framework that enables effective offline reinforcement learning with diffusion policies by deriving a KL-regularized objective over diffusion latents that optimizes using clean action-space value estimates without requiring backpropagation through the denoising process.

Original authors: Mahmoud Selim, Cristina Cipriani, Karl H. Johansson

Published 2026-09-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mahmoud Selim, Cristina Cipriani, Karl H. Johansson

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of artificial intelligence, there is a persistent challenge known as offline reinforcement learning. Imagine a student trying to learn how to play a complex video game, but they are forbidden from touching the controller. They can only watch hours of recorded footage of other players, some of whom were experts and some of whom made frequent mistakes. The goal is to learn a strategy that wins more often than the people in the videos, without ever risking a new move that might lead to a crash. This is difficult because the AI must learn from a fixed set of data without the safety net of trial and error. If the AI guesses a move that was never seen in the recordings, it might fail catastrophically, as it has no way to know if that move is actually good or just a dangerous hallucination.

To solve this, researchers have turned to a type of AI model called a diffusion policy. These models are exceptionally good at understanding complex patterns, such as the many different ways a human hand might grasp a tool or a robot might navigate a maze. They work by starting with a chaotic, noisy guess and slowly refining it, step by step, until it becomes a clear, usable action. However, a fundamental problem arises when trying to teach these models to win: the model makes its decisions in a hidden, noisy space, but the rewards it receives are based on the final, clean actions it actually takes. It is like trying to grade a student's essay by looking at their messy, crossed-out drafts rather than the final polished page. The feedback loop is broken, and the AI struggles to learn which of its hidden steps led to success.

A team of researchers has now bridged this gap with a new method called Noisy-Space Policy Gradient. Instead of trying to force the noisy, hidden steps to match the final rewards directly, they created a new way to assign value to those hidden steps. They realized that a single noisy guess does not correspond to just one final action, but rather to a whole cloud of possible actions that could emerge from it. By calculating the average reward of all the possible clean actions that could come from a specific noisy guess, they created a fair and accurate score for that guess. This score tells the AI exactly how good a particular hidden step is, based on the potential outcomes it could produce, rather than forcing it to commit to a single, possibly wrong, final answer.

The researchers tested this approach on a wide variety of tasks, ranging from simple robotic movements to complex visual puzzles. In these experiments, the new method consistently outperformed previous techniques, especially in difficult scenarios where the data was sparse or the tasks required long chains of precise actions. For instance, in tasks involving navigating large mazes or manipulating delicate objects, the new method achieved significantly higher success rates than earlier models. It proved particularly effective in situations where the AI had to choose between a common, safe action and a rare, high-reward action, a choice that often confused older models. By looking at the full range of possibilities rather than a single path, the AI learned to aim for the high-reward outcomes more reliably.

One of the most significant aspects of this work is how it handles the relationship between the AI's internal thinking process and the real world. Previous methods often tried to evaluate the AI's hidden steps directly, which led to confusion and instability. Others tried to simplify the AI's thinking process to make it easier to evaluate, but this often stripped away the very complexity that made the AI powerful in the first place. The new approach avoids these pitfalls by keeping the evaluation strictly in the world of real actions while allowing the AI to think in its complex, noisy way. It acts as a translator, ensuring that the feedback the AI receives is always grounded in reality, even though the AI is learning in a hidden space.

The results suggest that this method provides a solid foundation for training advanced AI systems on historical data. The researchers found that the method remained stable and efficient, requiring only a small number of calculations to estimate the value of each step. This efficiency is crucial, as it means the method can be applied to large, complex problems without becoming too slow to be practical. The team also explored how sensitive the system was to different settings, finding that it worked well across a broad range of conditions without needing constant, delicate tuning. This robustness makes it a promising tool for future applications in robotics and automation, where learning from past data is often the only option.

Ultimately, this work resolves a long-standing disconnect in how AI learns from experience. It demonstrates that by properly understanding the relationship between hidden thoughts and visible actions, we can teach complex systems to improve without ever stepping outside the bounds of their training data. The method does not rely on magic or guesswork; it relies on a clear, mathematical understanding of how noise transforms into action. As the field of artificial intelligence continues to evolve, approaches that can safely and effectively learn from the past will be essential for building systems that are both capable and reliable. This new technique offers a principled way forward, turning the messy process of learning from static data into a clear path toward better decision-making.

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

1. Problem Statement

Offline Reinforcement Learning (RL) aims to learn effective policies from fixed datasets without further environment interaction. A primary challenge in this setting is the "distributional shift" problem: policies that extrapolate beyond the support of the dataset often suffer from biased value estimates. To address the complex, multimodal action distributions found in offline data, expressive generative models like diffusion policies have emerged. These policies generate actions through an iterative denoising process, mapping a sequence of noisy latent variables to a final clean action.

However, integrating diffusion policies with value-based RL introduces a fundamental structural mismatch:

Optimization Space: Policy optimization occurs in a noisy latent space (the sequence of variables $x^{(k)}$ during the denoising process).
Value Definition: Reinforcement learning objectives (value functions, advantages, gradients) are defined over executed, clean actions ( $a \in \mathcal{A}$ ) that interact with the environment.

Existing approaches to bridge this gap suffer from significant drawbacks:

Action-Space Sampling: Treating diffusion models merely as samplers and reweighting actions based on an action-space critic bypasses latent-space guidance.
End-to-End Backpropagation: Optimizing policies by backpropagating value gradients through the entire denoising trajectory creates long stochastic computation graphs, leading to computational overhead and instability.
Direct Latent Evaluation: Injecting value guidance directly into noisy latents by evaluating Q-functions on intermediate variables ignores the fact that these variables are not the actions that determine rewards.

The paper identifies a gap in defining a value formulation that bridges noisy latent-space optimization and action-space returns in a principled, semantically coherent manner.

2. Methodology

The authors propose a framework centered on a Noisy-Space Action-Value Function and a derived Noisy-Space Policy Gradient (NSPG).

2.1 Noisy-Space Action-Value Function ( $Q_{noisy}$ )

The core theoretical contribution is the definition of a value function for noisy latent variables that does not assign intrinsic value to the noise itself but rather to the distribution of clean actions it induces.

Definition: For a state $s$ , diffusion step $k$ , and noisy latent $x^{(k)}$ , the noisy-space action-value function is defined as the expected action-space return induced by the diffusion policy conditioned on $(s, x^{(k)}, k)$ :
$Q_{noisy}^\pi(s, x^{(k)}, k) \triangleq \mathbb{E}_{a \sim \pi(\cdot|s, x^{(k)}, k)} [Q^\pi(s, a)]$
Semantic Correctness: The authors prove that $Q_{noisy}$ exactly equals the expected discounted return obtained when the policy is initialized at $(s, x^{(k)}, k)$ and executed in the environment. This ensures the value function is consistent with the environment interface, which only observes clean actions.
Surrogate Semantics: To make conditioning on $x^{(k)}$ well-defined without altering the policy's reverse-time dynamics, the authors introduce a surrogate joint distribution. This models the relationship between a clean action $a$ and a noisy latent $x^{(k)}$ via the forward diffusion process (corrupting a clean action). This allows the derivation of gradients without changing the policy's generative mechanism.

2.2 Noisy-Space Policy Gradient (NSPG)

Using the surrogate semantics, the authors derive a closed-form expression for the gradient of $Q_{noisy}$ with respect to the noisy latent $x^{(k)}$ .

Gradient Formulation:
$\nabla_{x^{(k)}} Q_{noisy}^\pi(s, x^{(k)}, k) = \mathbb{E}_{a \sim \pi} \left[ A_{noisy}^\pi(s, a, x^{(k)}, k) \nabla_{x^{(k)}} \log q_k(x^{(k)} | a) \right]$
Where:
- $A_{noisy}^\pi = Q^\pi(s, a) - Q_{noisy}^\pi(s, x^{(k)}, k)$ is the noisy-space advantage.
- $\nabla_{x^{(k)}} \log q_k(x^{(k)} | a)$ is the score function of the forward diffusion kernel.
Key Distinction: Unlike approaches that map a noisy latent to a single decoded action $\hat{a}$ and differentiate $Q(s, \hat{a})$ , NSPG differentiates the expected value over the distribution of compatible clean actions. This allows the gradient to reinforce directions in noisy space associated with high-value compatible actions while suppressing low-value ones, even if the most probable action is not the highest value.
No Backpropagation: Crucially, this gradient is computed using only clean action-space value estimates ( $Q^\pi$ ) and the forward diffusion kernel. It avoids backpropagation through the full denoising trajectory, significantly improving stability and efficiency.

2.3 Noisy-Space Actor-Critic (NSAC)

The authors instantiate this framework into NSAC, a practical offline RL algorithm:

Critic: Learns a conservative action-space Q-function ( $Q_\phi$ ) using offline data and a Lower-Confidence Bound (LCB) target to mitigate overestimation and distributional shift. The critic is never evaluated on noisy latents.
Actor: Updates the diffusion policy using a regression objective that combines the standard denoising loss (behavior cloning) with a value-guidance term derived from the NSPG.
- The actor loss includes a term: $\lambda \cdot \epsilon_\theta^\top \nabla_{x^{(k)}} Q_{noisy}$ .
- This allows the policy to improve in the latent space while remaining anchored to the data distribution.

3. Key Contributions

Identification of the Mismatch: The paper formally identifies the semantic gap between noisy latent optimization and action-space value functions in diffusion-based RL.
Theoretical Formulation: Introduction of $Q_{noisy}$ , a value function that assigns value to noisy latents strictly through the distribution of clean actions they generate, ensuring semantic consistency with the MDP.
NSPG Derivation: Derivation of a policy gradient that optimizes noisy latents using only clean action-space value estimates, eliminating the need for backpropagation through the denoising process.
NSAC Algorithm: Development of a stable actor-critic algorithm that implements NSPG, demonstrating that expressive diffusion policies can be trained effectively in offline settings.

4. Experimental Results

The authors evaluate NSAC on state-based D4RL benchmarks (locomotion, AntMaze, Adroit) and vision-based OGBench tasks.

Performance: NSAC achieves state-of-the-art or competitive performance across diverse environments.
- Locomotion: Outperforms or matches prior methods (e.g., CQL, IQL, Diffuser) on HalfCheetah, Hopper, and Walker2d. Notably, it achieves ~10% gains on HalfCheetah medium-replay.
- AntMaze: Demonstrates substantial improvements over existing methods, particularly in large maze environments where it outperforms the closest baseline (DAC) by over 50% relative performance.
- Adroit: Achieves the strongest performance on Pen-Cloned-v1 among evaluated methods.
- Vision (OGBench): Attains the highest aggregate score (271.5) on visual tasks, outperforming flow-based policies and DAC.
Ablation Studies:
- Guidance Strength ( $\lambda$ ): The method is robust across a range of guidance strengths, with intermediate values yielding optimal performance. Overly aggressive guidance degrades performance by pushing the policy out of the data support.
- Monte Carlo Approximation: The noisy-space expectation can be estimated efficiently with a small number of samples (e.g., 2-4), incurring minimal computational overhead compared to DAC.
- Pessimism: The algorithm remains robust across a broad range of pessimism coefficients.

5. Significance and Claims

The paper claims that NSPG provides a principled and effective basis for training diffusion policies in offline RL. By resolving the mismatch between latent-space optimization and action-space value, the method avoids the instability of end-to-end backpropagation and the semantic incoherence of evaluating critics on noisy variables.

The authors emphasize that their approach:

Preserves the expressivity of multi-step diffusion policies without sacrificing tractability.
Offers a modest but significant improvement in stability and performance, particularly in challenging, long-horizon, and multimodal settings where value estimation and policy optimization are tightly coupled.
Suggests that the framework could naturally extend to online and offline-to-online learning, though this is left for future work.

The work does not claim to solve all offline RL problems but specifically addresses the integration of diffusion policies with value-based objectives, offering a theoretically grounded alternative to heuristic sampling or unstable end-to-end training.