Improving Diffusion Planners by Self-Supervised Action Gating with Energies

Dit artikel introduceert SAGE, een zelftoezichtsmethode die diffusioplanners verbetert door dynamisch inconsistente trajecten te straffen via een energiegebaseerd herordeningssysteem, waardoor de robuustheid en prestaties zonder extra training of omgevingssimulaties worden verhoogd.

De kern

Stel je voor dat je een robot wilt leren een complexe taak uitvoeren, zoals een robotarm die een kopje thee zet of een vierpotige robot die door een doolhof loopt. Je hebt geen tijd om de robot duizenden keren te laten vallen en opnieuw te proberen (dat is te gevaarlijk en te duur). Dus geef je de robot een groot boek met verhalen van hoe een andere, slimme robot die taak eerder heeft gedaan. Dit noemen we "offline reinforcement learning".

De robot leest dit boek en probeert een nieuw plan te bedenken. Moderne robots gebruiken hiervoor een slimme techniek genaamd Diffusion Planning.

Het Probleem: De Dromer die Vergeet te Kijken naar de Grond

Stel je voor dat de robot een droommachine is. Hij droomt van 100 verschillende manieren om de taak te voltooien. Hij kijkt naar al deze dromen en kiest degene die er in de droom het mooist uitziet en het meeste "punten" oplevert.

Maar hier zit een addertje onder het gras:
Soms droomt de robot iets dat er op papier perfect uitziet (hij krijgt veel punten), maar dat in de echte wereld onmogelijk is.

• Voorbeeld: De robot droomt dat hij in één seconde van de ene kant van de kamer naar de andere springt. In zijn droom is dit een snelle, efficiënte route. Maar in de realiteit heeft hij geen superkrachten; hij zou tegen de muur aanvliegen en vallen.

De robot kiest deze "slechte droom" omdat hij alleen naar de score kijkt, niet naar of het haalbaar is. Hij begint de taak, maar faalt direct omdat hij een onrealistisch begin heeft gekozen.

De Oplossing: SAGE (De Slimme Portier)

De auteurs van dit paper hebben een oplossing bedacht genaamd SAGE (Self-supervised Action Gating with Energies). Je kunt SAGE zien als een slimme portier of een reality-check die tussen de droommachine en de robot staat.

Hier is hoe het werkt, stap voor stap:

1. De Droommachine (De Generator):
   De robot droomt nog steeds 100 verschillende plannen, net als voorheen. Hij doet dit nog steeds op dezelfde manier.

2. De Reality-Check (De Energie):
   Voordat de robot een plan mag uitvoeren, loopt SAGE langs de eerste paar stappen van elk plan. SAGE heeft een speciaal "gevoel" (een model) dat is getraind op het boek met de oude verhalen. Dit gevoel weet precies hoe de wereld werkt: "Als je hier een stap zet, is de volgende stap logisch?"

   • Als het plan logisch is (bijv. de robot loopt rustig naar de deur), geeft SAGE een lage energie (een groen lichtje).
   • Als het plan onlogisch is (bijv. de robot springt door de muur), geeft SAGE een hoge energie (een rood waarschuwingslichtje). Dit noemen ze "energie" omdat het voelt als een zware last die je moet dragen als het plan niet klopt.

3. De Gating (De Selectie):
   SAGE kijkt naar alle 100 dromen. Hij gooit direct de dromen weg met de hoogste energie (de onmogelijke dromen). Hij houdt alleen de plannen over die haalbaar zijn.
   Van de overgebleven, haalbare plannen, kiest de robot dan nog steeds degene uit die de meeste punten oplevert.

Waarom is dit zo slim?

• Geen Nieuw Boek Leren: De robot hoeft niet opnieuw te leren hoe hij moet dromen. SAGE is als een extra bril die je opzet nadat je het boek hebt gelezen. Je verandert het boek niet, je filtert alleen de slechte dromen eruit.
• Geen Probeer-en-Fout: SAGE heeft geen nieuwe experimenten nodig in de echte wereld. Hij leert puur uit het oude boek (de offline data) door te voorspellen wat er zou gebeuren als je een bepaalde stap zet.
• Veiligheid: Het voorkomt dat de robot in paniek raakt of vastloopt omdat hij een onmogelijke stap probeert.

De Metafoor van de Reisplanner

Stel je voor dat je een reisplanner gebruikt om een wandeltocht te maken.

