Technische Samenvatting: Trust Region Q-Adjoint Matching (TRQAM)

Probleemstelling

Het artikel adresseert de uitdaging van off-policy fine-tuning voor voorgeprogrammeerde flow policies. Hoewel flow-matching policies rijke, multi-modale actie-distributies bieden die superieur zijn aan conventionele unimodale Gaussische policies, maakt hun impliciete definitie via multi-stap denoising-processen gradient-based policy improvement instabiel. Directe backpropagatie door de sampling-keten is rekenkundig duur en vatbaar voor instabiliteit.

Bestaande oplossingen proberen dit te omzeilen door ofwel additieve residualen te leren (waarbij multi-stap dynamica wordt genegeerd) ofwel te opereren in de ruïtenspace (beperkt door de expressiviteit van de bevroren policy). Een recente aanpak, Q-learning met Adjoint Matching (QAM), herformuleert fine-tuning als een geheugenloos Stochastisch Optimaal Besturings (SOC) probleem, waarbij een geleerde criticus het sampling-proces leidt via adjoint matching. De auteurs identificeren echter een fundamentele kwetsbaarheid in QAM: in off-policy settings is de geleerde criticus onvermijdelijk imperfect. Kleine benaderingsfouten in de criticus worden exponentieel versterkt wanneer de policy-update vertrouwt op een vaste "temperatuur" (of inverse temperatuur $\beta$), wat leidt tot destructieve drift van de voorgeprogrammeerde prior en uiteindelijk modelinstorting. Deze instabiliteit blijft bestaan, zelfs met gradient clipping.

Methodologie: Trust Region Q-Adjoint Matching (TRQAM)

TRQAM stelt een stabiel off-policy fine-tuning-algoritme voor dat een trust-region mechanisme direct integreert in de SOC sampling-dynamica. De kerninnovatie is de adaptieve controle van de Kullback-Leibler (KL) divergentie in de padruimte tussen de gefinetunteerde policy en de voorgeprogrammeerde prior.

Belangrijke Theoretische Inzichten

1. Exponentiële Versterking van Fouten: De auteurs formaliseren (Lemma 1) dat adjoint matching met vaste temperatuur lijdt aan een exponentiële versterking van criticusfouten. Een vaste $\beta$ kan niet tegelijkertijd een betrouwbare criticus benutten en bescherming bieden tegen een onbetrouwbare.
2. $\lambda$ als Trust-Region Parameter: TRQAM introduceert een scalair parameter $\lambda$ die de diffusiecoëfficiënt in de SOC SDE schaalt met $\sqrt{\lambda}$.
3. Exacte KL Identiteit (Stelling 1): Met behulp van de stelling van Girsanov bewijzen de auteurs dat het schalen van de diffusie met $\sqrt{\lambda}$ de KL-divergentie in de padruimte maakt tot een exacte, gesloten-vorm functie van $\lambda$:
   $$D_{KL}(P_u \parallel P_{base}) = \mathbb{E}_{X \sim P_u} \left[ \frac{1}{2\lambda} \int_0^1 \|u(X_\tau, \tau)\|^2 d\tau \right]$$
   Dit vestigt een directe structurele link tussen de parameter $\lambda$ en de afwijking van de basispolicy.
4. Terminal KL Bound (Propositie 1): De KL-divergentie in de padruimte boven-begint de KL van de terminal actie-distributie, waardoor het garanderen dat het controleren van de KL in de padruimte de uiteindelijke policy-afwijking effectief beperkt.

Algoritmisch Kader

In plaats van een KL-term toe te voegen als externe loss-straf (die kan worden genegeerd door sterke criticus-gradiënten), internaliseert TRQAM $\lambda$ binnen de sampling-dynamica:

• Sampling: De gecontroleerde SDE gebruikt een diffusiecoëfficiënt $\sqrt{\lambda}\sigma(\tau)$. Het aanpassen van $\lambda$ herschikt de SDE zelf, waardoor de gecontroleerde dynamica dichter bij of verder van de basisdynamica wordt getrokken.
• Adaptatie: Het algoritme maakt gebruik van geprojecteerde duale daling om $\lambda$ dynamisch aan te passen. Het schat de gerealiseerde KL in de padruimte ($\hat{D}_n$) via een Monte Carlo-schader afgeleid van de gediscretiseerde OT-geheugenloze sampler.
• Update-regel: $\lambda$ wordt bijgewerkt om een doel-KL-budget $\epsilon_{KL}$ af te dwingen:
  $$\lambda_{n+1} \leftarrow \max\{0, \lambda_n + \eta_\lambda (D_n - \epsilon_{KL})\}$$
  Als de gerealiseerde KL het budget overschrijdt, neemt $\lambda$ toe, waardoor de dynamica conservatiever wordt; als deze onder het budget valt, neemt $\lambda$ af, waardoor agressievere verbetering mogelijk wordt.

