Technische Samenvatting: Score-Based One-step MeanFlow Policy Optimization (SOM)

1. Probleemstelling

Versterkingsleren (RL) in continue actieruimten heeft steeds vaker generatieve modellen, zoals diffusie en flow matching, aangenomen om de expressiviteitsbeperkingen van unimodale Gaussische beleidsstrategieën te overwinnen. Deze modellen kunnen complexe, multimodale actieverdelingen vastleggen, waardoor een rijker onderzoek en diverse optimale gedragingen mogelijk worden. Een kritieke bottleneck blijft echter bestaan: standaard diffusie- en flow-matching-beleidsstrategieën vereisen iteratieve, meestaps denoising om een enkele actie te genereren. Dit introduceert aanzienlijke computatiekosten bij inferentie, wat bijzonder schadelijk is voor online RL, waarbij beleidsinferentie continu wordt afgewisseld met datacollectie.

Hoewel MeanFlow als oplossing is voorgesteld om het samplingproces te reduceren tot één enkele netwerkbewerking door het leren van een gemiddeld snelheidsveld, wordt de directe toepassing ervan op online RL gehinderd door een fundamentele vereiste: het construeren van het doelsnelheidsveld vereist doorgaans steekproeven uit de doelverdeling (het optimale beleid). In online RL zijn deze grondwaarheidssteekproeven niet beschikbaar, aangezien de agent de doelverdeling gelijktijdig met het beleid moet leren. Bestaande benaderingen die proberen dit te omzeilen (bijvoorbeeld Best-of-N sampling) hebben vaak moeite met schaalbaarheid naarmate de actiedimensie toeneemt.

2. Methodologie

De auteurs stellen Score-Based One-step MeanFlow Policy Optimization (SOM) voor, een actor-critic-algoritme dat MeanFlow-training in volledig online RL mogelijk maakt zonder expliciete steekproeven uit de doelverdeling te vereisen. De kerninnovatie ligt in het construeren van het doelsnelheidsveld direct vanuit de Q-functie via schatting van de score.

2.1 Tijdsafhankelijke Energiefunctie

De methode behandelt de negatieve Q-functie als de energie van de doelverdeling. In een RL-kader met entropie-geregulariseerd volgen de optimale beleidsstrategieën een Boltzmann-verdeling:
$$\pi^*(a|s) \propto \exp(\alpha Q(s, a))$$
waarbij $\alpha$ een temperatuurparameter is. De auteurs definiëren een tijdsafhankelijke energiefunctie $Q_t(s, a_t)$ die deze Boltzmann-structuur behoudt over de diffusietijdstippen $t \in [0, 1]$. De score van de verdeling op elk tijdstip $t$ wordt afgeleid als:
$$\nabla_{a_t} \log p_t(a_t|s) = \nabla_{a_t} Q_t(s, a_t)$$
Cruciaal is dat, hoewel het berekenen van $Q_t$ in gesloten vorm onuitvoerbaar is, de methode alleen zijn gradiënt vereist, die benaderd kan worden.

2.2 Score-schatting via iDEM-aanpassing

Om de score $\nabla_{a_t} Q_t(s, a_t)$ te schatten zonder doelsteekproeven, past SOM de Iterated Denoising Energy Matching (iDEM)-schatter aan. In plaats van te samplingen uit de ware doelverdeling, samplet de methode kandidaat-acties uit een Gaussisch voorstel dat is gecentreerd rond de huidige verstoordetoestand $a_t$. Deze kandidaten worden gewogen op basis van hun Q-waarden om de posterior-verwachting te benaderen. De resulterende scoreschatting is:
$$\nabla_{a_t} Q_t(s, a_t) \approx \nabla_{a_t} \log \sum_{i=1}^K \exp(\alpha Q(s, a^{(i)}_{0|t}))$$
waarbij $a^{(i)}_{0|t}$ steekproeven zijn getrokken uit een Gaussische verdeling die is geconditioneerd op $a_t$. Dit beïnvloedt effectief de actievernieuwingen naar gebieden met hoge waarden, terwijl exploratie via het ruisproces behouden blijft.

2.3 Constructie van het Snelheidsveld

De geschatte score wordt geïntegreerd in een Probability Flow ODE (PF-ODE) om het doelsnelheidsveld af te leiden. Voor een Variance Preserving (VP) SDE wordt het snelheidsveld gedefinieerd als:
$$v_t(a_t|s) = -\frac{1}{2}\beta(t) \left( a_t + \nabla_{a_t} Q_t(s, a_t) \right)$$
Dit snelheidsveld transporteert ruis direct naar hoogwaardige modi die worden geïnduceerd door de huidige Q-functie. Het MeanFlow-beleid $u_\theta$ wordt vervolgens getraind om te voldoen aan de MeanFlow-identiteit met behulp van dit geconstrueerde doelsnelheidsveld, waarbij het verlies wordt geminimaliseerd:
$$\mathcal{L}(\theta) = \mathbb{E} \left[ \| u_\theta(a_t, r, t, s) - \text{sg}(u^{\text{new}}_{\text{tgt}}) \|^2_2 \right]$$
waarbij $u^{\text{new}}_{\text{tgt}}$ is afgeleid uit de PF-ODE-snelheid.

