Technische Samenvatting: Temporal-Distance JEPA (TD-JEPA)

Probleemstelling

Joint-Embedding Predictive Architectures (JEPAs) zijn uitgegroeid tot effectieve backbones voor latente wereldmodellen, waarbij ze leren om toekomstige representaties te voorspellen in plaats van pixels te reconstrueren. Er bestaat echter een structurele kloof tussen het trainingsdoel van JEPAs en de vereisten van Model Predictive Control (MPC).

• De Train-Plan Kloof: JEPA-training optimaliseert doorgaans korte-horizon latente voorspelling. In contrast hiermee vereist planning het rangschikken van meerstaps-geïmageerde toekomsten op basis van voortgang richting een doel.
• De Geometrische Beperking: Bestaande JEPA-planners (bijv. LeWM) vertrouwen op latente Euclidische afstand als een de facto planningskost. Hoewel deze geometrie in sommige contexten kan correleren met voortgang, is dit een bijproduct van representatie-leren in plaats van een signaal dat expliciet wordt gewonnen voor "stappen-naar-doel" of gerichte bereikbaarheid.
• De Uitdaging: Hoe kunnen beloningsvrije offline demonstratie-logs een temporele progressie-structuur opleveren die deze kloof verkleint? En wanneer moet deze gewonnen structuur direct worden ingezet als een planningskost versus het vormen van representaties voor geometrische planning?

Methodologie: TD-JEPA

De auteurs stellen Temporal-Distance JEPA (TD-JEPA) voor, die de LeWM encoder–predictor backbone behoudt maar deze uitbreidt met een mechanisme voor het gericht winnen van temporele kosten. De methode bestaat uit drie kerncomponenten:

1. Backbone en Voorspelling

TD-JEPA maakt gebruik van de standaard LeWM-architectuur: een actie-geconditioneerde encoder ($Enc_\theta$) en predictor ($Pred_\phi$) die observaties mappen naar latente toestanden en toekomstige embeddings voorspellen. Het behoudt de SIGReg regularisatie om representatie-instorting te voorkomen en de standaard één-stap voorspellingsverlies ($\mathcal{L}_{pred}$).

2. Gericht Winnen van Temporele Kosten

In plaats van uitsluitend te vertrouwen op Euclidische afstand, leert TD-JEPA een gerichte, niet-negatieve energiefunctie $d_\psi(z_s, z_g)$ die de kosten vertegenwoordigt om een doel-latente $z_g$ te bereiken vanaf een bron-latente $z_s$.

• Architectuur (MRN Head): De cost head gebruikt een Metric-Residual Network (MRN) parametrisering. Het deelt de kosten op in een symmetrische configuratieterm en een asymmetrische residu:
  $$d_\psi(z_s, z_g) = \|\phi_{sym}(z_s) - \phi_{sym}(z_g)\|_2 + \max_k \text{ReLU}(\phi_{asym, k}(z_s) - \phi_{asym, k}(z_g))$$
  Dit ontwerp zorgt ervoor dat de kosten niet-negatief zijn en inherent asymmetrisch ($d_\psi(z_s, z_g) \neq d_\psi(z_g, z_s)$), wat de gerichte aard van het bereiken van een doel weerspiegelt.
• Supervisie (Temporele Labels): Het model wordt getraind op offline demonstratie-logs zonder dichte beloningen.
  • Positieven: Voor indices $i < j$ binnen dezelfde traject, dient de temporele kloof $\tau(i, j) = j - i$ als een proxy voor stappen-naar-subdoel. Het model regresseert $d_\psi(z_i, z_j)$ naar $\tau(i, j)$.
  • Negatieven: Cross-traject paren fungeren als heuristische negatieven, die onterecht lage kosten bestraffen voor ongerelateerde staat-paren via een hinge loss.
• Rollout Consistentie: Een multi-step teacher-forcing loss ($\mathcal{L}_{roll}$) wordt toegevoegd om de geleerde kosten af te stemmen op de open-loop horizon die door de planner wordt gebruikt, waardoor cumulatieve fouten tijdens de $H$-stap rollout worden verminderd.

3. Planning en Implementatie

Tijdens de testfase zoekt de planner (met behulp van de Cross-Entropy Methode of iCEM) naar actiesequenties die de terminale energie minimaliseren na een $H$-stap latente rollout.

