Steering Your Diffusion Policy with Latent Space Reinforcement Learning

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

Lizenzen

Open Access
Paper: CC BY 4.0

Alles Folgende stützt sich auf das vollständige Paper unter seiner Lizenz CC BY 4.0; die Abbildung wird mit Nennung der Autoren gezeigt.

Kurz gesagt

DSRL verbessert eine eingefrorene Diffusions- oder Flow-Policy, indem es eine kleine RL-Policy trainiert, die deren Eingangsrauschen wählt, und passt reale Roboter und π0 in einigen Dutzend bis wenigen Hundert Episoden an.

Warum es wichtig ist

Ein Team kann eine funktionierende, per Behavioral Cloning gelernte Policy vor Ort mit einigen Dutzend realen Episoden und einer spärlichen Erfolgsbelohnung verbessern, ohne ihre Gewichte anzufassen; dadurch lassen sich auch Black-Box-Modelle oder nur per API zugängliche Modelle anpassen. In den π0-Experimenten trainierte DSRL auf einer RTX 3070 mit 8 GB, während das π0-Repository angibt, dass LoRA-Fine-Tuning mindestens 22,5 GB benötigt.

Problem

Behavioral-Cloning-Policies, heute oft Diffusion Policies, können nicht aus ihrer eigenen Einsatzerfahrung lernen, sodass die Behebung schwacher Leistung meist das Sammeln weiterer menschlicher Demonstrationen bedeutet. RL könnte sie autonom verbessern, doch das Fine-Tuning einer Diffusion Policy mit RL erfordert Backpropagation durch ihre mehrstufige Denoising-Kette, was rechenintensiv und numerisch instabil ist, und Standard-RL benötigt für reale Roboter zu viele Samples.

Lösungsansatz

DSRL behandelt die vortrainierte Diffusion Policy (mit deterministischem DDIM- oder Flow-Sampling) als Teil der Umgebung und lernt eine separate Latent-Noise-Policy, die das Anfangsrauschen wählt, statt es aus N(0, I) zu ziehen. Dadurch wird die Anpassung zu gewöhnlichem RL über einem 'Latent-Action-MDP' und benötigt nur Vorwärtsdurchläufe der Basis-Policy, nicht ihre Gewichte. Die Noise-Aliasing-Variante DSRL-NA trainiert einen Critic auf realen Aktionen und destilliert ihn in einen Critic im Rauschraum, sodass sie Offline-Daten nutzen und die vielen Rauschwerte ausschöpfen kann, die auf dieselbe Aktion abgebildet werden.

Wichtigste Erkenntnisse

Bei deterministischem DDIM- oder Flow-Sampling bestimmt das Eingangsrauschen die Aktion vollständig, sodass die Wahl des Rauschens die Policy steuert, ohne sie zu verändern.
Da die Basis-Policy nur Aktionen innerhalb der Verteilung ausgibt, liefert die Optimierung über ihr Rauschen die Konservativität des Offline-RL ohne expliziten Strafterm.
Viele Rauschwerte werden auf nahezu dieselbe Aktion abgebildet; die Destillation eines Critics im Aktionsraum in den Critic im Rauschraum nutzt dieses Aliasing aus und benötigte etwa halb so viele Samples wie einfaches SAC im Rauschraum.
Für den 1.600-dimensionalen, in Chunks gegliederten Aktionsraum von π0 ist die Wiederholung eines einzigen 32-dimensionalen Rauschvektors über den Chunk ausdrucksstark genug, um die Policy zu steuern.
Die Steuerung funktionierte über Basis-Policies unterschiedlicher Größe, Trainingsdauer und Datenqualität hinweg, solange die Basis-Policy ein gewisses aufgabenlösendes Verhalten zeigt.

Hauptergebnisse

