MAGE: Multi-scale Autoregressive Generation for Offline Reinforcement Learning

MAGE is een nieuwe offline versterkingsleermethode die gebruikmaakt van multi-schaal autoregressieve generatie en conditionele geleiding om coherente en controleerbare trajecten te produceren in langdurige taken met schaarse beloningen.

De kern

MAGE: De Slimme Architect voor Robot-Plannen

Stel je voor dat je een robot wilt leren om een complexe taak uit te voeren, zoals een robotarm die een pen vasthoudt en een tekening maakt, of een robot die door een enorm labyrint moet lopen om een schat te vinden. Het probleem is dat deze robots vaak alleen maar "kijken" naar oude video's van hoe het gedaan moet worden (offline leren), zonder dat ze zelf kunnen experimenteren.

De uitdaging? Veel van deze taken zijn langdurig en hebben weinig directe beloningen. Het is alsof je een lange reis maakt waarbij je pas aan het einde een prijs krijgt, en tussendoor geen enkel "goed gedaan"-signaal. Bestaande methoden raken hier vaak de draad kwijt; ze maken mooie kleine stapjes, maar vergeten het grote plaatje, waardoor de robot tegen de muur loopt of de verkeerde kant op gaat.

De auteurs van dit papier hebben MAGE bedacht. Laten we MAGE uitleggen met een paar creatieve vergelijkingen.

1. Het Probleem: De "Kijk-Door-de-Verrekijker"-Methode

Bestaande AI-methoden kijken vaak als een mens door een verrekijker: ze zien heel goed wat er direct voor hun neus gebeurt (de korte termijn), maar ze zien de horizon niet.

• Andere methoden (zoals Decision Transformer): Ze schrijven een verhaal zin voor zin. Ze weten wat ze net hebben geschreven, maar ze verliezen soms het hoofdidee van het verhaal uit het oog.
• Andere methoden (zoals Diffusion): Ze proberen het hele verhaal in één keer te "ontwarren" uit ruis. Dit werkt goed voor korte verhalen, maar bij lange, ingewikkelde verhalen worden de zinnen soms logisch, maar past het verhaal niet bij elkaar (lokaal logisch, globaal gek).

2. De Oplossing: MAGE als een "Meer-Schaalige Architect"

MAGE werkt niet zin voor zin, maar als een architect die een gebouw ontwerpt.

Stap 1: De Grove Schets (De "Vogelperspectief")
Stel je voor dat je een huis wilt bouwen. Je begint niet met het kiezen van de kleur van de tegels in de badkamer. Je begint met een schets op een vel papier: Waar komt de woonkamer? Waar de slaapkamer? Waar de trap?
MAGE doet dit eerst. Het maakt een grove, lange-termijn schets van de hele route. Dit is de "coarse" (grove) schaal. Het weet: "Eerst ga ik naar links, dan naar de schat, dan naar de uitgang."

Stap 2: Het Verfijnen (De "Microscoop")
Nu de schets er is, gaat MAGE het verfijnen.

• Schaal 2: Waar komen de muren precies?
• Schaal 3: Hoe ziet de vloer eruit?
• Schaal 4: Welke tegels gebruiken we?
  MAGE werkt autoregressief (stap voor stap), maar altijd van groot naar klein. Het gebruikt de grove schets om de volgende, fijnere details te bepalen. Hierdoor weet de robot altijd waar hij naartoe gaat (het grote doel), terwijl hij tegelijkertijd weet hoe hij zijn voet precies moet zetten (de kleine details).

3. De Magische "Kompass" (Conditionering)

Een ander probleem is dat robots soms vergeten waar ze begonnen zijn. Ze beginnen misschien met de verkeerde houding.
MAGE heeft een kompass dat continu wordt gecontroleerd.

• De robot krijgt een opdracht: "Ga naar de schat (doel) en begin hier (start)."
• MAGE gebruikt een speciale module (de decoder met adapter) die als een strenge leraar fungeert. Als de robot tijdens het plannen begint te afdwalen van de startpositie of het doel, corrigeert deze leraar het plan direct.
• Vergelijking: Het is alsof je een GPS hebt die niet alleen de route tekent, maar ook constant checkt: "Wacht, je bent nu 10 meter van de start weg, dat klopt niet. Pas je route aan."

4. Waarom werkt dit zo goed?

In de tests (zoals het laten schrijven van een robot met een pen of het vinden van een weg in een labyrint) bleek MAGE superieur:

• Bij lange taken: Waar andere robots de draad kwijtraken na 10 stappen, houdt MAGE het overzicht over 100 stappen.
• Bij weinig beloningen: Omdat MAGE eerst het "grote plaatje" (de schets) maakt, weet het dat het even moet doorwerken voordat het de prijs krijgt. Het raakt niet gefrustreerd.
• Snelheid: Ondanks dat het slim is, is het snel genoeg om in real-time te werken (zoals een mens die reageert).

Samenvattend

MAGE is als een meester-architect die een robot leert plannen.

1. Eerst tekent hij de hoofdlijnen van de reis (van start naar finish).
2. Dan verfijnt hij die lijnen stap voor stap tot gedetailleerde instructies.
3. Een strenge leraar (de conditionering) zorgt ervoor dat de robot nooit de startpositie of het doel uit het oog verliest.

Hierdoor kunnen robots nu complexe, lange taken uitvoeren die eerder te moeilijk waren, zonder dat ze zelf hoeven te experimenteren. Ze leren gewoon van de beste oude plannen, maar dan op een manier die zowel het grote plaatje als de kleine details perfect combineert.

Technische samenvatting

Hieronder volgt een gedetailleerde technische samenvatting van het paper "MAGE: Multi-Scale Autoregressive Generation for Offline Reinforcement Learning", geschreven in het Nederlands.

1. Probleemstelling

Offline Reinforcement Learning (RL) richt zich op het trainen van agenten uitsluitend op bestaande datasets, zonder verdere interactie met de omgeving. Hoewel generatieve modellen (zoals Diffusion-modellen en Transformers) succesvol zijn in het modelleren van complexe trajectverdelingen, kampen ze met aanzienlijke beperkingen bij lange-horizon taken met schaarse beloningen (sparse rewards).

De bestaande methoden ondervinden de volgende problemen:

• Gebrek aan lange-termijn bewustzijn: Methoden zoals Decision Transformers zijn vaak unidirectioneel en hebben moeite met het begrijpen van de globale context over lange tijdshorizons.
• Lokale generatie-bias: Diffusiemodellen genereren vaak trajecten die lokaal plausibel zijn, maar die globally inconsistent of oncoherent zijn, wat leidt tot falen in complexe planningstaken.
• Beperkte hiërarchische modellen: Bestaande hiërarchische methoden gebruiken vaak een starre tweelaagse structuur (hoog niveau voor subdoelen, laag niveau voor acties). Dit beperkt het vermogen om multi-schaal temporele afhankelijkheden volledig te modelleren en introduceert optimalisatieproblemen door het gezamenlijk trainen van twee afhankelijke beleidsfuncties.

2. Methodologie: MAGE

MAGE (Multi-scale Autoregressive GEneration) is een nieuw raamwerk voor offline RL dat trajecten genereert in een top-down, van grof naar fijn (coarse-to-fine) proces. Het doel is om zowel de globale structuur als de lokale dynamiek van een traject te modelleren.

Het systeem bestaat uit twee kerncomponenten:

A. Multi-Scale Traject Auto-Encoder (MTAE)

Dit component encodeert een traject $\tau$ in een hiërarchie van discrete latente representaties (tokens) op verschillende tijdschalen.

• Input: Het traject wordt gemodelleerd als een reeks paren van (toestand $s$, Return-to-Go $R$).
• VQ-VAE Architectuur: Het gebruikt een Vector Quantized Variational Autoencoder (VQ-VAE) met een gedeelde codebook.
• Multi-Schaal Encoding: Een traject wordt omgezet in een reeks token-kaarten $M = (m_1, m_2, ..., m_K)$, waarbij $m_1$ de meest grove (globale) structuur vastlegt en $m_K$ de fijnste details bevat.
• Residuale Encoding: Bij het coderen wordt de informatie van de grovere schaal afgetrokken van de input voordat de volgende, fijnere schaal wordt gecodeerd. Dit zorgt ervoor dat elke schaal de resterende details vastlegt.

B. Multi-Scale Condition-Guided Autoregressive Generator

Dit component genereert de trajectrepresentaties sequentieel van grof naar fijn.

• Autoregressieve Generatie: Een Transformer voorspelt de token-kaart $m_k$ voor schaal $k$, conditioned op de eerder gegenereerde grovere kaarten $m_{<k}$, de starttoestand $s_0$ en de gewenste Return-to-Go $R_0$.
• Formule: $p(m_1, ..., m_K | s_0, R_0) = \prod_{k=1}^K p(m_k | m_{<k}, s_0, R_0)$.
• Actie Bepaling: Na het genereren van de latente representaties $Z$, wordt een actie bepaald via een latente inverse dynamiek model. Dit model leert de actie direct uit de samengevoegde latente vector, wat efficiënter is dan het decoderen van het volledige traject eerst.

