Sintesi Tecnica: Temporal-Distance JEPA (TD-JEPA)

Definizione del Problema

Le Joint-Embedding Predictive Architectures (JEPA) sono emerse come backbone efficaci per i modelli di mondo latenti, apprendendo a predire rappresentazioni future piuttosto che ricostruire i pixel. Tuttavia, esiste un divario strutturale tra l'obiettivo di addestramento delle JEPA e i requisiti della Model Predictive Control (MPC).

• Il Divario Train-Plan (Addestramento-Pianificazione): L'addestramento JEPA tipicamente ottimizza la predizione latente a breve orizzonte. Al contrario, la pianificazione richiede di classificare futuri immaginati multi-step in base al progresso verso un obiettivo.
• La Limitazione Geometrica: Gli esistenti pianificatori JEPA (ad es., LeWM) si affidano alla distanza euclidea latente come costo di pianificazione de facto. Sebbene questa geometria possa correlare con il progresso in alcuni contesti, è un sottoprodotto dell'apprendimento delle rappresentazioni piuttosto che un segnale esplicitamente estratto per i "passaggi verso l'obiettivo" o la raggiungibilità diretta.
• La Sfida: Come possono i log di dimostrazioni offline privi di reward fornire una struttura di progresso temporale che riduca questo divario? Inoltre, quando questa struttura estratta debba essere utilizzata direttamente come costo di pianificazione rispetto all'uso per modellare le rappresentazioni per la pianificazione geometrica?

Metodologia: TD-JEPA

Gli autori propongono Temporal-Distance JEPA (TD-JEPA), che mantiene l'architettura encoder-predittore di LeWM ma la aumenta con un meccanismo di estrazione del costo temporale diretto. Il metodo consiste in tre componenti principali:

1. Backbone e Predizione

TD-JEPA utilizza l'architettura standard LeWM: un encoder condizionato dall'azione ($Enc_\theta$) e un predittore ($Pred_\phi$) che mappano le osservazioni in stati latenti e prevedono i futuri embedding. Preserva la regolarizzazione SIGReg per prevenire il collasso della rappresentazione e la standard perdita di predizione a singolo step ($\mathcal{L}_{pred}$).

2. Estrazione del Costo Temporale Diretto

Invece di fare affidamento esclusivamente sulla distanza euclidea, TD-JEPA apprende una funzione di energia diretta e non negativa $d_\psi(z_s, z_g)$ che rappresenta il costo per raggiungere un obiettivo latente $z_g$ da uno stato latente sorgente $z_s$.

• Architettura (Testa MRN): La testa del costo utilizza una parametrizzazione Metric-Residual Network (MRN). Essa decompone il costo in un termine di configurazione simmetrico e un residuo asimmetrico:
  $$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))$$
  Questo design garantisce che il costo sia non negativo e intrinsecamente asimmetrico ($d_\psi(z_s, z_g) \neq d_\psi(z_g, z_s)$), riflettendo la natura diretta del raggiungimento dell'obiettivo.
• Supervisione (Etichette Temporali): Il modello è addestrato su log di dimostrazioni offline senza reward dense.
  • Positivi: Per gli indici $i < j$ sulla stessa traiettoria, il gap temporale $\tau(i, j) = j - i$ funge da proxy per i passi verso il sotto-obiettivo. Il modello regredisce $d_\psi(z_i, z_j)$ verso $\tau(i, j)$.
  • Negativi: Le coppie tra traiettorie diverse agiscono come negativi euristici, penalizzando costi spurieamente bassi per coppie di stati non correlati tramite una hinge loss.
• Consistenza del Rollout: Una perdita di teacher-forcing multi-step ($\mathcal{L}_{roll}$) viene aggiunta per allineare il costo appreso con l'orizzonte open-loop utilizzato dal pianificatore, riducendo gli errori di accumulo durante il rollout di $H$ step.

3. Pianificazione e Deployment

Al tempo di test, il pianificatore (utilizzando Cross-Entropy Method o iCEM) cerca sequenze di azioni che minimizzino l'energia terminale dopo un rollout latente di $H$ step.