• Dual Deployment Strategie: Het artikel identificeert twee verschillende rollen voor het gewonnen signaal op basis van de taakstructuur:
  1. Topologie-gedomineerde Taken (bijv. Navigatie): De gewonnen gerichte kosten $d_\psi$ wordt direct ingezet als het planningsdoel.
  2. Contact-rijke Taken (bijv. Manipulatie): Dezelfde met temporele signalen getrainde checkpoint wordt gebruikt, maar de planner vertrouwt op latente Euclidische afstand ($\ell_2$). Hier fungeert de temporele supervisie als een representatie-signaal om de geometrische planning te vormen, wat de effectiviteit ervan verbetert.

Belangrijkste Bijdragen

1. Diagnose van de JEPA Train-Plan Kloof: De auteurs tonen aan dat voorspellende JEPAs dynamiek leren uit logs, maar dat planning vaak rust op ongetrainde embedding-geometrie. Ze laten zien dat het winnen van een gerichte temporele kostenfunctie uit logs de ranking mismatch aanzienlijk verkleint (Spearman $\rho$ stijgt van 0.65 voor Euclidische afstand naar 0.95 voor de gewonnen kosten op Push-T).
2. Temporal-Distance Mining Framework: TD-JEPA introduceert een methode om de structurele progressie van de tijd te extraheren uit beloningsvrije logs met behulp van gerichte temporele supervisie, cross-traject negatieven en horizon-gematchte rollout consistentie.
3. Co-design van Kosten en Implementatie: Het artikel levert bewijs dat de keuze van de planningskost (gewonnen temporeel versus geometrisch) moet worden co-ontworpen met de implementatie tijdens de planning. Temporele kosten blinken uit in topologische navigatie, terwijl geometrische kosten superieur blijven voor contact-rijke manipulatie, zelfs wanneer de onderliggende representatie met temporele signalen is getraind.

Experimentele Resultaten

De methode werd geëvalueerd in vier omgevingen: Two-Room (navigatie), Reacher (reiken), Push-T (manipulatie) en OGB-Cube (contact-rijke manipulatie). Evaluaties werden uitgevoerd onder een "locked" protocol (vaste seeds, checkpoints en episode manifests) om een eerlijke vergelijking met LeWM en de gelijktijdige RC-aux baseline te garanderen.

• Navigatie & Reaching: Het inzetten van de gewonnen kosten $d_\psi$ verbeterde de succesratio's aanzienlijk.
  • Two-Room: Bereikte 100.0% succes (vs. 97.4% voor LeWM).
  • Reacher: Bereikte 97.0% succes (vs. 96.0% voor LeWM).
• Manipulatie: Bij contact-rijke taken verbeterde de temporele training de latente representatie, wat de geometrische planning ten goede kwam.
  • OGB-Cube: Het gebruik van de met temporele signalen getrainde checkpoint met $\ell_2$ planning verhoogde het succes met 14.2 procentpunt over LeWM (82.2% vs. 68.0%).
  • Push-T: Verbeterde van 83.6% (LeWM) naar 86.0%.
• Ablatie-studies: Het verwijderen van de gerichte residu (symmetrische head), cross-traject negatieven of de rollout consistentie leidde in alle gevallen tot een verslechtering van de prestaties, wat de noodzaak van elk onderdeel bevestigt.
• Foutanalyse: Op Push-T presteerde pure temporele kosten ($d_\psi$) minder goed dan geometrische kosten ($\ell_2$), omdat fouten primair optraden tijdens de post-contact uitlijning, waarbij lokale geometrie kritischer is dan de ranking van temporele progressie.

Betekenis en Claims

Het artikel beweert dat TD-JEPA de train–plan kloof voor JEPA world-model planners verkleint door de temporele progressie-structuur te ontdekken in offline logs en de vorm van de kosten te co-ontwerpen met de implementatie tijdens de planning.

• Bescheiden Claims: De auteurs beweren niet dat temporele kosten universeel de geometrische kosten vervangen. In plaats daarvan pleiten ze voor een conditionele implementatiestrategie: gebruik de gewonnen gerichte kosten wanneer progressie topologisch is, en gebruik het als een representatie-signaal voor geometrische planning wanneer de contact-geometrie domineert.
• Empirische Validatie: Onder een "locked" evaluatie evenaart of overtreft TD-JEPA zowel LeWM als RC-aux in elke omgeving, wat aantoont dat reward-free temporele mining een levensvatbare weg is om latent MPC te verbeteren zonder dat daar beloningsmodellen of online interactie voor nodig zijn.

De code is beschikbaar via de verstrekte GitHub-repository.