Improving Diffusion Planners by Self-Supervised Action Gating with Energies

Das Paper stellt SAGE vor, eine selbstüberwachte Nachbewertungsmethode, die durch die Bestrafung dynamisch inkonsistenter Pläne mit einem latenten Konsistenzsignal die Robustheit und Leistung von Diffusionsplanern im Offline-Reinforcement-Learning ohne erneutes Training oder Umgebungsrollouts verbessert.

Das Wesentliche

Stell dir vor, du möchtest einen Roboter lernen lassen, durch ein Labyrinth zu laufen oder einen Stuhl zu schieben, ohne dass er dabei jemals selbst mit dir interagiert hat. Du gibst ihm nur ein riesiges Videoarchiv von jemand anderem, der das schon einmal getan hat. Das nennt man „Offline Reinforcement Learning".

Das Problem dabei ist: Der Roboter ist wie ein sehr kreativer, aber etwas naiver Traumtänzer. Er schaut sich die Videos an und versucht, neue, noch bessere Wege zu erfinden. Aber manchmal träumt er sich einen Weg aus, der auf dem Papier toll aussieht (vielleicht ist er der kürzeste), aber in der Realität sofort scheitert, weil er eine Wand durchquert oder gegen die Schwerkraft ankämpft. Er plant etwas, das unmöglich ist.

Die Forscher in diesem Papier haben eine Lösung namens SAGE entwickelt. Hier ist die Erklärung, wie das funktioniert, ganz einfach und mit ein paar Bildern im Kopf:

1. Das Problem: Der „Traum" vs. die „Realität"

Stell dir vor, der Roboter plant seine nächsten Schritte wie ein Regisseur, der viele verschiedene Film-Skripte entwirft.

• Der alte Weg: Der Regisseur schaut sich alle Skripte an und wählt den aus, der am meisten Action und Punkte verspricht. Aber manchmal ist dieser Skript so verrückt, dass der Schauspieler (der Roboter) beim ersten Schritt schon gegen die Wand läuft. Der Film bricht ab.
• Das Problem: Der Regisseur (der Planer) ist zu sehr darauf fixiert, wohin es gehen soll (die Punkte), und vergisst, ob es überhaupt möglich ist, dorthin zu kommen.

2. Die Lösung: SAGE – Der „Realitäts-Check"

SAGE ist wie ein strenger, aber fairer Korrektor, der hinzukommt, bevor der Roboter einen Schritt macht. Er schaut sich die Pläne nicht an, um zu sagen, ob sie „cool" sind, sondern ob sie „machbar" sind.

Der Name steht für Self-supervised Action Gating with Energies. Klingt kompliziert, ist aber eigentlich ganz einfach:

• Der Lehrer (JEPA): Zuerst lernt SAGE aus den alten Videos, wie sich die Welt normalerweise verhält. Es ist wie ein Lehrer, der dem Roboter beibringt: „Wenn du nach rechts drückst, bewegst du dich nach rechts. Wenn du gegen eine Wand drückst, passiert nichts." Der Lehrer schaut nicht auf Belohnungen, sondern nur auf die Physik der Dinge.
• Der Energie-Messwert: Wenn der Roboter einen neuen Plan (einen Film-Skript) entwirft, schaut SAGE sich die ersten paar Schritte an.
  • Wenn der Plan sagt: „Ich springe jetzt durch die Wand", sagt SAGE: „Autsch! Das kostet viel Energie (im Sinne von Anstrengung gegen die Gesetze der Physik). Das ist unmöglich."
  • Wenn der Plan sagt: „Ich gehe um die Ecke", sagt SAGE: „Gut, das kostet wenig Energie. Das ist machbar."

3. Wie es funktioniert: Der „Gating"-Effekt

Stell dir vor, der Roboter hat 50 verschiedene Pläne für die nächsten Schritte in der Hand.

