Evolving Diffusion and Flow Matching Policies for Online Reinforcement Learning

Dit paper introduceert GoRL, een framework dat de stabiliteit en expressiviteit van online versterkend leren verbetert door optimalisatie te ontkoppelen van generatie via een tractabele latente ruimte en een conditionele generatieve decoder, wat leidt tot superieure prestaties op complexe controle-taken.

De kern

Stel je voor dat je een robot wilt leren om te dansen. Je hebt twee manieren om dit te doen, maar beide hebben een groot probleem.

Het oude probleem: De "Strakke Dans" vs. De "Chaos"

1. De Strakke Dans (De oude methode):
   Stel je een robot voor die alleen maar rechte lijnen en simpele cirkels kan dansen. Dit is makkelijk te berekenen en stabiel. Hij zakt nooit in elkaar. Maar als de muziek plotseling vraagt om een ingewikkelde salto of een sprong naar links én rechts tegelijk, kan deze robot het niet. Hij probeert het midden te vinden, wat eruitziet als een rare, onnatuurlijke beweging. Hij mist de "expressiviteit".

2. De Chaos (De nieuwe, complexe methoden):
   Nu hebben we robots die kunnen dansen als echte mensen: ze kunnen springen, draaien en complexe bewegingen maken. Dit zijn de "generatieve" modellen (zoals Diffusie of Flow Matching). Ze zijn geweldig in het nabootsen van complexe bewegingen. Maar als je ze probeert te leren terwijl ze dansen (online leren), gaan ze vaak in de war. Ze proberen te veel dingen tegelijk, de berekeningen worden te zwaar, en ze vallen vaak om of leren niets meer. Het is alsof je een beginnende danser probeert te leren terwijl je hem blinddoekt en in een storm zet.

De Oplossing: GORL (De "Twee-Manen Dans")

De auteurs van dit papier hebben een slimme oplossing bedacht genaamd GORL. Ze lossen het probleem op door de robot in twee personen te splitsen: een Leraar en een Danser.

Hier is hoe het werkt, in een simpele analogie:

1. De Twee Personen

• De Leraar (De Latente Encoder): Dit is een simpele, slimme robot die alleen maar "ideeën" bedenkt. Hij werkt in een veilige, simpele wereld waar alles makkelijk te berekenen is. Hij weet precies welke ideeën goed werken en welke niet. Hij is stabiel en leert snel.
• De Danser (De Generatieve Decoder): Dit is de complexe, artistieke robot. Hij kan alles doen: springen, draaien, acrobatiek. Hij is echter niet goed in het leren van nieuwe dingen door zelf te proberen; hij is beter in het uitvoeren van wat hem wordt verteld.

2. De Dansstijl: "Wisselende Tijden"

In plaats dat de Danser en de Leraar tegelijkertijd proberen te leren (wat leidt tot chaos), doen ze het in fases:

• Fase 1: De Leraar denkt na.
  De Danser staat stil (of doet wat hij al kan). De Leraar probeert nieuwe ideeën uit in de veilige wereld. Hij leert: "Als ik dit idee denk, krijg ik een punt." Omdat hij in een simpele wereld werkt, leert hij snel en zonder te vallen.
• Fase 2: De Danser oefent.
  Nu is de Leraar klaar met denken. Hij geeft zijn beste ideeën aan de Danser. De Danser kijkt naar deze ideeën en zegt: "Oké, hoe vertaal ik dit idee naar een echte dansbeweging?" De Danser leert dan hoe hij die ideeën het mooist kan uitvoeren.
• Het Geheim: De "Vaste Anker"
  Dit is het slimste deel. Normaal gesproken zou de Danser proberen te leren van de ideeën die de Leraar nu heeft. Maar de Leraar verandert constant! Als de Danser leert van de Leraar die net verandert, raakt hij in de war (het "zelf-reconstructie" probleem).
  GORL lost dit op door de Danser te laten oefenen op een vaste, simpele lijst van ideeën (een "Gaussische prior"). De Danser leert dan: "Hoe vertaal ik deze simpele ideeën naar de beste dansbewegingen die ik net heb gezien?" Hierdoor wordt de Danser steeds beter, zonder in de war te raken door de veranderende Leraar.

3. Het Resultaat

Door deze cyclus te herhalen (Leraar leert -> Danser oefent op vaste basis -> Leraar leert weer), krijgen ze het beste van twee werelden:

• De stabiliteit van de simpele Leraar (geen vallende robots).
• De expressiviteit van de complexe Danser (mooie, ingewikkelde dansen).

