Technische Samenvatting: Schaalbare Maximum Entropy Reinforcement Learning voor Diffusiebeleid via Adjoint Matching

Probleemstelling

Diffusiebeleid (diffusion policies) zijn een krachtig paradigma geworden voor het representeren van complexe, multimodale actieverdelingen in Reinforcement Learning (RL). De toepassing ervan op online RL wordt echter beperkt door de uitdaging van schaalbare training in afwezigheid van grondwaarheid-data (ground-truth data). Standaard trainingsdoelstellingen zoals score matching of bridge matching zijn niet direct toepasbaar omdat online RL geen statische dataset van optimale monsters bezit.

Bestaande benaderingen voor het trainen van diffusiebeleid in online settings worden geconfronteerd met drie primaire beperkingen:

1. Reverse KL Divergentie: Methoden die de reverse KL-divergentie minimaliseren, vereisen backpropagation door het gehele diffusieproces. Dit veroorzaakt enorme geheugenoverhead en computationele inefficiëntie, wat opschaling naar grotere problemen verhindert.
2. Importance Weighted Matching: Benaderingen die importance sampling gebruiken om fouten tussen voorspelde en doelacties te minimaliseren, hebben moeite met het schalen naar hoogdimensionele actieruimtes.
3. Q-Score Methoden: Technieken die gebruikmaken van de gradiënt van de Q-functie ($\nabla_a Q$) vertrouwen vaak op overdempte Langevin-dynamica. Deze processen vereisen doorgaans afkaping na een eindig aantal stappen, wat weinig theoretische garanties biedt in het eindige tijdsregime, of vertrouwen op heuristische samplings-technieken die geen garanties bieden voor het genereren van monsters uit de ware doelverdeling.

Methodologie: Adjoint Matching Diffusion Policy (AMDP)

De auteurs stellen Adjoint Matching Diffusion Policy (AMDP) voor, een raamwerk dat de optimalisatie van diffusiebeleid voor Maximum Entropy RL (MaxEnt-RL) formuleert als een Stochastisch Optimaal Controle (SOC) probleem. De kerninnovatie is de adaptatie van Adjoint Matching (AM), specifiek Reciprocal Adjoint Matching (RAM), om simulatievrije training mogelijk te maken.

1. Theoretische Fundering

De methode formuleert de beleidsverbeteringsstap als het vinden van een driftfunctie $u$ zodanig dat de marginale verdeling van een diffusieproces op de terminale tijd $\tau=1$ overeenkomt met het optimale MaxEnt-beleid $\pi^*(a|s) \propto \exp(Q^\pi(s,a)/\alpha)$.

• Schrödinger Bridge: Het probleem wordt geformuleerd als een Schrödinger bridge-probleem, waarbij de kinetische kosten van de controle $u$ worden geminimaliseerd onder randvoorwaarden.
• Reciprocal Adjoint Matching: In plaats van door het diffusiepad te backpropageren (wat het opslaan van tussenliggende toestanden vereist), maakt AMDP gebruik van de reciproque projectie van het gecontroleerde proces op een ongecontroleerd referentieproces. Dit maakt het mogelijk om de loss te minimaliseren door tussenliggende toestanden $X_\tau$ direct te samplen uit de terminale staat $X_1$ met behulp van een conditioneel reciproque proces $\Pi^0_{\tau|1}$.
• Simulatievrije Training: Door $X_\tau$ te samplen uit de terminale staat $X_1$ (die in een replay buffer wordt opgeslagen), vermijdt het algoritme het simuleren van de voorwaartse SDE tijdens de gradiënt-update. Dit omzeilt de noodzaak van expliciete likelihood-schatting en kostbare backpropagation door de diffusieketen.

2. Belangrijke Algoritmische Componenten

• Actieruimte Adaptatie (Error Function Squashing):
  Om gebonden actieruimtes te hanteren (bijv. $A = [-1, 1]^d$), introduceren de auteurs een differentieerbare transformatie $f: \mathbb{R}^d \to [-1, 1]^d$. In tegenstelling tot de standaard elementgewijze $\tanh$-squashing, gebruikt AMDP een geschaalde error function (erf) squashing:
  $$f_{\text{erf}}(x) = (\text{erf}(k x_i))_{i=1}^d$$
  De schaalfactor $k$ wordt gekozen zodanig dat de determinant van de Jacobiaan van de transformatie exact overeenkomt met de Gaussische marginale dichtheid van het referentieproces. Deze specifieke keuze zorgt ervoor dat de log-ratio van de referentiedichtheid en de Jacobiaan-determinant een constante wordt, wat de regressiedoelstelling aanzienlijk vereenvoudigt en de numerieke stabiliteit verbetert.

• Beleiditeratie (Policy Iteration):

  • Evaluatie: Er wordt een gemodificeerde "zachte" Bellman-backup gebruikt. Omdat de marginale entropie van diffusiemodellen onhandelbaar is, stellen de auteurs een handelbare entropie-ondergrens voor die afgeleid is van de controlekosten, wat de convergentie-eigenschappen behoudt.
  • Verbetering: Het beleid wordt bijgewerkt door de AMDP-loss te minimaliseren, wat een fixed-point iteratie is. Het unieke vaste punt van deze loss komt overeen met de optimale controle $u^*$.

