Recovering Hidden Reward in Diffusion-Based Policies

1. Probleemstelling

Op diffusie gebaseerde beleidsstrategieën zijn een state-of-the-art paradigma geworden voor het leren van robotische manipulatievaardigheden uit expertdemonstraties, vanwege hun vermogen om complexe, multimodale actie-distributies te modelleren. Deze beleidsstrategieën vertrouwen echter doorgaans op Behavior Cloning (BC), wat leert om handelingen na te bootsen zonder expliciet de onderliggende intentie of beloningsstructuur te modelleren.

• Beperkingen van Standaard Diffusie: Ze missen een expliciete energievoorstelling, waardoor ze vatbaar zijn voor falen wanneer situaties tijdens het testen afwijken van de trainingsdistributie (Out-of-Distribution of OOD). Ze kunnen niet eenvoudig een beloningssignaal leveren voor downstream Reinforcement Learning (RL) zonder te resorteren tot instabiele adversarial training of computationeel duur Inverse Reinforcement Learning (IRL).
• Het Gat: Er is behoefte aan een raamwerk dat generatieve actie-modellering verenigt met beloningsherwinning, waardoor zowel hoogwaardige imitatie als de extractie van een geldig beloningssignaal voor beleidsverfijning mogelijk wordt, terwijl tegelijkertijd theoretische garanties op generalisatie worden gewaarborgd.

2. Methodologie: ENERGYFLOW

De auteurs stellen ENERGYFLOW voor, een raamwerk dat een scalar energiefunctie $E_\theta(s, a)$ parametrisert, waarvan de gradiënt dient als het denoising-veld voor actiegeneratie.

Kernarchitectuur

• Scalar Energieparametrisatie: In tegenstelling tot standaard diffusiebeleid dat direct een vectorwaarde ruisvoorspelling $\epsilon_\theta(s, a)$ regresseert, leert ENERGYFLOW een scalar energiefunctie $E_\theta(s, a, t)$.
• Conservatief Vectorveld: De scorefunctie (denoising-veld) wordt afgeleid via automatische differentiatie:
  $$S_\theta(a, s, t) = -\nabla_a E_\theta(a, s, t)$$
  Door constructie zorgt dit ervoor dat het vectorveld conservatief is (rotatievrij, $\nabla \times S = 0$), wat betekent dat het de gradiënt is van een scalar potentiaal. Dit dwingt integreerbaarheid af, waardoor ervoor wordt gezorgd dat actievoorkeuren transitief en fysiek realiseerbaar zijn.
• Training: Het model wordt getraind met Denoising Score Matching (DSM). Het doel minimaliseert het verschil tussen de geleerde gradiënt en de ware ruisrichting:
  $$\mathcal{L} = \mathbb{E}_{t, a_0, \epsilon} \left[ \sigma^2(t) \left\| -\nabla_{a_t} E_\theta(a_t, s, t) + \frac{\epsilon}{\sigma(t)} \right\|^2 \right]$$
• Inferentie: Acties worden gegenereerd door de probability-flow ODE te integreren met behulp van de gradiënt van de geleerde energiefunctie.

Beloningsextractie

Het raamwerk haalt een beloningssignaal direct uit het geleerde energielandschap. Om om te gaan met de staat-afhankelijke integratieconstante $c(s)$ (die ruwe energiewaarden ambigu maakt over verschillende staten heen), stellen de auteurs Centered Shaping voor:
$$\tilde{r}(a, s) = -\left( E_\theta(a, s) - \mathbb{E}_{a' \sim \mathcal{N}}[E_\theta(a', s)] \right)$$
Dit trekt een staat-afhankelijke baseline af (geschat via Monte Carlo-sampling), waardoor de beloning wordt genormaliseerd om alleen relatieve actievoorkeuren binnen een staat weer te geven.

3. Belangrijkste Theoretische Bijdragen

Het paper biedt rigoureuze theoretische bewijzen die het verband aantonen tussen diffusie-scorematching en beloningsherwinning:

