EXPO: Stable Reinforcement Learning with Expressive Policies

1. Probleemstelling

Het paper adresseert de uitdaging van het fine-tunen van expressieve policy-klassen (zoals Diffusion Policies en Flow-Matching policies) met behulp van Online Reinforcement Learning (RL) gegeven een offline dataset.

• Het Kernconflict: Hoewel expressieve policies uitstekend zijn in het vastleggen van complexe, multi-modale gedragsverdelingen via imitatielearning (IL), zijn ze moeilijk direct te optimaliseren met RL. In tegenstelling tot eenvoudige Gaussische policies worden expressieve policies geparametriseerd door lange denoising-ketens. Het terugpropageren van waarden-gradienten door deze ketens om policy-parameters bij te werken is rekenkundig duur en instabiel, wat vaak leidt tot trainingsschijven of slechte convergentie.
• Het Gat: Bestaande online RL-methoden zijn ontworpen voor eenvoudige Gaussische policies en slagen er niet in om het vermogen van vooraf getrainde expressieve modellen te benutten. Omgekeerd worstelen methoden die proberen expressieve policies te fine-tunen vaak met stabiliteit of vereisen complexe modificaties aan het denoising-proces.
• Doel: Ontwikkelen van een steekproef-efficiënt, stabiel online RL-algoritme dat vooraf getrainde expressieve policies kan fine-tunen met offline data zonder de parameters van de expressieve policy direct te optimaliseren tegen de waardenfunctie.

2. Methodologie: EXPO (Expressive Policy Optimization)

De auteurs stellen EXPO voor, een kader dat de expressieve policy ontkoppelt van directe waardenmaximalisatie. In plaats van gradienten terug te laten propageren door de expressieve policy, bouwt EXPO een "on-the-fly" (OTF) policy die Q-waarden maximaliseert met behulp van twee onderscheidende componenten:

A. Twee-componenten Policy Architectuur

1. Basis Expressieve Policy ($\pi_{base}$):
   • Rol: Genereert initiële actie-candidaten op basis van de gedragsverdeling.
   • Training: Getraind met een stabiel Imitatielearning (IL) doel (bijv. Denoising Diffusion Probabilistic Models). Het wordt nooit expliciet getraind om de Q-waarde te maximaliseren, waardoor de stabiliteit en het vermogen om complexe modi te vangen behouden blijven.
2. Lichtgewicht Edit Policy ($\pi_{edit}$):
   • Rol: Een kleine, Gaussische policy die een actie $a$ neemt die is bemonsterd uit $\pi_{base}$ en een residuele edit $\hat{a}$ output om een verfijnde actie te produceren $\tilde{a} = a + \hat{a}$.
   • Training: Getraind met standaard RL policy gradients (maximaliseren van $Q(s, a + \hat{a})$) met entropieregulering.
   • Beperking: De edit is beperkt (geschaald met een hyperparameter $\beta$) om dicht bij de originele actie te blijven, waardoor de edit policy een lokaal optimalisatieprobleem oplost in plaats van een globaal. Dit zorgt voor efficiënte en stabiele training.

B. On-the-Fly (OTF) Policy Extractie

Het daadwerkelijke beleid dat wordt gebruikt voor bemonstering en Temporal-Difference (TD) backup is geen enkel neuronaal netwerk, maar een dynamisch selectieproces:

1. Bemonstering: Voor een gegeven staat $s$ bemonstert het systeem $N$ acties uit $\pi_{base}$.
2. Bewerken: Elke gebemonsterde actie wordt verfijnd door $\pi_{edit}$.
3. Selectie: Het systeem evalueert de Q-waarde voor zowel de originele acties ($a_i$) als de bewerkte acties ($\tilde{a}_i$).
4. Maximalisatie: De actie met de hoogste Q-waarde onder alle kandidaten wordt geselecteerd voor uitvoering en gebruikt als doel voor de TD-backup.

C. Entropy Backup voor Beperkte Data Regimes

In scenario's met kleine offline datasets introduceren de auteurs een entropy backup. Omdat expressieve policies (zoals diffusie) geen gesloten vorm-entropie hebben, construeren ze een zachte bemonsteringsverdeling over de $N$ basis+bewerkte actie-candidaten, gewogen op basis van hun Q-waarden. Dit maakt het opnemen van een entropiebonus in de RL-verliesfunctie mogelijk om exploratie te stimuleren zonder analytische entropieberekeningen te vereisen.

