Traj-LeWM: Path-Aware World-Model Planning via Latent Trajectory Cost

Riepilogo Tecnico di Traj-LeWM: Pianificazione di Modelli di Mondo Sensibili al Percorso tramite Costo di Traiettoria Latente

1. Definizione del Problema

Il paper affronta due limitazioni specifiche di LeWM (Lightweight Visual World Model), un'architettura predittiva a embedding congiunto (JEPA) che apprende la dinamica latente dai pixel per eseguire una pianificazione condizionata all'obiettivo. Sebbene LeWM ottenga prestazioni elevate su compiti semplici, il suo tasso di successo diminuisce significativamente in ambienti più complessi (ad esempio, dal 96% su Push-T al 74% su OGBench-Cube). Gli autori identificano due cause profonde per questo degrado:

1. Mancanza di Supervisione a Livello di Traiettoria: L'obiettivo di addestramento di LeWM si basa sulla predizione del passo successivo (transizioni locali) e su un termine di regolarizzazione (SIGReg) per prevenire il collasso della rappresentazione. Nessuno dei due componenti richiede esplicitamente che l'encoder condiviso preservi informazioni sulla traiettoria completa rispetto all'obiettivo del compito. Di conseguenza, lo spazio latente non è addestrato direttamente per valutare la qualità di una completa traiettoria candidata.
2. Insensibilità alla Valutazione Basata Solo sull'Endpoint: Durante la pianificazione, LeWM classifica le sequenze di azioni candidate basandosi esclusivamente sulla distanza euclidea tra l'endpoint latente predetto e la rappresentazione dell'obiettivo. Questo approccio ignora le informazioni sul percorso intermedio. Due traiettorie con endpoint identici ma evoluzioni intermedie diverse (una riuscita e una fallita) ricevono lo stesso punteggio, portando a una selezione subottimale dei candidati quando le predizioni del modello divergono dagli esiti reali dell'esecuzione.

2. Metodologia: Traj-LeWM

Gli autori propongono Traj-LeWM, un framework che mantiene l'originale obiettivo di dinamica locale di LeWM e la valutazione dell'endpoint, introducendo un Costo di Traiettoria Latente (LTC) Condizionato all'Obiettivo. Il metodo opera in tre fasi:

A. Costo di Traiettoria Latente (LTC) Condizionato all'Obiettivo
L'innovazione principale è una funzione di costo scalare apprendibile, $S_\psi(\tau | g)$, che valuta una traiettoria latente completa $\tau$ dato un obiettivo $g$. A differenza della distanza dell'endpoint, l'LTC aggrega informazioni su tutto il percorso. È definito come:
$$S_\psi(\tau | g) = \frac{1}{T} \sum_{t=0}^{T-1} \ell_\psi(z_t, \Delta z_t, d^g_t, \eta_t)$$
Dove:

• $z_t$: Stato latente corrente.
• $\Delta z_t$: Evoluzione locale (differenza tra stati adiacenti).
• $d^g_t = z_g - z_t$: Quantità relativa all'obiettivo (condizionamento).
• $\eta_t$: Fase temporale normalizzata.
• $\ell_\psi$: Una funzione di contributo per passo non negativa e apprendibile.

Questa struttura permette al costo di variare in base a come uno stato viene raggiunto, non solo dove termina.

