PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Mahmoud Selim, Cristina Cipriani, Karl H. Johansson

arXiv:2609.06882v1cs.LGcs.AIcs.RO

TL;DR

Diffusion policies create a mismatch between noisy latent-space optimization and action-space values in offline RL. The paper resolves this with Qnoisy, NSPG, and a diffusion-compatible improvement objective, with NSAC showing strong performance across state-based and vision-based benchmarks. The work’s scope includes benchmark-version differences and leaves systematic online or offline-to-online evaluation for future work.

  • Problem

    Diffusion policies optimize through noisy latent variables while reinforcement-learning values are defined over executed actions, leaving a fundamental gap in value-based optimization.

  • Method

    The paper defines Qnoisy through expected clean-action returns, derives NSPG, and instantiates it as NSAC without backpropagating through the full denoising process.

  • Results

    NSAC achieves strong empirical performance across state-based D4RL and vision-based OGBench benchmarks, with especially clear gains in sparse-reward, long-horizon, and multimodal settings.

  • Takeaways & Limitations

    Noisy-space value objectives provide a principled basis for training expressive diffusion policies in offline reinforcement learning.

  • Takeaways & Limitations

    Some scores use different benchmark versions, and systematic extensions to online and offline-to-online learning remain future work.

Abstract

from arXiv · show

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/

1. Introduction

Diffusion policies represent multimodal actions, but offline RL must optimize noisy latent representations using values defined over executed actions. The paper introduces NSPG and NSAC to bridge this mismatch while avoiding full denoising backpropagation.

  • Offline RL must improve policies within dataset support because extrapolation beyond observed state–action pairs can produce severely biased values.
  • Diffusion policies represent rich, multimodal action distributions through iterative denoising across latent variables that are not directly observable by the environment.
  • Their latent-space optimization conflicts with RL value functions defined over executed, noise-free actions, creating a fundamental latent–action mismatch.
  • NSPG defines noisy-space values through expected action-space returns induced by clean actions, while querying critics only on executable actions.
  • NSPG avoids backpropagation through the full denoising process, and NSAC instantiates the objective as a stable noisy-latent actor–critic algorithm.
  • NSAC demonstrates consistent performance gains over prior methods on state-based D4RL and vision-based OGBench tasks.

2. Preliminaries

The preliminaries define offline RL, action-space value functions, and diffusion policies. Diffusion generation proceeds through noisy latent variables, but only the final clean action affects the environment.

  • An MDP comprises state and action spaces, transitions, rewards, an initial-state distribution, and a discount factor governing expected returns.
  • An action-value function assigns the expected discounted return to an executed action in a state, followed by the policy thereafter.
  • Offline RL learns from a fixed dataset collected by an unknown behavior policy without additional environment interaction.
  • Diffusion policies generate actions by sampling an initial prior and applying state-conditioned reverse-time denoising over multiple latent variables.
  • The denoising trajectory ends in a clean action, whereas intermediate noisy latents are not MDP actions and never directly affect rewards or transitions.

3. Noisy-Space Value Functions and Policy Optimization

The paper defines noisy-space values as conditional expected clean-action values, establishes their semantic correctness, and derives NSPG as a local score-weighted gradient for latent optimization.

  • 3.1. Noisy-Space Action-Value Function: Qnoisy assigns each noisy latent the expected action-space value over executable actions induced by the diffusion policy.
  • 3.1. Noisy-Space Action-Value Function: Because values are assigned through generated actions rather than latent variables themselves, Qnoisy provides a well-defined latent-space optimization objective.
  • 3.1. Noisy-Space Action-Value Function: The semantic correctness result identifies Qnoisy with the expected discounted return obtained by initializing the diffusion policy at a noisy latent and executing it in the environment.
  • 3.2. Surrogate Semantics for Noisy Latents: A surrogate noisy joint samples a clean action from the policy’s marginal distribution and corrupts it with the forward diffusion kernel, making conditioning on noisy latents well-defined.
  • 3.2. Surrogate Semantics for Noisy Latents: The surrogate semantics interface with value functions without constraining the diffusion policy or changing its reverse-time dynamics.
  • 3.3. Noisy-Space Policy Gradient (NSPG): NSPG provides a local noisy-space improvement direction that avoids backpropagation through the full denoising trajectory.
  • 3.3. Noisy-Space Policy Gradient (NSPG): Unlike decoded-action guidance, NSPG evaluates expected clean-action value over compatible actions rather than differentiating value after mapping each latent to one action.
  • 3.3. Noisy-Space Policy Gradient (NSPG): Its noisy-space advantage acts as a signed weight on the forward-diffusion score, amplifying high-value action contributions and suppressing below-average ones.