2.4 Implementatiedetails

• Normalisatie: De scoregradiënt wordt genormaliseerd naar een eenheidsnorm en herschaald met een instelbare coëfficiënt $w$ om training te stabiliseren over variërende beloningschalen.
• Critic-training: De critic wordt bijgewerkt met een distributionele Bellman-doelstelling met twin Q-netwerken.
• Best-of-N (BoN): Hoewel de primaire bijdrage de score-gebaseerde snelheid is, incorporeren de auteurs BoN-sampling tijdens de critic-bijwerkingen om training te versnellen door de kandidaat met de hoogste Q-waarde uit meerdere steekproeven te selecteren voor de schatting van de doelwaarde.

3. Belangrijkste Bijdragen

1. Identificatie van een Bottleneck: Het artikel identificeert het gebrek aan steekproeven uit de doelverdeling als de primaire barrière voor het toepassen van MeanFlow in online RL.
2. Score-gebaseerde Snelheidsconstructie: Het stelt een nieuwe methode voor om het doelsnelheidsveld direct te construeren vanuit Q-functie-gradiënten, waardoor de noodzaak van expliciete doelsteekproeven wordt omzeild.
3. SOM-algoritme: De introductie van SOM, een actor-critic-algoritme dat MeanFlow integreert met online RL, wat expressief beleidsleren mogelijk maakt met enkelstaps actiesampling.
4. Empirische Validatie: Aantonen dat SOM state-of-the-art prestaties behaalt op locomotietaken, terwijl het tegelijkertijd de trainings- en inferentietijd aanzienlijk reduceert in vergelijking met iteratieve diffusie- en flow-matching-baselines.

4. Experimentele Resultaten

De auteurs hebben SOM geëvalueerd op negen MuJoCo v4 continue controletaken (waaronder Ant, HalfCheetah, Hopper, Humanoid, Walker2d, Swimmer, Pusher, Reacher en InvertedDoublePendulum).

• Prestaties: SOM behaalde de hoogste gemiddelde return op vier van de vijf locomotie-benchmarks en de op één na hoogste op Hopper-v4. De verbetering was het meest uitgesproken op Humanoid-v4 (hoge actiedimensie), waar SOM de op één na beste methode met een duidelijke marge versloeg, wat suggereert dat de op gradiënten gebaseerde aanpak gunstig schaalt met de dimensie.
• Efficiëntie: Zoals getoond in Tabel 2, bereikt SOM een enkelstaps inferentietijd van 0,218 ms op HalfCheetah-v4, vergeleken met 0,6–1,5 ms voor iteratieve baselines (die 10–100 denoising-stappen vereisen). Ook de trainingstijd is aanzienlijk gereduceerd.
• Beleidsanalyse: Visualisaties (t-SNE en Q-landschappen) geven aan dat SOM actiesteekproeven strak concentreert rond hoogwaardige Q-modi, terwijl concurrerende enkelstaps methoden (zoals MFP) de neiging hebben in te storten tot enkele modi of diffuus te verdelen.
• Robuustheid: Experimenten met verstoord Q-functies toonden aan dat de door Gaussische ruis gesmeerde scoreschatting in SOM robuust blijft, waarbij grove critic-fouten effectief worden gemiddeld terwijl de lokale structuur nabij hoogwaardige modi behouden blijft.

5. Betekenis en Beweringen

Het artikel beweert dat SOM de spanning oplost tussen de expressiviteit van generatieve beleidsstrategieën en de computatiebeperkingen van online RL. Door het doelsnelheidsveld direct af te leiden uit de Q-functie, stelt SOM MeanFlow-beleidsstrategieën in staat om in volledig online omgevingen te worden getraind waar doelsteekproeven niet beschikbaar zijn.

De auteurs benadrukken dat SOM state-of-the-art prestaties bereikt met één generatiestap, wat een praktische oplossing biedt voor real-time controletoepassingen waar inferentielatentie kritiek is. Ze merken op dat hoewel de methode momenteel beperkt is tot continue controletaken (MuJoCo) en afhankelijk is van de kwaliteit van de geleerde critic, het een belangrijke stap voorwaarts is naar efficiënte, expressieve beleidsoptimalisatie. Het werk suggereert dat toekomstige richtingen het uitbreiden van SOM naar hogedimensionale of gedeeltelijk waarneembare omgevingen (bijvoorbeeld visuele controle) en het integreren van geavanceerdere onzekerheidsbewuste critics omvatten.