Lagrangian Perturbation Diffusion Steering: Latent Reinforcement Learning for Generative Policies

Behavior cloning with high-capacity generative policies achieves strong imitation performance, but is often limited by demonstration coverage and distribution shift. Direct reinforcement learning fine-tuning can improve performance, but updating large action decoders is frequently unstable and sample inefficient. We propose Lagrangian Perturbation Diffusion Steering (LP-DS), a lightweight adaptation method that improves a frozen generative policy by learning a compact noise-space perturbation before decoding. LP-DS optimizes this perturbation with a Lagrangian trust-region objective, improving downstream value while constraining deviation from the latent prior. Across RoboMimic manipulation, OpenAI Gym locomotion, and Adroit dexterous manipulation benchmarks, LP-DS improves sample efficiency, success, and return while maintaining higher action-space entropy than unconstrained noise-space steering, with return improvements of up to 25\% over prior baselines. Additional evaluations with flow-matching backbones, a large vision-language-action model, and physical Franka deployment show that LP-DS is not limited to compact diffusion policies or simulated benchmarks. Project page: https://sites.google.com/view/lp-ds/home.

Robots are increasingly trained by watching demonstrations, but these demonstrations cannot cover every possible situation the robot may face. As a result, a robot may work well in familiar settings but fail when objects, positions, or conditions change. This paper introduces LP-DS, a method for improving an already trained robot policy without changing the large model that produces the robot’s actions. Instead of retraining the whole policy, LP-DS makes small, controlled adjustments to the random input used by the policy when generating actions. These adjustments help the robot choose actions that lead to better outcomes while keeping the behavior close to what the original policy learned from demonstrations. This is important because overly large changes can make the robot behave unpredictably or lose useful alternative behaviors. We test LP-DS in simulation and on real Franka robot tasks. Across these settings, it improves task success while better preserving diverse ways of solving the task.