4. Learning Diffusion Policies with Noisy-Space Actor–Critic (NSAC)

NSAC performs policy improvement in diffusion’s noisy latent space while estimating value strictly on executable actions. It combines KL-regularized score-based improvement, a diffusion-compatible actor regression loss, and pessimistic action-space critic learning.

  • NSAC overview: NSAC alternates noisy-space diffusion actor improvement with conservative action-space critic learning from offline data.The critic remains confined to executable actions, while its information reaches the actor through Qnoisy and its gradient.
  • Actor: Noisy-Space Policy Improvement: Directly optimizing the KL-regularized policy objective is infeasible for diffusion policies because their policy densities are intractable.The method therefore rewrites policy improvement using score functions.
  • Actor: Noisy-Space Policy Improvement: NSAC extends score-based policy improvement to noisy perturbations induced by the forward diffusion process.Its value term uses Qnoisy rather than applying an action-value function directly to noisy latents, preserving the score-function structure.
  • Actor: Noisy-Space Policy Improvement: The noisy-space policy improvement objective admits a diffusion-compatible regression form, avoiding backpropagation through the full denoising process.The actor loss includes guidance strength λ, which controls value-guidance influence relative to denoising.
  • Critic: Pessimistic Action-Space Value Learning: The critic is an action-space Q-function trained with pessimistic value estimation and a lower-confidence-bound target.An ensemble uses expectation and variance, with ρ > 0 controlling the pessimism level.
  • NSAC algorithm: Algorithm 1 samples offline transitions, diffusion steps, and noise; updates the critic; computes ∇x(k)Qnoisy; and updates the diffusion actor.This procedure returns a diffusion policy after alternating critic and actor updates.

5. Prior Work

Prior offline-RL work regularizes policies, estimates values conservatively, or constrains policy updates. Diffusion-policy methods either use generative models as samplers, backpropagate through generation, or simplify the policy class at a cost to expressiveness.

  • Offline Reinforcement Learning: Offline-RL policy-regularization methods constrain improvement toward behavior policies using generative, divergence-based, or KL-regularized approaches.Examples include BCQ, BEAR, BRAC, and TD3+BC.
  • Offline Reinforcement Learning: Conservative value-based methods address overestimation on out-of-distribution actions through pessimistic objectives, asymmetric losses, or in-sample maximization.Representative methods include CQL, IQL, IDQL, and IVR.
  • Offline Reinforcement Learning: Trust-region and KL-regularized methods constrain the policy improvement step to mitigate error amplification from bootstrapped value estimates.This line of work motivates adopting KL-regularized policy iteration.
  • Diffusion and Flow Policies for Offline Reinforcement Learning: Some diffusion methods use generative models primarily as behavior samplers or planners rather than directly optimizing them with reinforcement-learning objectives.Diffuser is an example of return-conditioned planning with guided sampling.
  • Diffusion and Flow Policies for Offline Reinforcement Learning: Other methods directly optimize diffusion or flow policies with value-based objectives, often requiring backpropagation through the generative process.This adds computational overhead and potential instability.
  • Diffusion and Flow Policies for Offline Reinforcement Learning: One-step flow policies improve tractability by avoiding iterative denoising but trade off part of the expressive capacity motivating multi-step diffusion policies.FQL exemplifies this policy-class simplification.

6. Experiments

