Steering Your Diffusion Policy with Latent Space Reinforcement Learning

Authors: Andrew Wagenmaker, Mitsuhiko Nakamoto, Yunchu Zhang, Seohong Park, Waleed Yagoub, Anusha Nagabandi, Abhishek Gupta, Sergey Levine · CoRL 2025

Licences

Open Access
Paper: CC BY 4.0

Everything below draws on the full paper under its CC BY 4.0 licence; its figure is credited to the authors.

TL;DR

DSRL improves a frozen diffusion or flow policy by training a small RL policy that picks its input noise, adapting real robots and π0 in tens to a few hundred episodes.

Why it matters

A team can improve a working behaviour-cloned policy in place, with a few dozen real episodes and a sparse success reward, without touching its weights, which also makes black-box or API-only models adaptable. In the π0 experiments DSRL trained on an 8 GB RTX 3070, while the π0 repository states LoRA fine-tuning needs at least 22.5 GB.

Problem

Behavioural-cloning policies, now often diffusion policies, cannot learn from their own deployment experience, so fixing weak performance usually means collecting more human demonstrations. RL could improve them autonomously, but fine-tuning a diffusion policy with RL needs back-propagation through its multi-step denoising chain, which is computationally heavy and numerically unstable, and standard RL needs too many samples for real robots.

Proposed solution

DSRL treats the pretrained diffusion policy (with deterministic DDIM or flow sampling) as part of the environment and learns a separate latent-noise policy that chooses the initial noise instead of drawing it from N(0, I). This turns adaptation into ordinary RL over a 'latent-action MDP' and needs only forward passes of the base policy, not its weights. The noise-aliased variant, DSRL-NA, trains a critic on real actions and distils it into a noise-space critic, so it can use offline data and exploit the many noise values that map to the same action.

Key insights

With deterministic DDIM or flow sampling the input noise fully determines the action, so choosing the noise steers the policy without changing it.
Because the base policy only outputs in-distribution actions, optimising over its noise gives offline RL conservatism without an explicit penalty term.
Many noise values map to nearly the same action; distilling an action-space critic into the noise-space critic exploits this aliasing and needed about half the samples of plain SAC on the noise space.
For π0's 1,600-dimensional chunked action space, repeating one 32-dimensional noise vector across the chunk is expressive enough to steer the policy.
Steering worked across base policies of different sizes, training lengths and data quality, as long as the base policy shows some task-solving behaviour.

Main results

Online adaptation on Robomimic and OpenAI Gym reached near-optimal performance with at least 5-10x fewer samples than DPPO, IDQL, DQL, DIPO and QSM.
On 10 OGBench offline tasks DSRL was state-of-the-art on about half and improved substantially over the unsteered flow policy.
Real Franka single-task cube-into-bowl: success rose from 20% to nearly 100% in 3,500 online steps (about 40 episodes); RLPD, with or without human interventions, did not solve the task.
Real WidowX 250 with a multi-task policy trained on Bridge V2: significant gains on pick-and-place, drawer closing and block stacking within 100-150 episodes, while RLPD did not learn.
π0 on LIBERO: success from about 20% to about 100% after about 10,000 online samples; ResiP and V-GPS did not improve it significantly.
Real Franka with the public π0-DROID checkpoint: improved toaster and spoon-on-plate tasks in 80 and 65 episodes (about 11,000 and 10,000 steps), which the authors call the first real-world RL fine-tuning of π0.

Limitations

Exploration is bounded by the base policy: there is no guarantee that every diffusion policy is steerable, and a very narrow action distribution may leave too few options (authors).
There is no way yet to quantify in advance how much a given base policy can be improved (authors).
Like any RL method it needs reward signals, online rollouts for the online variant, and resets in the real world, beyond what imitation learning requires (authors).
Real-world tasks use hand-defined 0-1 success rewards, and evaluations are small (10-20 trials per task on the WidowX).

Field impact

Finalist for the CoRL 2025 Best Paper / Best Student Paper awards, with 242 citations so far. The paper links a project website (diffusion-steering.github.io) and builds on public π0 checkpoints; the licensed PMLR copy does not name a code repository.

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.

Open Source Robotics