Riepilogo Tecnico: AdaJEPA: Un Modello del Mondo Latente Adattivo

Definizione del Problema

I modelli del mondo latente, in particolare le Joint-Embedding Predictive Architectures (JEPA), consentono una pianificazione efficiente prevedendo stati futuri in uno spazio latente compatto. Questi modelli vengono tipicamente addestrati su traiettorie offline prive di reward e accoppiati con il Model Predictive Control (MPC) per compiti condizionati da un obiettivo. Tuttavia, esiste un limite critico: il modello del mondo è solitamente congelato dopo l'addestramento.

In fase di deployment, se le previsioni del modello diventano imprecise a causa di shift di distribuzione al tempo di test (ad esempio, cambiamenti visivi come rumore o illuminazione, o cambiamenti dinamici come attrito e massa), l'MPC ottimizza un obiettivo errato. Azioni che appaiono efficaci nella "immaginazione latente" possono fallire nell'ambiente reale. Piccoli errori a un singolo step si accumulano lungo l'orizzonte di pianificazione, portando a un sostanziale degrado delle prestazioni. Gli approcci attuali richiedono spesso dimostrazioni esperte aggiuntive, dati di fine-tuning per il dominio target o complessi loop esterni per affrontare questo problema, ostacolando l'applicazione nel mondo reale.

Metodologia: AdaJEPA

Gli autori propongono AdaJEPA, un modello del mondo latente adattivo che esegue l'adattamento al tempo di test all'interno del loop chiuso dell'MPC. Invece di trattare il modello come statico, AdaJEPA ricalibra continuamente il modello del mondo utilizzando segnali auto-supervisionati derivati dalle interazioni stesse dell'agente.

Il Ciclo Plan–Execute–Adapt–Replan (Pianifica–Esegui–Adatta–Ripianifica)

Il meccanismo centrale opera come segue ad ogni step dell'MPC:

1. Plan (Pianifica): L'agente pianifica una sequenza di azioni utilizzando il modello del mondo attuale (potenzialmente obsoleto) per minimizzare un costo latente di raggiungimento dell'obiettivo.
2. Execute (Esegui): L'agente esegue il primo chunk di azioni ($a_t$) nell'ambiente.
3. Observe (Osserva): L'agente riceve l'osservazione successiva ($o_{t+1}$).
4. Adapt (Adatta): La transizione $\{o_t, a_t, o_{t+1}\}$ viene aggiunta a un buffer di replay. I parametri del modello vengono aggiornati tramite un piccolo numero di step di gradiente (spesso solo uno) per minimizzare l'errore di predizione latente su questa nuova transizione.
5. Replan (Ripianifica): Il modello aggiornato viene immediatamente utilizzato per la successiva iterazione di pianificazione.

Componenti Tecnici

• Architettura: Il modello del mondo sottostante consiste in un encoder sensoriale ($\mathcal{E}_s$), un encoder di azione ($\mathcal{E}_a$) e un predittore ($f_\theta$). È addestrato offline per prevedere rappresentazioni latenti future ($z_{t+1}$) piuttosto che ricostruire i pixel.
• Perdita di Adattamento (Adaptation Loss): L'adattamento utilizza lo stesso obiettivo di predizione latente auto-supervisionato del pre-addestramento:
  $$\mathcal{L}_{ada}(\mathcal{B}) = \frac{1}{|\mathcal{B}|} \sum_{(o_i, a_i, o_{i+1}) \in \mathcal{B}} \ell(f_\theta(z_i, \mathcal{E}_a(a_i)), \text{sg}(z_{i+1}))$$
  dove $\mathcal{B}$ è un buffer di transizioni recenti, $\ell$ è una perdita di predizione latente (ad esempio, MSE) e $\text{sg}(\cdot)$ è un operatore stop-gradient per prevenire il collasso.
