Technische Samenvatting: WAM-RL

Probleemstelling

Recente World-Action (WA) modellen hebben een sterke generalisatie en data-efficiëntie aangetoond bij het leren van robotbeleid door toekomstige observaties en acties gezamenlijk te modelleren via video-gebaseerde generatieve modellen. Echter, bestaande WA-modellen leunen zwaar op supervised learning van expert-trajecten. Deze afhankelijkheid legt twee fundamentele beperkingen op:

1. Distributionele Beperkingen: Het geleerde beleid wordt beperkt door de ondersteuning (support) van de trainingsdata, waardoor het moeite heeft met het verwerven van fijnmazige manipulatievaardigheden die buiten de demonstratieverdeling vallen.
2. Gebrek aan Continue Verbetering: De modellen kunnen niet continu verbeteren door interactie met de echte wereld, wat hun vermogen belemmert om zich aan te passen aan nieuwe scenario's of fouten tijdens de uitvoering te corrigeren.

Hoewel Reinforcement Learning (RL) een natuurlijke oplossing is, is het toepassen ervan op WA-modellen niet triviaal. In tegenstelling tot standaard Vision-Language-Action (VLA) modellen, waarbij de visuele representatie vaak vaststaat, bestaan WA-modellen uit twee nauw gekoppelde componenten: een world model (dat toekomstige voorspellingen genereert) en een action model (dat voorspellingen naar acties vertaalt). Het naïef toepassen van RL of online fine-tuning kan leiden tot distributieverschuivingen in de latente ruimte van het world model, wat instabiliteit en prestatievermindering veroorzaakt omdat het action model diep afhankelijk is van deze latente ruimte.

Methodologie: WAM-RL

De auteurs stellen WAM-RL voor, een framework dat de gezamenlijke optimalisatie van het world model en het action model mogelijk maakt door middel van online interactie. De kerninzicht is dat het world model de voorspellende structuur en planningscapaciteit biedt, terwijl het action model primair fungeert als een vertaler die latente voorspellingen omzet in uitvoerbare acties.

Het framework hanteert een tweeledig optimalisatieschema:

1. Online Video SFT met KL-regularisatie (World Model)

Om het world model te verfijnen, maken de auteurs gebruik van trajecten die verzameld zijn tijdens succesvolle interacties voor online zelf-gesuperviseerde fine-tuning (SFT).

• Doel: Het world model wordt getraind om toekomstige observaties te voorspellen met behulp van een standaard video-modelleringsobjectief (bijv. flow matching).
• Stabilisatie: Om te voorkomen dat de latente feature-distributie drastisch verschuift (wat het action model zou breken), wordt een KL-regularisatieterm geïntroduceerd. Deze dwingt de geüpdatete latente features van het world model om dicht bij die van een bevroren kopie van het vooraf getrainde model te blijven.
• Implementatie: De auteurs benaderen de latente distributies als Gaussische verdelingen, waarbij een exponential moving average (EMA) wordt gebruikt om de covariantie van het huidige model te schatten terwijl de referenticovariantie vast blijft. Dit maakt geleidelijke adaptatie mogelijk terwijl de geometrie die de actor verwacht behouden blijft.

2. Reconstructie-gebaseerde RL (Action Model)

Het action model wordt geoptimaliseerd met behulp van RL om ervoor te zorgen dat uitgevoerde trajecten de voorspellingen van het world model getrouw realiseren.

• Beloningssignaal: In plaats van ijle (sparse) taakbeloningen, definiëren de auteurs een reconstructie-gebaseerde dichte beloning. Deze beloning meet de consistentie (gelijkenis) tussen de verbeelde toekomstige observaties (voorspeld door het world model) en de werkelijke observaties verkregen door de actie in de omgeving uit te voeren.
• Gelijkenismetrieken: De auteurs onderzoeken diverse gelijkenisfuncties, waaronder pixel-niveau MSE, optische flow-consistentie en semantische feature-gelijkenis (DINOv2, V-JEPA2).
• Optimalisatie: De actor wordt bijgewerkt via policy gradient-methoden met behulp van deze reconstructiebeloningen, wat de kloof tussen de verbeelding van het world model en de echte uitvoering effectief overbrugt.

