Coupled Local and Global World Models for Efficient First Order RL

Technische Samenvatting: Gekoppelde Lokale en Globale Wereldmodellen voor Efficiënte First Order RL

Probleemstelling

Reinforcement Learning (RL) heeft aanzienlijk succes bereikt in locomotie-taken via geparalleliseerde GPU-simulatoren, maar het stuit op substantiële barrières in perceptuele controle en manipulatie. Twee primaire problemen belemmeren de voortgang:

1. Sample-efficiëntie: Praktische online leerprocessen direct vanuit interacties met echte robots zijn vaak onhaalbaar vanwege de hoge kosten van dataverzameling.
2. Simulatiegetrouwheid: Het simuleren van manipulatie-taken is inherent complexer dan locomotie. De staat van de omgeving verandert dynamisch met de acties van de agent (bijv. vervormbare objecten, granulaire dynamica), wat een oneindige variabiliteit introduceert die huidige fysica-simulatoren moeilijk met voldoende getrouwheid kunnen modelleren voor sim-to-real transfer.

Hoewel het leren van agenten binnen differentieerbare wereldmodellen met First-Order Gradient (FoG) methoden een pad biedt om de sample-efficiëntie te verbeteren, kampen bestaande benaderingen met een schaalbaarheidsbottleneck. Grootschalige wereldmodellen (bijv. diffusiemodellen) die nodig zijn voor hoogwaardige real-world dynamica zijn computationeel duur om doorheen te differentiëren, wat standaard backpropagation onhandelbaar maakt voor het trainen van FoG-policies. Omgekeerd missen lichtgewicht modellen die geschikt zijn voor differentiatie vaak de expressiviteit om complexe real-world dynamica te vangen.

Methodologie

Het artikel stelt een nieuw framework voor dat de voorwaartse simulatie ontkoppelt van de achterwaartse differentiatie om efficiënte First-Order Model-Based RL (FoG-MBRL) mogelijk te maken met grootschalige wereldmodellen. De kernarchitectuur bestaat uit twee gekoppelde componenten:

1. Globaal Wereldmodel (Forward Pass):

   • Een grootschalig, offline getraind model (ofwel een diffusie-gebaseerd model zoals DIAMOND, of een transformer-gebaseerd latent diffusiemodel zoals DreamerV4) fungeert als een "globale" simulator.
   • Dit model opereert in pixelruimte (of een hoogdimensionale latente ruimte) om hoogwaardige voorwaartse trajecten ($s^{\rightarrow}_{t+1}$) te genereren die complexe visuele en contact-dynamica accuraat vastleggen.
   • Het is getraind op ongestructureerde, hoog-entropische "play data" verzameld van echte robots, waardoor de noodzaak voor handmatig ontworpen fysica-simulatoren wordt geëlimineerd.

2. Lokaal Wereldmodel (Backward Pass):

   • Een lichtgewicht Recurrent State Space Model (RSSM) fungeert als een "lokale" surrogaat in een laag-dimensionale latente ruimte.
   • Dit model benadert de lokale dynamica en levert de gradiënten die nodig zijn voor policy-optimalisatie.
   • Cruciaal is dat de gradiënten worden berekend door de Jacobiaan van het lokale model te evalueren op de door het globale model gegenereerde state-punten ($s^{\rightarrow}_{t+1}$), in plaats van op de eigen voorspellingen van het lokale model. Deze ontkoppeling stelt het systeem in staat om de nauwkeurigheid van het globale model te benutten voor trajectgeneratie, terwijl het lokale model wordt gebruikt voor computationeel hanteerbare gradiëntschatting.

3. Trainingsprotocol:

   • Dataverzameling: Ongestructureerde "play data" wordt verzameld op echte robots zonder strikte taaksegmentatie.
   • Beloningsmodellering: Afhankelijk van de taak worden beloningssignalen afgeleid van contrastieve energie-functies (getraind op play/demonstratie data) of lichtgewicht classificatiekoppen die taakrelevante gedragingen identificeren.
   • Optimalisatie: De policy wordt volledig geoptimaliseerd binnen het wereldmodel met behulp van het Decoupled forward-backward Model-based policy Optimization (DMO) algoritme. Het lokale model wordt continu bijgetraind tijdens de RL-training om de nauwkeurigheid te handhaven ten opzichte van de huidige actie-distributie van de policy.

