MODIP: Efficient Model-Based Optimization for Diffusion Policies

Technische Samenvatting: MODIP – Efficiënte Model-Based Optimalisatie voor Diffusiebeleid

1. Probleemstelling

Diffusiebeleid (Diffusion Policies, DPs) is een krachtige representatie geworden voor robotbesturing, met name in imitatie-leren (Behavioral Cloning, BC), vanwege het vermogen om complexe, multimodale actieverdelingen te modelleren. Er bestaat echter een fundamentele beperking: DPs die uitsluitend getraind zijn op offline demonstratiedata, hebben de neiging om het gedrag van de dataset te reproduceren zonder de capaciteit om daar voorbij te gaan.

Hoewel Reinforcement Learning (RL) een pad biedt om vooraf getrainde beleidsregels te verbeteren, is het direct finetunen van DPs met RL uitdagend. In tegenstelling tot standaard Gaussische beleidsregels genereren DPs acties via een iteratief meerstaps denoising-proces. Dit maakt beleidsoptimalisatie complex, wat vaak leidt tot instabiliteit en hoge computationele kosten tijdens training en inferentie. Bestaande methoden die proberen direct RL fine-tuning toe te passen (bijv. Diffusion-QL, DPPO), worstelen met deze inefficiënties. Het artikel stelt dat een effectievere aanpak is om directe RL-optimalisatie van de diffusie-objectief te vermijden en in plaats daarvan model-gebaseerde planning te gebruiken om verbeterde trajecten te genereren voor supervised distillatie.

2. Methodologie: MODIP

MODIP (Model-Based Optimization for Diffusion Policies) is een framework ontworpen voor offline-naar-online fine-tuning van DPs. Het werkt door een latent wereldmodel te gebruiken om een Model Predictive Control (MPC) planner te sturen, die hoogwaardige trajecten genereert. Deze trajecten worden vervolgens gebruikt als supervised targets om de DP te finetunen via standaard Behavioral Cloning (BC), waardoor een iteratieve verbeteringslus ontstaat zonder de kern van het denoising-objectief van de DP te wijzigen.

Het framework bestaat uit vijf geleerde componenten:

1. Encoder ($h_\eta$): Brengt observaties in kaart naar latente toestanden.
2. Latent Dynamics Model ($f_\theta$): Voorspelt toekomstige latente toestanden gegeven huidige toestanden en acties.
3. Reward Model ($\hat{R}_\xi$): Scoort korte-horizon rollouts in de verbeelding.
4. Terminal Value Function ($V_\psi$): Schat de lange-horizon returns in voorbij de MPC-planninghorizon.
5. Diffusion Policy ($\pi_\phi$): Dient als een expressieve actiesequentie-prior voor de planner.

Belangrijke Technische Mechanismen:

• MPC-naar-Policy Distillatie: In plaats van de DP direct te optimaliseren met een RL-loss, gebruikt MODIP een DP-gestuurde MPC-planner om acties te selecteren tijdens online interactie. De planner verfijnt kandidaat-actiesequenties in de latente ruimte. De resulterende verbeterde trajecten worden opgeslagen in een online buffer, gemengd met offline data, en gebruikt om de DP bij te werken met de standaard supervised denoising loss. Dit stelt het beleid in staat om te verbeteren door een betere datadistributie, terwijl de stabiliteit van supervised learning behouden blijft.
• Hybride Planning met Terminal State Value: De planner gebruikt Model Predictive Path Integral (MPPI) controle. Een cruciale efficiëntie-innovatie is het gebruik van een terminale state value $V(s)$ in plaats van een beleidsafhankelijke state-action waarde $Q(s, \pi(s))$.
  • Standaard Aanpak: Het gebruik van $Q(s, \pi(s))$ vereist het uitvoeren van het volledige meerstaps denoising-proces van de DP bij de terminale staat voor elke kandidaat-traject, wat computationeel duur is.
  • MODIP Aanpak: Het gebruik van $V(s)$ maakt een enkele forward pass mogelijk om lange-horizon returns te schatten, wat de inferentietijd aanzienlijk vermindert.
