PaperlayerView exact revision on arXiv

Source-linked AI summary

Steering Your Diffusion Policy with Latent Space Reinforcement Learning

Andrew Wagenmaker, Mitsuhiko Nakamoto, Yunchu Zhang, Seohong Park, Waleed Yagoub, Anusha Nagabandi, Abhishek Gupta, Sergey Levine

arXiv:2506.15799v2cs.ROcs.LG

TL;DR

Pretrained behavioral-cloning policies cannot directly exploit deployment experience, while conventional RL adaptation is often too sample-intensive for real-world robotics. The paper introduces DSRL, which runs RL over a diffusion policy’s latent-noise space using black-box access rather than changing policy weights. DSRL enables effective, sample-efficient improvement across simulated and real-world settings, including generalist policies.

  • Problem

    Behavioral-cloning policies cannot directly use deployment experience, and RL finetuning of pretrained robotic policies remains costly in computation, time, and samples.

  • Method

    DSRL adapts diffusion policies by optimizing their input-noise distribution through RL, treating latent noise as actions while leaving the base policy weights unchanged.

  • Results

    DSRL improves diffusion policies across simulated and real-world tasks, including success gains from 20% to 90% with fewer than 50 online episodes and improvement of π0.

  • Takeaways & Limitations

    DSRL provides a black-box, sample-efficient route for autonomous real-world improvement of pretrained diffusion and generalist robot policies.

  • Takeaways & Limitations

    DSRL’s exploration is constrained by the underlying diffusion policy, and the paper does not guarantee that all such policies are steerable or exploratory enough.

Abstract

from arXiv · show

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.

1 Introduction

Behavioral cloning policies cannot directly use deployment experience, while RL adaptation is often too sample-intensive for real-world robotics. DSRL adapts diffusion policies by optimizing their input-noise distribution rather than modifying policy weights, enabling black-box, sample-efficient improvement.

  • Behavioral cloning translates offline demonstrations into effective policies but cannot directly use experience collected during deployment.
  • RL can improve pretrained policies, but existing finetuning remains expensive in computation, time, and samples for real-world robotic control.
  • DSRL runs RL over a diffusion policy’s input-noise distribution, steering behavior without modifying the pretrained policy’s weights.
  • DSRL requires only forward denoising passes and black-box policy access, supporting online adaptation without direct access to policy weights.
  • 20% to 90% success improvement can occur with less than 50 episodes of online interaction.
  • DSRL also improves pretrained generalist robot policies, including π0, during real-world deployment.

2 Related Work

Related work spans behavioral cloning with diffusion and other generative policies, reinforcement learning methods built around diffusion models, and optimization of diffusion initial noise. DSRL applies latent-noise optimization to robotic policy adaptation.

  • Behavioral cloning has produced strong robotic policies using diffusion, flow, and autoregressive generative model parameterizations.
  • Diffusion-based RL methods include value-weighted offline training, direct reward optimization, and action rejection sampling.
  • Prior generative-modeling work optimizes initial diffusion noise for image-quality gains or diffusion-guidance effects, whereas DSRL targets robotic control adaptation.

3 Preliminaries

The paper frames diffusion policies as generative action models and formulates adaptation as reinforcement learning over their latent-noise inputs. Deterministic denoising is central to treating noise selection as action selection.

  • Markov decision processes: An MDP comprises states, actions, transitions, initial-state distribution, rewards, and a discount factor.
  • Markov decision processes: RL seeks a policy that maximizes expected discounted return under the MDP dynamics and reward function.
  • Diffusion models and flow matching: Diffusion models transform an easy-to-sample distribution into a complex target distribution through a learned denoising process.
  • Diffusion models and flow matching: The reverse diffusion process starts from xT ∼N(0, I) and repeatedly applies denoising updates to generate x0.
  • Diffusion models and flow matching: DDIM sampling is deterministic given initial noise and can reduce denoising steps, which is useful when robots need actions generated quickly.
  • Diffusion models and flow matching: Flow matching similarly samples tractable noise and deterministically processes it into a sample from a goal distribution.
  • Behavioral cloning with diffusion and flow policies: For diffusion or flow BC policies, the denoising process maps state and latent noise to actions; adaptation seeks reward maximization while requiring the policy to be steerable.
  • Problem setting: The approach assumes a pretrained diffusion policy and considers adapting it using online or offline experience.

4 Diffusion Steering via Reinforcement Learning

