Recurrent Action Transformer with Memory

Deze paper introduceert RATE, een nieuw transformer-architectuur voor offline versterkingsleren die een recurrente geheugenmechanisme integreert om effectieve besluitvorming op lange termijn in deels waarneembare omgevingen mogelijk te maken, wat leidt tot aanzienlijke prestatieverbeteringen op geheugenintensieve taken zonder in te leveren op standaard benchmarks.

De kern

Stel je voor dat je een speler bent in een heel moeilijk computerspel. Je moet een doolhof doorlopen, maar er is een probleem: het spel is zo lang dat je niet alles in één keer kunt onthouden. Je moet een hint zien die je 100 stappen geleden kreeg, om nu de juiste afslag te nemen.

Dit is precies het probleem waar kunstmatige intelligentie (AI) vaak tegenaan loopt. De slimste modellen van vandaag, de zogenaamde Transformers (denk aan de technologie achter ChatGPT), zijn fantastisch in het begrijpen van lange zinnen of verhalen. Maar in een spel of een complexe taak worden ze snel "kortzichtig". Ze kunnen maar een beperkt stukje van het verleden onthouden. Zodra iets te lang geleden is gebeurd, is het voor hen alsof het nooit heeft bestaan.

De auteurs van dit paper hebben een oplossing bedacht: RATE (Recurrent Action Transformer with Memory). Laten we uitleggen hoe dit werkt met een paar simpele vergelijkingen.

1. Het probleem: De "Gouden Kooi"

Stel je een standaard Transformer voor als een student die een heel dik boek leest, maar die maar 10 pagina's tegelijk op zijn bureau kan hebben. Als hij pagina 100 moet raadplegen terwijl hij op pagina 10 zit, kan hij dat niet. Hij moet het hele boek opnieuw lezen (wat te lang duurt) of hij raakt de informatie kwijt. In een spel betekent dit: als de hint voor de afslag 500 stappen geleden was, vergeet de AI het en loopt hij de verkeerde kant op.

2. De oplossing: RATE en het "Magische Notitieblok"

RATE is als diezelfde student, maar dan met een magisch notitieblok (de "Memory") en een slimme poortwachter (de "Memory Retention Valve").

Hier is hoe RATE werkt, stap voor stap:

• Deel het verhaal op: In plaats van het hele boek (het spel) in één keer te proberen te lezen, deelt RATE het verhaal op in kleine hoofdstukken (segmenten).
• Het Magische Notitieblok: Aan het begin van elk hoofdstuk pakt RATE een notitieblok. Hierin schrijft hij de belangrijkste dingen op die hij tot nu toe heeft gezien.
  • Vergelijking: Stel je voor dat je door een lange gang loopt. Je ziet een rode pijl aan het begin. Je loopt 100 meter, draait om een hoek, en ziet nog steeds de rode pijl in je hoofd. Dat is het notitieblok.
• De Poortwachter (MRV): Dit is het slimste onderdeel. Normaal gesproken zou je bij het schrijven van een nieuw hoofdstuk je oude notities kunnen overschrijven met nieuwe, onbelangrijke praatjes.
  • De Poortwachter kijkt echter kritisch: "Is dit nieuwe ding belangrijk? Of moeten we de oude rode pijl bewaren?"
  • Als het nieuwe ding niet belangrijk is, laat de poortwachter de oude, cruciale informatie (zoals de rode pijl) veilig in het notitieblok staan. Hij blokkeert het overschrijven. Dit heet in het paper de Memory Retention Valve.

3. Waarom is dit zo goed?

In de paper testen ze RATE op verschillende "doolhoven":

• Het T-Maze: Een agent moet aan het begin een hint zien (links of rechts) en die onthouden tot hij aan het einde van een zeer lange gang is.
  • Resultaat: De oude modellen (zoals de standaard Transformer) vergeten de hint zodra hij uit beeld is. RATE onthoudt het perfect, zelfs als de gang 9.600 stappen lang is (veel langer dan wat andere modellen kunnen).
• ViZDoom (Een 3D spel): Hier moet je onthouden welke kleur zuil je zag, om later de juiste munten te verzamelen.
  • Resultaat: RATE is veel consistenter en maakt minder fouten dan de concurrenten.

4. Is het alleen maar voor moeilijke spellen?