• Strategia di Deployment Duale: Il paper identifica due ruoli distinti per il segnale estratto in base alla struttura del task:
  1. Task Dominati dalla Topologia (es. Navigazione): Il costo diretto estratto $d_\psi$ è distribuito direttamente come obiettivo di pianificazione.
  2. Task Ricchi di Contatto (es. Manipolazione): Viene utilizzato lo stesso checkpoint addestrato temporalmente, ma il pianificatore si affida alla distanza euclidea latente ($\ell_2$). In questo caso, la supervisione temporale agisce come un segnale di rappresentazione per modellare lo spazio latente, migliorando l'efficacia della pianificazione geometrica.

Contributi Chiave

1. Diagnosi del Divario Train-Plan nelle JEPA: Gli autori dimostrano che le JEPA predittive apprendono la dinamica dai log, ma la pianificazione spesso si basa sulla geometria degli embedding non addestrata. Mostrano che l'estrazione di un costo temporale diretto dai log riduce significativamente questo mismatch di ranking (lo Spearman $\rho$ aumenta da 0.65 per la distanza euclidea a 0.95 per il costo estratto su Push-T).
2. Framework di Estrazione della Distanza Temporale: TD-JEPA introduce un metodo per estrarre la struttura del progresso temporale da log reward-free utilizzando la supervisione temporale diretta, negativi tra traiettorie e consistenza del rollout accoppiata all'orizzonte.
3. Co-Design di Costo e Deployment: Il paper fornisce prove che la scelta del costo di pianificazione (temporale estratto vs geometrico) debba essere co-progettata con la struttura del task. I costi temporali eccellono nella navigazione topologica, mentre i costi geometrici rimangono superiori per la manipolazione ricca di contatto, anche quando la rappresentazione sottostante è addestrata con segnali temporali.

Risultati Sperimentali

Il metodo è stato valutato su quattro ambienti: Two-Room (navigazione), Reacher (raggiungimento), Push-T (manipolazione) e OGB-Cube (manipolazione ricca di contatto). Le valutazioni sono state condotte sotto un protocollo "locked" (seed, checkpoint e manifest degli episodi fissi) per garantire un confronto equo con LeWM e la baseline concorrente RC-aux.

• Navigazione e Raggiungimento: L'uso del costo estratto $d_\psi$ ha migliorato significativamente i tassi di successo.
  • Two-Room: Ha raggiunto il 100.0% di successo (rispetto al 97.4% di LeWM).
  • Reacher: Ha raggiunto il 97.0% di successo (rispetto al 96.0% di LeWM).
• Manipolazione: Nei task ricchi di contatto, l'addestramento temporale ha migliorato la rappresentazione latente, beneficiando la pianificazione geometrica.
  • OGB-Cube: L'uso del checkpoint addestrato temporalmente con pianificazione $\ell_2$ ha aumentato il successo di 14.2 punti percentuali rispetto a LeWM (82.2% vs 68.0%).
  • Push-T: Migliorato dall'83.6% (LeWM) all'86.0%.
• Studi di Ablazione: Rimuendo il residuo diretto (testa simmetrica), i negativi tra traiettorie o la consistenza del rollout, si è verificato un degrado delle prestazioni, confermando la necessità di ciascun componente.
• Analisi dei Fallimenti: Su Push-T, il costo temporale puro ($d_\psi$) è stato meno performante del costo geometrico ($\ell_2$) perché i fallimenti avvenivano principalmente durante l'allineamento post-contatto, dove la geometria locale è più critica del ranking del progresso temporale.

Significato e Rivendicazioni

Il paper sostiene che TD-JEPA riduce il divario train-plan per i pianificatori world-model JEPA scoprendo la struttura del progresso temporale nei log offline e co-progettando la forma del costo con il deployment al tempo di pianificazione.

• Rivendicazioni Modeste: Gli autori non sostengono che i costi temporali sostituiscano universalmente i costi geometrici. Sostengono invece una strategia di deployment condizionale: usare il costo diretto estratto quando il progresso è topologico, e usarlo come segnale di rappresentazione per la pianificazione geometrica quando la geometria del contatto domina.
• Validazione Empirica: Sotto valutazione "locked", TD-JEPA eguaglia o supera sia LeWM che RC-aux in ogni ambiente, dimostrando che l'estrazione temporale reward-free è una via percorribile per migliorare l'MPC latente senza richiedere modelli di reward o interazione online.

Il codice è disponibile nel repository GitHub fornito.