Enhancing Policy Learning with World-Action Model

Dit paper introduceert het World-Action Model (WAM), een actie-geregulariseerd wereldmodel dat door het integreren van een inverse dynamica-doelstelling de representaties verbetert, wat leidt tot aanzienlijk hogere succespercentages en een efficiëntere training voor robotbeheersingstaken in vergelijking met bestaande methoden zoals DreamerV2.

De kern

Stel je voor dat je een robot wilt leren om complexe taken te doen, zoals een la openen of een lampje aan- en uitzetten. Normaal gesproken moet je die robot duizenden keren laten oefenen in de echte wereld. Dat is traag, duur en soms gevaarlijk als de robot iets breekt.

Om dit op te lossen, gebruiken wetenschappers een "Wereldmodel". Dit is als een droomwereld in het hoofd van de robot. De robot leert hoe de wereld werkt door in zijn hoofd te "dromen" over wat er gaat gebeuren, in plaats van alles in het echt te proberen.

Dit artikel introduceert een nieuwe, slimme versie van zo'n droomwereld, genaamd WAM (World-Action Model). Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De "Zieke" Droom

Stel je voor dat je een filmkijker bent die alleen kijkt naar de beelden op het scherm, maar nooit luistert naar de geluiden of de dialogen.

• De oude methode (DreamerV2): De robot leert alleen te voorspellen hoe het beeld er morgenuit zal zien als hij een knop indrukt. Hij leert dus alleen wat hij ziet, maar niet precies waarom het gebeurt.
• Het gevolg: De robot leert een vaag beeld van de wereld. Hij weet dat een la open gaat, maar hij begrijpt niet precies welke beweging van zijn hand dat veroorzaakte. Het is alsof hij een film kijkt met de geluidsdemping aan: hij ziet de actie, maar mist de context.

2. De Oplossing: De "Actieve" Droom (WAM)

De auteurs van dit artikel zeggen: "Laten we de robot niet alleen laten kijken, maar ook laten luisteren naar zijn eigen bewegingen."

Ze voegen een extra hersenstreek toe aan het model: een Inverse Dynamics Head.

• De Analogie: Stel je voor dat je een dansleraar bent. De oude robot keek alleen naar de danspasjes van de leraar en probeerde ze na te doen. De nieuwe robot (WAM) kijkt ook naar de spieren van de leraar en vraagt zich af: "Welke beweging heb ik precies gemaakt om die danspas te krijgen?"
• Hoe het werkt: WAM leert niet alleen te voorspellen wat er gaat gebeuren, maar ook welke actie er precies nodig was om die verandering te veroorzaken. Het model moet dus twee dingen tegelijk doen:
  1. Zeggen hoe de wereld er morgen uitziet.
  2. Zeggen welke beweging ik heb gemaakt om daar te komen.

3. Waarom is dit zo krachtig?

Door deze extra taak te geven, wordt het "geheugen" van de robot veel scherper.

• De "Filter"-effect: Omdat de robot moet uitleggen welke beweging hij deed, moet hij in zijn hoofd (in de latente ruimte) de belangrijke details bewaren. Hij leert onderscheid te maken tussen dingen die belangrijk zijn voor zijn beweging (zoals de positie van een handgreep) en dingen die irrelevant zijn (zoals een vliegje dat over de muur kruipt).
• Het Resultaat: De robot bouwt een veel nauwkeuriger "simulatie" van de wereld op. Het is alsof hij van een wazige tekening overschakelt naar een 3D-animatie met perfecte fysica.

4. De Resultaten: Sneller en Beter

De onderzoekers testten dit op de CALVIN-benchmark, een reeks van 8 robot-taken (zoals laden openen en schuiven).

