Steering Your Diffusion Policy with Latent Space Reinforcement Learning

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.

• The paper proposes Dsrl, which steers diffusion policies by modifying the latent noise distribution using reinforcement learning, enabling efficient policy adaptation.
• The method is sample-efficient and works with black-box access, avoiding modifications to policy weights while leveraging noise aliasing for improved learning.
• Experimental results on simulated and real-world robotic tasks show that Dsrl significantly outperforms traditional behavioral cloning and other RL variants.

Steering Your Diffusion Policy with Latent Space Reinforcement Learning

Introduction

The paper "Steering Your Diffusion Policy with Latent Space Reinforcement Learning" (2506.15799) addresses the challenge of adapting behavioral cloning (BC) policies for robotic control learned from human demonstrations. Traditional BC approaches do not utilize experience gathered during deployment to refine policy behavior, necessitating further demonstrations, which are costly and time-consuming. Reinforcement learning (RL) offers autonomous online policy improvement but often demands numerous samples. The authors propose an efficient approach called Diffusion Steering via Reinforcement Learning (Dsrl), exploiting the latent-noise space of diffusion policies for fast, effective policy adaptation.

Figure 1: The Dsrl framework modifies the initial noise distribution with an RL-trained policy to enhance real-world policy adaptation.

Methodology

Diffusion Steering Approach

The Dsrl technique modifies the initial noise distribution used by diffusion models in contrast to traditional methods that adjust policy weights. This reimagining of the diffusion process involves altering the input noise, which determines policy behavior. The latent-noise space policy, trained via RL, efficiently steers the diffusion policy to yield desired actions. Crucially, Dsrl operates with black-box access without requiring modification of the underlying policy weights, facilitating practical application and enabling adaptation through API access.

Figure 2: Dsrl's adaptation capabilities demonstrated on OpenAI Gym tasks.

Latent-Noise Space Optimization

Dsrl formalizes the latent-noise space optimization by leveraging RL techniques. It conceptualizes the noise space as a transformed MDP, allowing for standard policy optimization using RL algorithms. This approach circumvents the computational challenges of back-propagating through multi-step denoising, enabling sample-efficient training of auxiliary latent-space policies.

Noise Aliasing Strategy

Dsrl integrates a noise aliasing strategy that enhances learning efficiency. This approach exploits similarities in the denoised output of different latent noise values, thereby reducing exploration requirements. The actor-critic framework incorporates these dynamics, offering a robust pathway to offline and online learning.

Figure 3: Dsrl significantly improves the performance of a multi-task BC policy.

Experimental Results

Simulated Benchmarks

Dsrl was evaluated on OpenAI Gym environments and the Robomimic benchmark, demonstrating significant sample efficiency compared to existing methods. Dsrl efficiently steered diffusion policies to optimal behavior, outperforming variants such as Dipo and Idql.

Figure 4: Dsrl enables successful adaptation of diffusion policies in Robomimic tasks.

Real-World Applications

In practical settings, Dsrl demonstrated substantial improvements in robotic tasks, refining both single-task and multi-task policies with limited online interaction. Its adaptation capabilities on the Bridge V2 dataset with the WidowX robot and $\pi_0$ generalist policies highlight its real-world applicability.

Figure 5: Real-world demonstration of Dsrl adaptation capabilities.

Implications and Future Directions

Dsrl's framework presents a practical approach for real-world robotic policy adaptation, offering significant improvements in sample efficiency and ease of implementation. Future work could explore Dsrl's application beyond robotics, such as in image generation or protein modeling, leveraging its latent-space optimization potential.

Additionally, investigations into the theoretical underpinnings of diffusion steering could provide insights into the expressivity and structure of noise space, enhancing understanding and application across domains.

Conclusion

The paper introduces Dsrl as a novel approach for adapting BC-trained diffusion policies using latent-space RL. Dsrl showcases robust improvements in sample efficiency for policy refinement, making it a promising tool for autonomous robotic learning and potentially other domains where diffusion models are prevalent.