Lagrangian Perturbation Diffusion Steering: Latent Reinforcement Learning for Generative Policies¶

Conference: ICML 2026
arXiv: 2606.01151
Code: https://sites.google.com/view/lp-ds/home (project page)
Area: Reinforcement Learning / Generative Policies / Trust Region Methods
Keywords: Diffusion Policy, Latent RL, Trust Region, Lagrangian, Mode Collapse

TL;DR¶

LP-DS treats frozen diffusion or flow-matching policies as black-box decoders \(\Phi(s,w)\) and learns a state-conditioned residual only on the initial noise \(w=\epsilon+\Delta_\theta(s)\). By applying a Lagrangian trust region constraint \(\mathbb{E}_s[\|\Delta_\theta(s)\|_2^2]\le\delta\) to bound the perturbation magnitude, it enables sample-efficient online RL fine-tuning while preserving multimodal priors. It demonstrates superior stability over DSRL and DPPO on RoboMimic, Gym, Adroit, and LIBERO, with returns increasing by up to +25%.

Background & Motivation¶

Background: High-capacity generative policies (e.g., Diffusion Policy, Flow Matching \(\pi_0\) series) have become the mainstream paradigm for imitation learning in continuous control and manipulation due to their ability to represent multimodal action distributions.

Limitations of Prior Work: Pure imitation learning is limited by demonstration coverage and distribution shift, necessitating RL fine-tuning. However, directly updating a massive diffusion or flow-matching decoder is sample-inefficient and unstable due to the long-chain denoising or ODE integration gradients. Recent work like DSRL moves RL into the latent noise space (black-box decoding), but it replaces the pre-trained prior with a new latent policy. This leads to two failure modes: (i) latent variables drift out of the \(\mathcal{N}(0,I)\) decoder training support, triggering off-manifold behavior; (ii) the multimodal prior collapses into a single mode.

Key Challenge: Decoders are trained on \(\mathcal{N}(0,I)\), but RL value gradients tend to push latent variables toward increasingly extreme high-value regions—creating a direct conflict between "improving returns" and "staying within the decoder's training support." Figure 2 in the paper provides evidence: in DSRL on HalfCheetah, the latent variable magnitude \(\|w\|\) grows monotonically until decoding fails, causing the success rate to drop to zero.

Goal: Achieve online RL improvements without modifying a single decoder weight; employ an explicit mechanism to "anchor" back to the prior by controlling the magnitude of latent perturbations; and provide an interpretable "knob" for users to balance task reward against multimodal preservation.

Key Insight: Re-formulate latent space RL as constrained optimization. Instead of learning a new latent policy, the method learns a residual \(\Delta_\theta(s)\) and incorporates the "perturbation magnitude \(\approx\) approximate KL divergence from the prior" as a hard constraint via a Lagrangian formulation.

Core Idea: Use \(w=\epsilon+\Delta_\theta(s)\) combined with a trust region constraint \(\mathbb{E}_s[\|\Delta_\theta(s)\|_2^2]\le\delta\) and a dual variable \(\alpha\) updated via projected gradients. This creates an intrinsic mechanism where "exceeding the trust region leads to automatic tightening" and "staying within it allows for relaxation."

Method¶

Overall Architecture¶

LP-DS formulations the frozen generative policy as a black-box decoder \(\Phi:\mathcal{S}\times\mathcal{W}\to\mathcal{A}\). At each state \(s\), baseline noise \(\epsilon\sim\mathcal{N}(0,I)\) is sampled, and a small MLP \(\Delta_\theta(s)\) generates the latent query \(w\); the agent interacts with the environment using \(a=\Phi(s,w)\). Value learning adopts a "dual critic" structure: the action-side critic \(Q_\psi^\mathcal{A}(s,a)\) follows standard TD learning, while the latent-side critic \(Q_\phi^\mathcal{W}(s,w)\) is obtained via distillation of "\(Q^\mathcal{A}\circ\Phi\)" over the baseline noise. The actor update follows the latent critic to avoid backpropagating through the decoder. Only \(\Delta_\theta, Q_\psi^\mathcal{A}, Q_\phi^\mathcal{W}, \text{ and } \alpha\) are updated; the decoder remains frozen.

Key Designs¶

