Next Embedding Prediction Makes World Models Stronger

Dit paper introduceert NE-Dreamer, een decoder-vrije MBRL-agent die next-embedding predictie met een temporele transformer gebruikt om effectieve state-representaties te leren zonder reconstructieverlies, wat leidt tot verbeterde prestaties in complexe, deels waarneembare omgevingen.

De kern

Stel je voor dat je een agent (een digitale robot) wilt leren een complex doolhof navigeren. De robot kan alleen kijken, maar hij ziet niet alles tegelijk; hij krijgt telkens een nieuwe foto van zijn omgeving. Dit is wat we "gedeeltelijke waarneembaarheid" noemen.

Het oude probleem was: hoe leer je deze robot om te onthouden waar hij was, zonder dat hij elke foto perfect moet kunnen natekenen?

Hier komt NE-Dreamer om de hoek kijken. Het is een nieuwe manier om robots slim te maken, en het werkt als volgt:

1. De Oude Manier: De "Fotograaf"

Stel je voor dat de robot vroeger moest leren door te proberen elke foto die hij zag, exact na te tekenen.

• Het probleem: De robot besteedde veel tijd en energie aan het onthouden van details die niet belangrijk waren, zoals de textuur van het tapijt of de kleur van de muur. Hij werd een uitstekende fotograaf, maar een slechte strateeg. Hij wist niet goed wat er volgende zou gebeuren, omdat hij te veel focus had op het huidige plaatje.

2. De Nieuwe Manier: De "Profeet" (NE-Dreamer)

NE-Dreamer stopt met het tekenen van foto's. In plaats daarvan leert hij een nieuwe vaardigheid: voorspellen.

Stel je voor dat je een detective bent die een verhaal schrijft.

• De oude detective schreef elke zin perfect uit, inclusief elke beschrijving van de kleding van de personages (reconstructie).
• De nieuwe detective (NE-Dreamer) kijkt naar wat er nu gebeurt en probeert te voorspellen: "Wat zal de volgende zin in het verhaal zijn?"

Hij doet dit niet door woorden te raden, maar door te raden wat de essentie van de volgende scène is.

Hoe werkt het precies? (De Analogie van de Voorspeller)

