Technische Samenvatting: Terug naar Parsimonious Latents: Het Leren van Taakgerichte Wereldmodellen vanuit Visuele Fundamenten

Probleemstelling

Wereldmodellen stellen agenten in staat om toekomstige dynamiek te voorspellen, geconditioneerd op acties, een vermogen dat centraal staat in planning en controle. Echter, bestaande benaderingen staan voor een fundamentele afweging in de latente representatie:

1. Pixel-gebaseerde modellen (bijv. Dreamer, MuZero) leren dynamiek direct vanuit pixels of compacte latente ruimtes die worden gevormd door beloningssignalen. Hoewel effectief, gaan ze vaak gepaard met hoge rekenkosten of vereisen ze online interactie/beloningssupervisie om de latente ruimte vorm te geven.
2. Foundation embedding-modellen (bijv. DINO-WM, V-JEPA) maken gebruik van bevroren, vooraf getrainde visuele encoders om direct te voorspellen in een hoogdimensionale semantische ruimte. Hoewel deze rijke semantische priors en zero-shot generalisatie bieden, bevatten de embeddings overmatige taak-irrelevante details (texturen, verlichting, achtergrond) die geen invloed hebben op de controle.

In beloningsvrije offline settings, waar agenten moeten leren uit vaste trajecten zonder beloningssupervisie, worden deze problemen verergerd. Het gebruik van ruwe foundation embeddings als de toestandsruimte dwingt het model om capaciteit te verspillen aan irrelevante variaties, wat leidt tot een slechte modelleer-efficiëntie, verminderde controle-interpreteerbaarheid en potentiële representatie-inzakking, met name in hoogdimensionale actieruimtes (bijv. 7-DoF robotarmen). De kernuitdaging is het extraheren van een compacte, taak-toereikende latente representatie uit deze rijke maar ruisachtige foundation embeddings zonder online interactie.

Methodologie: TC-WM

De auteurs stellen TC-WM (Task-Centric World Model) voor, een kader dat vooraf getrainde foundation embeddings niet behandelt als de definitieve toestandsruimte, maar als een "semantisch steiger" waaruit taakrelevante dynamiek wordt gedistilleerd.

Kernarchitectuur

TC-WM werkt op een pijplijn van visuele embeddings ($x_t$) en proprioceptieve signalen ($s^p_t$):

1. Gecombineerde Embedding: Een bevroren visuele encoder (bijv. DINOv2) beeldt afbeeldingen af op visuele kenmerken, terwijl een trainbare lineaire laag proprioceptie encodeert. Deze worden geconcateneerd om een gecombineerde embedding $x_t$ te vormen.
2. Latente Encoder (Projectie): Een lineaire encoder projecteert de hoogdimensionale gecombineerde embedding naar een compacte latente staat $z_t$. Cruciaal is dat $z_t$ wordt gefactoriseerd in twee deelruimtes:
   • $z^s_t$ (Staat-gerelateerd): Uitgelijnd met de fysieke staat.
   • $z^c_t$ (Complementair): Behoudt residuële semantische structuur (uiterlijk, achtergrond) om reconstructie-voorwaarden te voldoen.
3. Latente Dynamiek: Een op Transformer gebaseerde predictor ($f_\theta$) modelleert temporele transities in de compacte latente ruimte ($\hat{z}_{t+1} = f_\theta(z_{t-H:t}, a_{t-H:t})$).
4. Taakgerichte Uitlijning: Het model hanteert een contrastief doel (InfoNCE) om de $z^s_t$-deelruimte uit te lijnen met de proprioceptieve staat $s^p_t$. Dit dwingt de encoder actief om taakrelevante fysieke factoren te scheiden van storende variaties.
5. Embedding Decoder: Een lineaire decoder reconstrueert de oorspronkelijke visuele embedding vanuit $z_t$. Dit fungeert als een regularisator, waarbij wordt gegarandeerd dat de compacte latente staat verankerd blijft in het foundation-variëteit en degeneratieve oplossingen worden voorkomen.

Trainingsdoel

De totale loss combineert vier componenten:

• Latente Dynamiek Loss ($L^z_{dyn}$): Minimaliseert voorspelfout in de latente ruimte.
• TC Dynamiek Loss ($L^s_{dyn}$): Voorspelt toekomstige proprioceptie direct uit het verleden om taakgerichte dynamiek af te dwingen.
• Uitlijningsloss ($L_{align}$): Contrastieve loss die $z^s_t$ uitlijnt met $s^p_t$.
• Reconstructieloss ($L_{rec}$): Reconstructeert de visuele embedding vanuit $z_t$ om visuele structuur te behouden.

Planning op Testtijd