Waarom is dit belangrijk?
In de echte wereld (zoals robots die lopen of auto's die rijden) zijn situaties vaak complex. Soms moet je links springen, soms rechts, en soms rechtuit. Simpele robots (de "Strakke Dans") kunnen dat niet goed. Complexe robots (de "Chaos") vallen om als je ze probeert te leren.

GORL laat zien dat je een robot kunt bouwen die stabiel blijft terwijl hij steeds beter en complexer wordt. Op een heel moeilijke taak (HopperStand, waarbij een robot op één poot moet staan) deed deze methode het meer dan 3 keer zo goed als de beste bestaande methoden.

Kort samengevat:
GORL is als het hebben van een coach (de Leraar) die veilig in een kantoor zit en strategieën bedenkt, en een atleet (de Danser) die op het veld oefent. De coach bedenkt de strategie, de atleet oefent die strategie in, en dan wisselen ze van rol. Zo wordt de atleet steeds sterker zonder dat de coach in paniek raakt, en zonder dat de atleet de weg kwijtraakt.

Technische samenvatting

Probleemstelling: De Spanning tussen Stabiliteit en Expressiviteit

In online Reinforcement Learning (RL) voor continue controle bestaat er een fundamenteel conflict tussen optimalisatiestabiliteit en representatiekracht:

1. Stabiliteit vs. Expressiviteit: Traditionele methoden (zoals PPO en SAC) gebruiken vaak eenvoudige, unimodale verdelingen (zoals Gaussische verdelingen) voor het beleid (policy). Deze zijn stabiel omdat ze analytische likelihoods en gladde gradiënten bieden. Echter, in complexe omgevingen zijn de optimale acties vaak multimodaal (er zijn meerdere verschillende, goed werkende actiestrategieën). Een unimodale verdeling kan deze niet goed modelleren; het "dekken" van meerdere modi leidt tot het toekennen van waarschijnlijkheid aan lage-reward gebieden tussen de modi, wat resulteert in suboptimale prestaties.
2. De Valstrik van Generatieve Modellen: Om dit op te lossen, gebruiken recente werken expressieve generatieve modellen zoals Diffusiemodellen en Flow Matching. Deze kunnen complexe multimodale verdelingen modelleren. In online RL zijn deze echter vaak instabiel omdat:
   • De likelihoods (kansen) vaak onberekenbaar (intractable) zijn.
   • Gradiënten moeten worden teruggepropageerd (backpropagation) door lange generatieve steekproefketens (bijv. tientallen denoising-stappen of ODE-oplossers).
   • Dit leidt tot hoge variantie en instabiliteit wanneer de data-distributie continu verschuift (non-stationarity) tijdens het leren.

Bestaande oplossingen maken vaak compromissen: ze bevriezen de generator (beperkt expressiviteit) of gebruiken surrogate-objectieven die kunnen divergeren van de werkelijke likelihood.

Methodologie: GORL (Generative Online Reinforcement Learning)

De auteurs introduceren GORL, een raamwerk dat dit probleem oplost door een structureel principe toe te passen: ontkoppeling van optimalisatie en generatie.

Kernidee:
In plaats van het beleid direct te optimaliseren in de actieruimte (waar de generatieve decoder zit), wordt de optimalisatie beperkt tot een tractable latente ruimte. De expressieve generatie wordt overgelaten aan een decoder die los van de RL-gradiënten wordt getraind.

De Architectuur:
Het beleid wordt gefactoriseerd in twee componenten:

1. Encoder (Latent Policy) $\pi_\theta(\varepsilon | s)$: Een tractable policy (bijv. een Gaussische verdeling) die een latent variabele $\varepsilon$ genereert op basis van de staat $s$. Deze wordt geoptimaliseerd met standaard RL-algoritmen (zoals PPO).
2. Decoder (Generative Decoder) $g_\phi(s, \varepsilon)$: Een conditioneel generatief model (Diffusie of Flow Matching) dat de latent variabele $\varepsilon$ omzet in een actie $a$. Deze heeft geen expliciete likelihood nodig voor de RL-stap.

Het Twee-Tijdschaal Alternend Trainingsproces:
GORL gebruikt een cyclus van twee fasen om stabiliteit en expressiviteit te combineren:

• Fase 1: Encoder Optimalisatie (RL-stap):

  • De decoder $g_\phi$ wordt bevroren.
  • De encoder $\pi_\theta$ wordt geoptimaliseerd in de latente ruimte om de verwachte beloning te maximaliseren. Omdat de decoder vaststaat, gedraagt het zich als een deterministische dynamica.
  • Hierdoor blijven de RL-gradiënten stabiel en tractable (geen backprop door de generatieve keten).

• Fase 2: Decoder Verfijning (Supervised Learning-stap):

  • De encoder $\pi_\theta$ wordt bevroren.
  • De decoder $g_\phi$ wordt bijgewerkt via gesuperviseerde generatieve training op recente rollouts.
  • Cruciaal Innovatiepunt (Fixed-Prior Anchoring): In plaats van de decoder te trainen op latente waarden gegenereerd door de huidige encoder (wat zou leiden tot een "self-reconstruction" lus en weinig winst), worden de invoer voor de decoder getrokken uit een vaste prior (bijv. een standaard Gaussische verdeling $N(0, I)$).
  • Dit forceert de decoder om de verbeteringen die de encoder heeft gevonden te consolideren in een sterkere generatieve kaart, zonder vast te komen zitten in de drift van de encoder.

• Stadia-gewijze Re-initialisatie: Aan het begin van elke nieuwe trainingsfase wordt de encoder opnieuw geïnitieerd naar de prior. Dit voorkomt dat de encoder "vastloopt" in een slechte configuratie ten opzichte van de vernieuwde decoder, en zorgt voor een stabiele start voor de volgende cyclus van verbetering.

Belangrijkste Bijdragen

1. Analyse van Instabiliteit: Een theoretische analyse (in Appendix A) die aantoont waarom generatieve beleidsmodellen (Diffusie/Flow) inherent instabiel zijn in online RL vanwege onberekenbare likelihoods en gradiënten die door diepe steekproefketens moeten.
2. GORL Raamwerk: Het introduceren van een algoritme-onafhankelijk raamwerk dat optimalisatie en generatie structureel ontkoppelt. Dit combineert de stabiliteit van latente RL met de expressiviteit van generatieve modellen.
3. Theoretische Garantie: Bewijzen dat gradiënten in de latente ruimte onbevooroordeelde schatters zijn voor het totale beleid en dat beperkte divergentie in de latente ruimte leidt tot een beperkte prestatiedaling in de actieruimte.
4. Empirisch Succes: Demonstratie dat GORL consistent beter presteert dan zowel unimodale baselines (Gaussian PPO) als recente generatieve baselines (FPO, DPPO) op diverse continue controle-taken.

Resultaten

De auteurs hebben GORL getest op zes taken uit de DMControl Suite (o.a. CheetahRun, WalkerWalk, HopperStand).

• Prestaties: GORL (zowel met Flow Matching als Diffusie) behaalde hogere uiteindelijke returns dan alle baselines.
• HopperStand (Kritieke Test): Op deze uitdagende taak (balanceren van een hopper) behaalde GORL een episodische return van > 870. Dit is meer dan 3 keer zo hoog als de sterkste baseline (die rond de 286 bleef steken).
• Stabiliteit: Terwijl generatieve baselines zoals FPO vaak instabiel waren (prestaties daalden halverwege het trainen en herstelden niet), bleef GORL stabiel leren.
• Multimodaliteit: Visuele analyse van de actie-distributies toont aan dat GORL succesvol evolueert van een unimodale naar een bimodale verdeling (twee duidelijke pieken voor verschillende stabiele strategieën), terwijl Gaussische PPO beperkt blijft tot één modus.
• Ablatie Studies:
  • Zonder "Fixed-Prior Anchoring" (d.w.z. trainen op de driftende encoder) stort het systeem in.
  • Zonder "Stage-wise Re-initialization" daalt de prestatie drastisch bij het overstappen naar een nieuwe decoder-fase.

Betekenis en Conclusie

GORL biedt een praktische weg naar het creëren van RL-beleidsmodellen die zowel stabiel zijn (zodat ze kunnen leren in real-time omgevingen) als hoog expressief zijn (zodat ze complexe, multimodale strategieën kunnen leren).

De paper lost een langdurig probleem op in het veld: hoe maak je gebruik van de kracht van moderne generatieve AI (Diffusie/Flow) binnen de strenge eisen van online Reinforcement Learning zonder instabiliteit. Door de "ontkoppeling" en het gebruik van een vaste prior voor de decoder, creëren de auteurs een virtuele cyclus waarin het beleid en de generatieve capaciteit elkaar wederzijds versterken. Dit opent de deur voor robuustere en slimmere robots en besturingssystemen die complexe, niet-lineaire taken kunnen aanpakken.