1. Der alte Weg: Er nimmt den Plan mit den meisten Punkten.
2. Der neue Weg (mit SAGE):
   • SAGE wirft zuerst alle Pläne weg, die physikalisch unmöglich sind (die „Wand-durchqueren"-Pläne). Das nennt man „Gating" (wie ein Tor, das nur für machbare Pläne offen bleibt).
   • Von den übrig gebliebenen, machbaren Plänen wählt er dann den aus, der die meisten Punkte bringt.

Warum ist das so genial?

• Kein neues Training nötig: Der Roboter muss nicht neu lernen. SAGE ist wie ein Add-on, das man einfach dazwischen schaltet, wenn der Roboter plant.
• Kein Experimentieren: SAGE braucht keine neuen Versuche in der echten Welt. Es lernt alles nur aus den alten Videos.
• Es rettet den Roboter vor sich selbst: Es verhindert, dass der Roboter in eine Sackgasse läuft, nur weil der Plan auf dem Papier gut aussah.

Zusammenfassung in einem Satz

SAGE ist wie ein Sparringspartner, der dem Roboter zuruft: „Hey, dein Plan sieht toll aus, aber du wirst dabei gegen die Wand laufen! Nimm lieber den anderen Plan, der zwar etwas weniger Punkte verspricht, aber wenigstens funktioniert."

Dadurch werden Roboter, die nur aus alten Daten lernen, viel sicherer, robuster und können auch in schwierigen Umgebungen (wie Küchen oder Labyrinthen) viel besser arbeiten, ohne ständig stecken zu bleiben.

Technische Zusammenfassung

1. Problemstellung

Diffusionsplaner (Diffusion Planners) haben sich als leistungsstarke Methode für das Offline-Reinforcement-Learning (RL) etabliert. Sie modellieren die Verteilung von Aktionssequenzen oder Trajektorien und generieren Kandidaten durch iteratives Denoising. Ein zentrales Problem bei diesen Ansätzen ist jedoch die lokal inkonsistente Dynamik:

• Der Fehlermodus: Diffusionsplaner generieren oft viele Kandidaten-Trajektorien, die von einem gelernten Wertmodell (Critic) hoch bewertet werden, weil sie einen hohen erwarteten Return versprechen. Diese Trajektorien können jedoch lokal mit den tatsächlichen Umgebungs-Dynamiken unvereinbar sein (d. h., die ersten Schritte sind aus dem aktuellen Zustand heraus nicht ausführbar).
• Die Folge: Wenn ein Agent einen solchen Plan ausführt, scheitert er sofort, da die ersten Aktionen nicht realisierbar sind. Dies führt zu einer „zerbrechlichen" (brittle) Ausführung, insbesondere bei langen Planungshorizonten.
• Herausforderung: Bestehende Methoden versuchen oft, diese Probleme durch Guidance während des Generierungsprozesses oder durch zusätzliche Verifikationsmodelle zu lösen. Dies erhöht jedoch die Trainingskomplexität, erfordert oft negative Samples oder Umgebungs-Interaktionen und kann die ursprüngliche Verteilung der Trajektorien verzerren.

2. Methodik: SAGE (Self-supervised Action Gating with Energies)

SAGE ist ein Inferenzzeit-Modul, das Diffusionsplaner ergänzt, ohne deren Generator oder Critic neu zu trainieren. Es trennt explizit die Bewertung der Wertigkeit (Value) von der Durchführbarkeit (Feasibility).

Der Ansatz besteht aus zwei trainierten Komponenten, die rein offline auf dem Datensatz lernen:

A. Lernprozess (Offline Training)

1. JEPA-Encoder (Joint-Embedding Predictive Architecture):
   • Ein Encoder wird auf reinen Zustandssequenzen trainiert, um eine latente Repräsentation zu lernen, die konsistent mit den Datendynamiken ist.
   • Es wird ein Masked Prediction Ansatz verwendet: Der Encoder versucht, die latente Repräsentation zukünftiger Zustände (Targets) aus einem maskierten Kontextfenster vorherzusagen.
   • Ein Exponential Moving Average (EMA) Teacher stabilisiert das Training.
2. Aktions-konditionierter latenter Prädiktor:
   • Ein zweites Modell (ein Transformer) lernt im eingefrorenen latenten Raum, wie sich Zustände über kurze Zeithorizonte ändern, gegeben eine Aktionssequenz.
   • Es wird mit drei Zielen trainiert:
     • Teacher-forced Loss: Vorhersage des nächsten latenten Zustands unter Verwendung der wahren Historie.
     • Rollout Loss: Konsistenz bei autoregressiver Anwendung über einen kurzen Horizont.
     • Action-Usage Hinge: Eine Regularisierung, die sicherstellt, dass der Prädiktor tatsächlich auf die Aktionen reagiert (durch Permutation der Aktionen im Batch und Bestrafung von Vorhersagen, die auch bei falschen Aktionen korrekt bleiben).

B. Inferenzprozess (Runtime)

SAGE greift in den Entscheidungsprozess ein, nachdem der Diffusionsplaner Kandidaten-Trajektorien generiert hat:

1. Energie-Bewertung: Für jeden generierten Kandidaten wird die „latente Konsistenz-Energie" berechnet. Dies ist der Fehler zwischen der vom Prädiktor vorhergesagten latenten Entwicklung und der tatsächlich im Kandidaten geplanten Entwicklung über einen kurzen Präfix (z. B. die ersten $K$ Schritte).
   • Niedrige Energie: Der Plan ist dynamisch konsistent mit dem Datensatz (durchführbar).
   • Hohe Energie: Der Plan enthält inkonsistente Übergänge (nicht durchführbar).
2. Gating und Re-Ranking:
   • Kandidaten mit zu hoher Energie werden gefiltert (Gating).
   • Die verbleibenden Kandidaten werden neu sortiert, indem die ursprüngliche Wert-Score ($J$) um einen Energie-Strafterm ($\lambda \cdot E$) reduziert wird.
   • Die Aktion des bestbewerteten verbleibenden Kandidaten wird ausgeführt.

3. Schlüsselbeiträge

• Trennung von Wert und Durchführbarkeit: SAGE führt ein separates, selbstüberwachtes Signal für die lokale Durchführbarkeit ein, anstatt dies implizit in den Critic zu integrieren.
• Kein Re-Training: Das Modul ist modular und kann in bestehende Diffusionsplaner-Pipelines integriert werden, ohne den Generator oder den Critic neu zu trainieren.
• Keine Umgebungs-Interaktion: Das Training erfolgt ausschließlich auf dem Offline-Datensatz; es werden keine Rollouts oder zusätzlichen Umgebungs-Abfragen benötigt.
• Skalierbarkeit: Durch den Verzicht auf negative Sampling-Strategien oder komplexe Guidance-Mechanismen während des Trainings ist die Methode gut skalierbar auf große und diverse Datensätze.

4. Ergebnisse

Die Methode wurde auf dem D4RL-Benchmark in drei Domänen evaluiert:

1. Lokomotion (MuJoCo): Halbe Zebra, Hopper, Walker2d.
2. Manipulation (Franka Kitchen): Komplexe Aufgaben mit mehreren Schritten.
3. Navigation (AntMaze, Maze2D): Langfristige Planung in Labyrinthen.

Ergebnisse:

• Leistungssteigerung: SAGE verbessert konsistent die Leistung bestehender State-of-the-Art-Planer (insbesondere des „Generate-and-Rank"-Ansatzes DV).
  • Im Durchschnitt über alle MuJoCo-Umgebungen stieg die Leistung von 82,9 (DV) auf 84,4.
  • In der Manipulationsdomäne (Kitchen) wurde eine deutliche Verbesserung von 81,8 auf 85,6 erzielt.
  • In der Navigation (AntMaze) stieg der Score von 81,6 auf 84,5.
• Robustheit: SAGE reduziert signifikant das Auftreten von „Wall-Crossing" (Durchschreiten von Wänden) und anderen dynamisch unmöglichen Trajektorien, wie in den Visualisierungen (Abb. 3) gezeigt.
• Statistische Signifikanz: Die Verbesserungen sind statistisch signifikant (p-Werte < 0,05), mit Ausnahme von Maze2D, wo die Ergebnisse bereits nahe am Optimum lagen (Ceiling-Effekt).
• Overhead: Der zusätzliche Rechenaufwand für das Gating beträgt nur ca. 6,8 % der Inferenzzeit, was als sehr gering eingestuft wird.

5. Bedeutung und Fazit

SAGE adressiert eine fundamentale Schwäche von Diffusionsplanern: die Tendenz, attraktive, aber physikalisch unmögliche Pläne zu generieren. Indem es ein selbstüberwachtes, reward-freies Signal nutzt, um die lokale Konsistenz von Plänen zu bewerten, ermöglicht es eine zuverlässigere Ausführung ohne Kompromisse bei der Generierungsfähigkeit des Modells.

Die Arbeit zeigt, dass die Trennung von Generierung (was ist möglich?) und Selektion (was ist gut?) durch ein zusätzliches, leichtgewichtiges Gating-Modul effektiver ist als das Versuchen, beides in einem einzigen Modell zu vereinen. Dies bietet einen praktischen Weg, um Offline-RL-Systeme robuster und zuverlässiger zu machen, insbesondere in sicherheitskritischen Anwendungen wie der Robotik.