• Aggiornamenti dei Parametri: Per garantire l'efficienza, l'adattamento è leggero. Gli autori aggiornano solo un piccolo sottoinsieme di parametri (ad esempio, gli strati finali dell'encoder e del predittore) invece dell'intera rete.
• Strategie del Buffer: Il buffer di replay $\mathcal{B}$ può essere gestito tramite "recent-N" (mantenendo le ultime $N$ transizioni) o "hard-N" (mantenendo le transizioni con i maggiori errori di predizione).

Contributi Chiave

1. Primo Adattamento JEPA durante la Pianificazione: Per quanto ne sappiano gli autori, questo è il primo lavoro a adattare un modello del mondo JEPA durante la fase di pianificazione, invece di fare affidamento esclusivamente sul pre-training o sul fine-tuning offline.
2. Auto-Supervisione in Loop Chiuso: Il metodo utilizza le conseguenze immediate delle azioni eseguite come segnale auto-supervisionato, eliminando la necessità di etichette di reward, dimostrazioni esperte o fasi separate di raccolta dati.
3. Efficienza di Campionamento e Computazionale: Il framework ottiene guadagni significativi nelle prestazioni con un carico computazionale minimo. In molti esperimenti, un singolo step di gradiente per ogni step di ripianificazione dell'MPC è sufficiente.
4. Robustezza agli Shift di Distribuzione: L'approccio dimostra robustezza rispetto a shift visivi (rumore, sfocatura, cambiamenti di colore), shift dinamici (massa, smorzamento) e shift strutturali (layout di labirinti non visti).

Risultati Sperimentali

Gli autori hanno valutato AdaJEPA sui benchmark PushT (spinta di blocchi) e PointMaze (navigazione) sotto vari shift di distribuzione tra training e test.

• Performance In-Distribution: Anche su compiti simili ai dati di addestramento, AdaJEPA migliora i tassi di successo (ad esempio, un guadagno >20% su forme specifiche di PushObj) o mantiene le prestazioni di base senza degradazione, provando che è sicuro applicarlo in-distribution.
• Performance Out-of-Distribution (OOD):
  • Shift di Forma: Su forme di blocchi non viste, AdaJEPA ha quasi raddoppiato il tasso di successo della pianificazione rispetto al modello congelato.
  • Shift Visivi: Il metodo ha mostrato chiari guadagni sotto sfocatura gaussiana, rumore sale e pepe e illuminazione scura.
  • Shift di Dinamica e Layout: In PointMaze, AdaJEPA è riuscito a recuperare dai fallimenti causati da bassa massa o alto smorzamento e ha migliorato significativamente i tassi di successo su layout di labirinti non visti (ad esempio, un guadagno del +25.3% su specifici shift di layout).
• Latenza: Il passaggio di adattamento introduce una latenza trascurabile (circa 0.01–0.03s per ripianificazione) perché aggiorna solo un piccolo sottoinsieme di parametri. Inoltre, recuperando più velocemente dagli errori di predizione, l'agente spesso raggiunge l'obiettivo in meno ripianificazioni totali, riducendo il tempo di esecuzione complessivo.
• Generalizzazione alla Scarsità di Dati: Nei regimi con pochi dati, l'adattamento al tempo di test ha compensato la limitata copertura di addestramento. Ad esempio, con dati di addestramento limitati, AdaJEPA ha superato un modello congelato addestrato con 16 volte più traiettorie.

Significato e Rivendicazioni

L'articolo sostiene che i modelli del mondo latenti non dovrebbero rimanere fissi dopo l'addestramento, ma dovrebbero continuare a migliorare dall'esperienza incontrata durante il deployment.

• Ispirazione Biologica: L'approccio è fondato sull'adattamento sensomotorio biologico, in cui i sistemi (come il cervelletto) regolano il comportamento in base a input e dinamiche variabili.
• Complemento alla Scalabilità dei Dati: Mentre dati di addestramento più grandi e diversificati migliorano la generalizzazione, AdaJEPA fornisce una via efficiente in termini di campionamento per raffinare i modelli durante il deployment, agendo come complemento alla scalabilità dei dati.
• Direzione Futura: Gli autori suggeriscono che il passo successivo sia combinare questo adattamento leggero al tempo di test con l'apprendimento continuo e attivo per espandere la copertura del modello del mondo nel tempo, abilitando una percezione e una pianificazione più resilienti in ambienti mutevoli.

Gli autori sottolineano che il loro metodo è agnostico rispetto all'implementazione specifica del modello del mondo, mostrando miglioramenti coerenti attraverso diverse architetture (ResNet, Transformer, basate su DINO) e diversi planner (Gradient Descent, CEM).