De geleerde taakgerichte latente ruimte ondersteunt diverse planningsstrategieën:

• Laagdimensionale controle (Navigatie): Model Predictive Control (MPC) met behulp van de Cross-Entropy Method (CEM) om actiesequenties in de latente ruimte te optimaliseren.
• Hoogdimensionale manipulatie: Latent Diffusion Planning (LDP) gecombineerd met een Inverse Dynamics Model (IDM) om uitvoerbare acties te genereren vanuit latente trajecten.

Belangrijkste Bijdragen

1. Taakgerichte Decompositie: Het artikel introduceert een methode om hoogdimensionale foundation embeddings te decomponeren in een compacte, taak-toereikende latente ruimte door fysieke staten-dynamiek expliciet te scheiden van semantische storende variabelen.
2. Actieve Identificeerbaarheid: In tegenstelling tot eerdere werken die passief aannemen dat mechanismediversiteit zorgt voor identificeerbaarheid, gebruikt TC-WM contrastieve uitlijning met proprioceptie om actief onderscheidende conditionals voor verschillende fysieke staten te induceren, wat theoretisch de herwinning van taakgerichte factoren garandeert tot een affiene transformatie.
3. Beloningsvrije Offline Learning: Het kader bereikt hoogwaardige wereldmodellering en planning in offline settings zonder beloningssignalen, uitsluitend vertrouwend op visuele waarnemingen en proprioceptieve zij-informatie.
4. Theoretische Garantie: De auteurs leveren een stelling die bewijst dat onder milde aannames (contextuele dynamiek, latent-geobserveerde variabiliteit en voldoende uitlijning) de geleerde compacte latente staat identificeerbaar is en de ware fysieke staat herwint tot een herparameterisering.

Experimentele Resultaten

De auteurs evalueren TC-WM over 9 offline benchmarks (Maze, Wall, Push-T, Lift, Can, Square, Reacher, Cheetah, Hopper) tegen baselines waaronder TD-MPC2, DreamerV3, MuZero en DINO-WM.

• Kwaliteit Wereldmodel: TC-WM bereikt de laagste voorspelfout (MSE) voor latente rollouts over bijna alle taken, met name in manipulatie-omgevingen waar accurate dynamiek op lange horizon kritiek is. Hoewel reconstructie in de beeldruimte concurrerend is, ligt de primaire winst in de stabiliteit en nauwkeurigheid van de latente dynamiek.
• Planningsprestaties:
  • Bij navigatie (laagdimensionale acties) komt TC-WM overeen met de prestaties van sterke baselines zoals TD-MPC2.
  • Bij manipulatie (hoogdimensionale, contactrijke acties) presteert TC-WM aanzienlijk beter dan alle baselines. Bijvoorbeeld, op de Robomimic "Lift"-taak bereikt TC-WM een slagingspercentage van 54% vergeleken met 29% voor DINO-WM en 31% voor TD-MPC2.
• Representatie-analyse:
  • Lineair Probing: Lineaire regressors getraind op de latente ruimte van TC-WM herstellen proprioceptieve en objectstaten met hogere $R^2$-scores dan baselines, wat bevestigt dat de latente staat fysiek onderbouwd is.
  • Anti-Inzakking: TC-WM demonstreert een beter gebruik van de latente ruimte (hogere effectieve rang) in vergelijking met DINO-WM, dat lijdt aan representatie-inzakking bij direct gebruik van ruwe embeddings.
  • Robuustheid: TC-WM behoudt prestaties onder visuele verstoringen (Gaussisch ruis, kleurjitter) beter dan DINO-WM, met kleinere dalingen in SSIM.

Betekenis en Claims

Het artikel claimt dat foundation embeddings het beste worden benut als een semantisch steiger in plaats van een directe toestandsrepresentatie. Door deze embeddings te projecteren naar een taakgerichte deelruimte die is uitgelijnd met fysieke signalen, kunnen agenten het volgende bereiken:

• Parsimonious Latents: Een reductie in dimensionaliteit die irrelevante visuele details filtert terwijl de voor controle noodzakelijke dynamiek behouden blijft.
• Verbeterde Controle: Verbeterde planningsprestaties in complexe, hoogdimensionale omgevingen waar standaard foundation-modellen falen door ruis en irrelevante variaties.
• Theoretische Soundheid: Een formele rechtvaardiging voor waarom lineaire projecties en contrastieve uitlijning voldoende zijn om taakrelevante factoren te herwinnen uit bevroren visuele encoders.

De auteurs positioneren TC-WM als een brug tussen de generalisatiekracht van foundation-modellen en de controleerbaarheid die vereist is voor fysieke agenten, specifiek gericht op het overbruggen van de kloof in beloningsvrije offline leerscenario's.