Latent Space Residual Perturbation:
- Function: Compresses "policy refinement via RL" into a learnable state-conditioned offset on \(\mathcal{N}(0,I)\), strictly recovering BC behavior when initialized at \(\Delta_\theta(\cdot)\approx 0\).
- Mechanism: Replaces DSRL’s design of "directly learning \(w\sim\pi_\theta^\mathcal{W}(\cdot\mid s)\)" with \(w=\epsilon+\Delta_\theta(s)\) where \(\epsilon\sim\mathcal{N}(0,I)\). The offset acts on the ODE integration starting point (\(x_T\) for diffusion or \(x_n\) for flow). Combined with deterministic DDIM/flow decoding, this is equivalent to a lightweight translation of the generative distribution anchored to the prior.
- Design Motivation: The prior is the true carrier of the policy's multimodal structure. Treating it as a reference rather than a baseline to be replaced explicitly protects the ability to cover multiple behavioral modes during improvement.
Lagrangian Trust Region Constraint:
- Function: Controls the latent perturbation magnitude via an interpretable single knob \(\delta\), ensuring it remains near the \(\mathcal{N}(0,I)\) support to avoid off-manifold decoding and early mode collapse.
- Mechanism: Using the closed-form approximation \(D_{\mathrm{KL}}(q_\theta(\cdot\mid s)\|p_0)\approx\frac{1}{2}\|\Delta_\theta(s)\|_2^2\), the objective is written as \(\max_\theta\mathbb{E}[Q^\mathcal{W}(s,\epsilon+\Delta_\theta(s))]\) s.t. \(\mathbb{E}_s\|\Delta_\theta(s)\|_2^2\le\delta\). This is dualized as \(\mathcal{L}(\theta,\alpha)=\mathbb{E}[Q^\mathcal{W}(s,w)-\alpha(\|\Delta_\theta(s)\|_2^2-\delta)]\). \(\theta\) is updated via gradient ascent on the primal problem, and \(\alpha\) follows projected dual ascent: \(\alpha\leftarrow[\alpha+\eta_\alpha\mathbb{E}_s(\|\Delta_\theta(s)\|_2^2-\delta)]_+\).
- Design Motivation: Creates an automatic regulation loop—if the perturbation exceeds \(\delta\), \(\alpha\) is pushed higher, forcing the actor to be conservative; if the perturbation is small, \(\alpha\) decreases, allowing for exploration.
Dual Critics and Latent Distillation:
- Function: Avoids the instability of backpropagating through the decoder to the actor while still being driven by real environment rewards.
- Mechanism: The action critic \(Q_\psi^\mathcal{A}(s,a)\) is learned via TD \(y=r+\gamma\bar Q^\mathcal{A}(s',a')\), where \(a'=\Phi(w';s')\) and \(w'=\epsilon'+\Delta_{\theta'}(s')\). The latent critic \(Q_\phi^\mathcal{W}(s,w)\) is distilled via \(\mathcal{L}_\phi=\mathbb{E}_{s,\epsilon}[(Q^\mathcal{W}_\phi(s,\epsilon)-Q^\mathcal{A}_\psi(s,\Phi(\epsilon;s)))^2]\) on the baseline noise distribution. The actor only computes gradients against \(Q^\mathcal{W}\), removing the requirement for a differentiable decoder.
- Design Motivation: Explicitly decouples "value learning" and "gradient pathways" at the decoder boundary, preserving the real reward signal while keeping large diffusion/flow matching decoders as read-only.

Loss & Training¶

The single loop involves: 1 environment step → 1 \(Q^\mathcal{A}\) TD update → 1 \(Q^\mathcal{W}\) distillation update → 1 actor primal update + 1 \(\alpha\) projected dual update. \(\delta\) is set to 0.35 in most experiments (0.5 for Hopper, 0.10 for Lift, 0.66 for Pen, etc.). ODE/DDIM decoding is used with an action chunk size \(T_a=8\).

Key Experimental Results¶

Main Results¶

Cross-domain comparison (summarized from Figure 3, average of 6 seeds, units in success rate/return):

Domain	Task	LP-DS	DSRL	DPPO	IDQL/DQL	Remarks
RoboMimic	Square	Near highest, fastest convergence	Slow convergence	Medium	Low	Strong advantage in precision-sensitive tasks
Gym Control	Walker2D-v2	≈5000	≈4000 (Strongest baseline)	—	—	+25% improvement in return
Adroit	Pen/Hammer/Door/Relocate	Best overall	Second	Slightly worse	Poor	Best in both success rate and return for dexterous tasks
LIBERO-90	Cream Cheese	Significant improvement over frozen π0	—	—	—	Leverages large VLA by training only the perturbation
Franka Real	Pick-and-Place	33/40	—	—	18/40 (Frozen baseline)	Perturbation trained in sim, deployed directly
Franka Real	Mug Hanging	17/20	—	—	11/20 (Frozen baseline)	Same as above

Ablation Study¶

Configuration	Pen Success EMA	k-NN Action Entropy	Explanation
Full LP-DS	Highest	High	Benefits of Trust Region + Lagrangian
w/o Lagrangian	Medium → Unstable	Monotonically decreasing	Collapses to single mode without auto-tightening
w/o Lag. & noise bound	Lowest, sharp oscillation	Extremely low	Latent variables drift out of \(\mathcal{N}(0,I)\) support
DSRL	Low	Lowest	Direct replacement of prior causes early collapse
LP-DS-A (Action residual)	Early plateau	—	Modifying noise is much more expressive than action correction

Key Findings¶

\(\delta\) as a Multimodal Knob: In a symmetric four-mode toy environment, \(\delta=0.01\) maintains coverage of all four modes, \(\delta=0.05\) makes trajectories more direct but remains multimodal, and \(\delta=0.1\) collapses to a single target. DSRL collapses to one mode immediately.
Robustness: Performance is relatively insensitive to \(\delta\) within the 0.1–0.66 range. The authors describe it as a "coarse knob" rather than a fragile hyperparameter, proving the engineering friendliness of the trust region design.
Latent vs. Action Residuals: Latent space residuals significantly outperform action space residuals. For high-capacity decoders, "shifting the starting point \(w\) to choose a different ODE path" is far more informative than "correcting the endpoint \(a\) locally."
Sim-to-Real: Physical Franka experiments demonstrate that training the perturbation module in simulation and deploying it on a frozen hardware-resident policy works, showing robustness to sim-to-real gaps.

Highlights & Insights¶

Using \(\frac{1}{2}\|\Delta\|^2\) to approximate the KL trust region is an undervalued practical simplification. For "baseline + shift" Gaussians, the mean shift squared is the dominant term. Using this as a constraint allows the Lagrangian derivation to close in a single line, making it extremely simple to implement.
The dual-critic + distillation architecture to keep gradients out of the decoder is a clean solution for the instability of diffusion backpropagation. This is especially vital for large VLAs like \(\pi_0\), where keeping the computation graph for backpropagation is often infeasible.
The projected dual update for \(\alpha\) automatically achieves "constraint adaptation," allowing the same hyperparameters to work across RoboMimic, Gym, and Adroit without major tuning, which is a clear advantage over static weight KL regularization.

Limitations & Future Work¶

The trust region approximation \(\frac{1}{2}\|\Delta\|^2 \approx \text{KL}\) assumes small residuals and an \(\mathcal{N}(0,I)\) base. This approximation may be biased when the decoder prior is non-isotropic (e.g., conditional flow), a scenario the authors did not test.
The paper acknowledges that long-horizon, partially observable scenes were not systematically covered; future work involves making the trust region adaptive across states or time.
Physical experiments were limited to medium difficulty (pick-and-place, mug hanging), lacking high-contact or long-chain dexterous tasks, which limits the strength of the transferability conclusions.

vs. DSRL: DSRL learns \(\pi_\theta^\mathcal{W}(w\mid s)\) to replace the prior; LP-DS learns a residual \(\Delta_\theta(s)\) and explicitly constrains its magnitude. The difference in collapse vs. restraint is clearly visible in the paper’s visualizations.
vs. DPPO: DPPO uses policy gradients to fine-tune all decoder parameters. LP-DS only updates a lightweight perturbation MLP, leading to significantly better sample efficiency and stability, though its upper bound is tied to the quality of the prior.
vs. IDQL / DQL: These offline-to-online methods treat diffusion as a parameterization for the actor and still update the main network. LP-DS's "read-only decoder + latent RL" is a more lightweight, deployment-friendly route.
vs. Noise Optimization in Vision (ReNO, Noise Hypernetworks): Those methods focus on single-image generation quality. LP-DS adopts the idea but replaces the target with long-horizon return and an explicit trust region, representing a complete adaptation of latent optimization from generation to decision-making.

Rating¶

Novelty: ⭐⭐⭐⭐ Mapping the KL trust region to latent RL via residual approximation is elegant, though the individual components (latent RL, trust regions, distillation) are established.
Experimental Thoroughness: ⭐⭐⭐⭐⭐ Covers 4 simulation domains, large VLA backbones, and dual real-world Franka tasks, addressing concerns about backbone dependence and real-world applicability.
Writing Quality: ⭐⭐⭐⭐ Algorithm 1 and the derivations are compact and clear. Toy environment visualizations provide intuitive insights into the tradeoff between multimodality and performance.
Value: ⭐⭐⭐⭐ Provides a ready-to-use pipeline for scenarios where one needs to "boost" a large generative policy with RL without touching its weights, which is particularly useful for \(\pi_0\)-style foundation models.