Wasserstein Proximal Policy Gradient

Deze paper introduceert de Wasserstein Proximal Policy Gradient (WPPG), een methode voor entropie-geregulariseerd reinforcement learning die gebruikmaakt van Wasserstein-geometrie om een globale lineaire convergentie te garanderen en direct toepasbaar is op expressieve impliciete stochastische beleidsstrategieën zonder de log-dichtheid of zijn gradiënt te hoeven evalueren.

De kern

Stel je voor dat je een robot wilt leren lopen. De robot moet beslissen welke bewegingen hij maakt (bijvoorbeeld: "til je linkervoet op" of "duw met je rechterhand"). Dit noemen we Versterkend Leren (Reinforcement Learning). De robot probeert dingen, krijgt beloningen als hij goed doet, en past zijn strategie aan om in de toekomst meer punten te scoren.

Deze paper introduceert een nieuwe, slimme manier om die robot-strategie te verbeteren, genaamd WPPG (Wasserstein Proximal Policy Gradient). Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De "Gokker" vs. De "Rekenaar"

In de oude methoden (zoals PPO of SAC) moet de robot een strategie hebben die hij precies kan uitleggen. Hij moet kunnen zeggen: "De kans dat ik mijn voet 10 cm optil is 30%." Dit heet een expliciete strategie. Het probleem is dat dit vaak te simpel is. De echte wereld is complex, en soms wil je een strategie die heel gekke, creatieve bewegingen maakt die je niet makkelijk in een simpele formule kunt gieten.

De nieuwe methode (WPPG) werkt met impliciete strategieën. Dit is alsof je de robot een "magische doos" geeft. Je zegt: "Hier is een willekeurig getal (een latent variabele), en hier is je huidige positie. Wat doe je?" De robot doet iets, maar jij weet niet precies hoe hij dat getal omzet in een beweging. Je kunt de kansverdeling niet makkelijk uitrekenen, maar je kunt wel heel goed samples (voorbeelden) genereren. Dit is krachtiger, maar heel moeilijk om te optimaliseren met de oude methoden.

2. De Oplossing: Twee Soorten "Stapjes"

De auteurs zeggen: "Laten we de robot niet dwingen om een formule te kennen. Laten we hem gewoon laten bewegen." Ze gebruiken een wiskundig concept uit de Wasserstein-geometrie.

Stel je voor dat de robot-strategie een klontje klei is op een tafel.

• Oude methode (KL-divergentie): Dit is alsof je de klei probeert te vervormen door te kijken naar de "informatie" in de vorm. Het is lastig als de klei een vreemde vorm heeft.
• Nieuwe methode (Wasserstein): Dit is alsof je de klei fysiek schuift over de tafel. Je duwt de klei in de richting waar de beloning het hoogst is.

De WPPG-methode doet dit in twee stappen, net als het bakken van brood:

1. Stap 1: De Duw (Optimal Transport)
   De robot kijkt waar de beloning (de "Q-waarde") het hoogst is. Hij duwt zijn huidige bewegingen een beetje in die richting. Het is alsof je een bal een duw geeft zodat hij naar beneden rolt in een dal.
2. Stap 2: De Warmte (Gaussian Noise)
   Als je alleen duwt, wordt de robot te zeker van zichzelf en stopt hij met verkennen (hij blijft in één klein punt hangen). Daarom voegen ze een beetje "ruis" of "warmte" toe. Dit is alsof je de klei een beetje laat smelten en weer laten stollen. Hierdoor blijft de robot een beetje willekeurig en blijft hij nieuwe dingen uitproberen. In de wiskunde heet dit het toevoegen van Gaussisch ruis.

3. Waarom is dit zo cool?

Het grootste voordeel is dat je niet hoeft te weten hoe de robot denkt.

• Bij oude methoden moest je de robot vragen: "Wat is de kans dat je dit doet?" (De log-dichtheid). Als je een complexe "impliciete" robot hebt, kan hij dat antwoord vaak niet geven.
• Bij WPPG hoef je dat niet te vragen. Je zegt gewoon: "Hier is een willekeurige input, hier is de beloning, en hier is de richting waarin je moet bewegen." Je hoeft alleen maar te kijken hoe de robot reageert op een kleine verandering.

