Diffusion policies improve offline reinforcement learning with noisy action values

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

Machine LearningArtificial IntelligenceRobotics

Summary

In robotics and AI, teaching machines to make decisions in complex environments is hard, especially when learning from past experiences without active exploration. This paper presents a new way to improve decision-making policies that generate actions using a process called diffusion, which involves adding and removing noise. The authors introduce a technique to better estimate values of actions by considering the noisy steps involved, enabling more effective learning without needing to backtrack through complex noise removal steps. Their tests on simulated control tasks and visual benchmarks show this method helps train better policies for machines learning from previously collected data.

What this means in practice

  • For robotics engineers: Create more reliable control policies for robots using offline data without requiring complex noise gradient calculations.
  • For autonomous vehicle developers: Develop better decision-making systems that learn from recorded driving data by handling noisy action representations efficiently.

Authors

Mahmoud Selim, Cristina Cipriani, Karl H. Johansson

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/
View PDFOpen arXiv