Technische Samenvatting: Dueling World Models

Probleemstelling

Latente wereldmodellen, in het bijzonder die welke de Joint-Embedding Predictive Architecture (JEPA) volgen, plannen door toekomstige staat-embeddings te voorspellen op basis van huidige staten en acties. Echter, wanneer een scène beweging bevat waar de agent geen controle over heeft (distractoren), lijden deze modellen aan een "stille" modus van falen: de voorspellingen voor verschillende acties worden ononderscheidbaar, zelfs terwijl de algehele trainingsverlies (reconstructie- of voorspellingsfout) blijft verbeteren.

Bestaande remedies proberen controleerbare dynamiek te scheiden van exogene distractoren door externe mechanismen toe te voegen, zoals reconstructieverliezen, taakbeloningen, hulpdoelen voor inverse dynamica of aparte netwerkvertakkingen. Deze benaderingen vereisen vaak specifieke aannames over het distractor-proces of extra supervisie. De auteurs identificeren dat in standaard monolithische voorspellers $P(z, a)$ het actie-kanaal verstrengeld raakt met de distractor-dynamica, tenzij dit expliciet wordt gegrond door ware actie-labels of wordt gezuiverd via segmentatiemaskers. Cruciaal is dat standaard modelselectie op basis van validatieverlies vaak faalt om deze "actie-blindheid" te detecteren, waarbij actief checkpoints worden geselecteerd waar het actie-signaal is ingestort.

Methodologie: AD-JEPA

Het artikel stelt AD-JEPA (Action-Decomposed JEPA) voor, een minimale herstructurering van actie-geconditioneerde latente voorspellers die het actie-kanaal isoleert zonder dat er beloningen, reconstructie of distractor-specifieke hulpverliezen nodig zijn.

Kernmechanisme: Action-Mean Centering

De methode deelt de voorspeller op in twee stromen:

1. Passieve Stroom ($B(z)$): Voorspelt hoe de scène evolueert wanneer de actie gemarginaliseerd wordt.
2. Gecentreerde Actiestroom ($u(z, a)$): Representeert de actie-specifieke offset, berekend door de gemiddelde voorspelling over acties af te trekken van de ruwe actie-offset.

De update-regel wordt gedefinieerd als:
$$\hat{z}_{t+1} = B(z_t) + \underbrace{[\Delta(z_t, a_t) - \bar{\Delta}(z_t)]}_{u(z_t, a_t)}$$
Waarbij $\bar{\Delta}(z) = \mathbb{E}_{a' \sim \nu}[\Delta(z, a')]$ de gemiddelde offset is over een voorstelverdeling $\nu$ van acties.

Theoretische Fundering

Deze benadering past de Dueling DQN decompositie ($Q = V + (A - \bar{A})$) aan van scalaire waarden naar vector-waardige latente dynamica. De auteurs bewijzen dat deze aftrekking een algebraïsche identiteit is in plaats van een geleerde neiging:

• Common-Mode Cancellatie: Elke variatie die de actie-offsets identiek binnenkomt over alle acties heen (common-mode variatie), inclus ook actie-onafhankelijke distractoren, wordt exact geannuleerd.
• Exactheid: De annulering geldt zowel voor discrete actiesets (via exacte enumeratie) als voor continue actiesets (via Monte-Carlo sampling), mits het gemiddelde wordt geschat over een eindige set kandidaten.
• Geen Exogeniteitsaanname: De methode vereist geen proces-niveau exogeniteitsaannames (dat wil zeggen, dat distractoren onafhankelijk zijn van de policy van de agent op procesniveau); het vereist alleen dat de bijdrage van de distractor aan de transitie common-mode is over de acties.

Trainingsdoel

Het model wordt getraind met een standaard JEPA-voorspellingsverlies tegen een stop-gradient EMA target, aangevuld met:

• InfoNCE ($L_{act}$): Moedigt de voorspelling van de uitgevoerde actie aan om het dichtst bij de doelwaarde te liggen tussen de contrafeitelijke acties.
• Compactheidsstraf ($L_{off}$): Ontmoedigt het actie-kanaal om de staat-persistentie te absorberen.
• VICReg Regularisatie: Voorkomt representatie-instorting.

