Technische Samenvatting: XP-JEPA: Cross-Predictieve Fysische Gronding voor Voorspelbare Latente Dynamica

1. Probleemstelling

Latente wereldmodellen maken planning mogelijk door te zoeken in de representatieruimte: een agent codeert de huidige scène, rolt kandidaat-acties door geleerde dynamica heen en selecteert acties waarvan de voorspelde uitkomsten een doel benaderen. Deze aanpak rust op representaties met betrouwbare actie-geconditioneerde dynamica.

Echter, zelf-predictieve modellen (waarbij de encoder en de predictor gezamenlijk worden geoptimaliseerd) lijden vaak aan een specifieke foutmodus: de representatie en de predictor co-adapteren naar latente transities die intern voorspelbaar zijn, maar slechts zwak geconstrueerd worden door de onderliggende fysieke evolutie van de scène. Zoals in het artikel wordt geïllustreerd, kunnen twee verschillende fysieke toestanden (bijv. een blok vastgehouden door een grijper versus een blok rustend op een tafel langs dezelfde camerastraal) projecteren naar bijna identieke beelden. Een zelf-predictief doel vereist niet dat hun latente transities zich onderscheiden, zelfs als ze onder dezelfde actie op tegenovergestelde wijze evolueren. Deze ambiguïteit leidt tot "rollout-drift", waarbij voorspelde toekomsten significant afwijken van de fysieke realiteit, wat de controleprestaties ondermijnt.

Hoewel geprivilegieerde fysieke toestanden (zoals objectposes, snelheden) een externe referentie bieden om deze ambiguïteit te verminderen, gebruiken bestaande methoden deze vaak enkel om de inhoud van de representatie te superviseren (via distillatie of staat-regressie) of het downstream gedrag. Het artikel stelt dat het decodeerbaar maken van fysieke toestanden uit een enkele latente snapshot niet noodzakelijkerwijs beperkt hoe die latenten in de loop van de tijd evolueren onder een actie. De kernuitdaging is om geprivilegieerde trajecten te gebruiken om de predictieve transitie zelf te superviseren, waardoor de geleerde dynamica geworteld wordt in de fysieke realiteit zonder dat hiervoor geprivilegieerde inputs nodig zijn tijdens de testfase.

2. Methodologie: XP-JEPA

De auteurs introduceren Cross-Predictive JEPA (XP-JEPA), een framework dat visuele observaties en geprivilegieerde fysieke toestanden behandelt als twee weergaven van hetzelfde traject. In tegen tegenstelling tot methoden die een vaste fysieke target distilleren in een visuele representatie, koppelt XP-JEPA de evolutie van beide modaliteiten via een gedeelde predictor.

Kernarchitectuur en Doelstelling

XP-JEPA bouwt voort op de LeWM (Joint-Embedding Predictive Architecture) baseline, die gezamenlijk een visuele encoder $f_\theta$ en een actie-geconditioneerde predictor $g_\psi$ leert. XP-JEPA breidt dit uit met een fysieke tak:

1. Dubbele Encoders: Een visuele encoder $f_\theta$ verwerkt observaties $o_t$, en een fysieke encoder $h_\phi$ verwerkt geprivilegieerde toestanden $s_t$.
2. Gedeelde Predictor: Een enkele actie-geconditioneerde predictor $g_\psi$ brengt geschiedenissen van beide modaliteiten ($Z^o_t$ of $Z^s_t$) voort in de tijd.
3. Cross-Predictieve Loss: Het model optimaliseert een verliesfunctie die vier termen bevat:
   • Binnen-modaliteit termen: Visuele zelf-predictie ($g_\psi(Z^o_t, a_t) \to z^o_{t+1}$) en fysieke zelf-predictie ($g_\psi(Z^s_t, a_t) \to z^s_{t+1}$).
   • Cross-modale termen: Een visuele geschiedenis voorspelt de toekomstige fysieke representatie ($g_\psi(Z^o_t, a_t) \to z^s_{t+1}$), en een fysieke geschiedenis voorspelt de toekomstige visuele representatie ($g_\psi(Z^s_t, a_t) \to z^o_{t+1}$).
     De verliesfunctie wordt gedefinieerd als:
     $$\mathcal{L}_{\text{XP-JEPA}} = \sum_{m,n \in \{o,s\}} \| g_\psi(Z^m_t, a_t) - z^n_{t+1} \|^2_2 + \beta [\Omega(z^o) + \Omega(z^s)]$$
     waarbij $\Omega$ een isotropie-regularisator (SIGReg) is om representatie-instorting te voorkomen.

Unified Privileged-State Interface

Om dit toepasbaar te maken over heterogene taken (bijv. grijpen, duwen, stapelen), ontwerpen de auteurs een verenigde rigid-body schema. Scènes worden gerepresenteerd door een set tokens:

• Object Tokens: Bevatten canonieke geometrie (gepoolde Signed Distance Field), half-extents, pose (rotatie en positie).
• Effector Token: End-effector rotatie, positie en vingeropening.
• Table Token: Een geleerd token dat de set compleet maakt.
  Scènes worden gepadded naar een vast aantal object-slots (bijv. zes), waarbij ongebruikte slots worden gemaskeerd. Dit maakt het mogelijk voor een enkele fysieke encoder om diverse configuraties te verwerken zonder taakspecifieke labels.

Training versus Implementatie