1. Score-Beloning Equivalentie: Onder de aanname van Maximum Entropy Optimaliteit bewijst het paper dat de scorefunctie die door diffusiemodellen wordt geleerd evenredig is met de gradiënt van de zachte Q-functie van de expert ($\nabla_a Q^* \propto S^*$). Scorematching herstelt dus impliciet de beloningsgradiënt zonder adversarial min-max optimalisatie.
2. Complexiteidsreductie via Conservatieve Beperkingen: De auteurs bewijzen dat het beperken van de hypothese-ruimte tot conservatieve velden (gradiënten van scalairen) de Rademacher-complexiteit strikt verlaagt in vergelijking met onbeperkte vectorvelden. Specifiek schaalt de complexiteitsgrens voor conservatieve velden met de gladheid van de representatie ($L$), terwijl onbeperkte velden schalen met de outputdimensie ($\sqrt{d}$).
3. OOD Generalisatiegrenzen: Door gebruik te maken van de gereduceerde complexiteit, leidt het paper strakkere generalisatiegrenzen af voor Out-of-Distribution (OOD) scenario's. De conservatieve beperking fungeert als een geometrische regularisator, waardoor het model wordt voorkomen dat het willekeurige vectorrichtingen leert en wordt gedwongen om de vorm van het energielandschap te extrapoleren, wat leidt tot betere robuustheid onder distributieveranderingen.
4. Identificeerbaarheid en Foutpropagatie: Het paper stelt vast dat hoewel de absolute Q-waarde niet uniek kan worden geïdentificeerd (vanwege de onbekende staat-afhankelijke constante), de actie-rangschikking binnen een staat exact is. Bovendien wordt de propagatie van score-schatfouten naar actievoorkeuren begrensd, waarbij wordt aangetoond dat fouten lineair schalen met de afstand tussen acties (Lipschitz-continuïteit).

4. Experimentele Resultaten

De auteurs hebben ENERGYFLOW geëvalueerd op RoboMimic en Meta-World benchmarks, evenals op real-world robottaken.

• Imitatieprestaties (RQ1): ENERGYFLOW behaalde State-of-the-Art (SOTA) prestaties.
  • RoboMimic: 93,8% gemiddelde succesratio (vs. 91,2% voor Diffusion Policy).
  • Meta-World: 92,5% gemiddelde succesratio (vs. 90,7% voor Diffusion Policy).
  • Het presteerde aanzienlijk beter dan andere op energie gebaseerde methoden (zoals Implicit BC) en IRL-baselines.
• Real-world Implementatie (RQ2): Geïmplementeerd op een AGIBOT G1 robot voor "Bottle" en "Drawer" taken. ENERGYFLOW behaalde een 100% succesratio met soepele trajecten en minder aarzelingen in vergelijking met baselines.
• Kwaliteit van Beloning voor Downstream RL (RQ3): Wanneer gebruikt als beloningssignaal om een Soft Actor-Critic (SAC) agent te trainen:
  • De "Centered Energy"-beloning stelde SAC in staat om bijna-orakel prestaties te bereiken, aanzienlijk beter dan spaarzame beloningen en ruwe energiebeloningen.
  • Het bood dichte, stap-voor-stap begeleiding die vroege trainingsplateaus voorkwam.
• OOD Generalisatie (RQ4): Bij toenemende verstoringen van de beginpositie behield ENERGYFLOW robuuste succesratio's waar onbeperkte diffusie- en flow-beleidsstrategieën aanzienlijk degradeerden. Dit bevestigt het theoretische voordeel van de conservatieve beperking.
• Efficiëntie (RQ6): ENERGYFLOW behaalde inferentiesnelheden vergelijkbaar met niet-energie-gebaseerde Flow Policies (~10ms latentie) terwijl het de voordelen van een expliciete energiefunctie bood, waardoor de trage bemonsteringstijden van traditionele Energy-Based Models (EBMs) werden vermeden.

5. Betekenis en Impact

• Vereniging van Imitatie en RL: ENERGYFLOW overbrugt de kloof tussen Behavior Cloning en Reinforcement Learning. Het staat toe dat een beleidsstrategie via imitatie wordt getraind terwijl tegelijkertijd een geldig beloningssignaal wordt blootgelegd voor verdere verfijning, waardoor de noodzaak voor handmatig beloningsontwerp of instabiele adversarial IRL wordt geëlimineerd.
• Theoretische Garanties voor Diffusie: Het werk biedt de eerste theoretische rechtvaardiging voor waarom het afdwingen van een conservatief veld (integreerbaarheid) in diffusiebeleid voordelig is, door dit te koppelen aan gereduceerde hypothese-complexiteit en verbeterde OOD-generalisatie.
• Praktisch Nut: De methode is computationeel efficiënt, schaalbaar naar hoog-dimensionale actieruimtes en effectief in real-world robotische omgevingen, en biedt een robuuste oplossing voor het leren van complexe manipulatievaardigheden uit beperkte demonstraties.

Kortom, ENERGYFLOW demonstreert dat de structurele beperkingen die nodig zijn voor geldige beloningsextractie (conservatieve velden) tegelijkertijd dienen als krachtige inductieve biases die de generalisatie en robuustheid van op diffusie gebaseerde beleidsstrategieën versterken.