Opvallend genoeg rapporteren de auteurs dat het toevoegen van een geleerd gating-mechanisme om dimensies te selecteren inert was, wat degradeerde naar uniforme down-scaling, wat bevestigt dat centrering alleen de werkende mechanisme is.

Belangrijkste Bijdragen

1. Action-Mean Centering: Een een-regels herstructurering van actie-geconditioneerde voorspellers die het actie-kanaal common-mode invariant maakt. De auteurs leveren bewijzen (Proposities 1 en 2) die de exacte annulering van common-mode componenten tonen voor zowel discrete als gesamplede actiesets.
2. Empirische Validatie: De methode werd gevalideerd over vier verschillende settings:
   • Gridworld (FourRooms): Met rollende binaire distractor-cellen.
   • Synthetische Generatoren: Met bekende controleerbare en nuisance factoren.
   • Continuous Control (DMC): Met pixel-niveau bewegende occluders.
   • Atari (Freeway): Met natuurlijke verkeers-distractoren.
     In alle gevallen behield het gecentreerde kanaal het effect van de agent terwijl het de nuisance-lekkage verminderde tot ononderscheidbaar van nul, terwijl verstrengelde voorspellers faalden.
3. Readout-Only Toepasbaarheid: Omdat het mechanisme een eigenschap is van de readout en niet van de architectuur, kan het post-hoc worden toegepast op bevroren, vooraf getrainde wereldmodellen (specifiek RePo en TIA). De auteurs demonstreerden dat het aftrekken van het actie-gemiddelde van de ruwe readouts van deze bevroren modellen succesvol een actie-kanaal herstelde dat anders latent en niet gerouteerd was.
4. Gedefinieerde Grens: Het artikel identificeert expliciet de beperking: de methode faalt wanneer distractor-dynamica actie-gecorreleerd zijn (bijv. reafferente distractoren of agenten die reageren op de gebruiker). In dergelijke gevallen is de distractor niet common-mode en wordt deze niet geannuleerd.

Resultaten

• Herstel vs. Afwijzing: In synthetische tests bereikte het gecentreerde kanaal een hoge herwinning van de controleerbare factor ($R^2 \approx 0.87\text{--}0.93$) terwijl het de nuisance factor afwees ($R^2 \approx -0.00$). Niet-gecentreerde varianten herstelden het signaal maar lekten de nuisance zwaar.
• Controleprestaties: In de gridworld behield een model-predictive controller met het gecentreerde kanaal een hoge doelbereikingssucces ($\approx 0.92$), zelfs met 30 distractor-cellen, terwijl de standaard voorspeller instortte naar $\approx 0.57$.
• Post-Hoc Succes: Het toepassen van de aftrekking op bevroren RePo en TIA modellen herstelde actie-kanalen met $R^2$-waarden tussen 0.19 en 0.52, terwijl de ruwe actie-delta readouts nabij nul waren.
• Falen: De methode faalde correct (instorten naar kansniveau) in actie-gecorreleerde stress-tests, wat overeenkomt met de theoretische grens.

Betekenis en Claims

Het artikel claimt dat centrering een universele plug-in is voor actie-geconditioneerde wereldmodellen. De betekenis ligt in:

• Eenvoud: Het bereikt distractor-afwijzing zonder de complexiteit van reconstructie, reward-gebaseerde factorisatie of hulp-inverse-dynamica verliezen.
• Structurele Garantie: De scheiding is een algebraïsche eigenschap van de parametrisering, niet een geleerde neiging afhankelijk van specifieke trainingssignalen.
• Retrofitting Mogelijkheid: Het stelt gebruikers in staat om bestaande, bevroren wereldmodellen tijdens inferentie te "ont-distracteren" zonder hertraining, waardoor actie-kanalen naar boven komen die voorheen verborgen waren door verstrengelde dynamica.

De auteurs zijn bescheiden over de schaalbaarheid van controle: hoewel het kanaal controle mogelijk maakt in gridworlds, faalde een vooraf geregistreerde poging om dit te schalen naar DMC continue controle-taken door niet beter te presteren dan random baselines. Zij schrijven dit toe aan model-accuraatheid bottlenecks in de planner in plaats van een defect in het kanaal zelf. De primaire bijdrage is dus op representatie-niveau: het bieden van een schoon, common-mode-invariant actie-kanaal waar verstrengelde voorspellers "actie-blind" gaan.