Steering Your Diffusion Policy with Latent Space Reinforcement Learning

TL;DR

This work introduces Diffusion Steering via Reinforcement Learning (Dsrl), a method to adapt pretrained diffusion-based behavioral-cloning policies by guiding their latent-noise inputs rather than updating weights. By formulating a transformed latent-action MDP and employing an efficient noise-aliasing RL scheme, Dsrl achieves fast, sample-efficient real-world policy improvement with black-box access, applicable to online, offline, and offline-to-online settings. The approach yields state-of-the-art or near-state-of-the-art performance across simulated benchmarks, real-world robotics, and steering of generalist policies, often requiring far fewer online samples than baselines. These results highlight Dsrl’s practical potential for autonomous, data-efficient refinement of diffusion-based robotic policies in complex, real-world tasks.

Abstract

Robotic control policies learned from human demonstrations have achieved impressive results in many real-world applications. However, in scenarios where initial performance is not satisfactory, as is often the case in novel open-world settings, such behavioral cloning (BC)-learned policies typically require collecting additional human demonstrations to further improve their behavior -- an expensive and time-consuming process. In contrast, reinforcement learning (RL) holds the promise of enabling autonomous online policy improvement, but often falls short of achieving this due to the large number of samples it typically requires. In this work we take steps towards enabling fast autonomous adaptation of BC-trained policies via efficient real-world RL. Focusing in particular on diffusion policies -- a state-of-the-art BC methodology -- we propose diffusion steering via reinforcement learning (DSRL): adapting the BC policy by running RL over its latent-noise space. We show that DSRL is highly sample efficient, requires only black-box access to the BC policy, and enables effective real-world autonomous policy improvement. Furthermore, DSRL avoids many of the challenges associated with finetuning diffusion policies, obviating the need to modify the weights of the base policy at all. We demonstrate DSRL on simulated benchmarks, real-world robotic tasks, and for adapting pretrained generalist policies, illustrating its sample efficiency and effective performance at real-world policy improvement.

Figures (22)

• Figure 1: Overview of our proposed approach, Diffusion Steering via Reinforcement Learning (Dsrl). Standard deployment of a BC-trained diffusion policy $\pi_{\mathrm{dp}}$ first samples noise $\bm{w} \sim \mathcal{N}(0,I)$ that is then denoised through the reverse diffusion process to produce an action $\bm{a}$. We propose modifying the initial distribution of $\bm{w}$ with an RL-trained latent-noise space policy $\pi^{ \mathcal{W}}$, that instead of choosing $\bm{w} \sim \mathcal{N}(0,I)$, chooses $\bm{w}$ to steer the distribution of actions produced by $\pi_{\mathrm{dp}}$ in a desirable way, enabling highly sample efficient real-world adaptation of robot policies.
• Figure 2: By selecting which point in the latent-noise space to denoise, we can steer the action produced by $\pi_{\mathrm{dp}}$ to a desired mode. At the same time, different points in the latent-noise space may in some cases be denoised to the same point in the original action space, allowing one to infer the behavior of noise actions $\bm{w}_2$ and $\bm{w}_3$ at $\bm{s}$ given only observation $(\bm{s},\bm{a}')$, forming the basis of our noise aliasing approach (\ref{['sec:noise_aliasing']}).
• Figure 3: Dsrl enables online adaptation of pretrained diffusion policies on OpenAI Gymbrockman2016openai.
• Figure 4: Dsrl enables online adaptation of pretrained diffusion policies on Robomimicrobomimic2021.
• Figure 5: Dsrl can make effective use of offline data to speed up online learning.