Technische Samenvatting: FlowMPC

Probleemstelling

Flow Matching (FM) is een krachtige methode geworden voor behavior cloning in multimodale actieruimtes, met name voor robotmanipulatie waarbij meerdere verschillende actietrajecten dezelfde taak kunnen oplossen. FM is echter fundamenteel een imitation learning-methode die getraind is om expertgedrag te matchen, in plaats van direct de verwachte return te maximaliseren. Daarom zijn FM-policies beperkt door de kwaliteit en de dekking van de demonstratiegegevens. Ze hebben vaak moeite met het herstellen van kleine fouten, het stapelen van fouten (compounding mistakes), of staten die niet goed vertegenwoordigd zijn in de expertdataset.

Hoewel recent werk heeft geprobeerd dit aan te pakken door FM-policies te verfijnen met reinforcement learning (RL), vereisen dergelijke benaderingen het wijzigen van de FM-trainingsdoelstelling zelf. Dit artikel onderzoekt een andere vraag: Kan een geleerd wereldmodel een FM-policy verbeteren zonder de FM-trainingsdoelstelling te wijzigen? Het doel is om te bepalen of modelgebaseerde planning een aanvulling kan vormen op flow-gebaseerde imitatie-policies tijdens de testtijd.

Methodologie: FlowMPC Framework

Het voorgestelde framework, FlowMPC, integreert een via imitation learning geleerde FM-policy met een geleerd wereldmodel binnen een Model Predictive Path Integral (MPPI) planning-loop. Het kernidee is om de FM-policy te gebruiken als een gestructureerd voorstelmechanisme voor het genereren van kandidaat-actiesequenties, die vervolgens worden geëvalueerd en verfijnd door het wereldmodel.

1. Trainingsparadigma

De trainingspijplijn bestaat uit vijf stadia:

1. Expert Data Collectie: Een SAC-agent met geprivilegieerde staatinformatie genereert expertdemonstraties.
2. Traject Collectie: Volledige trajecten (RGB, proprioceptie, acties, beloningen) worden verzameld.
3. FM Policy Training: Een FM-policy wordt getraind op deze expert-trajecten om een multimodale actie-prior te leren.
4. Visuele Preprocessing: Expert-trajecten worden opnieuw verzameld waarbij RGB-observaties zijn voorbewerkt door de visuele encoder van de FM-policy om compacte visuele latents ($z_{viz,t}$) te produceren.
5. Wereldmodel Training: Het wereldmodel wordt online getraind met behulp van een replay buffer die vooraf is geladen met expertdata, waarbij de voorbewerkte visuele latents gecombineerd worden met de proprioceptieve staat ($x_t = [s_t, z_{viz,t}]$).

2. Architecturale Aanpassingen aan TD-MPC2

Het framework bouwt voort op TD-MPC2, maar vervangt de geleerde actor door de via imitation learning geleerde FM-policy. Verschillende kritieke aanpassingen zijn gedaan om deze hybride vorm te accommoderen:

• Behavior-Cloning Action Predictor ($\pi_{BC}$): Een dedicated netwerk brengt de huidige modelstaat in kaart naar een actie. Dit wordt gebruikt voor terminale waarde bootstrapping tijdens planning-rollouts om ervoor te zorgen dat de waardefunctie acties evalueert die binnen de distributie vallen, wat het planningsproces stabiliseert.
• Unified Representation Interface: Het wereldmodel consumeert de FM visuele latent direct geconcateneerd met de proprioceptieve staat, in plaats van aparte encoders te leren. Dit stelt het wereldmodel in staat om de bevroren visuele kenmerken van de FM-policy te benutten.
• Vereenvoudigde Doelstellingen: Gezien de binaire beloningsstructuur van de taken, werden de complexe gediscretiseerde symlog-getransformeerde beloningen van standaard TD-MPC2 vervangen door binaire cross-entropie voor beloningen en standaard gemiddelde kwadratische fout voor Q-waarde regressie.
• Dual World Model Varianten:
  • State-Space Variant: Voorspelt direct in de ruimte van geconcateneerde staat en visuele latent. Gebruikt voor de PickCube taak ($H=1$).
  • Latent-Space Variant: Volgt de standaard TD-MPC2 met een geleerde latente representatie en multi-step rollouts. Gebruikt voor de complexere PickSingleYCB taak ($H=3$).

3. Inferentie en Planning

Tijdens de testtijd werkt de planner als volgt:

1. Voorstel: Gegeven de huidige observatie, samplet de planner $K$ kandidaat-actietrajecten. Een subset ($K_\pi$) wordt voorgesteld door de FM-policy, terwijl de rest willekeurig wordt gesampled.
2. Rollout & Scoring: Elk traject wordt gerold in het geleerde wereldmodel. Trajecten worden gescoord door de voorspelde beloningen op te tellen met een terminale waarde-inschatting ($\hat{Q}$) afgeleid van de behavior-cloning predictor.
3. Verfijning: MPPI herweegt en verfijnt de trajecten op basis van deze scores. De eerste actie van de geoptimaliseerde sequentie wordt uitgevoerd.
4. Efficiëntie: De FM visuele encoding wordt één keer per observatie berekend en hergebruikt over alle samples en integratiestappen.

Belangrijkste Resultaten

De aanpak werd geëvalueerd op twee ManiSkill manipulatietaken: PickCube-v1 en PickSingleYCB-v1.

• End-of-Episode Success: De meest significante winst werd waargenomen in het behouden van succes aan het einde van de episode.
  • PickCube: Verbeterde van 93.14% (alleen FM) naar 97.44% (FlowMPC).
  • PickSingleYCB: Verbeterde van 56.81% (alleen FM) naar 66.41% (FlowMPC).
• Anytime Success: Verbeteringen werden ook gezien in het bereiken van een succesvolle staat op enig moment, hoewel de winst kleiner was.
  • PickCube: 95.78% $\to$ 98.68%.
  • PickSingleYCB: 68.77% $\to$ 69.78%.

De resultaten geven aan dat het wereldmodel bijzonder effectief is bij het corrigeren van kleine actiefouten en het behouden van stabiele controle nabij de taakvoltooiing, in plaats van alleen te helpen de agent een tijdelijke successtaat te laten bereiken.

Betekenis en Claims

Het artikel stelt dat wereldmodel-gebaseerde planning effectief een flow-gebaseerde imitatie-policy kan aanvullen zonder een nieuwe FM-trainingsdoelstelling te vereisen.

• Hybride Aanpak: Het werk demonstreert een levensvatbare hybride tussen imitation learning en model-based RL. Het gebruikt imitatie om een sterke, multimodale actie-prior te leren en planning om beslissingen te verfijnen tijdens de inferentietijd.
• Geen Her-training Vereist: De verbetering komt voort uit een betere benutting van de policy tijdens de testtijd, niet uit het hertrainen van de policy met een nieuwe RL-doelstelling.
• Bescheiden Omvang: De auteurs geven expliciet aan dat de resultaten beperkt zijn tot twee taken en dat de optimale wereldmodel-ontwerp (state-space vs. latent-space, planning horizon) varieerde tussen taken. Ze claimen geen universeel optimale configuratie, maar suggereren dat de combinatie van expressieve generatieve policies en lookahead planning een veelbelovende richting is voor robuuste beleidsverbetering.
• Prestatie Baseline: De methode verbeterde ten opzichte van een sterke FM-baseline (die al dicht bij de TD-MPC2 prestaties lag op één taak), wat suggereert dat de planner een sterke policy "verscherpt" in plaats van slechts een zwakke policy te redden.