Belangrijkste Bijdragen

• Ontkoppeld FoG-MBRL Framework: De auteurs introduceren een methode die hoogwaardige voorwaartse simulatie (via grote diffusie/flow modellen) scheidt van efficiënte achterwaartse differentiatie (via lichtgewicht RSSM's). Dit lost de trade-off op tussen modelcomplexiteit en de kosten van gradiëntberekening.
• Simulator-vrij Leren van Real-World Data: De aanpak traint agenten volledig vanaf nul met real-world robotdata, waarbij de noodzaak voor fysica-simulatoren wordt omzeild en de sim-to-real kloof direct wordt aangepakt.
• Zero-Shot Deployment: Het framework demonstreert het vermogen om policies te trainen in het wereldmodel en deze vervolgens zonder verdere fine-tuning of behavior cloning guidance in te zetten op echte robots.
• Robuustheid tegen Complexe Dynamica: De methode gaat succesvol om met taken die partiële observabiliteit, complexe contact-dynamica en een hoge mate van vrijheidsgraden (bijv. humanoid manipulatie) bevatten.

Experimentele Resultaten

De auteurs evalueerden hun aanpak op drie real-world robotische taken van toenemende complexiteit:

1. Push-T: Een tafelmanipulatie-taak waarbij een T-vormig object wordt geduwd.
2. Ego-Centric Push Cube: Een vierpotige robot die een kubus naar een doel duwt, waarbij loop-dynamica en partiële occlusie een rol spelen.
3. Humanoid Grasp and Lift: Een hoog-DoF humanoïde die een doos grijpt en optilt met een behendige hand.

Belangrijkste Bevindingen:

• Sample- en Tijd-efficiëntie: De voorgestelde DMO-methode presteerde aanzienlijk beter dan PPO (Proximal Policy Optimization) in zowel sample-efficiëntie als wall-clock tijd over alle taken. PPO had moeite met convergentie of vereiste excessieve samples, waarbij het vaak in suboptimale lokale oplossingen terechtkwam.
• Noodzaak van Globale Modellen: Ablatie-studies ("No Diffusion"), waarbij het globale diffusiemodel werd vervangen door een standaard RSSM voor de voorwaartse passes, resulteerden in volledig falen van de taak (0% succespercentage), wat bevestigt dat hoogwaardige voorwaartse modellen essentieel zijn voor complexe manipulatie.
• Superioriteit boven Behavior Cloning: De geleerde policies presteerden beter dan Behavior Cloning baselines (Action Chunking Transformer), wat het vermogen aantoont om langetermijnstrategieën te ontdekken en regio's van de state-space te beheren die niet goed gedekt werden door demonstraties.
• Zero-Shot Succes: Policies die volledig in het wereldmodel zijn getraind, behaalden hoge succespercentages (bijv. 9/10 op Push-T, 8/10 op Humanoid Grasp) wanneer ze op echte robots werden ingezet zonder aanvullende training in de echte wereld.

Betekenis en Claims

Het artikel claimt dat het leren van agenten volledig binnen data-gedreven wereldmodellen een haalbaar en veelbelovend pad is voor het oplossen van moeilijk te modelleren taken in de beeldruimte. Door de eisen voor voorwaartse nauwkeurigheid en achterwaartse berekenbaarheid te ontkoppelen, laten de auteurs zien dat grootschalige, computationeel zware wereldmodellen effectief kunnen worden gebruikt voor efficiënte policy-optimalisatie.

Dit werk positioneert zich als een oplossing voor de afhankelijkheid van handmatig ontworpen fysica-simulatoren, en biedt een methode die schaalt naar complexe, ongestructureerde omgevingen via real-world data acquisitie. De auteurs merken op dat hoewel hun huidige implementatie het trainen van globale modellen vanaf nul vereist en vertrouwt op handmatig gelabelde beloningen voor bepaalde taken, het framework een veilige proxy biedt voor real-robot training, wat schaalbare policy-optimalisatie mogelijk maakt zonder het risico op hardwarebeschadiging. Toekomstig werk wordt geïdentificeerd als het benutten van internet-schaal pre-trained videomodellen om de datavereisten verder te verlagen en beloningslabeling te automatiseren.