Diffusion steering changes which latent noise is denoised, thereby selecting actions without altering the base policy. DSRL recasts this selection as RL in a latent-action MDP and uses noise aliasing to improve sample efficiency, including in offline learning.

  • Diffusion steering: Changing the input-noise distribution can steer a diffusion policy toward desired actions without modifying its weights or postprocessing outputs.
  • 4.1 Diffusion Steering as Latent-Noise Space Policy Optimization: DSRL treats latent noise as the action space of a transformed MDP, with the diffusion policy’s action mapping encapsulated inside the environment.
  • 4.1 Diffusion Steering as Latent-Noise Space Policy Optimization: The transformed latent-action MDP can be interacted with as a black box, allowing standard RL algorithms to optimize latent-noise actions.
  • Computational efficiency: DSRL avoids backpropagation through multi-step denoising, which can be computationally intensive and numerically unstable.
  • Computational efficiency: Black-box access enables adaptation when only forward policy queries or API access are available, including for proprietary models.
  • Noise aliasing: Different latent-noise actions can map to the same observed action, allowing noise aliasing to infer unexecuted latent actions from executed behavior.
  • Noise-Aliased DSRL: DSRL-NA uses action-space and latent-noise critics, transferring value information from observed actions to corresponding latent-noise actions.
  • Noise-Aliased DSRL: In offline learning, DSRL-NA avoids unseen-action value queries because the diffusion policy maps latent noise to in-distribution actions, providing conservatism without an explicit penalty.

5 Experiments

The experiments evaluate DSRL across online, offline, offline-to-online, real-world, and generalist-policy settings. Across these settings, DSRL adapts diffusion and flow-based policies with strong sample efficiency, including substantial real-world improvements.

  • Experimental scope: DSRL is evaluated in simulated online, offline, and offline-to-online settings, as well as real-world robotic control and generalist-policy adaptation.The experiments use DSRL-NA and SAC-based DSRL across these settings.
  • Offline adaptation: DSRL achieves state-of-the-art performance on approximately half of 10 OGBench offline tasks and substantially improves over the unsteered diffusion policy.The offline experiments apply DSRL-NA to the same offline data used to train the frozen diffusion policy.
  • Real-world adaptation: Using 100-150 online episodes, DSRL significantly improves a real-world multi-task policy on Pick-and-Place, Drawer Closing, and Block Stacking, while RLPD does not learn.Performance is averaged over 20 evaluations for Pick-and-Place and 10 for the other tasks.
  • Generalist policies: On Libero, DSRL improves π0 from ≈20% to ≈100% after ≈10,000 online samples, unlike RESIP and V-GPS.The evaluation also includes the challenging bimanual AlohaTransferCube task.
  • Generalist policies: DSRL significantly improves π0 on two real-world tasks using approximately 10,000-11,000 total steps, providing a successful real-world RL-based adaptation demonstration.The tasks are turning on a toaster and placing a spoon on a plate.

6 Discussion

DSRL addresses the persistent challenge of efficiently and stably improving diffusion policies after deployment. The discussion identifies broader directions involving alternative policy inputs, other diffusion-model domains, and theoretical analysis.

  • DSRL enables efficient reinforcement-learning-based improvement of diffusion policies in simulated and real-world settings.
  • Follow-up directions: The paper focuses on optimizing input noise, leaving observation or prompt modification as a direction for future policy improvement.
  • Follow-up directions: The underlying latent-noise optimization principle may extend beyond control to image generation or protein modeling.
  • Follow-up directions: Theoretical analysis remains open regarding diffusion steering’s expressivity, noise-space structure, and ease of learning policies over that space.

Contributions

The contribution statement credits the team with leading algorithm development, experiments across simulated and real-world settings, and the paper’s primary writing.

  • Andrew Wagenmaker led project coordination, core algorithm development, primary paper writing, and several simulation, ablation, and real-world experiments.
  • Mitsuhiko Nakamoto led the π0 experiments, while Yunchu Zhang led real-world single-task diffusion-policy experiments.
  • Seohong Park led offline experiments and contributed to algorithm development.

Limitations

DSRL’s effectiveness depends on the underlying diffusion policy’s exploration and requires standard RL resources such as rewards, rollouts, and real-world resets.

  • DSRL’s exploration capabilities are determined by the underlying diffusion policy, with no guarantee that every policy is steerable or explores adequately.
  • A highly concentrated action distribution may provide too few options for improvement and limits advance estimates of achievable gains.
  • Online DSRL requires reward signals, online rollouts, and resets when learning in the real world.

