Accurate and Efficient World Modeling with Masked Latent Transformers

Technische Samenvatting: Nauwkeurige en Efficiënte Wereldmodellering met Masked Latent Transformers

Probleemstelling
Model-gebaseerde reinforcement learning (RL) algoritmen, in het bijzonder de Dreamer-serie, hebben aanzienlijk succes geboekt door agenten te trainen op gesimuleerde trajecten binnen een latente wereldmodel. Echter, er bestaat een kritische beperking: de gecomprimeerde aard van de latente ruimte in standaardbenaderingen (bijv. DreamerV3) leidt vaak tot het verlies van cruciale omgevingsdetails, wat de prestaties van de agent in complexe omgevingen zoals Crafter verslechtert. Daartegenover staan recente methoden die streven naar hogere nauwkeurigheid, zoals $\Delta$-IRIS en DIAMOND, die vaak vereisen dat agenten direct worden getraind vanaf gereconstrueerde pixelafbeeldingen. Deze aanpak vermindert de trainingsefficiëntie en voorkomt dat de agent kan profiteren van de rijke interne representaties (bijv. langetermijngeheugen) die door het wereldmodel zijn geleerd. De uitdaging is om een wereldmodelleringsbenadering te ontwikkelen die tegelijkertijd nauwkeurig is in het vastleggen van fijnmazige details en efficiënt is in training en inferentie.

Methodologie: EMERALD
De auteurs stellen EMERALD (Efficient MaskEd latent tRAnsformer worLD model) voor, een Transformer-gebaseerd RL-algoritme dat is ontworpen om de afweging tussen nauwkeurigheid en efficiëntie aan te pakken. De architectuur bestaat uit drie primaire componenten:

1. Ruimtelijke Latente Toestand met MaskGIT: In tegen tegenstelling tot eerdere methoden die observaties comprimeren tot een enkele vector, gebruikt EMERALD een ruimtelijke latente toestand ($z_t$). Deze toestand wordt gerepresenteerd als een grid van tokens in plaats van een platte vector. Om deze tokens efficiënt te decoderen, maakt het model gebruik van MaskGIT (Masked Generative Image Transformer). In plaats van autoregressieve sequentiële decodering, wat traag is, gebruikt MaskGIT parallelle decodering volgens een schema. Tijdens de training worden tokens gemaskeerd volgens een cosinus-schema, en het model voorspelt ze parallel. Tijdens de verbeelding (inferentie) verfijnt het model voorspellingen over meerdere stappen (bijv. $S=3$) om coherentie te waarborgen en hallucinaties te voorkomen.
2. Temporele Transformer Netwerk: Het model integreert een Temporal Masked Transformer (TSSM) om langetermijnafhankelijkheden te beheren. Het verwerkt de ruimtelijke latente toestanden om een temporele verborgen toestand ($h_t$) te genereren met behulp van gemaskeerde zelf-aandacht (self-attention). Dit vervangt de Recurrent State-Space Models (RSSM) met GRU's die in DreamerV3 worden gebruikt, waardoor betere modellering van historische context mogelijk is zonder de bottlenecks van recurrente verwerking.
3. Latente Ruimte Actor-Critic Leren: De agent (actor- en critic-netwerken) wordt volledig getraind binnen de latente ruimte met behulp van de interne representaties van het wereldmodel. Hierdoor kan de agent profiteren van het door het wereldmodel geleerde geheugen en de ruimtelijke kenmerken, in tegenstelling tot pixelgebaseerde benaderingen waarbij de agent representaties opnieuw moet leren van gereconstrueerde afbeeldingen.

Het trainingsdoel minimaliseert een samengestelde verliesfunctie inclusief:

• Masking Verlies ($L_{mask}$): KL-divergentie tussen de doel-latente toestand en de voorspelde gemaskeerde tokens.
• Dynamics Verlies ($L_{dyn}$): Regularisatie om ervoor te zorgen dat de encoder-representaties voorspelbaar zijn vanuit de verborgen toestand.
• Reconstructie Verlies ($L_{rec}$): Pixel-niveau reconstructie van de invoerobservatie.
• Reward en Continuation Verliezen ($L_{rew}, L_{con}$): Voorspellen van omgevingsbeloningen en episode-continuatievlaggen.

Kernbijdragen

• Nieuw Architectuur: Introductie van een wereldmodel dat ruimelijke latente toestanden combineert met MaskGIT-voorspellingen, wat hoogwaardige trajectgeneratie mogelijk maakt zonder de computationele kosten van autoregressieve decodering.
• Efficiëntie en Nauwkeurigheid: De methode bereikt bijna perfecte beeldreconstructie (zoals aangetoond door lage pixelfouten en het behoud van details zoals diamanten en skeletpijlen) terwijl de hoge trainingsefficiëntie behouden blijft door de agent in de latente ruimte te houden.
• State-of-the-Art Prestaties: De aanpak zet nieuwe benchmarks op de Crafter omgeving, waarbij het vorige model-gebaseerde methoden en menselijke expertprestaties overtreft.

Resultaten
Geëvalueerd op de Crafter benchmark met een budget van 10 miljoen omgevingsstappen:

• Behaalde Score: EMERALD behaalde een score van 58,1%, waarmee het de vorige beste (DreamerV3 XL met 39,6%) overtrof en de prestaties van menselijke experts overtrof (50,5%).
• Behaalde Dekking (Achievement Coverage): De agent ontgrendelde succesvol alle 22 prestaties ten minste één keer tijdens de evaluatie, wat duidt op robuuste langetermijnplanning en perceptie.
• Trainingsefficiëntie: EMERALD verzamelde 27 Frames Per Second (FPS) op een RTX 3090, wat aanzienlijk sneller is dan $\Delta$-IRIS (12 FPS) en DreamerV3 XL (23 FPS), terwijl het hen in score overtreft.
• Ablatie-studies: Experimenten bevestigden dat de combinatie van ruimtelijke latente toestanden en de Transformer-architectuur cruciaal is. Het verwijderen van ruimtelijke latenten of het gebruiken van een enkele decoderingstap verminderde de prestaties en voorspellingsnauwkeurigheid aanzienlijk.

Betekenis en Claims
Het artikel claimt dat EMERALD een belangrijke vooruitgang vertegenwoordigt in model-gebaseerde RL door de spanning tussen de nauwkeurigheid van het wereldmodel en de efficiëntie van de training op te lossen. Door MaskGIT te gebruiken voor parallelle decodering in een ruimtelijke latente ruimte, stelt de methode agenten in staat om cruciale omgevingsdetails waar te nemen die voorheen verloren gingen in gecomprimeerde vectorrepresentaties. De auteurs benadrukken dat dit de eerste methode is die de prestaties van menselijke experts overtreft op de Crafter benchmark binnen 10 miljoen stappen. Verder demonstreert het werk dat Transformer-gebaseerde architecturen met gemaskeerde aandacht effectief langetermijnafhankelijkheden en ruimtelijke relaties kunnen modelleren, wat een schaalbare alternatief biedt voor recurrente modellen voor complexe, geheugenintensieve omgevingen. De auteurs positioneren dit werk als een stap naar meer capabele autonome agenten in complexe domeinen zoals Minecraft.