• Training: Zowel de encoders als de gedeelde predictor worden gezamenlijk getraind. De fysieke tak fungeert als een tweede online predictieve weergave, die de visuele dynamica beperkt.
• Implementatie: De fysieke encoder $h_\phi$ en de geprivilegieerde staat $s_t$ worden weggegooid. Het geïmplementeerde model bestaat alleen uit de visuele encoder en de predictor. Het ontvangt dezelfde visuele observaties als de baseline, zonder geprivilegieerde inputs tijdens de inferentie. Planning vindt plaats door kandidaat-actiesequenties uit te rollen in de geleerde latente dynamica en deze te scoren tegen een gecodeerd doel.

3. Belangrijkste Bijdragen

1. Cross-Predictive JEPA: Een nieuwe architectuur die latente transities fundeert in geprivilegieerde fysieke trajecten door visuele en fysieke geschiedenissen door een gedeelde predictor te laten lopen en elke voorspelling te matchen met de toekomsten van beide modaliteiten. Dit koppelt pointwise correspondentie met actie-geconditioneerde predictie.
2. Unified Privileged-State Interface: Een rigid-body schema dat een enkele fysieke encoder in staat stelt om over diverse object-interactieconfiguraties te opereren zonder categorie-labels of taakspecifieke state-velden.
3. Scheiding van Eigenschappen: Het werk maakt expliciet onderscheid tussen decodeerbaarheid (het herstellen van fysieke variabelen uit een snapshot), voorspelbaarheid (betrouwbare actie-geconditioneerde predictie over de tijd) en planning nut. Het toont aan dat een hoge decodeerbaarheid niet garandeert dat er ook een hoge voorspelbaarheid of controle succes optreedt.

4. Experimentele Resultaten

De auteurs evalueren XP-JEPA op een multi-task suite (gebaseerd op Meta-World, 22 configuraties) en vier single-task omgevingen (Push-T, OGBench-Block, Two-Room, Reacher).

Voorspelbaarheid (Forecastability)

• Rollout-drift: XP-JEPA vermindert de relatieve rollout-drift van een nieuw aangepaste predictor aanzienlijk. Op de multi-task suite daalt de drift van 0.361 (Visuele baseline) naar 0.104.
• Generalisatie: Deze verbetering is consistent over alle single-task omgevingen (bijv. in Two-Room daalt de drift van 0.503 naar 0.221).
• Ablatie: Het verwijderen van cross-predictie (ALIGN-ONLY) vermindert de drift minder effectief (0.158), terwijl het verwijderen van de gedeelde predictor (CROSS-ONLY) een lage drift behoudt (0.101), wat aangeeft dat cross-modale predictie de primaire drijfveer is voor voorspelbaarheid.

Controleprestaties

• Succespercentages: Op de multi-task suite verhoogt XP-JEPA het gemiddelde controle succes van 53.6% naar 78.2% (+24.6 procentpunten).
• Taakafhankelijkheid: Hoewel de voorspelbaarheid universeel verbetert, variëren de winst in controle per taak. Zo presteert XP-JEPA bijvoorbeeld aanzienlijk beter dan de baseline in Two-Room en OGBench-Block, maar laat het slechts marginale winst zien in Push-T en Reacher.
• Vergelijking met Regressie: Een baseline die een uitgebreide pose-gebaseerde state target regresseert vanuit de visuele latent (REGRESS) bereikt een hoge positie-decodeerbaarheid ($R^2 = 0.991$) maar faalt in het verbeteren van de voorspelbaarheid (drift 0.373) of de controle (53.7% succes). Dit bevestigt dat de inhoud van een snapshot alleen niet zorgt voor een betrouwbare rollout-predictie.

Inzichten uit Ablaties

• Correspondentie versus Predictie: Juiste visuele-fysieke correspondentie (het matchen van het juiste traject) is cruciaal voor controle. De SHUFFLE ablatie (gebruikmakend van dezelfde architectuur maar met mismatchende trajecten) behoudt een lage drift (0.167) maar laat de controle instorten naar 16.6%, wat bewijst dat voorspelbare dynamica onvoldoende zijn als ze niet de onderscheidende kenmerken behouden die relevant zijn voor actie-selectie.
• Predictor Delen: Het enkel delen van de predictor (SHARE-ONLY) levert weinig winst op (56.7% succes), wat suggereert dat de koppeling van de modaliteiten belangrijker is dan de gedeelde gewichten.

5. Betekenis en Claims

Het artikel claimt dat cross-predictieve fysieke gronding meer voorspelbare latente dynamica produceert voor rollout-gebaseerde controle zonder dat hiervoor tijdens de testfase geprivilegieerde inputs nodig zijn.

• Mechanisme: De verbetering komt niet voort uit het simpelweg coderen van meer fysieke informatie, maar uit het beperken van de predictieve evolutie van de representatie met behulp van corresponderende fysieke trajecten.
• Onderscheid: Het werk verheldert dat "decodeerbaarheid" (kunnen we de staat lezen?) en "voorspelbaarheid" (kunnen we de toekomst voorspellen?) verschillende eigenschappen zijn. Een model kan zeer decodeerbaar zijn maar slecht voorspelbaar, wat leidt tot slechte controle.
• Beperkingen: De auteurs merken bescheiden op dat de resultaten beperkt zijn tot simulatie, gepaard gaan met de noodzaak van gekoppelde geprivilegieerde trajecten tijdens de training, en vertrouwen op een vast rigid-body state schema. Het uitbreiden hiervan naar real-world observaties en bredere fysieke interacties blijft toekomstig werk.

Samenvattend toont XP-JEPA aan dat door de geprivilegieerde staat te behandelen als een tweede predictieve weergave en cross-modale consistentie in de transitiemodellen af te dwingen, men visuele wereldmodellen kan leren die robuust zijn tegen fysieke ambiguïteiten, wat leidt tot aanzienlijk verbeterde planning en controleprestaties.