Learning Abstract World Models with a Group-Structured Latent Space

Technische Samenvatting: Het Leren van Abstracte Wereldmodellen met een Groepsgestructureerde Latente Ruimte

Probleemstelling

Het leren van betekenisvolle abstracte modellen van Markov-beslissingsprocessen (MDP's) is essentieel voor het verbeteren van generalisatie vanuit beperkte data in versterkend leren (RL). Hoewel bestaande methoden voor het leren van wereldmodellen succesvol hoogdimensionale data inlagere dimensionale abstracte ruimtes embedden om de sample-efficiëntie te verbeteren, vereisen state-of-the-art benaderingen vaak grote trainingsdatasets om prestaties te bereiken die superieur zijn aan die van mensen. Een belangrijke beperking is het ontbreken van integratie van bekende geometrische structuren (symmetrieën) die inherent zijn aan de dynamiek van de omgeving. Zonder dergelijke priors kunnen agenten moeite hebben om te generaliseren naar ongezette scenario's, of ze hebben overmatige data nodig om eenvoudige symmetrieën te leren, zoals rotatie- of translatie-invarianties.

Methodologie

De auteurs stellen een kader voor dat geometrische priors oplegt aan de laagdimensionale representatiemanifold van een geleerd transitie-model. De kernmethodologie omvat:

1. Groepsgestructureerde Latente Ruimte: In plaats van de abstracte toestandsruimte $Z$ te behandelen als een standaard Euclidische vectorruimte, modelleren de auteurs deze als een quotiëntruimte (bijvoorbeeld $\mathbb{R}/k\mathbb{Z}$) of een product van dergelijke ruimtes. Deze keuze weerspiegelt de onderliggende symmetriegroepen van de omgeving (bijvoorbeeld cyclische groepen voor rotaties).
2. Groepsacties als Transitiedynamiek: Het transitie-model $\tau: Z \times A \to Z$ wordt beperkt tot werken via een groepsactie. Specifiek wordt de volgende latente toestand voorspeld als de huidige toestand gecombineerd met een geleerd transitievector via een groepsoperatie $\oplus$:
   $$\hat{z}_{t+1} = z_t \oplus \Delta(z_t, a; \theta_{trans})$$
   Voor quotiëntruimtes zoals $\mathbb{R}/k\mathbb{Z}$ wordt deze operatie geïmplementeerd via modulair rekenen, wat op natuurlijke wijze periodieke randvoorwaarden codeert.
3. Ontkoppelde Representaties: Om omstandigheden te hanteren die zowel symmetrische (gestructureerde) als niet-symmetrische (ongestructureerde) kenmerken bevatten, bevordert het kader ontkoppeling. Een sparseregularisatie wordt toegepast op de transitievector $\Delta$, waardoor kenmerken die invariant moeten blijven onder een specifieke actie gedwongen worden om nul te zijn. Dit wordt bereikt door specifieke coördinaten van de transitievector te maskeren op basis van een mapping $\sigma(a)$ die definieert welke latente dimensies worden beïnvloed door actie $a$.
4. Leerdoel: Het model wordt getraind met behulp van een zelftoezichtend contrastief leerdoel (InfoNCE-verlies) om voorspelde transities af te stemmen op de grondwaarheid, gecombineerd met een verlies voor het voorspellen van beloningen en een volumeregularisatieverlies om de abstracte ruimte compact en interpreteerbaar te houden. In tegenstelling tot eerdere benaderingen die equivariantie afdwingen via gewicht-gekoppelde netwerkarchitecturen (wat computatiever duur kan zijn), dwingt deze methode structuur direct af binnen de geometrie van de latente ruimte.

Belangrijkste Bijdragen

• Geometrische Priors in de Latente Ruimte: Het artikel toont aan hoe bekende symmetrische structuren (zoals rotatie- en translatiesymmetrieën) kunnen worden geïncorporeerd door geschikte latente ruimtes (quotiëntruimtes) en bijbehorende groepsacties te selecteren, in plaats van de neurale netwerkarchitectuur te wijzigen.
• Hanteren van Gemengde Symmetrieën: Het kader leert succesvol representaties die gestructureerde (symmetrische) en ongestructureerde kenmerken combineren, waardoor het model kan schalen naar complexe omgevingen zoals 3D-weergaven vanuit het eerste gezicht, waar niet alle dynamiek symmetrisch is.
• Verbeterde Generalisatie: De aanpak leidt tot eenvoudigere, meer ontkoppelde representaties die beter generaliseren naar ongezette transities en minder trainingsvoorbeelden vereisen in vergelijking met volledig ongestructureerde basismethoden.
• RL-prestaties Downstream: De geleerde abstracte modellen dienen als effectieve priors voor downstream RL-taken, wat de sample-efficiëntie in regimes met weinig data aanzienlijk verbetert.

Experimentele Resultaten

De auteurs evalueerden hun aanpak op verschillende omgevingen:

• Speelgoed-MDP's (Passage & Torus): In een eenvoudige cyclische gridwereld en een torus-MDP (product van twee cyclische groepen) voorspelde het groepsgestructureerde model nauwkeurig ongezette transities waar ongestructureerde basismethoden faalden. Het model herstelde de juiste topologie (bijvoorbeeld het mappen van de latente ruimte naar een torus in $\mathbb{R}^3$) met slechts 2 latente dimensies.
• MiniGrid (Top-Down View): De methode ontkoppelde succesvol oriëntatie (gemodelleerd als een cyclische groep) van positie (gemodelleerd als Euclidische ruimte), wat resulteerde in een representatie als productruimte $\mathbb{R}/2\pi\mathbb{Z} \times \mathbb{R}^2$.
• VizDoom (First-Person 3D View): In een hoogdimensionale omgeving met RGB-invoer mapte het model complexe visuele data naar een 3-dimensionale latente ruimte. Het legde correct de cyclische aard van rotatie en de translatie-aard van beweging vast, en presteerde beter dan basismethoden op generalisatiemetrics (Hits@k, Mean Reciprocal Rank) en downstream RL-taken.
• Kwantitatieve Metrics: In alle omgevingen presteerde het model "AWM + Geometrische Priors" significant beter dan ongestructureerde basismethoden (AWM zonder priors) en andere methoden (PRAE, Rotatiematrix) in generalisatie naar ongezette transities, met name in regimes met weinig data. In downstream RL-taken (Torus, MiniGrid, VizDoom) behaalden agenten die de geometrische prior gebruikten hogere cumulatieve beloningen met minder trainingsstappen.

Betekenis en Claims

Het artikel beweert dat door het benutten van geometrische priors via de keuze van latente ruimte en groepsacties, het mogelijk is om meer sample-efficiënte en generaliseerbare wereldmodellen te leren zonder de computatiever overhead van equivariante neurale netwerken. De auteurs benadrukken dat hun methode een flexibele combinatie van symmetrische en niet-symmetrische kenmerken mogelijk maakt, waardoor deze toepasbaar is op complexe, real-world scenario's zoals 3D-games vanuit het eerste gezicht. Het werk suggereert dat het respecteren van de onderliggende symmetrie van de omgeving in de abstracte representatie leidt tot "eenvoudigere en meer ontkoppelde representaties", wat op zijn beurt betere voorspelling en planning faciliteert. De auteurs leveren volledige code om reproduceerbaarheid te garanderen.