Dit maakt het mogelijk om veel complexere, krachtigere robots te bouwen die niet vastzitten in simpele wiskundige formules.

4. Het Resultaat: Sneller en Slimmer

De auteurs hebben bewezen dat deze methode wiskundig gegarandeerd naar de beste oplossing convergeert (het "globale minimum" bereikt). In de praktijk hebben ze het getest op zware taken, zoals het laten lopen van een humanoïde robot (een robot met een menselijk lichaam).

• De vergelijking: Ze hebben WPPG vergeleken met de beste bestaande methoden (zoals SAC en PPO).
• De uitkomst: De nieuwe methode (WPPG) deed het net zo goed als de beste bestaande methoden, maar de versie met de "impliciete" robot (WPPG-I) deed het zelfs beter dan alle anderen. De robot leerde sneller en bereikte hogere scores.

Samenvattend in één zin:

Deze paper introduceert een slimme manier om robots te leren bewegen door hun strategieën te "schuiven" in de richting van de beloning en ze tegelijkertijd een beetje "wilde" te houden met ruis, zonder dat je hoeft te weten hoe hun interne logica precies werkt. Hierdoor kunnen ze complexere en betere bewegingen leren dan ooit tevoren.

Technische samenvatting

Probleemstelling

Reinforcement Learning (RL) met continue actie-ruimtes maakt vaak gebruik van policy gradient-methoden. Traditionele methoden, zoals TRPO en PPO, baseren zich op de Euclidische geometrie van de parameterruimte of op informatie-geometrie (via Kullback-Leibler divergentie). Deze benaderingen hebben echter beperkingen:

1. Afhankelijkheid van dichtheidsfuncties: Ze vereisen dat de log-dichtheid ($\log \pi(a|s)$) en de gradiënt daarvan expliciet berekenbaar zijn. Dit maakt ze ongeschikt voor complexe, "impliciete" stochastische beleidsfuncties (zoals pushforward maps) waarbij de dichtheid niet in gesloten vorm bestaat.
2. Geometrie van de actie-ruimte: KL-gebaseerde methoden behandelen acties vaak als onafhankelijke categorieën en negeren de onderliggende meetkunde van de continue actie-ruimte.
3. Theoretische lacunes: Hoewel er eerdere werk is over Wasserstein-gradienten in RL, ontbreekt het aan sterke convergentiegaranties voor continue actie-ruimtes, vooral voor parametrische beleidsfuncties die verder gaan dan deeltjesbenaderingen.

Het doel van dit artikel is een nieuwe methode te ontwikkelen die de voordelen van de Wasserstein-geometrie benut (respect voor de actie-ruimte) zonder de noodzaak van expliciete dichtheidsberekeningen, en dit te onderbouwen met globale convergentiebewijzen.

Methodologie: Wasserstein Proximal Policy Gradient (WPPG)

De auteurs introduceren WPPG, een algoritme dat beleidsupdates formuleert in de ruimte van kansverdelingen onder de 2-Wasserstein-metriek ($W_2$), gecombineerd met entropieregulering.

Kernconcepten:

1. Wasserstein Proximal Update: In plaats van een KL-straf, wordt een proximal update uitgevoerd in de $W_2$-ruimte. Dit maximaliseert de verwachte waarde (Q-waarde) onder een straal van de huidige beleidsverdeling, gemeten in $W_2$.
2. Operator-Splitting Schem: De auteurs ontleden de complexe proximal update in twee afzonderlijke stappen via een operator-splitting (Lie-Trotter):
   • Stap 1: Optimal Transport (Drift): Een update die de acties verschuift in de richting van de gradiënt van de Q-functie ($\nabla_a Q$) om de verwachte beloning te verhogen.
   • Stap 2: Heat Flow (Entropie): Een stap die Gaussisch ruis injecteert om de entropieregulering te modelleren. Wiskundig komt dit overeen met het convolueren van de verdeling met een Gaussische kern.