Belangrijkste Bijdragen

1. Eerste RL-framework voor WA-modellen: WAM-RL is het eerste framework dat reinforcement learning integreert in het World-Action paradigma, specifiek door de gezamenlijke optimalisatie van zowel het world model als het action model mogelijk te maken.
2. Stabiele Trainingsstrategie: De auteurs stellen een nieuwe trainingsstrategie voor die online video SFT met KL-regularisatie combineert voor het world model en reconstructie-gebaseerde RL voor de actor. Dit ontwerp faciliteert de co-evolutie van beide componenten zonder de latente ruimte te destabiliseren.
3. Empirische Inzichten in Long-Horizon Taken: Het werk biedt cruciale inzichten in de beperkingen van alleen-actor optimalisatie. De auteurs laten zien dat het optimaliseren van enkel de actor weliswaar verbeteringen oplevert op korte-horizon taken, maar faalt in het bieden van significante winst op lange-horizon taken vanwege cumulatieve voorspellingsfouten. Gezamenlijke optimalisatie blijkt cruciaal voor complexe, lange-horizon prestaties.

Experimentele Resultaten

De methode werd geëvalueerd op de LIBERO (object-gecentreerde manipulatie) en RLBench (multi-step vaardigheden) benchmarks.

• Prestatiewinst: WAM-RL presteerde consequent beter dan zowel de vooraf getrainde baseline (Base) als de alleen-actor RL (πRL).
  • Op LIBERO-Object verbeterde de succesratio van 68% (Base) naar 82% (WAM-RL), wat de alleen-actor aanpak (78%) aanzienlijk overtreft.
  • Op RLBench (Water Plants) verbeterde WAM-RL de succesratio van 19% naar 22%, terwijl de alleen-actor methode (πRL) niet boven het basismodel (18%) uitkwam.
• Gedragsanalyse:
  • Herstelgedrag (Recovery Behaviors): De online adaptatie van het world model maakte de generatie van meer realistische herstelgedragingen mogelijk (bijv. het opnieuw proberen van mislukte grijpbewegingen) in plaats van het voorspellen van onmiddellijk falen.
  • Ablatie op Beloningen: De studie toonde aan dat hoewel optische flow het meest discriminerende beloningssignaal bood tussen succes en falen, pixel-niveau MSE de beste downstream prestaties opleverde. De auteurs hypothetiseren dat dit komt omdat pixel MSE beter aansluit bij het trainingsdoel van het world model en een sterkere regularisatie afdwingt tegen out-of-distribution acties.
  • Ablatie op Video SFT: Modellen die getraind zijn zonder online video SFT slaagden er niet in om te herstellen van initiële fouten en weken snel af naar out-of-distribution gedrag. In contrast hiermee leerden modellen met video SFT om falen te voorzien en corrigerende acties te genereren binnen een enkele open-loop chunk.

Betekenis en Claims

Het artikel claimt dat de effectiviteit van World-Action modellen fundamenteel wordt beperkt door de kwaliteit van het world model. Hoewel alleen-actor optimalisatie korte-horizon taken kan verbeteren, kan het de cumulatieve voorspellingsfouten in lange-horizon settings niet aanpakken.

De betekenis van WAM-RL ligt in het aantonen dat het gezamenlijk verbeteren van het world model en het action model essentieel is voor robuuste, lange-horizon besluitvorming. Door het world model toe te staan om door online interactie te leren van faal-dynamiek en herstelstrategieën, kan het systeem afwijkingen anticiperen en corrigeren, wat leidt tot een robuustere uitvoering.

De auteurs erkennen bescheiden hun beperkingen, waarbij zij opmerken dat KL-regularisatie de mate van adaptatie van het world model beperkt, wat potentieel de capaciteit tot uitbreiding buiten de vooraf getrainde distributie beperkt. Verder kan de afhankelijkheid van reconstructie-gebaseerde beloningen (vaak gebaseerd op vooraf getrainde representaties) een beperkt contrast bieden tussen succes en falen, wat suggereert dat er in toekomstig werk behoefte is aan meer expressieve, taakbewuste beloningsformuleringen.