Advantage-Guided Diffusion for Model-Based Reinforcement Learning

Dit paper introduceert AGD-MBRL, een methode die de voordeelgeschattingen van een agent gebruikt om het reverse diffusion-proces te sturen, waardoor model-gebaseerde versterkingslering op korte horizon myopie overwint en aanzienlijk betere sample-efficiëntie en prestaties bereikt dan bestaande baselines.

De kern

Stel je voor dat je een robot wilt leren een complexe dans uitvoeren, zoals het rennen over een ongelijk terrein of het balanceren op één been. Dit is wat Reinforcement Learning (RL) doet: een agent leert door te proberen, te vallen en te leren van zijn fouten.

Maar er is een probleem: als je de robot alleen maar in de echte wereld laat oefenen, kost dat veel tijd en energie (en in het geval van een echte robot, veel geld en risico). Daarom gebruiken wetenschappers Model-Based Reinforcement Learning (MBRL). In plaats van direct in de echte wereld te oefenen, bouwt de robot eerst een virtuele wereld (een "droom") en oefent daar.

Het oude probleem: De "Stapelende Foutjes"

Vroeger maakten deze virtuele werelden een simpele fout: ze voorspelden de volgende stap, en gebruikten die voorspelling om de stap daarna te voorspellen.

• Analogie: Stel je voor dat je een verhaal vertelt aan een vriend, die het verhaal aan een ander vertelt, en die weer aan een derde. Na een paar mensen is het verhaal compleet anders dan het origineel.
• In de robotwereld noemen we dit compounding errors. Als de robot één klein foutje maakt in zijn droom, wordt dat foutje bij elke volgende stap groter, tot zijn droom volledig onherkenbaar en nutteloos is.

De nieuwe oplossing: Diffusie (Het "Denoising" Proces)

Recente technologie, genaamd Diffusiemodellen, lost dit op. In plaats van stap voor stap te voorspellen, "droomt" de robot het hele verhaal in één keer.

• Analogie: Stel je voor dat je een foto van een landschap hebt die volledig met ruis (witte vlekjes) is bedekt. Een diffusiemodel is als een kunstenaar die langzaam, stap voor stap, de ruis wegpoetst totdat het hele landschap helder verschijnt. De kunstenaar kijkt niet naar één boom en zegt "daarna komt een boom", maar ziet het hele landschap tegelijk en maakt het helder. Dit voorkomt dat foutjes zich stapelen.

Het nieuwe probleem: Korte Visie (Myopie)

Maar er zit nog een addertje onder het gras. Soms is de "droom" van de robot te kort. Hij ziet misschien alleen de eerste paar seconden van de dans.

• Het oude systeem: Als de robot alleen kijkt naar de beloning (punten) die hij in die korte droom krijgt, kan hij in de valkuil trappen.
• Analogie: Stel je bent een wandelaar die een berg wil beklimmen. Je kijkt alleen naar de eerste 10 meter. Je ziet een pad dat direct omhoog gaat (veel punten nu), maar dat pad leidt naar een afgrond. Een ander pad gaat eerst een beetje omlaag (minder punten nu), maar leidt naar de top. Een robot met "korte visie" kiest het eerste pad en valt de afgrond in, omdat hij niet ziet wat er na die 10 meter gebeurt.

De oplossing van dit papier: De "Advantage" Gids