• Beleid-onafhankelijk Critic Leren: Om de trainingskosten verder te verlagen, ontkoppelt MODIP het critic-leren van het huidige beleid. Standaard actor-critic methoden vereisen vaak het bevragen van het huidige beleid om Temporal Difference (TD) targets te berekenen (bijv. $Q(s', \pi(s'))$). Voor DPs is dit kostbaar. MODIP vormt TD-targets met state-value schattingen ($V(s')$) van target-netwerken, waardoor herhaalde DP denoising-iteraties tijdens critic-updates worden vermeden.

Trainingsprocedure:

1. Offline Pretraining: De DP en alle MPC-componenten (dynamics, reward, value) worden getraind op de offline dataset.
2. Online Fine-tuning: De agent gebruikt de DP-gestuurde MPC-planner om te interageren met de omgeving. Trajecten worden opgeslagen in een online buffer. Updates worden uitgevoerd op minibatches die offline en online data mengen, met een tijdvariërende ratio die geleidelijk de focus verschuift naar online, door de planner verbeterde data.

3. Belangrijkste Bijdragen

Het artikel schetst vier primaire bijdragen:

1. Hybride DP-gestuurde MPC Planning: Introductie van een planner die een DP gebruikt als een multimodale actiesequentie-prior, waarbij voorstellen worden verfijnd via latente MPC.
2. MPC-naar-Policy Distillatie: Een methode om DPs te finetunen met supervised learning op trajecten die zijn verbeterd door de MPC-planner, waardoor directe RL-optimalisatie van het diffusie-objectief wordt vermeden.
3. Efficiënte Value-gebaseerde Traject Scoring: De vervanging van de standaard terminale estimate $Q(s, \pi(s))$ door een terminale state value $V(s)$, wat kostbare DP-queries bij terminale staten elimineert en de planning-tijd inferentie-kosten vermindert.
4. Beleid-onafhankelijke Critic Targets: Een mechanisme om de waarde-lering te ontkoppelen van de huidige DP door state-value schattingen te gebruiken voor TD-targets, wat de computationele kosten van critic-updates en de trainingstijd vermindert.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op drie benchmark-families: D4RL MuJoCo (locomotie), D4RL Kitchen (long-horizon manipulatie) en RoboMimic (state-based manipulatie).

• Prestaties vs. Directe RL Fine-tuning: MODIP presteert consistent beter dan of gelijk aan directe RL fine-tuning methoden (Diffusion-QL, DPPO, DSRL, PA-RL) over de meeste taken. Met name op de HalfCheetah taak presteren MODIP en PA-RL significant beter dan andere DP-baselines, toegeschreven aan het vermogen van MPC om bij elke stap te herplannen en fouten effectiever te herstellen dan methoden die lange actie-chunks uitvoeren.
• Prestaties vs. Model-Based Baselines: MODIP presteert beter dan TD-MPC2 (een sterke hybride baseline met een Gaussische policy prior) op alle MuJoCo taken, wat suggereert dat de expressieve DP prior een superieur startpunt biedt voor hybride planning vergeleken met standaard MLP policies.
• Ablatie Studies:
  • MPC-naar-Policy Distillatie: Het verwijderen van de MPC-verfijning (Policy-only) verslechtert de prestaties op de HalfCheetah en Kitchen taken aanzienlijk, wat de waarde van model-gebaseerde verfijning bevestigt.
  • Terminale Value ($V$ vs. $Q$): Het gebruik van $V(s)$ in plaats van $Q(s, \pi(s))$ verbeterde de prestaties over alle taken en verminderde de inferentietijd met 2.90x (van ~85s naar ~29s per 1000 stappen).
  • Critic Learning: Beleid-onafhankelijke targets verminderden de totale trainingstijd met 1.6x (van ~40h naar ~25h) door dure DP-queries tijdens critic-updates te vermijden.
  • Prior Parameterisatie: Het gebruik van een DP prior presteerde aanzienlijk beter dan een Gaussische MLP policy prior, met name op de Hopper en Kitchen taken.

5. Betekenis en Claims

Het artikel claimt dat MODIP erin slaagt de kloof te overbruggen tussen de expressieve kracht van Diffusiebeleid en de sample-efficiëntie van Model-Based RL. Door de MPC-planner te behanden als een indirecte beleidsverbeteringsoperator, behoudt MODIP de stabiliteit en eenvoud van supervised learning, terwijl het de mogelijkheid biedt om de beperkingen van de initiële offline dataset te overstijgen.

De auteurs benadrukken dat de specifieke ontwerpkeuzes — het gebruik van een terminale state value $V(s)$ en beleid-onafhankelijke critic targets — niet louter optimalisaties zijn, maar essentieel zijn om DP-gestuurde planning computationeel haalbaar te maken. Het werk demonstreert dat het combineren van een multimodale prior met model-gebaseerde planning superieure resultaten oplevert vergeleken met zowel directe RL fine-tuning van DPs als traditionele model-gebaseerde methoden met Gaussische policies.

Beperkingen: De auteurs merken op dat hoewel diffusiebeleid veelzijdig is voor het leren van pixels, alle benchmarks in deze studie gebruikmaakten van real-valued vector state informatie. Het beoordelen van het potentieel van MODIP om direct van afbeeldingen te leren, wordt geïdentificeerd als een noodzakelijke volgende stap voor toekomstig werk.