B. Supervisione delle Preferenze della Traiettoria
Per addestrare l'LTC e modellare l'encoder condiviso, gli autori introducono un meccanismo di apprendimento basato sulle preferenze utilizzando una perdita logistica a coppie. Il sistema impara ad assegnare costi inferiori alle traiettorie "preferite" (corrispondenti all'esperto) e costi superiori alle traiettorie "negative". Vengono utilizzati due tipi di esempi negativi:

1. Preferenze Sintetiche:
   • Disallineamento dell'Obiettivo: Una traiettoria esperta accoppiata con un obiettivo spostato e non correlato.
   • Perturbazione Preservante l'Endpoint: Una traiettoria esperta con rumore aggiunto agli stati intermedi (mantenendo fissi inizio e fine). Ciò costringe il modello a distinguere tra diversi percorsi intermedi anche se gli endpoint sono identici.
2. Preferenze da Fallimento in Anello Chiuso:
   • L'attuale modello esegue la pianificazione utilizzando solo la valutazione dell'endpoint (senza LTC) ed esegue il piano nell'ambiente.
   • Le esecuzioni fallite vengono estratte e accoppiate con traiettorie esperte che hanno lo stesso inizio e obiettivo.
   • Questi fallimenti nel mondo reale forniscono esempi negativi che la metrica basata solo sull'endpoint non è riuscita a identificare, calibrando l'LTC rispetto ai modelli di fallimento specifici del modello.

C. Addestramento Congiunto e Pianificazione

• Addestramento: La perdita totale combina la perdita di predizione originale di LeWM ($L_{LeWM}$), la perdita di preferenza sintetica ($L_{path}$) e la perdita estratta dai fallimenti ($L_{mined}$). Il ramo sintetico aggiorna sia l'LTC che l'encoder condiviso, mentre il ramo dei fallimenti aggiorna solo l'LTC (tramite gradienti staccati) per evitare di confondere l'apprendimento della dinamica del predittore.
• Pianificazione: I candidati sono classificati usando un Punteggio Congiunto:
  $$C_\lambda(\hat{\tau}, g) = C_0(\hat{\tau}, g) + \lambda \tilde{S}_\psi(\hat{\tau} | g)$$
  Dove $C_0$ è la distanza originale dell'endpoint e $\tilde{S}_\psi$ è un punteggio LTC calibrato. Ciò consente al pianificatore di utilizzare le informazioni sul percorso intermedio quando le predizioni dell'endpoint sono ambigue o imprecise.

3. Contributi Chiave

1. Costo di Traiettoria Latente (LTC): Introduzione di un costo apprendibile e condizionato all'obiettivo che valuta traiettorie predette complete, aggregando l'evoluzione latente relativa all'obiettivo per fornire un segnale di pianificazione sensibile al percorso.
2. Apprendimento delle Preferenze della Traiettoria: Sviluppo di un meccanismo che utilizza negativi sintetici (disallineamenti dell'obiettivo e perturbazioni intermedie) e fallimenti estratti da anello chiuso per addestrare l'LTC e modellare la rappresentazione condivisa senza alterare l'obiettivo originale del predittore.
3. Valutazione Completa: Test estesi su quattro compiti simulati (Push-T, OGBench-Cube, Reacher, Two-Room) e su un robot fisico Franka FR3, dimostrando miglioramenti rispetto al baseline.

4. Risultati Sperimentali

Risultati di Simulazione:
Traj-LeWM supera LeWM in tutti e quattro i compiti, con miglioramenti del tasso di successo di:

• Push-T: +3% (96% $\to$ 99%)
• OGBench-Cube: +14% (74% $\to$ 88%)
• Reacher: +7% (86% $\to$ 93%)
• Two-Room: +7% (87% $\to$ 94%)

Valutazione su Robot Fisico:
Su un Franka FR3 fisico con 20 compiti start-goal, Traj-LeWM ha ottenuto un tasso di successo del 70% (14/20) rispetto al 50% (10/20) di LeWM. I fallimenti dovuti a collisioni sono diminuiti dal 20% al 10%.

Esperimenti Controllati e Ablazioni:

• Modellazione della Rappresentazione (Gap 1): Anche quando l'LTC viene rimosso al tempo di test, l'encoder condiviso di Traj-LeWM (addestrato con preferenze di traiettoria) produce punteggi dell'endpoint con un'AUC più alta nel distinguere successo vs fallimento rispetto a LeWM.
• Consapevolezza del Percorso (Gap 2): Quando le predizioni dell'endpoint sono imprecise, l'aggiunta dell'LTC migliora l'ordinamento corretto dei candidati. Nei test "controllati dall'endpoint" (dove gli stati iniziale e finale sono identici), l'LTC distingue con successo le traiettorie di successo da quelle fallite in oltre il 94% dei casi su Reacher, provando che utilizza le informazioni sul percorso intermedio.
• Complementarità: Gli studi di ablazione mostrano che la supervisione della traiettoria (addestramento) e la classificazione esplicita dell'LTC (pianificazione) sono complementari. Il punteggio congiunto fornisce i risultati migliori, indicando che non sono intercambiabili.

5. Significato e Affermazioni

Il paper afferma che Traj-LeWM affronta con successo le limitazioni dei modelli di mondo basati solo sulla dinamica locale integrando informazioni a livello di traiettoria sia nell'apprendimento della rappresentazione che nella classificazione dei candidati.

• Affermazioni Modeste: Gli autori dichiarano esplicitamente che i risultati sul robot fisico sono "prove preliminari della fattibilità su robot reale piuttosto che un vantaggio prestazionale conclusivo" a causa del numero limitato di compiti accoppiati (20).
• Insight Centrale: Il lavoro dimostra che la sola distanza dell'endpoint è insufficiente per una pianificazione affidabile in ambienti complessi. Apprendendo un costo su completi rollout latenti, il modello può sfruttare le prove del percorso intermedio per classificare i candidati in modo più efficace, specialmente quando le predizioni dell'endpoint sono incerte.
• Efficienza: Il metodo aggiunge meno di 1 milione di parametri al modello, mantenendo un'architettura leggera paragonabile all'originale LeWM pur aumentando significativamente le prestazioni in compiti impegnativi di manipolazione e navigazione.