Nee! De auteurs laten zien dat RATE ook heel goed is in "normale" taken waar je geen lange termijn geheugen voor nodig hebt (zoals het spelen van Atari-spelletjes of het laten lopen van een robot).

• Vergelijking: Het is alsof je een auto hebt met een superkrachtige navigatie voor lange reizen, maar die ook perfect werkt als je gewoon naar de supermarkt rijdt. Je hoeft de navigatie niet uit te schakelen; hij helpt gewoon mee zonder in de weg te zitten.

Samenvatting in één zin

RATE is een slimme AI-architectuur die lange verhalen (of spelrondes) opdeelt in stukjes, en gebruikmaakt van een slimme poortwachter om ervoor te zorgen dat de belangrijkste herinneringen uit het verleden nooit worden gewist, zelfs niet als het verhaal duizenden stappen lang is.

Dit maakt het mogelijk voor robots en AI-agenten om taken te volbrengen die veel tijd en geduld vereisen, zonder dat ze de draad kwijtraken.

Technische samenvatting

1. Het Probleem

In het veld van Offline Reinforcement Learning (RL) hebben Transformers (zoals de Decision Transformer) bewezen zeer effectief te zijn door agent-trajecten te behandelen als sequenties. Echter, in gedeeltelijk waarneembare omgevingen (POMDP's) en taken met lange tijdshorizons, ondervinden standaard Transformers fundamentele beperkingen:

• Kwadratische complexiteit: De zelf-attentie (self-attention) mechanismen hebben een complexiteit van $O(N^2)$, wat de contextlengte beperkt.
• Verlies van informatie: In POMDP's is het cruciaal om informatie uit het verleden (bijv. een hint aan het begin van een episode) te onthouden om later de juiste beslissing te nemen. Zodra deze informatie de vaste contextvenster verlaat, kunnen standaard Transformers deze niet meer benaderen.
• Schaalbaarheid: Bestaande oplossingen om de context te verlengen (zoals sparse attention) leiden vaak tot trainingsinstabiliteit of generaliseren slecht buiten NLP-taken.

2. Methodologie: RATE

De auteurs stellen RATE (Recurrent Action Transformer with Memory) voor, een nieuwe architectuur die de kracht van Transformers combineert met een recurrente geheugenmechanisme om lange-termijn afhankelijkheden te modelleren zonder de contextlengte te hoeven vergroten.

Kerncomponenten van RATE:

1. Segmentatie en Recurrente Verwerking:

   • Het traject wordt opgedeeld in $N$ segmenten van lengte $K$.
   • In plaats van het hele traject in één keer te verwerken, wordt het sequentieel verwerkt.
   • Elk segment $S_n$ wordt voorafgegaan en gevolgd door geheugen-embeddings ($M_n$). De prefix zorgt voor "lees-toegang" (het segment kan naar het verleden kijken), en de suffix zorgt voor "schrijftoegang" (het segment kan het geheugen updaten).

2. Geheugen-embeddings en Caching:

   • Learned Memory Embeddings: Een set van learnable tokens die belangrijke historische informatie vasthouden.
   • Cached Hidden States: Net als bij Transformer-XL worden de verborgen staten van eerdere segmenten opgeslagen en hergebruikt als context voor het volgende segment, wat de effectieve horizon verlengt.

3. Memory Retention Valve (MRV):

   • Dit is het meest innovatieve onderdeel. Een naïef doorgeven van geheugen kan leiden tot het overschrijven van belangrijke informatie of het accumuleren van fouten.
   • De MRV is een cross-attention module die bepaalt welke informatie uit het nieuwe geheugen ($M_{n+1}$) behouden moet blijven en wat moet worden vervangen door het oude geheugen ($M_n$).
   • Het werkt als een "klep" die selectief informatie doorlaat, gebaseerd op de relevantie van de nieuwe data ten opzichte van het bestaande geheugen.
   • Theoretische garantie: De auteurs bewijzen een stelling (Theorem 1) die aantoont dat de MRV zorgt voor een ondergrens aan geheugenbehoud, waardoor catastrofaal vergeten wordt voorkomen.

Werkingsprincipe:
Het model encodeert observaties, acties en returns-to-go. Voor elk segment wordt het transformer-model aangeroepen met het segment en de huidige geheugen-embeddings. De output bevat een nieuwe set geheugen-embeddings, die vervolgens door de MRV wordt gefilterd voordat deze naar het volgende segment wordt doorgegeven.

3. Belangrijkste Bijdragen

1. Nieuwe Architectuur (RATE): Een transformer voor offline RL die drie complementaire mechanismen combineert: (i) geleerde geheugen-embeddings, (ii) caching van verborgen staten, en (iii) de MRV voor selectieve informatiestroom.
2. Uitgebreide Evaluatie: RATE is getest op een breed scala aan geheugen-intensieve taken (ViZDoom-Two-Colors, T-Maze, Memory Maze, Minigrid-Memory, POPGym) en standaard benchmarks (Atari, MuJoCo).
3. Theoretische Analyse: Een wiskundige onderbouwing van de MRV die garandeert dat belangrijke informatie niet verloren gaat tijdens updates, zelfs bij zeer lange sequenties.
4. Generalisatie: Het bewijs dat RATE niet alleen interpolatie binnen de trainingslengte beheerst, maar ook uitstekend extrapolatie toont naar veel langere inferentielengtes (tot 9600 stappen), waar andere modellen falen.

4. Resultaten

De experimentele resultaten tonen aan dat RATE significant beter presteert dan bestaande baselines in geheugen-afhankelijke omgevingen:

• Geheugen-intensieve taken:

  • T-Maze: RATE bereikt een succesratio van 100% over een breed scala aan corridorlengtes (tot 9600 stappen), terwijl modellen zoals Decision Transformer (DT) en RMT (Recurrent Memory Transformer) snel degradëren tot willekeurige prestaties (50%) zodra de hint buiten het contextvenster valt.
  • ViZDoom-Two-Colors: RATE behaalt de hoogste totale beloning en de laagste onbalans tussen rode en groene pilaren, wat aantoont dat het de kleurhint van het begin van de episode consistent onthoudt. DT faalt zodra de hint uit beeld verdwijnt.
  • POPGym: Op een suite van 48 deels waarneembare taken behaalt RATE de hoogste score (9.54), en is het enige model dat een positieve score behaalt op de specifieke "memory puzzle" taken (0.45), terwijl alle andere modellen negatief scoren.
  • Minigrid-Memory & Memory Maze: RATE toont sterke interpolatie en extrapolatie op onbekende kaartgroottes en lange trajecten.

• Standaard RL-benchmarks (Atari & MuJoCo):

  • RATE presteert competitief met, en vaak beter dan, gespecialiseerde offline RL-algoritmen (zoals CQL, DT, Mamba) op standaard MDP-taken. Dit bewijst dat de geheugenmechanismen de prestaties in niet-geheugen-taken niet belemmeren.

• Ablatie Studies:

  • Het verwijderen van de MRV leidt tot een drastische daling in prestaties op lange trajecten, wat aantoont dat de klep essentieel is voor het behoud van informatie.
  • Het corrupten van geheugen-embeddings tijdens inferentie bevestigt dat deze tokens de primaire drager zijn van de cruciale hint-informatie.

5. Betekenis en Impact

Dit paper is significant omdat het een oplossing biedt voor een van de grootste uitdagingen in offline RL: het modelleren van lange-termijn afhankelijkheden in gedeeltelijk waarneembare omgevingen zonder de rekenkosten van self-attention exponentieel te laten toenemen.

• Unificatie: RATE fungeert als een universele architectuur die zowel korte als lange tijdshorizons effectief aanpakt, waardoor de noodzaak voor specifieke modellen per taaktype afneemt.
• Efficiëntie: Door trajecten te segmenteren en geheugen te gebruiken, kan RATE veel langere sequenties verwerken dan een standaard Transformer met een vast contextvenster, terwijl het GPU-geheugenverbruik lager blijft dan bij het verwerken van volledige lange trajecten in één keer.
• Toekomstperspectief: De resultaten onderstrepen dat geïntegreerde geheugenmechanismen (zoals MRV) essentieel zijn voor het oplossen van complexe POMDP's en stellen een nieuwe standaard voor voor architecturen die gericht zijn op lange-termijn besluitvorming.

Kortom, RATE bewijst dat het combineren van transformer-attentie met een slim, selectief recurrente geheugenmechanisme de sleutel is tot robuust offline leren in complexe, langdurige omgevingen.