C. Condition-Guided Refinement

Om ervoor te zorgen dat het gegenereerde traject exact begint bij de gewenste starttoestand $s_0$ en voldoet aan de doelstelling $R_0$, introduceert MAGE een condition-guided adapter in de decoder.

• Een extra verliesfunctie ($L_{cond}$) straft afwijkingen van de initiële conditie $(s_0, R_0)$ na het decoderen.
• Dit zorgt voor precieze controle over kortetermijngedrag en voorkomt dat het traject afwijkt van de beoogde route.

3. Belangrijkste Bijdragen

1. Multi-Schaal Autoregressieve Generatie: MAGE is de eerste offline RL-methode die een multi-schaal autoregressieve aanpak (geïnspireerd door Visual Autoregressive modellen) toepast op trajectgeneratie, waardoor zowel lange-termijn afhankelijkheden als korte-termijn details effectief worden gemodelleerd.
2. Unified Policy: In tegenstelling tot traditionele hiërarchische methoden die meerdere beleidsfuncties vereisen, leert MAGE één enkel beleidsmodel dat werkt over alle tijdschalen, wat de training stabiliseert en de coherentie verbetert.
3. Condition-Guided Decoder: De introductie van een adapter-module voor conditionele verfijning lost het probleem op van trajecten die wel globaal correct zijn maar lokaal afwijken van de startconditie.
4. Efficiëntie: Ondanks de complexiteit van de multi-schaal structuur, behoudt MAGE een hoge inferentiesnelheid, wat het geschikt maakt voor real-time toepassingen.

4. Resultaten

MAGE is geëvalueerd op vijf offline RL-benchmarks tegen 15 baseline-algoritmen (inclusief Decision Transformer, Diffusion-QL, HDMI, HD, etc.).

• Prestaties bij Schaarse Beloningen: MAGE behaalt state-of-the-art resultaten, vooral in lange-horizon taken zoals Adroit (dexterous manipulation), Franka Kitchen (compositional tasks) en AntMaze/Maze2D (navigatie).
  • In de Adroit omgevingen (bijv. Pen, Door, Hammer) overtreft MAGE alle baselines aanzienlijk, zelfs op datasets met menselijke of gekloonde data.
  • In navigatietaken (AntMaze) behaalt MAGE de hoogste scores op bijna alle datasets, terwijl andere methoden vaak vastlopen in muren of falen in het bereiken van het doel.
• Dense Rewards: MAGE presteert ook competitief in omgevingen met dichte beloningen (Gym locomotion tasks), wat aantoont dat de methode algemeen toepasbaar is.
• Ablatie Studies:
  • Het gebruik van meerdere tijdschalen ($K$) verbetert de prestaties tot een optimaal punt (rond $K=8$).
  • Het modelleren van $(R, S)$ (Return en State) werkt beter dan het modelleren van alleen acties of alleen toestanden.
  • De conditionele loss ($L_{cond}$) is cruciaal om afwijkingen in het begin van het traject te voorkomen.
• Snelheid: MAGE is aanzienlijk sneller dan diffusie-gebaseerde methoden (ongeveer 50x sneller dan HD en 80x sneller dan DD), met een inferentietijd van ongeveer 27 ms per stap, wat binnen de eisen voor robotbesturing valt.

5. Betekenis en Conclusie

MAGE biedt een krachtig nieuw paradigma voor offline Reinforcement Learning door het probleem van lange-horizon planning te benaderen via multi-schaal autoregressieve generatie. Door de globale structuur eerst te schetsen en deze vervolgens iteratief te verfijnen, slaagt MAGE erin om coherentie en controle te combineren waar eerdere generatieve modellen tekortschoten.

De methode lost het fundamentele dilemma op tussen het modelleren van lange-termijn afhankelijkheden en het behouden van lokale precisie. De resultaten suggereren dat MAGE een robuuste en efficiënte oplossing biedt voor complexe sequentiële besluitvormingstaken in de echte wereld, zoals robotica en medische planning, waar schaarse beloningen en lange planningshorizons de norm zijn. De openbaarmaking van de broncode draagt bij aan de reproduceerbaarheid en verdere ontwikkeling in het veld.