A Additional Ablations of Diffusion Steering

The ablations examine denoising, policy-data quality, policy diversity, algorithm choice, and implementation settings affecting DSRL’s convergence and sample efficiency.

  • Denoising steps: DSRL’s denoising-step sensitivity is evaluated separately for training and inference step counts.
  • Algorithm comparison: DSRL-NA requires fewer samples than DSRL-SAC on Robomimic Square, while the same trend is observed less strongly on Can.
  • Base-policy quality: The quality of πdp affects convergence speed and final solution quality, although DSRL remains surprisingly robust to variation in base-policy performance.
  • Data diversity: Training πdp on all data splits does not improve DSRL performance over training on a narrower dataset, despite producing more diverse behaviors.
  • Design choices: Recommended design choices include UTD of at least 20, layer normalization, large actor and critic networks, and bounded latent-action magnitude.
  • Experimental setup: Experiments use DSRL-NA except on Robomimic Transport, where DSRL-SAC is used, with parallel environments and replay-buffer initialization from πdp rollouts.

C.4.1 Single-Task Diffusion Policy

The single-task Franka experiments initialize DSRL from ten diffusion-policy rollouts and use a binary success criterion for cube pick-and-place.

  • Setup: DSRL initializes its replay buffer with data from 10 episodes generated by πdp using w ∼ N(0, I).The setup trains πdp with a self-collected dataset and uses a CNN architecture.
  • Evaluation: Cube pick-and-place episodes are successful when the cube is dropped into the bowl, using a 0-1 success reward.
  • Setup: The single-task experiments use the Franka robotic setup shown in Figure 18.

C.4.2 Multi-Task Diffusion Policy

The multi-task WidowX experiments use a DiT diffusion policy trained on Bridge v2, initialize learning from 20 policy rollouts, and evaluate three tasks with binary success rewards.

  • Policy and initialization: The WidowX diffusion policy πdp uses the DiT architecture and is trained on the Bridge v2 dataset with image goal conditioning.Deployment conditions πdp on a goal image of the desired final position.
  • Policy and initialization: DSRL and RLPD initialize their replay buffers from 20 episodes of πdp sampled with w ∼ N(0, I).RLPD additionally combines initial-rollout data with online data, while DSRL uses its own hyperparameters.
  • Evaluation: The multi-task experiments evaluate pick-and-place, drawer closing, and block stacking on WidowX using 0-1 success rewards.Success requires task-specific completion conditions, including object placement and release for the manipulation tasks.
  • Configuration: The experiments use the hyperparameter configurations listed for DSRL, πdp, and DiT πdp in Tables 10, 8, and 9.
  • Setup: The multi-task WidowX setups are illustrated in Figure 19.

C.5 Details of π0 Experiments (Section 5.5)

The π0 experiments study simulation and real-world tasks while steering a pretrained policy in latent noise space. Large action chunks are handled with repeated single-step noise, and DSRL training requires only an 8GB GPU locally.

  • Tasks: The π0 experiments study Libero and AlohaTransferCube simulation tasks, including online fine-tuning of specified manipulation behaviors.The Libero benchmark contains 130 tasks across separate suites.
  • Tasks: The simulation and real-world π0 task setups are shown in Figures 20 and 21.
  • Latent-noise formulation: π0 uses 50-step action chunks with 32-dimensional noise and actions per step, creating a total dimensionality of 1,600.This high-dimensional noise space makes naively training the actor and critic challenging.
  • Latent-noise formulation: DSRL-SAC avoids the full chunked noise space by training a single-step actor and repeating one sampled noise vector across the action-chunk axis.The repeated-noise formulation remains expressive enough to steer π0 while making SAC training more efficient and stable.
  • Real-world experiments: The real-world π0 experiments use the π0-DROID checkpoint and apply the same large-action-chunk handling method, with an action chunk of size 10.
  • Compute: π0 inference runs on a remote policy server, while DSRL training runs locally on an NVIDIA GeForce RTX 3070 GPU with 8GB of VRAM.The setup uses substantially less compute than the cited π0 LoRA and full-finetuning requirements.
  • Configuration: The π0 experiment hyperparameters are provided in Table 11.
Loading 2506.15799v2…