3. Belangrijkste Bijdragen

1. Stabiele Waardenmaximalisatiestrategie: Het paper introduceert een nieuwe aanpak om directe gradientenpropagatie door expressieve denoising-ketens te vermijden. Door waardenmaximalisatie uit te besteden aan een lichtgewicht edit policy en een selectiemechanisme, bereikt EXPO stabiliteit waar directe optimalisatie faalt.
2. Steekproef-efficiëntie: EXPO bereikt een 2-3x verbetering in steekproef-efficiëntie vergeleken met eerdere state-of-the-art methoden in zowel online RL als offline-naar-online fine-tuning settings.
3. Policy-agnostisch: Het kader is agnostisch ten opzichte van de specifieke parametrisering van de basis policy (werkt met Diffusion, Flow-Matching, etc.) en vereist geen vooraf getrainde waardenfunctie, waardoor het toepasbaar is op elke vooraf getrainde robotische policy.
4. On-the-Fly TD Backup: De auteurs tonen aan dat het gebruik van de waardenmaximaliserende actie voor het TD-doel (niet alleen voor bemonstering) cruciaal is voor leersnelheid, en zo effectief de kloof overbrugt tussen SARSA-achtige en Q-learning-achtige updates in de context van expressieve policies.

4. Experimentele Resultaten

De methode werd geëvalueerd op 12 uitdagende continue controle-taken verspreid over 4 domeinen (AntMaze, Adroit, Robomimic, MimicGen), allemaal met spaarzame beloningen.

• Online RL (Geen Pre-training): EXPO presteerde aanzienlijk beter dan baselines (waaronder RLPD, IDQL en QSM) in steekproef-efficiëntie, en convergeerde vaak sneller, zelfs zonder gebruik te maken van de offline dataset voor pre-training.
• Offline-naar-Online Fine-tuning:
  • EXPO slaagde erin om vooraf getrainde policies te fine-tunen zonder de prestatie-inzakking die vaak wordt gezien bij andere methoden (bijv. DAC stortte in, IDQL stagneerde).
  • Het behield hoge prestaties op fijnmazige manipulatie-taken (bijv. een naald inslaan, blokken stapelen) waar andere methoden moeite hadden om verder te gaan dan de initiële imitatielearning-prestaties.
• Ablatie Studies:
  • OTF TD Backup: Het verwijderen van de waardenmaximaliserende selectie voor het TD-doel (alleen bemonstering gebruiken) verlaagde de prestaties drastisch, wat het belang van Q-waarde-afstemming in de backup bevestigt.
  • Actie Edits: Het verwijderen van de edit policy resulteerde in slechte prestaties, vooral in taken die exploratie vereisten buiten de smalle gedragsverdeling.
  • Kwaliteit van Offline Data: De prestaties correleerden sterk met de kwaliteit van de offline dataset. Echter, de variant met entropy backup liet EXPO toe om zelfs te slagen met zeer slechte offline data (waar imitatielearning alleen <10% succes behaalde).

5. Betekenis

EXPO vertegenwoordigt een belangrijke stap voorwaarts in Robotica en RL door de bottleneck van "expressieve policy fine-tuning" op te lossen.

• Praktische Impact: Het maakt de implementatie van grote, vooraf getrainde generatieve modellen (die uitstekend zijn in het nabootsen van menselijke data) mogelijk in real-world omgevingen waar ze moeten aanpassen en verbeteren via online interactie.
• Theoretisch Inzicht: Het daagt de noodzaak uit om door complexe generatieve ketens te backpropageren voor RL, en laat zien dat een "lokaal edit + globaal selectie"-strategie stabieler en efficiënter is.
• Toekomstige Richting: Het werk opent de deur voor het fine-tunen van elke vooraf getrainde generatieve policy (diffusie, flow, etc.) in online settings, wat potentieel de adoptie van foundation-modellen in robotica kan versnellen.

Beperkingen: De methode vereist het bemonsteren van meerdere acties voor elke batch-update, wat rekenkundig duur is. Bovendien gaat het uit van een redelijke prior (offline dataset of vooraf getrainde policy); het is niet ontworpen voor leren vanaf nul zonder voorafgaande kennis.