• Trust-Region Updates:
  Om stabiliteit te garanderen in online settings waar de doelverdeling snel verschuift, wordt een trust-region constraint toegevoegd. Deze straft de afwijking van de nieuwe drift $u$ van de vorige iteratie $u_{\text{old}}$ af, wat agressieve updates voorkomt die kunnen leiden tot distributie-instorting of Q-functie divergentie. Dit wordt opgelost via een gerelaxeerde Lagrangian-formaliteit.

Belangrijkste Bijdragen

1. Schaalbaar Trainingsalgoritme: AMDP introduceert een zeer efficiënt algoritme dat diffusiebeleid optimaliseert zonder door het diffusieproces te backpropageren, wat de geheugenvereisten drastisch vermindert.
2. Simulatievrije Optimalisatie: Door gebruik te maken van reciprocal adjoint matching, maakt AMDP optimalisatie mogelijk met enkel de terminale acties die in een replay buffer zijn opgeslagen, wat meerdere gradiënt-updates per sample mogelijk maakt zonder de SDE opnieuw te simuleren.
3. Numeriek Stabiele Actie-Squashing: Het voorstel van error function squashing biedt een stabiel alternatief voor $\tanh$, waardoor de trainingsoelstelling wordt vereenvoudigd door specifieke termen in de loss-functie te elimineren.
4. Stabiele Beleidsupdates: De integratie van een trust-region loss zorgt voor stabiele beleidsiteratie in online settings zonder de theoretische garanties van het adjoint matching-raamwerk aan te tasten.

Experimentele Resultaten

De auteurs evalueerden AMDP op 63 omgevingen binnen MuJoCo Playground, ManiSkill en HumanoidBench in de on-policy setting, en op 7 hoogdimensionele hond- en humanoïde omgevingen (DMC suite) in de off-policy setting.

• On-Policy Prestaties:

  • AMDP evenaart of overtreft sterke Gaussische baselines (REPPO, PPO, SPO) en state-of-the-art diffusiemethoden (DIME, DPPO, FPO).
  • In complexe, hoogdimensionele taken (HumanoidBench) presteert AMDP aanzienlijk beter dan alle baselines, zelfs zonder Best-of-N (BoN) sampling.
  • AMDP convergeert sneller dan REPPO in sample-efficiëntie op MuJoCo DMC en humanoïde locomotie-taken.
  • Wanneer gecombineerd met BoN sampling (het selecteren van de beste van 16 monsters), bereikt AMDP superieure prestaties over alle benchmarks.

• Off-Policy Prestaties:

  • Op DMC hond- en humanoïde taken evenaart AMDP de prestaties van DIME (die reverse KL optimaliseert) en convergeert het iets sneller, wat de concurrentiekracht van de adjoint matching-doelstelling in off-policy settings aantoont.

• Computationele Efficiëntie:

  • Wall-clock Tijd: AMDP brengt slechts een kleine constante overhead aan (~10%) met zich mee vergeleken met zeer efficiënte Gaussische methoden zoals REPPO, ongeacht het aantal diffusiestappen.
  • Schaalbaarheid: In contrast hiermee zijn methoden die gebruikmaken van reverse KL backpropagation 10x langzamer met 16 diffusiestappen en tot wel 72-80x langzamer met 128 stappen.
  • Model Schaling: De auteurs slaagden erin om een groot diffusiebeleid (106M parameters, ~300x groter dan standaard modellen) te trainen met AMDP, wat hun vermogen aantoont om te schalen naar complexe modellen waar backpropagation-gebaseerde methoden zouden falen vanwege geheugenbeperkingen.

Betekenis en Claims

Het artikel claimt dat AMDP een theoretisch gefundeerd, computationeel efficiënt en robuust pad biedt voor het inzetten van expressieve generatieve beleids in real-world RL-toepassingen.

• Overwinnen van Geheugen-Bottlenecks: Door de noodzaak van backpropagation door de diffusieketen te elimineren, maakt AMDP training van diffusiebeleid levensvatbaar voor online RL-scenario's waar geheugen- en computationele middelen beperkt zijn.
• Theoretische Garanties: In tegenstelling tot heuristische Q-score methoden, is AMDP afgeleid van Stochastic Optimal Control met fixed-point garanties, wat verzekert dat het beleid convergeert naar de optimale MaxEnt-verdeling.
• Praktische Levensvatbaarheid: De combinatie van simulatievrije training, numerieke stabiliteit via erf-squashing en trust-region updates adresseert de primaire barrières (schaalbaarheid, stabiliteit en efficiëntie) die de adoptie van diffusiebeleid in online RL voorheen beperkten.

De auteurs concluderen dat hun aanpak de computationele overhead aanzienlijk vermindert terwijl het competitieve prestaties behoudt, waardoor diffusiebeleid een meer levensvatbaar alternatief maakt voor traditionele Gaussische beleid voor complexe continue controle-taken.