Technische Samenvatting: Sensorimotor World Models (SMWM)

Probleemstelling

Het artikel behandelt de uitdaging van het leren van latente wereldmodellen uit hoogdimensionale observaties (bijv. pixels) zonder beloningssignalen of taaklabels. Hoewel Joint Embedding Predictive Architectures (JEPA's) een veelbelovend kader bieden door toekomstige toestanden direct in een embedding-ruimte te voorspellen (om pixelreconstructie te vermijden), lijden ze aan representatie-instorting (representation collapse). Wanneer een encoder en een voorwaarts dynamica-model gezamenlijk worden getraind om de voorspellingsfout te minimaliseren, kan de encoder leren om alle observaties naar een constante embedding te mappen. Deze triviale oplossing minimaliseert de voorwaartse loss, maar maakt het model nutteloos voor downstream controle of planning.

Bestaande oplossingen om instorting te voorkomen, vertrouwen vaak op complexe mechanismen zoals:

• Het bevriezen van vooraf getrainde encoders (bijv. DINO-WM).
• Het toevoegen van distributionele regularisaties (bijv. het matchen van embeddings met een isotrope Gaussische verdeling via SIGReg).
• Het gebruik van Exponential Moving Average (EMA) target-netwerken met stop-gradients.
• Het hanteren van multi-term loss-functies met variantie-covariantie beperkingen.

De auteurs stellen dat deze benaderingen mogelijk niet volledig adresseren wat een nuttige wereldmodel-representatie moet behouden. Gebaseerd op het concept van "perceptie voor actie" en enactieve perceptie, betogen de auteurs dat representaties gevormd moeten worden door hun relevantie voor acties, waarbij onderscheid wordt behouden dat cruciaf is voor controleerbare dynamica, terwijl oncontroleerbare afleiders (distractors) worden weggefilterd.

Methodologie: Sensorimotor World Model (SMWM)

De auteurs stellen SMWM voor, een latent wereldmodel dat end-to-end wordt getraind vanaf offline, beloningsvrije trajecten $(o_t, a_t, o_{t+1})$. De kerninnovatie is het gebruik van inverse dynamics regularisatie als het enige mechanisme om representatie-instorting te voorkomen.

Architectuur

Het model bestaat uit drie componenten die gezamenlijk worden getraind:

1. Encoder ($f_\theta$): Brengt observaties $o_t$ in kaart naar een compacte latente embedding $z_t$.
2. Forward Dynamics Model ($g_\phi$): Voorspelt de volgende latente staat $\hat{z}_{t+1}$ gegeven de huidige staat $z_t$ en actie $a_t$.
3. Inverse Dynamics Model ($h_\psi$): Voorspelt de actie $\hat{a}_t$ die is uitgevoerd tussen twee opeenvolgende embeddings $z_t$ en $z_{t+1}$.

Trainingsdoel (Training Objective)

Het model wordt geoptimaliseerd door een gezamenlijke loss-functie te minimaliseren:
$$\mathcal{L} = \mathcal{L}_{fwd} + \lambda \mathcal{L}_{inv}$$

• Forward Loss ($\mathcal{L}_{fwd}$): Mean Squared Error (MSE) tussen de voorspelde volgende embedding $\hat{z}_{t+1} = g_\phi(z_t, a_t)$ en de werkelijke gecodeerde volgende staat $z_{t+1} = f_\theta(o_{t+1})$.
• Inverse Loss ($\mathcal{L}_{inv}$): MSE tussen de voorspelde actie $\hat{a}_t = h_\psi(z_t, z_{t+1})$ en de grondwaarheid actie $a_t$.
• Regularisatie Gewicht ($\lambda$): Een enkele hyperparameter die de sterkte van het inverse dynamics signaal regelt.

Werkingsmechanisme

De inverse loss fungeert als een anti-collapse mechanisme. Als de encoder instort naar een constante embedding, ontvangt het inverse model identieke inputs voor alle transities en kan het slechts een constante actie voorspellen (het gemiddelde van de actieverdeling). Om de inverse loss onder dit risico van een constante voorspeller te verlagen, moet de encoder informatie over de actie $a_t$ behouden die aanwezig is in de transitie $(o_t, o_{t+1})$. Dit dwingt de representatie om actie-relevante informatie (controleerbare vrijheidsgraden) te behouden, terwijl het op natuurlijke wijze oncontroleerbare afleiders (exogene ruis) filtert die de actieherstel niet beïnvloeden.

In tegenstelling tot distributionele regularisaties die de geometrie van de embedding-ruimte voorschrijven (bijv. het afdwingen van een Gaussische verdeling), verankert inverse dynamics de representatie aan een taak-gegrond kwantiteit (de actie) zonder specifieke geometrische priors op te leggen.

Belangrijkste Bijdragen

1. SMWM Architectuur: Een simpel, end-to-end trainbaar sensorimotor wereldmodel dat inverse dynamics gebruikt als het enige anti-collapse mechanisme, waarbij slechts één extra hyperparameter ($\lambda$) nodig is.
2. Representatie Eigenschappen: Empirisch bewijs dat de geleerde representaties niet alleen niet-gecollabeerd zijn, maar specifiek de controleerbare vrijheidsgraden volgen, ruimtelijke geometrie (zowel Cartesiaans als periodiek) behouden en oncontroleerbare afleiders wegfilteren.
3. Downstream Prestaties: Demonstratie van competitieve planning-prestaties over 2D en 3D controle-taken, wat aantoont dat de structurele voordelen van de representatie vertalen naar effectieve doel-geconditioneerde planning.

Experimentele Resultaten

Gecontroleerde Testomgevingen (Dot World & Sprite World)

• Herstel van Intrinsieke Dimensie: In een "dot world" omgeving slaagde het model erin de ware 2D intrinsieke dimensie van de staat-ruimte te herstellen uit pixels alleen, zelfs wanneer deze ingebed was in een 6ada-dimensionale latente ruimte.
• Filtering van Afleiders: Wanneer oncontroleerbare "afleider"-objecten (die willekeurig bewogen) werden geïntroduceerd, negeerde de encoder deze. De effectieve latente dimensie kwam overeen met de controleerbare vrijheidsgraden, niet met de totale observatie-dimensie.
• Controle-afhankelijke Perceptie: In een sprite world experiment leerde het model pose-variabelen (positie, oriëntatie) alleen te behouden wanneer deze controleerbaar waren. Wanneer de oriëntatie oncontroleerbaar was, middelde de representatie dit uit, wat resulteerde in een wazige reconstructie van de oriëntatie.

Planning Benchmarks

De auteurs evalueerden SMWM op vier offline controle-taken: TwoRoom (2D navigatie), Reacher (continue controle), Push-T (2D manipulatie) en OGBench-Cube (3D manipulatie).

• Baselines: Vergeleken met SIGReg (een sterke JEPA-baseline met Gaussische regularisaties), Forward-only (geen inverse loss) en Random actie-sampling.
• Prestaties:
  • SMWM kwam overeen met de SIGReg prestaties op de drie 2D-taken.
  • SMWM overtrof SIGReg op de 3D OGBench-Cube taak (84% succes vs. 59%).
  • Beide geregulariseerde methoden presteerden significant beter dan de Forward-only ablatie, wat de noodzaak van een anti-collapse mechanisme bevestigt.
• Latente Geometrie: Probing-experimenten toonden aan dat SMWM embeddings lage-dimensionale subruimtes vormen die de fysieke staat-geometrie weerspiegelen (bijv. cilindrische structuren voor periodieke gewrichtshoeken in Reacher, Cartesiaanse kaarten voor navigatie).

Betekenis en Claims

Het artikel claimt dat SMWM een geprincipeerde en eenvoudige instantiatie biedt van "perceptie voor actie" in latent wereldmodel leren. Door de latente staten te dwingen voldoende informatie te behouden om de acties van de agent te herstellen, bias het model de representatie op natuurlijke wijze richting de controleerbare vrijheidsgraden van de omgeving.

Belangrijke conclusies zijn onder meer:

• Eenvoud: Inverse dynamics alleen is voldoende om de end-to-end JEPA training te stabiliseren zonder reconstructie, beloningssignalen, bevroren encoders of complexe distributionele priors.
• Structurele Bias: De methode induceert representaties die interpreteerbaar zijn en in lijn zijn met de causale structuur van de omgeving (acties en hun effecten), in plaats van alleen bij correlaties.
• Praktisch Nut: De geleerde representaties ondersteunen effectieve latente planning, met name in complexe 3D-omgevingen waar andere geregulariseerde baselines moeite hebben.

De auteurs erkennen beperkingen, zoals de aanname dat acties herstelbaar zijn uit opeenvolgende observaties en de huidige onmogelijkheid om grootheden te vangen die niet identificeerbaar zijn vanuit enkelvoudige frames (zoals snelheid) zonder observatiegeschiedenissen. Ze positioneren SMWM echter als een robuuste, beloningsvrije aanpak voor het leren van actie-gerichte wereldmodellen uit offline data.