• Minder oefenen: De robot had 8,7 keer minder trainingstijd nodig om even goed te worden als de oude methode.
• Beter presteren:
  • Alleen maar kijken en nabootsen (Behavioral Cloning): De oude robot slaagde in 46% van de gevallen, de nieuwe WAM-robot in 62%.
  • Na extra oefenen in de droomwereld (PPO): De oude robot kwam uit op 80%, terwijl de nieuwe WAM-robot 93% haalde. Bij twee taken haalde hij zelfs 100% succes!

Samenvatting in één zin

WAM is als een robot die niet alleen naar een film kijkt, maar ook de regisseur is die precies begrijpt welke knoppen hij moet indrukken om de film te laten verlopen; hierdoor leert hij sneller, maakt hij minder fouten en wordt hij een veel betere "dromer" voor zijn toekomstige acties.

Kortom: Door de robot te dwingen te begrijpen waarom de wereld verandert (door zijn eigen acties), wordt zijn droomwereld realistischer en wordt hij een veel slimmere robot in de echte wereld.

Technische samenvatting

Probleemstelling

Wereldmodellen (world models) zijn cruciaal voor het leren van omgevingsdynamica in de robotica, omdat agents hierdoor kunnen plannen en gedrag kunnen verfijnen via "imaginatie" (rollouts) in een latente ruimte, zonder kostbare fysieke interacties. Echter, conventionele wereldmodellen (zoals DreamerV2) worden uitsluitend getraind om toekomstige observaties (beelden) te voorspellen op basis van het verleden en acties.

De kern van het probleem is een asymmetrie in de training:

1. Het model leert latent staten ($z_t$) die geoptimaliseerd zijn voor pixelreconstructie.
2. Deze latent staten worden direct gebruikt als input voor downstream beleidsmodellen (zoals diffusie-beleidsmodellen).
3. Omdat de trainingssignalen alleen uit beeldreconstructie komen, worden de latent representaties niet expliciet gedwongen om actie-relevante structuur te coderen. Hierdoor kunnen belangrijke informatie over hoe de omgeving reageert op acties verloren gaan of niet optimaal worden vastgelegd, wat de prestaties van het beleidsleren beperkt.

Bestaande oplossingen die acties en beelden gezamenlijk genereren, vereisen vaak grote fundamentele modellen en een volledige herschrijving van de architectuur, wat niet altijd efficiënt is.

Methodologie: Het World-Action Model (WAM)

De auteurs stellen WAM voor, een lichtgewicht uitbreiding op het bestaande DreamerV2-architectuur. In plaats van de architectuur te veranderen, wordt de trainingsdoelstelling aangepast om een inverse dynamiek-head toe te voegen.

Kerncomponenten:

1. Architectuur: WAM bouwt voort op de Recurrent State-Space Model (RSSM) van DreamerV2. Een dual-stream CNN-coder verwerkt statische en grijper-camera beelden en fuseert deze met proprioceptieve staten.
2. Inverse Dynamiek Head: Een extra component ($M_{action}$) voorspelt de actie ($\hat{a}_t$) die heeft geleid tot een overgang tussen twee opeenvolgende encoder-embeddings ($e_t$ en $e_{t+1}$).
   • Formule: $\hat{a}_t = \psi([e_t; e_{t+1}])$, waarbij $\psi$ een MLP is.
3. Trainingsdoelstelling: Het totale verlies ($L_{WAM}$) combineert drie componenten:
   • $L_{KL}$: Kullback-Leibler-divergentie tussen posterior en prior (regulatie).
   • $L_{recon}$: Reconstructie van de toekomstige observatie (beeld).
   • $L_{action}$: Voorspelling van de actie (inverse dynamiek).
   • Doel: $L_{WAM} = \lambda_{KL} L_{KL} + \lambda_{img} L_{recon} + \lambda_{act} L_{action}$.

Het "Cascading Effect":
Een cruciaal ontwerpaspect is dat de actie-head werkt op de encoder-embeddings ($e_t$) en niet op de RSSM-features ($f_t$). Omdat de RSSM-toestand $h_t$ de vorige actie al bevat, zou actievoorspelling daar triviaal zijn. Door de encoder direct te regulariseren, wordt de "actie-bewuste" structuur doorgegeven:

• Het beïnvloedt de posterior $z_t$.
• Via de KL-loss wordt deze structuur doorgegeven naar de prior $\hat{z}_t$.
• Uiteindelijk genereren de geimagineerde rollouts in de latente ruimte features die rijker zijn aan actie-informatie voor het downstream beleid.

Beleidsleren (Policy Learning):
Het getrainde WAM fungeert als een bevroren simulator voor twee fasen:

1. Behavioral Cloning (BC): Een diffusie-beleid (DiffusionMLP) wordt getraind om acties te genereren op basis van de WAM-features uit expert-demonstraties.
2. Offline Fine-tuning (PPO): Het beleid wordt verder verfijnd binnen de latente ruimte van het bevroren WAM-model via PPO (Proximal Policy Optimization), zonder fysieke interactie. Een reward-classifier schat de beloning in de latente ruimte.

Belangrijkste Bijdragen

1. WAM Introductie: Een nieuwe wereldmodel-variant die een inverse dynamiek-head toevoegt aan DreamerV2, waardoor latent representaties expliciet worden geregulariseerd voor actie-relevantie.
2. Verbeterde Generatiekwaliteit: WAM presteert beter dan DreamerV2 op videovoorspellingsmetrieken (PSNR, SSIM, LPIPS, FVD) op het CALVIN-benchmark, ondanks het gebruik van 8,7 keer minder trainingsstappen.
3. Superieur Beleidsleren: De verbeterde representaties leiden tot significante prestatieverbeteringen in downstream taken, zowel bij imitatie-leren (BC) als bij model-based versterkend leren (PPO), zonder de beleidsarchitectuur te wijzigen.

Resultaten

De evaluatie vond plaats op het CALVIN-benchmark (8 manipulatietaakken met een Franka Emika Panda robot).

1. Kwaliteit van Wereldmodel (Imaginatie):

• WAM overtreft DreamerV2 op alle vier de video-metrieken.
• Visuele rollouts tonen realistischere toekomstige staten, betere behoud van objectvormen en minder kleurdrift.

2. Behavioral Cloning (BC) Resultaten:

• Gemiddelde Succes率: WAM bereikte 61,7% succes, vergeleken met 45,8% voor de DiWA-baseline (DreamerV2).
• De grootste winsten werden geboekt bij taken met gearticuleerde objecten (bijv. "close drawer": +31,1%, "move slider right": +31,1%), wat aantoont dat WAM fijne ruimtelijke cues beter vastlegt.

3. Model-based PPO Fine-tuning Resultaten:

• Na 800 iteraties van fine-tuning bereikte WAM een gemiddelde succes率 van 92,8%, tegenover 79,8% voor de baseline.
• Twee taken ("turn on lightbulb" en "turn off led") bereikten 100% succes.
• Efficiëntie: WAM bereikte de prestaties van de baseline met 8,7 keer minder trainingsstappen voor het wereldmodel.

Significantie

Dit paper demonstreert dat het expliciet modelleren van de relatie tussen staten en acties (via inverse dynamiek) binnen een wereldmodel fundamenteel de kwaliteit van de geleerde representaties verbetert.

• Efficiëntie: Het toont aan dat men geen enorme foundation-modellen hoeft te bouwen om betere wereldmodellen te krijgen; een kleine aanpassing in de trainingsdoelstelling volstaat.
• Generalisatie: De methode werkt effectief voor zowel offline imitatie-leren als online versterkend leren in een gesimuleerde omgeving.
• Praktische Toepassing: Voor robotica betekent dit dat agents sneller en betrouwbaarder complexe manipulatietaakken kunnen leren met minder data, wat een grote stap is richting efficiëntere en robuustere autonome systemen.