ScoRe-Flow: Complete Distributional Control via Score-Based Reinforcement Learning for Flow Matching

Titel

ScoRe-Flow: Complete Distributional Control via Score-Based Reinforcement Learning for Flow Matching

1. Het Probleem

Flow Matching (FM) is een krachtige techniek voor robotbesturing die snellere actiegeneratie biedt dan traditionele diffusiemodellen, door actiepaden te modelleren via deterministische gewone differentiaalvergelijkingen (ODE's). Echter, FM-beleid dat uitsluitend via imitatielearning (IL) is getraind, heeft twee fundamentele beperkingen:

1. Suboptimale prestaties: Het beleid kan de kwaliteit van de demonstratie-data niet overtreffen en lijdt aan covariaatverschuiving.
2. Deterministische dynamiek: Omdat FM ODE's gebruikt, zijn de overgangen deterministisch (Dirac-maten). Dit maakt de waarschijnlijkheidsdichtheid van acties ongedefinieerd, wat het direct toepassen van Reinforcement Learning (RL) methoden zoals Policy Gradient (bijv. PPO) onmogelijk maakt, aangezien deze een tractabele log-likelihood vereisen.

Bestaande oplossingen om FM te fine-tunen met RL (zoals ReinFlow) converteren de deterministische stroom naar een Stochastische Differentiaalvergelijking (SDE) door ruis toe te voegen. Hoewel dit exploratie mogelijk maakt, beperkt deze "noise-only" aanpak de controle: de ruis perturbeert alleen de positie, maar modificeert niet de onderliggende dynamiek (de drift) van de stroom. Dit kan leiden tot inefficiëntie, vooral wanneer de demonstraties al sterke priors bieden.

2. Methodologie: ScoRe-Flow

De auteurs introduceren ScoRe-Flow, een methode die RL-finetuning mogelijk maakt door decoupled control (ontkoppelde controle) over de gemiddelde waarde (drift) en variantie van stochastische overgangen.

Kernconcepten:

• Score-Drift Dualiteit: De auteurs maken gebruik van de bekende dualiteit tussen de snelheidsvector (velocity field) van een lineaire Flow Matching-paad en de scorefunctie (de gradiënt van de log-dichtheid). Voor lineaire paden kan de score $s_t(a)$ in gesloten vorm worden berekend uit de reeds getrainde snelheidsvector $v_\theta$:
  $$s_t(a) = \frac{t \cdot v_\theta(t, a, s) - a}{1 - t}$$
  Dit elimineert de noodzaak voor extra netwerken om de score te schatten.
• Drift Modulatie: In plaats van alleen ruis toe te voegen, moduleren ze de drift-term van de SDE met de scorefunctie. De score wijst naar gebieden met hoge waarschijnlijkheid, waardoor het beleid stochastische trajecten kan sturen naar de "data manifold" (de ruimte van goede acties).
• Decoupled Control: ScoRe-Flow introduceert twee leerbare componenten:
  1. Een score-scheduler ($\alpha_\psi$) die de sterkte van de drift-correctie bepaalt.
  2. Een variantie-predictor ($\sigma_\phi$) die de grootte van de stochastische exploratie onafhankelijk bepaalt.
     Dit stelt het beleid in staat om de richting (via de score) en de intensiteit (via de variantie) van exploratie onafhankelijk van elkaar aan te passen.

Formulering:
De discrete updateregel voor de actie $a_{k+1}$ wordt:
$$a_{k+1} = a_k + \left[ v_\theta + \alpha_{scaled}(t) \cdot s_t(a_k) \right] \Delta t + \sigma_\phi \sqrt{\Delta t} \epsilon_k$$
Waarbij $\alpha_{scaled}(t)$ een tijdsafhankelijke afname bevat om numerieke stabiliteit te garanderen (voorkomend dat de term divergeert bij $t \to 1$).

3. Belangrijkste Bijdragen

1. Inzicht in Dynamiek: De auteurs tonen aan dat ruis-only controle de dynamiek van de stroom niet beïnvloedt. Ze introduceren het gebruik van de scorefunctie als een principieel mechanisme om de drift te sturen en exploratie naar hoog-waarschijnlijke gebieden te leiden.
2. ScoRe-Flow Architectuur: Een nieuwe RL-finetuning-methode die gesloten-vorm scores gebruikt voor drift-modulatie zonder extra netwerken, gecombineerd met een leerbare variantie-predictor. Dit resulteert in de eerste methode die volledige, ongekoppelde controle over gemiddelde en variantie biedt binnen Flow Matching.
3. Efficiëntie en Stabiliteit: De methode vereist geen extra parameters voor de scoreberekening en biedt aanzienlijke snelheidswinsten ten opzichte van bestaande methoden.

4. Resultaten

De methode werd geëvalueerd op drie benchmarks: D4RL (locomotie), Robomimic (visuele manipulatie) en Franka Kitchen (meerdere taken met schaarse beloningen).

• Convergentiesnelheid: ScoRe-Flow convergeert 2.4x sneller dan de huidige state-of-the-art (ReinFlow) op D4RL-taken bij een gelijk aantal denoising-stappen.
• Prestaties:
  • Op Robomimic (PickPlaceCan, Square, Transport) behaalde ScoRe-Flow tot 5.4% hogere success rates dan de beste baselines.
  • Op Franka Kitchen (een complexe, schaarse beloning omgeving) overtrof het alle baselines, met name op de "Partial" dataset (3.8 vs 3.4 sub-taken voltooid).
  • Op Humanoid-v3 (D4RL) behaalde het een nieuwe state-of-the-art score van $5100 \pm 47$.
• Efficiëntie: Door het gebruik van Flow Matching (4-10x minder stappen dan diffusiemodellen) is ScoRe-Flow structureel veel sneller dan diffusie-gebaseerde RL-methoden (zoals DPPO), met een snelheidswinst van tot 21.9x in wandelklok-tijd.
• Stabiliteit: In tegenstelling tot ReinFlow, dat hoge variantie en instabiel gedrag vertoont, toont ScoRe-Flow stabiele leercurves dankzij de score-gestuurde drift.

5. Betekenis en Impact

ScoRe-Flow is een doorbraak in het gebied van robotbesturing met generatieve modellen. Het lost het fundamentele probleem op van het toepassen van RL op Flow Matching door de deterministische aard van ODE's te overbruggen met een wiskundig elegante, maar praktische SDE-conversie.

De belangrijkste implicaties zijn:

• Superieure Sample-efficiëntie: Robotbeleid kan sneller leren dan de demonstraties, wat essentieel is voor complexe taken waar menselijke demonstraties suboptimaal zijn.
• Geen extra overhead: Het gebruik van de gesloten-vorm score betekent dat de methode geen zware extra netwerken vereist, waardoor deze schaalbaar is voor grote robotica-systemen.
• Robuustheid: De ontkoppelde controle van drift en variantie maakt het beleid robuuster tegenover instabiele exploratie, wat cruciaal is voor real-world robottoepassingen met schaarse beloningen.

Samenvattend biedt ScoRe-Flow een complete oplossing voor distributiecontrole in Flow Matching, wat leidt tot snellere training, hogere prestaties en meer stabiel gedrag in complexe robotica-taken.