1. Geen Tekenen, Alleen Voorspellen:
   De robot kijkt naar zijn verleden (een reeks foto's) en vraagt zich af: "Als ik nu deze actie doe, hoe zal de 'geest' van de volgende foto eruitzien?" Hij probeert niet de foto zelf te maken, maar de samenvatting (de embedding) van de volgende foto.

2. De Temporele Transformer (De "Tijdmachine"):
   De robot heeft een speciaal breinonderdeel (een 'temporal transformer') dat als een tijdmachine werkt. Hij kijkt niet alleen naar het nu, maar houdt een draad vast door de tijd heen. Hij zegt: "Ik zie nu een deur, en ik heb gisteren een sleutel gezien. Dus voorspel ik dat de volgende 'geest' van de wereld een open deur zal zijn."

3. De Check (De "Barlow Twins"):
   Als de robot zijn voorspelling doet, kijkt hij of het klopt met de werkelijkheid. Maar hij doet dit op een slimme manier: hij zorgt dat zijn voorspelling stabiel en nuttig is. Hij zorgt ervoor dat hij niet zomaar willekeurige dingen gaat voorspellen, maar dat zijn voorspellingen logisch aansluiten op wat er echt gebeurt.

Waarom is dit zo goed?

• Geen Afleiding: Omdat de robot niet hoeft na te tekenen hoe het tapijt eruitziet, heeft hij meer "breinruimte" om na te denken over wat er gaat gebeuren. Hij focust op wat belangrijk is voor het doel (bijvoorbeeld: waar is de sleutel?).
• Beter Onthouden: In doolhoven (zoals de DMLab-taken in het papier) moet je dingen onthouden die je een tijdje geleden hebt gezien. Omdat NE-Dreamer altijd probeert te voorspellen wat er volgende komt, is hij van nature beter in het onthouden van de context. Hij bouwt een coherent verhaal op in plaats van losse foto's.
• Sneller en Sterker: De resultaten tonen aan dat deze robot veel beter presteert in moeilijke, gedeeltelijk onzichtbare omgevingen dan de oude modellen, terwijl hij op simpele taken net zo goed blijft werken.

Samenvattend in één zin:

In plaats van een robot te maken die een perfecte fotograaf is (die elke foto na tekent), hebben we een robot gemaakt die een voorspeller is: hij leert de wereld te begrijpen door te voorspellen wat er als volgende gaat gebeuren, waardoor hij veel slimmer en beter in staat is om complexe doolhoven te navigeren.

Technische samenvatting

Probleemstelling

Modelgebaseerd versterkend leren (MBRL) in gedeeltelijk waarneembare, hoogdimensionale omgevingen (zoals pixelgebaseerde beelden) staat voor een fundamenteel uitdaging: het leren van een compacte, latente staat die lange-termijn voorspellingen en controle ondersteunt.

• Huidige aanpak: De dominante methode (zoals in Dreamer) gebruikt een pixel-decoder om de waarneming te reconstrueren. Dit levert rijke features op, maar introduceert een zware generatieve last. De decoder kan capaciteit verspillen aan visuele details die irrelevant zijn voor de taak (zoals texturen of achtergronden).
• Decoder-vrije methoden: Deze verwijderen de pixel-decoder om de pijplijn te vereenvoudigen. Echter, veel bestaande decoder-vrije doelen focussen op instantane overeenstemming (overeenkomst op hetzelfde tijdstip). In gedeeltelijk waarneembare omgevingen is dit onvoldoende; de representatie moet voorspellend zijn over de tijd. Zonder expliciete temporele beperkingen kan het leren "drijven" of instorten, wat leidt tot zwakke lange-termijn structuren, vooral in taken die geheugen en ruimtelijk redeneren vereisen.

Methodologie: NE-Dreamer

Het paper introduceert NE-Dreamer, een decoder-vrije MBRL-agent die de pixel-reconstructie vervangt door next-embedding prediction (voorspelling van de volgende embedding).

Kerncomponenten:

1. Latente Wereldmodel (RSSM): NE-Dreamer behoudt de Recurrent State-Space Model (RSSM) architectuur van Dreamer (met een deterministische recurrente staat $h_t$ en een stochastische latente staat $z_t$).
2. Vervanging van de Decoder: In plaats van een pixel-decoder te trainen om $x_t$ te reconstrueren, voorspelt het model de volgende encoder-embedding ($\hat{e}_{t+1}$) op basis van de geschiedenis tot tijdstip $t$.
3. Causale Temporal Transformer: Een lichtgewicht causale transformer wordt geïntegreerd om de geschiedenis van latenten en acties te verwerken en de embedding voor het volgende tijdstip te voorspellen.
4. Voorspellende Alignering (Barlow Twins):
   • Het model voorspelt $\hat{e}_{t+1}$ en aligneert dit met de werkelijke volgende embedding $e_{t+1}$ (via de encoder van de volgende observatie).
   • Er wordt een stop-gradient toegepast op het doel ($e_{t+1}$) om instabiliteit te voorkomen.
   • Als verliesfunctie wordt de Barlow Twins loss gebruikt. Deze straalt redundantie uit (vermindert correlaties tussen verschillende dimensies) en bevordert invariantie (verhoogt de diagonale correlaties), maar toegepast op voorspelling van de toekomst in plaats van overeenstemming op hetzelfde moment.
5. Actor-Critic: De agent leert een beleid en waardenfunctie in de latente ruimte via "imaginatie" (rollouts), net als in DreamerV3, maar gebruikt de nieuwe representatie.

Belangrijkste Bijdragen

1. Nieuw Doel: Voorstellen van een decoder-vrije wereldmodel-doelstelling gebaseerd op next-embedding prediction, die expliciet temporele voorspelbaarheid in de geleerde representatie afdwingt.
2. Architecturale Integratie: Het integreren van een causale temporal transformer in een Dreamer-stijl MBRL-pijplijn om voorspellingen van de geschiedenis naar de toekomst te maken binnen standaard RSSM-training.
3. Empirische Validatie: Uitgebreide evaluatie op DeepMind Control Suite (DMC) en DeepMind Lab (DMLab), waarbij NE-Dreamer presteert op of boven het niveau van DreamerV3 en andere toonaangevende agents.
4. Ablatie-studies: Het isoleren van de winstbronnen, wat aantoont dat de verbetering komt door predictive sequence modeling (transformer + verschuiving van het doel naar de volgende stap) en niet door reconstructie of extra regularisatie.

Resultaten

De resultaten worden gepresenteerd in twee hoofdgebieden:

• DMLab Rooms (Geheugen & Navigatie):

  • Op deze uitdagende taken, waar succes afhangt van het onthouden van informatie over lange tijdsperiodes, behaalt NE-Dreamer substantiële winsten.
  • Het overtreft zowel sterke decoder-gebaseerde baselines (DreamerV3) als decoder-vrije baselines (R2-Dreamer, DreamerPro).
  • De verbetering is het grootst in taken zoals "Rooms Collect" en "Rooms Watermaze", waar agenten complexe ruimtelijke lay-outs moeten onthouden.
  • Ablaties: Het verwijderen van de transformer of het terugdraaien van het doel naar "same-step matching" zorgt voor een instorting van de prestaties, wat bewijst dat temporele voorspelling cruciaal is.

• DeepMind Control Suite (DMC):

  • Op standaard continue controle-taken (waar de omgeving vaak volledig waarneembaar is of minder geheugen vereist) presteert NE-Dreamer op gelijke hoogte met DreamerV3 en andere baselines.
  • Dit bevestigt dat het verwijderen van reconstructie geen prestatieverlies veroorzaakt in standaard scenario's.

• Representatie Diagnostics:

  • Post-hoc reconstructie-experimenten tonen aan dat de latenten van NE-Dreamer temporeel consistent zijn. In tegenstelling tot andere methoden, waar objecten of attributen kunnen verdwijnen in de latente staat (drift), behoudt NE-Dreamer de identiteit van objecten en ruimtelijke lay-outs consistent over de tijd.

Betekenis en Conclusie

Het paper stelt dat next-embedding prediction met een causale temporal transformer een effectief en schaalbaar kader is voor MBRL in complexe, gedeeltelijk waarneembare omgevingen.

• Verschuiving in Paradigma: Het bewijst dat pixel-reconstructie niet noodzakelijk is voor het leren van robuuste wereldmodellen. In plaats daarvan kan het direct voorspellen van de toekomstige representatie (in plaats van het heden te reconstrueren) leiden tot betere lange-termijn planning.
• Efficiëntie: De methode bereikt superieure prestaties zonder de rekenlast van een pixel-decoder of zware data-augmentatie.
• Toekomst: Hoewel het huidige werk focust op omgevingen waar lange-termijn structuur belangrijker is dan fijne visuele details, opent dit de weg voor verdere onderzoek naar voorspellende, decoder-vrije wereldmodellen in nog complexere domeinen.

Kortom, NE-Dreamer demonstreert dat het afdwingen van temporele coherentie via voorspelling van de volgende embedding een krachtigere leerstrategie is dan het simpelweg reconstrueren van de huidige waarneming.