Die Online-Anpassung auf Robomimic und OpenAI Gym erreichte nahezu optimale Leistung mit mindestens 5- bis 10-mal weniger Samples als DPPO, IDQL, DQL, DIPO und QSM.
Auf 10 OGBench-Offline-Aufgaben lag DSRL bei etwa der Hälfte auf dem Stand der Technik und verbesserte sich deutlich gegenüber der ungesteuerten Flow-Policy.
Realer Franka, Einzelaufgabe Würfel in Schüssel: Der Erfolg stieg in 3.500 Online-Schritten (etwa 40 Episoden) von 20 % auf nahezu 100 %; RLPD löste die Aufgabe weder mit noch ohne menschliche Eingriffe.
Realer WidowX 250 mit einer auf Bridge V2 trainierten Multi-Task-Policy: deutliche Zugewinne bei Pick-and-Place, Schubladeschließen und Blockstapeln innerhalb von 100–150 Episoden, während RLPD nicht lernte.
π0 auf LIBERO: Erfolg von etwa 20 % auf etwa 100 % nach etwa 10.000 Online-Samples; ResiP und V-GPS verbesserten es nicht wesentlich.
Realer Franka mit dem öffentlichen π0-DROID-Checkpoint: verbesserte Toaster- und Löffel-auf-Teller-Aufgaben in 80 bzw. 65 Episoden (etwa 11.000 bzw. 10.000 Schritte), was die Autoren als erstes RL-Fine-Tuning von π0 in der realen Welt bezeichnen.

Einschränkungen

Die Exploration ist durch die Basis-Policy begrenzt: Es gibt keine Garantie, dass jede Diffusion Policy steuerbar ist, und eine sehr schmale Aktionsverteilung kann zu wenige Optionen lassen (Autoren).
Es gibt noch keine Möglichkeit, im Voraus zu quantifizieren, wie stark sich eine gegebene Basis-Policy verbessern lässt (Autoren).
Wie jedes RL-Verfahren benötigt es Belohnungssignale, bei der Online-Variante Online-Rollouts und Resets in der realen Welt, über das hinaus, was Imitation Learning erfordert (Autoren).
Reale Aufgaben verwenden von Hand definierte 0-1-Erfolgsbelohnungen, und die Evaluierungen sind klein (10–20 Versuche pro Aufgabe auf dem WidowX).

Wirkung auf das Fachgebiet

Finalist für die CoRL 2025 Best Paper / Best Student Paper Awards, bisher 242 Zitierungen. Das Paper verlinkt eine Projektwebsite (diffusion-steering.github.io) und baut auf öffentlichen π0-Checkpoints auf; die lizenzierte PMLR-Fassung nennt kein Code-Repository.

Zusammenfassung

Aus menschlichen Demonstrationen gelernte Robotersteuerungspolicies haben in vielen realen Anwendungen beeindruckende Ergebnisse erzielt. In Szenarien, in denen die anfängliche Leistung nicht zufriedenstellend ist, wie es in neuartigen Open-World-Umgebungen häufig der Fall ist, erfordern solche per Behavioral Cloning (BC) gelernten Policies jedoch typischerweise das Sammeln zusätzlicher menschlicher Demonstrationen, um ihr Verhalten weiter zu verbessern – ein teurer und zeitaufwendiger Prozess. Im Gegensatz dazu verspricht Reinforcement Learning (RL), eine autonome Online-Verbesserung von Policies zu ermöglichen, bleibt aber wegen der großen Zahl an Samples, die es typischerweise benötigt, oft dahinter zurück. In dieser Arbeit unternehmen wir Schritte hin zu einer schnellen autonomen Anpassung BC-trainierter Policies durch effizientes RL in der realen Welt. Mit besonderem Fokus auf Diffusion Policies – eine BC-Methodik auf dem Stand der Technik – schlagen wir Diffusion Steering via Reinforcement Learning (DSRL) vor: die Anpassung der BC-Policy durch Ausführen von RL über ihrem Latent-Noise-Raum. Wir zeigen, dass DSRL hochgradig sample-effizient ist, nur Black-Box-Zugriff auf die BC-Policy benötigt und eine wirksame autonome Policy-Verbesserung in der realen Welt ermöglicht. Darüber hinaus vermeidet DSRL viele der Herausforderungen, die mit dem Fine-Tuning von Diffusion Policies verbunden sind, und macht es überflüssig, die Gewichte der Basis-Policy überhaupt zu verändern. Wir demonstrieren DSRL auf simulierten Benchmarks, realen Roboteraufgaben und bei der Anpassung vortrainierter generalistischer Policies und veranschaulichen so seine Sample-Effizienz und wirksame Leistung bei der Policy-Verbesserung in der realen Welt.

Open Source Robotics