Steering Your Diffusion Policy with Latent Space Reinforcement Learning

TL;DR

This paper introduces DSRL, a method for efficiently adapting diffusion-based robotic control policies through reinforcement learning in the latent space, enabling fast autonomous policy improvement without modifying the original policy weights.

Contribution

The paper presents diffusion steering via reinforcement learning (DSRL), a novel approach that enables sample-efficient, real-world policy adaptation by operating over the latent noise space of diffusion policies.

Findings

DSRL achieves high sample efficiency in simulated and real-world tasks.

It enables autonomous policy improvement without modifying base policy weights.

Demonstrates effectiveness in adapting pretrained generalist policies.

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.