De auteurs van dit paper (Daniele Foffano en collega's) hebben een slimme oplossing bedacht: Advantage-Guided Diffusion (AGD-MBRL).

In plaats van de robot alleen te laten kijken naar de punten die hij nu krijgt, geven ze hem een kompas dat hem vertelt welke keuze op de lange termijn het beste is. Dit komas heet de Advantage-functie.

• Analogie: Stel je voor dat de robot een ervaren trainer heeft. De trainer zegt niet: "Kijk, hier krijg je een snoepje!" (dat is de korte beloning). De trainer zegt: "Als je hier naar links gaat, krijg je nu misschien minder snoepjes, maar over 10 minuten sta je op de top van de berg en krijg je een hele zak snoepjes."
• De robot gebruikt dit advies om zijn "droom" (de diffusie) te sturen. Hij gaat niet zomaar willekeurig dromen, maar hij droomt specifiek scenario's waarin hij op de lange termijn wint.

Hoe werkt het precies? (De twee manieren)

De auteurs hebben twee manieren bedacht om dit komas te gebruiken:

1. SAG (De Voorzichtige Trainer): Deze methode is als een trainer die zegt: "Als een keuze duidelijk beter is, doe het dan, maar wees niet te overtuigd als je niet zeker bent." Het is veilig en stabiel.
2. EAG (De Enthousiaste Trainer): Deze methode is als een trainer die roept: "Als iets er goed uitziet, ga daar dan direct voor!" Het is krachtiger en kan sneller leren, maar vereist dat de trainer (het model) zijn zaakjes goed heeft.

Wat leverde dit op?

De robot met dit nieuwe komas (AGD-MBRL) heeft getest op verschillende moeilijke taken (zoals rennen, springen en balanceren).

• Resultaat: Hij leerde sneller en beter dan eerdere methoden.
• Vergelijking: Terwijl andere robots soms vastliepen in slechte strategieën (door hun korte visie), vond deze robot de beste routes naar de top. In sommige gevallen was hij twee keer zo efficiënt als de beste bestaande methoden.

Samenvatting in één zin

Deze paper introduceert een slimme manier om robots te laten dromen van de toekomst, waarbij ze niet alleen kijken naar de punten die ze nu krijgen, maar gebruikmaken van een slim kompas dat hen wijst naar de beste lange-termijn resultaten, waardoor ze veel sneller en slimmer leren.

Technische samenvatting

Titel: Advantage-Guided Diffusion for Model-Based Reinforcement Learning (AGD-MBRL)

Auteurs: Daniele Foffano, Arvid Eriksson, David Broman, Karl H. Johansson, Alexandre Proutiere (KTH Royal Institute of Technology).

---

1. Het Probleem

Model-based Reinforcement Learning (MBRL) probeert de sample-efficiëntie te verbeteren door een wereldmodel te leren en daarbinnen te plannen. Traditionele autoregressieve wereldmodellen (die de volgende toestand voorspellen op basis van de huidige) lijden echter onder opstapeling van fouten (compounding errors): kleine voorspellingsfouten in één stap versterken elkaar over een lange horizon, wat leidt tot degradatie van de prestaties.

Diffusiemodellen bieden een oplossing door trajecten (reeksen van toestanden en acties) gezamenlijk te genereren in plaats van stap-voor-stap, waardoor opstapeling van fouten wordt geminimaliseerd. Echter, bestaande geleide diffusiemethoden (guided diffusion) voor RL hebben twee belangrijke beperkingen:

1. Beleid-alleen (Policy-only): Methoden zoals PolyGRAD sturen het generatieproces alleen op basis van het huidige beleid, maar negeren waardevolle informatie over de waarde (value) van toestanden.
2. Beloningsgebaseerd (Reward-based) en Myopisch: Methoden zoals Diffuser sturen het proces aan op basis van de cumulatieve beloning binnen het gegenereerde venster. Dit is kortzichtige (myopische) planning, vooral wanneer het diffusie-horizon (H) kort is ten opzichte van de werkelijke taakhorizon. Het model ziet dan de langetermijnwaarde van toekomstige toestanden (buiten het venster) niet, wat kan leiden tot het kiezen van suboptimale trajecten die op korte termijn goed lijken maar op lange termijn slecht zijn.

2. Methodologie: AGD-MBRL

De auteurs introduceren Advantage-Guided Diffusion for MBRL (AGD-MBRL). In plaats van alleen te kijken naar de directe beloning of het huidige beleid, gebruiken ze de advantage-functie ($A^\pi(s, a)$) van het RL-agent om het reverse diffusieproces te sturen. De advantage-functie, gedefinieerd als $Q^\pi(s, a) - V^\pi(s)$, geeft aan hoe goed een actie is ten opzichte van het gemiddelde, en houdt rekening met zowel huidige als toekomstige beloningen (langetermijnwaarde).

De methode introduceert twee specifieke geleidingsmechanismen (guides):

1. Sigmoid Advantage Guidance (SAG):

   • Modelleert de optimaliteit van een stap als een Bernoulli-probabiliteit met een sigmoidfunctie: $p(O_t=1|s_t, a_t) = \sigma(A^\pi(s_t, a_t))$.
   • Dit is een conservatieve aanpak. De sigmoidfunctie is begrensd (tussen 0 en 1), wat over-schattingen van de advantage door het RL-agent compenseert. Het geeft een voorkeur voor hoge advantage-waarden, maar dempt extreme waarden.

2. Exponential Advantage Guidance (EAG):

   • Gebruikt een energie-gebaseerde aanpak waarbij de waarschijnlijkheid exponentieel toeneemt met de cumulatieve advantage: $p(\tau) \propto \exp(\sum A^\pi(s_t, a_t))$.
   • Dit is een agressievere aanpak die het generatieproces sterk stuurt naar gebieden met hoge verwachte verbetering. Het convergeert sneller als de advantage-schatting goed is, maar is gevoeliger voor over- of onderschatting.

Wiskundige Fundament:
De auteurs bewijzen dat het sturen van het diffusiemodel met SAG of EAG wiskundig equivalent is aan het uitvoeren van hergewogen steekproefneming (reweighted sampling) van trajecten gegenereerd door een verbeterd beleid ($\pi'$). De gewichten voor trajecten met toestanden met een hoge advantage nemen toe. Onder standaard aannames impliceert dit dat het gegenereerde beleid een hogere waarde heeft dan het oorspronkelijke beleid.

Implementatie:

• De methode integreert naadloos met PolyGRAD-achtige architecturen.
• De geleiding wordt alleen toegepast op de state-componenten van het reverse diffusieproces.
• De actiegeneratie blijft beleids-geconditioneerd (via de PolyGRAD guide), wat zorgt voor consistentie tussen gegenereerde toestanden en acties.
• Er is geen wijziging nodig in het trainingsdoelwit van het diffusiemodel zelf; de geleiding gebeurt alleen tijdens het genereren (sampling).

3. Belangrijkste Bijdragen

1. Identificatie van Myopie: De auteurs tonen aan waarom beloningsgebaseerde geleiding bij korte horizons leidt tot myopische data-generatie en waarom de advantage-functie dit probleem oplost door langetermijnwaarde te incorporeren.
2. Theoretisch Bewijs: Ze leveren formele bewijzen (Propositie V.1 en V.2) dat geleiding via SAG of EAG leidt tot steekproefneming uit een beleid met een hogere verwachte waarde ($J(\pi') \geq J(\pi)$).
3. Empirische Validatie: De methode wordt getest op MuJoCo-taken (HalfCheetah, Hopper, Walker2D, Reacher) en overtreft bestaande methoden in sample-efficiëntie en finale prestaties.

4. Resultaten

De experimenten werden uitgevoerd op 1,5 miljoen omgevingstappen en vergeleken met:

• PolyGRAD: De basis architectuur (beleid-geleid).
• Online Diffuser: Een aanpassing van Diffuser (beloningsgeleid).
• Model-vrije baselines: PPO en TRPO.

Kernbevindingen:

• Algemene Superioriteit: AGD-MBRL (zowel met SAG als EAG) presteert over het algemeen beter dan alle andere baselines, vaak met een marge van 2x in vergelijking met model-vrije methoden.
• HalfCheetah: EAG presteert hier significant beter dan SAG, wat suggereert dat in omgevingen waar de waarde-functie makkelijk te schatten is, de agressieve exponentiële geleiding snel tot optimale oplossingen leidt.
• Walker2D: Hier presteert SAG beter dan EAG in de vroege trainingsfasen. De conservatieve aard van de sigmoidfunctie blijkt robuuster in omgevingen waar de optimale waarde-functie moeilijker te benaderen is.
• Stabiliteit: AGD-MBRL toont stabielere leercurves met minder frequentie en magnitude van prestatieretrogressies (regressies) vergeleken met PolyGRAD en Online Diffuser.
• Hopper: Alle methoden convergeren hier naar vergelijkbare prestaties, wat wijst op een makkelijke omgeving, maar AGD-MBRL blijft competitief.

5. Betekenis en Toekomstperspectief

De paper toont aan dat advantage-aware guidance een eenvoudige maar krachtige remedie is voor het probleem van kortzichtige planning in diffusie-model MBRL. Door de advantage-functie te gebruiken, kunnen agenten trajecten genereren die niet alleen op korte termijn belonend lijken, maar die ook leiden tot een betere langetermijnstrategie.

Beperkingen en Toekomst:

• De grootste beperking blijft de rekenkosten van diffusiemodellen, omdat het genereren van trajecten veel iteratieve stappen vereist.
• Toekomstig werk richt zich op het versnellen van de generatie, bijvoorbeeld door generatie in een latente ruimte of het gebruik van flow matching.

Conclusie: AGD-MBRL biedt een principieel verbeterde manier om wereldmodellen te gebruiken voor planning, waarbij de synthetische data die door het model wordt gegenereerd direct wordt geoptimaliseerd voor langetermijn succes in plaats van alleen korte-termijn beloning.