Belangrijkste Bijdragen

1. Identificatie van Kwetsbaarheid: Het artikel identificeert en formaliseert de exponentiële versterking van criticusfouten bij adjoint matching met vaste temperatuur, en toont aan dat dit leidt tot modelinstorting op benchmarks zoals Robomimic.
2. Theoretische Link: Het bewijst dat het schalen van de diffusiecoëfficiënt met $\sqrt{\lambda}$ een exacte identiteit creëert tussen de besturingskosten en de KL in de padruimte, waardoor $\lambda$ een principieel trust-region-parameter wordt.
3. Algoritme-ontwerp: TRQAM wordt voorgesteld als een methode die de trust-region-parameter internaliseert in de SOC-dynamica en deze aanpast via duale daling, waarbij de KL-bound op het sampling-niveau wordt afgedwongen in plaats van als een zachte loss-straf.
4. Empirische Superioriteit: Uitgebreide experimenten tonen aan dat TRQAM stabiele convergentie bereikt waar eerdere methoden falen en state-of-the-art baselines overtreft.

Experimentele Resultaten

De auteurs evalueren TRQAM op 50 taken uit de OGBench-benchmark en de Robomimic-manipulatiebenchmark.

• Offline RL-prestaties: Op OGBench bereikt TRQAM een algehele successrate van 68%, wat aanzienlijk beter is dan de sterkste baseline (DSRL met 46%) en de eerdere adjoint-matching-methoden (QAM met 35%, QAM-E met 45%). De winsten zijn vooral opvallend in taken met lange horizon en combinatorische taken.
• Stabiliteit: Op Robomimic vertonen methoden met vaste temperatuur (QAM, QAM-E) exploderende adjoint-lossen (tot $10^{20}$) en successraten die instorten tot bijna nul over verschillende hyperparameter-instellingen. Daarentegen blijft TRQAM stabiel over alle geteste budgetten.
• Validatie van het Mechanisme: Ablatiestudies bevestigen dat:
  • Adaptatie noodzakelijk is: Adaptieve KL-budgetten presteren beter dan vaste budgetten.
  • Internalisatie noodzakelijk is: Het internaliseren van $\lambda$ in de SDE (TRQAM) houdt de KL-budget nauwkeurig bij, terwijl externe KL-regulatie toelaat dat de gerealiseerde KL ver afwijkt van het doel, wat leidt tot prestatiedaling.
• Sensitiviteit: Het doel-KL-budget $\epsilon_{KL}$ fungeert als een voorspelbare regelaar; strakkere budgetten leveren over het algemeen betere prestaties op, hoewel de optimale waarde systematisch varieert met de taakstructuur (bijvoorbeeld zijn grotere budgetten gunstig voor taken met grotere toestandenruimtes zoals puzzle-4x4).

Betekenis en Claims

Het artikel claimt dat TRQAM een principieel, stabiel alternatief biedt voor bestaande off-policy fine-tuning-methoden voor flow policies. Door de trust-region-beperking te internaliseren in de sampling-dynamica via de stelling van Girsanov, vermijdt TRQAM de "destructieve drift" die wordt veroorzaakt door imperfecte critics in methoden zoals QAM.

De auteurs positioneren TRQAM als het off-policy equivalent van Trust Region Policy Optimization (TRPO) en Proximal Policy Optimization (PPO) voor on-policy RL. Zij betogen dat voor voorgeprogrammeerde flow policies, die waardevolle gedrags-priors vormen, het handhaven van een gecontroleerde afwijking via een exacte KL-identiteit cruciaal is om catastrofaal vergeten en modelinstorting te voorkomen. De methode wordt gepresenteerd als een robuuste oplossing voor zowel offline RL als de overgang naar online fine-tuning, en biedt een betrouwbare mechanisme om voorgeprogrammeerde vaardigheden te benutten zonder stabiliteit op te offeren.