• De oude manier: De planner kijkt naar alle mogelijke routes. Hij ziet een route die heel kort is en snel naar de top leidt (hoge score). Hij kiest deze. Maar als je er gaat lopen, blijkt die route over een steile afgrond te gaan. Je valt.
• De SAGE-methode: De planner kijkt ook naar die snelle route. Maar SAGE, de slimme portier, zegt: "Wacht even, ik heb gekeken naar de kaarten van eerdere wandelaars. Deze route is fysiek onmogelijk om te lopen zonder te vallen." SAGE blokkeert die route. De planner kiest dan de volgende beste route die wel haalbaar is.

Conclusie

SAGE maakt robots slimmer en veiliger door een extra laag van realiteit toe te voegen aan hun dromen. Het zorgt ervoor dat ze niet alleen plannen maken die er goed uitzien, maar plannen die ze ook daadwerkelijk kunnen uitvoeren. Het is een simpele, maar krachtige manier om robots minder "brittle" (kwetsbaar) te maken en ze beter te laten presteren in de echte wereld.

Technische samenvatting

1. Het Probleem

Offline Reinforcement Learning (RL) met behulp van diffusieplanners is een krachtige methode voor sequentiële besluitvorming. Deze planners genereren veel mogelijke toekomstige trajecten (state-action sequenties) via een diffusiemodel en selecteren vervolgens de beste optie op basis van een geleerde waarderingsfunctie (value function).

De kern van het probleem dat dit artikel adresseert, is de kwetsbaarheid bij de uitvoering:

• Lokale inconsistentie: Een diffusieplanner kan trajecten genereren die een hoge score krijgen volgens de waarde-functie (belangrijk voor lange termijn beloning), maar die lokaal onuitvoerbaar zijn. Dit betekent dat de eerste stappen van het geplande traject niet fysiek haalbaar zijn vanuit de huidige staat, gezien de dynamiek van de omgeving.
• Fragiele uitvoering: Wanneer een agent een dergelijk "onrealistisch" prefix kiest, faalt de uitvoering direct of leidt dit tot een cascade van fouten bij herplanning.
• Beperkingen van bestaande oplossingen: Bestaande methoden proberen dit op te lossen door de generatie te sturen (guidance) of extra verifiers te gebruiken. Dit verhoogt echter de trainingscomplexiteit, vereist vaak extra modellen, of kan de diversiteit van de gegenereerde trajecten onnodig beperken.

2. Methodologie: SAGE

De auteurs stellen Self-supervised Action Gating with Energies (SAGE) voor. Dit is een inference-time (tijdens het testen) methode die de bestaande diffusieplanner niet opnieuw traint, maar de selectie van acties verbetert.

SAGE splitst het beslissingsproces op in twee signalen: Waarde (lange termijn beloning) en Haalbaarheid (lokale dynamische consistentie).

De Architectuur van SAGE

SAGE bestaat uit twee componenten die puur offline worden getraind:

1. JEPA-Encoder (Joint-Embedding Predictive Architecture):

   • Een encoder wordt getraind op offline state-sequenties om een latente representatie te leren die de dynamiek van de dataset consistent weergeeft.
   • Dit wordt gedaan via self-supervision: het model voorspelt de latente embedding van een toekomstige staat (target) op basis van een gemaskeerd contextvenster.
   • Er wordt gebruik gemaakt van een EMA (Exponential Moving Average) teacher voor stabiliteit.

2. Actie-geconditioneerde Latente Predictor:

   • In een tweede fase wordt een predictor getraind in de bevroren latente ruimte van de JEPA.
   • Deze predictor leert korte-termijn overgangen te modelleren: gegeven een latente staat $z_t$ en een actie $a_t$, voorspelt het de volgende latente staat $\hat{z}_{t+1}$.
   • Het trainingsdoel omvat drie termen:
     • Teacher-forced loss: Nauwkeurige voorspelling van de volgende stap.
     • Rollout loss: Consistentie bij autoregressieve toepassing over een korte horizon.
     • Action-usage hinge: Een regularisatieterm die voorkomt dat de predictor acties negeert (door permutatie van acties binnen een batch en het straffen van lage fouten bij verkeerde acties).

Inference (Tijdens het testen)

Wanneer de diffusieplanner $C$ kandidaat-trajecten genereert, past SAGE het volgende proces toe:

1. Energieberekening: Voor elke kandidaat wordt een "energie"-score berekend op basis van de voorspellingsfout van de latente predictor over de eerste $K$ stappen van het traject.
   • $E(\hat{\tau}) = \frac{1}{K} \sum || f_\eta(z_{t+k}, a_{t+k}) - z_{t+k+1} ||_1$
   • Een lage energie betekent dat het traject consistent is met de dataset-dynamiek (haalbaar). Een hoge energie wijst op dynamische inconsistentie (onhaalbaar).
2. Gating en Herordening:
   • De kandidaten met de hoogste energie (onhaalbaar) worden gefilterd (bijv. de bovenste 20% verwijderd).
   • Van de resterende kandidaten wordt de beste gekozen door de oorspronkelijke waarde-score ($J$) te combineren met een straffende term voor de energie:
     $$i^* \in \arg \max_i (J(\hat{\tau}_i) - \lambda E(\hat{\tau}_i))$$
   • Hierbij is $\lambda$ een hyperparameter die de weging bepaalt.

3. Belangrijkste Bijdragen

• Scheiding van Waarde en Haalbaarheid: SAGE introduceert een expliciet, zelftoezichtend signaal voor haalbaarheid dat losstaat van de waarde-functie. Dit lost het spanningsveld op tussen het maximaliseren van beloning (extrapolatie) en het respecteren van dynamische beperkingen (conservatisme).
• Modulair en Traineringsvrij: SAGE vereist geen hertraining van de diffusiegenerator of de criticus. Het kan worden ingebouwd in bestaande pijplijnen die kandidaten genereren en scoren.
• Self-Supervised Learning: Het maakt gebruik van JEPA-architecturen en self-supervision, waardoor het geen negatieve steekproeven of interactie met de omgeving nodig heeft tijdens het trainen.
• Efficiëntie: De methode voegt slechts een kleine rekentijd toe (ongeveer 6-8% overhead) omdat de evaluatie beperkt blijft tot een korte prefix van het traject.

4. Resultaten

De auteurs evalueren SAGE op diverse benchmarks uit de D4RL-suite, waaronder locomotie (MuJoCo), navigatie (AntMaze, Maze2D) en manipulatie (Kitchen).

• Prestatieverbetering: SAGE verbetert de prestaties van state-of-the-art diffusieplanners (zoals DV en Diffuser) consistent.
  • Bijvoorbeeld: Op de MuJoCo-locomotie-benchmarks steeg het gemiddelde van 82.9 (DV) naar 84.4.
  • Op manipulatie-taken (Kitchen) steeg de score van 73.6 naar 74.5 (Mixed) en van 90.0 naar 96.6 (Partial).
• Robuustheid: SAGE vermindert het aantal "catastrofische fouten" (zoals doorlopen van muren of vallen) aanzienlijk zonder de diversiteit van de trajecten te vernietigen.
• Validatie van het Energie-signaal: Experimenten met kunstmatig beschadigde trajecten (waarbij acties worden geschud) tonen aan dat de berekende energie een scherpe piek vertoont op het moment van inconsistentie. De AUROC voor het detecteren van deze fouten ligt tussen 0.94 en 0.99, wat aantoont dat het signaal zeer betrouwbaar is.
• Statistische Significantie: De verbeteringen zijn statistisch significant (p-waarde < 0.05) in de meeste domeinen, met uitzondering van Maze2D waar de prestaties al dicht bij het optimum zaten (ceiling effect).

5. Betekenis en Impact

Dit werk biedt een praktische en schaalbare oplossing voor een fundamenteel probleem in offline RL: het vertrouwen op generatieve modellen die soms "droomtrajecten" genereren die in de realiteit onmogelijk zijn.

• Paradigmaverschuiving: In plaats van de generatie te beperken (wat vaak leidt tot minder diverse of suboptimale oplossingen), focust SAGE op het filteren en herordenen van de gegenereerde opties. Dit behoudt de kracht van de generator terwijl de uitvoerbaarheid wordt gewaarborgd.
• Toepasbaarheid: Omdat het een "plug-in" module is die geen extra interactie met de omgeving vereist, is het direct toepasbaar op bestaande systemen in robotica en AI.
• Toekomstige Richting: Het succes van SAGE onderstreept het belang van het scheiden van verschillende beslissingssignalen (waarde vs. haalbaarheid) in complexe planningstaken en opent de deur voor meer zelftoezichtende mechanismen in veilige RL.

Kortom, SAGE maakt diffusieplanners betrouwbaarder en robuuster door een laag van "realiteitscheck" toe te voegen die dynamische inconsistenties detecteert voordat de agent een actie uitvoert.