3. Omgaan met Impliciete Beleidsfuncties:
   • Het algoritme is ontworpen voor impliciete beleidsfuncties, gedefinieerd als $a = g_\theta(s, Z)$, waarbij $Z$ een eenvoudige verdeling is (bijv. standaardnormaal).
   • Door de entropie-stap te implementeren als het toevoegen van ruis aan de output van de generator ($g_\theta$), en de drift-stap te optimaliseren via de actie-gradiënt van de critic, hoeft het algoritme geen log-dichtheid of score-functie van het beleid te evalueren.
   • Dit maakt de methode direct toepasbaar op expressieve generatieve modellen (zoals diepe neural networks) die complexe, multimodale actie-distributies kunnen modelleren.

Convergentieanalyse:
De auteurs bewijzen een lineaire globale convergentiesnelheid voor WPPG.

• Ze gebruiken de $T_2$-transport-informatie-ongelijkheid (een verband tussen KL-divergentie en $W_2$) in plaats van de gebruikelijke "three-point identity" die bij KL-gebaseerde methoden (zoals Mirror Descent) wordt gebruikt.
• De analyse geldt zowel voor exacte Q-waarden als voor benaderde Q-waarden (Actor-Critic setting) met gecontroleerde schattingsfouten.

Belangrijkste Bijdragen

1. Nieuw Update-Schema: Een operator-splitting methode die optimal transport en Gaussische ruisinjectie combineert, waardoor beleidsupdates mogelijk zijn zonder expliciete dichtheidsberekeningen.
2. Theoretische Garantie: Het eerste bewijs van lineaire globale convergentie voor een Wasserstein-gebaseerde policy gradient methode in continue actie-ruimtes met entropieregulering, inclusief Actor-Critic implementaties.
3. Toepasbaarheid op Impliciete Modellen: De methode maakt het gebruik van krachtige, impliciete beleidsfuncties (pushforward maps) mogelijk, wat een groot voordeel is ten opzichte van bestaande methoden zoals SAC die beperkt zijn tot families met tractabele dichtheden (bijv. Gaussisch).
4. Empirische Prestaties: WPPG presteert competitief en vaak superieur op standaard benchmarks.

Experimentele Resultaten

De auteurs evalueren WPPG op MuJoCo-benchmarks (o.a. Hopper, Walker2d, Humanoid) en vergelijken het met PPO, SAC en WPO (Wasserstein Policy Optimization).

• WPPG (Expliciet): Met een Gaussisch MLP-beleid presteert het vergelijkbaar met SAC, wat aantoont dat de Wasserstein-geometrie net zo effectief is als de KL-geometrie.
• WPPG-I (Impliciet): Dit is de variant met een impliciete beleidsfunctie (gebaseerd op een latent variable model). WPPG-I overtreft consistent alle baselines (PPO, SAC, WPO) op bijna alle taken.
  • Het bereikt hogere returns, vooral op moeilijke, hoog-dimensionele taken zoals Humanoid.
  • Het bewijst dat de framework natuurlijk uitbreidt naar rijkere architecturen die complexere actie-distributies kunnen leren.
• Ablatie-studies:
  • De grootte van de geinjecteerde ruis ($\tau$) is cruciaal voor de exploratie-exploitatie balans.
  • De dimensie van de latente variabele in WPPG-I moet voldoende groot zijn voor exploratie, maar niet te groot om het leren te vertragen.
  • Het gebruik van "Double-Q" (twee critic netwerken) verbetert de stabiliteit en prestaties aanzienlijk.

Significantie en Impact

Dit werk is een belangrijke stap in de theorie en praktijk van Reinforcement Learning:

• Theoretisch: Het vult een gat in de literatuur door globale convergentie te garanderen voor Wasserstein-methoden in continue ruimtes, wat eerder een open vraag was. Het introduceert nieuwe wiskundige technieken (gebaseerd op $T_2$-ongelijkheden) die alternatief zijn voor de standaard KL-analyse.
• Praktisch: Het opent de deur voor het gebruik van zeer expressieve, impliciete beleidsfuncties in RL. Omdat veel moderne generatieve modellen (zoals Normalizing Flows of Diffusion Models) geen eenvoudige log-dichtheid hebben, biedt WPPG een oplossing om deze modellen direct te gebruiken voor RL zonder complexe aanpassingen.
• Toekomst: De methode suggereert dat de geometrie van de actie-ruimte een fundamentele rol speelt in het leren van complexe gedragspatronen, en dat het benutten van deze geometrie (via Wasserstein) superieure resultaten kan opleveren ten opzichte van traditionele KL-benaderingen.