Experiments evaluate NSAC against offline RL and diffusion-policy baselines on state-based and vision-based tasks, then test guidance and Monte Carlo sampling choices. NSAC performs strongly across benchmarks, with especially large gains on challenging AntMaze tasks and robust behavior under moderate guidance and few samples.

  • Benchmark setup: NSAC is evaluated against broad offline RL, Gaussian-policy, diffusion-policy, and flow-policy baselines across locomotion, AntMaze, Adroit, and visual OGBench tasks.DAC is the closest score-based diffusion actor–critic baseline because it also avoids denoising backpropagation but evaluates values directly on noisy latents.
  • State-based results: NSAC achieves strong and consistent performance, outperforming or matching prior methods across diverse state-based offline RL environments.The evaluation includes locomotion, AntMaze, and Adroit tasks using average normalized scores.
  • State-based results: Approximately 10% gains over prior methods on HalfCheetah medium and medium-replay show strong performance in long-horizon, action-sensitive locomotion.NSAC also achieves the strongest performance among evaluated methods on Pen-Cloned-v1.
  • State-based results: Over 50% relative improvement on AntMaze-large-play and AntMaze-large-diverse marks NSAC’s largest reported state-based gains over DAC.These tasks combine sparse rewards with numerous suboptimal trajectories and large maze environments.
  • Vision-based results: From 253.2 to 271.5, NSAC improves the strongest prior aggregate score on four vision-based OGBench tasks.It achieves the best or near-best performance on all four tasks, with particularly clear gains on visual cube and 4x4 puzzle tasks.
  • Ablation studies: NSAC remains stable with few Monte Carlo samples, and increasing the sample count produces diminishing returns.The noisy-space expectation can therefore be estimated efficiently without expensive sampling procedures.
  • Ablation studies: Intermediate guidance values λ ∈{1, 3, 10} yield strongest performance, while excessively strong guidance degrades learning.Moderate guidance balances policy improvement with behavioral regularization and avoids driving the policy away from offline-data support.

7. Discussion and Conclusion

The paper frames NSPG as a principled way to optimize diffusion policies through values assigned to noisy latents via induced clean-action distributions. NSAC performs strongly on state-based and vision-based benchmarks, while extensions beyond offline RL remain future work.

  • Discussion and conclusion: NSPG assigns noisy latents values through the distribution of clean actions they induce, optimizing expected action-space returns while querying critics only on executable actions.NSAC builds on this formulation for state-based and vision-based offline RL.
  • Discussion and conclusion: NSAC achieves strong empirical performance across state-based and vision-based benchmarks, especially in sparse-reward, long-horizon, and multimodal settings.The conclusion presents these settings as the clearest areas of benefit.
  • Discussion and conclusion: Systematic evaluation of NSPG for online and offline-to-online learning is left to future work.The paper describes these directions as potential extensions beyond its offline RL setting.

Impact Statement

The work targets more reliable and data-efficient diffusion-policy learning in offline RL using standard simulated benchmarks. It identifies no direct negative societal consequences requiring specific discussion beyond broader reinforcement-learning and autonomous-decision-making considerations.

  • Impact statement: The proposed methods target more reliable and data-efficient policy learning in offline reinforcement learning.They are evaluated on standard simulated benchmarks.
  • Impact statement: The paper identifies no direct negative societal consequences requiring specific discussion beyond broader considerations for reinforcement learning and autonomous decision-making.

A.1. Proof of Proposition 3.2

The proof shows that conditioning on a noisy latent affects the environment only through the induced distribution of the executed first action. Averaging the standard action-value over that conditional action distribution establishes the noisy-space value identity.

  • Assumptions: Under the environment assumption, transitions and rewards depend on executed actions, not directly on noisy latents.This lets the proof isolate the noisy latent’s effect to the first action distribution.
  • Trajectory construction: The proof fixes a state and noisy latent, samples the first action conditionally, then follows the policy’s marginal behavior thereafter.Subsequent states and actions are generated through the ordinary environment kernel and state-conditional policy marginals.
  • Value decomposition: Conditioning additionally on the executed first action yields the standard action-value function.Once that action is fixed, the later trajectory distribution no longer depends directly on the noisy latent.
  • Conclusion: Taking expectation over the conditional first-action distribution and applying total expectation proves the proposition.The final step integrates the ordinary action-value over actions induced by the noisy latent.
  • Gradient derivation: The theorem’s gradient derivation differentiates the noisy-space value under integrability and differentiation-under-the-integral assumptions.The score-function identity is then expressed using the forward diffusion kernel and its marginal likelihood.
  • Implementation: NSPG estimates the noisy-space gradient without differentiating through the reverse diffusion trajectory or evaluating critics on noisy latents.The implementation uses Monte Carlo samples of executed actions and averages score-weighted advantages.

B.2. Evaluation Protocol

The evaluation uses a fixed, seed-averaged actor–critic protocol that separates noisy-latent optimization from clean action-value estimation. NSPG estimates noisy-space values by reverse-diffusion Monte Carlo while retaining diffusion denoising and behavior-cloning objectives.

  • Evaluation protocol: Policies are evaluated every 10,000 gradient steps over 10 episodes, with final scores averaged over the last 50,000 steps and all seeds.Training uses fixed budgets without early-stopping selection.
  • Architecture: NSPG maintains a strict separation between latent-space policy optimization and action-space value estimation.The critic ensemble operates on executed actions, while diffusion components guide actor updates in noisy space.
  • Architecture: The diffusion actor uses a DDPM denoiser with T = 5 steps and a variance-preserving noise schedule.The actor predicts forward-process noise from state, noisy latent, and timestep inputs.
  • Noisy-space estimation: NSPG estimates Qnoisy by reverse-decoding multiple executed actions, evaluating each with the target critic, and averaging their action-space values.The resulting gradient uses score-weighted noisy-space advantages averaged across Monte Carlo samples.
  • Objective: The gradient estimator avoids reverse-process backpropagation and noisy-latent critic evaluation while preserving action-space value semantics.The actor combines standard denoising behavior cloning with soft noisy-space value guidance.
  • Stabilization: Critic targets use pessimistic lower-confidence-bound aggregation, while target networks use exponential moving averages for stability.These choices address overestimation and distributional shift in offline training.
  • Diagnostic visualization: Figure 4 compares NSPG’s expected clean-action gradient with a decoded baseline across denoising steps in a multimodal toy action distribution.The illustration shows gradient fields and ascent traces from noisy to clean latents; its behavior is diagnostic rather than a convergence claim.

C.1. Training Dynamics

NSAC shows stable learning across 15 environments, including challenging AntMaze and dexterous manipulation tasks. Performance remains robust across pessimism levels and improves with small critic ensembles before saturating.

  • Overall dynamics: Across 15 evaluation environments, NSAC exhibits stable training and consistent convergence, though harder tasks learn more slowly.The reported environments include locomotion, AntMaze, and Adroit tasks.
  • Pessimism sensitivity: Moderate pessimism coefficients provide the best balance between avoiding overestimation and preserving learning signal.Both excessively small and excessively large coefficients can degrade performance.
  • Critic ensemble size: Performance generally improves with small critic ensembles but saturates beyond a modest ensemble size.This supports ensemble-based pessimism without requiring large ensembles.

C.4. Training Time and Computational Cost.

NSAC has computational cost comparable to or lower than DAC under matched implementation settings. Increasing Monte Carlo samples raises cost while retaining similar wall-clock efficiency at the reported sample counts.

  • Computational cost: With two Monte Carlo samples, NSAC requires 3.2 ms per update and 0.9× DAC’s wall-clock time.The comparison uses matched architecture, diffusion horizon, and hardware.
  • Computational cost: With four Monte Carlo samples, NSAC requires 3.4 ms per update and 0.98× DAC’s relative wall-clock time.Additional samples increase cost but remain comparable to DAC in the reported comparison.
  • Related evaluations: The accompanying curves report normalized return across environments, while separate ablations vary pessimism and critic ensemble size.These figures provide training and sensitivity context rather than direct timing measurements.
Loading 2609.06882v1…