LeWAM: modellazione end-to-end del mondo e delle azioni con le JEPA

Abstract degli autori

Le architetture predittive a embedding congiunto (JEPA) apprendono modelli latenti compatti del mondo senza ricostruire i pixel. I sistemi JEPA esistenti addestrano in genere le rappresentazioni visive e le dinamiche condizionate dalle azioni prima di aggiungere una policy di controllo o un risolutore. Presentiamo LeWAM, una JEPA end-to-end per il controllo, che combina: (1) un obiettivo di addestramento congiunto per la previsione del successivo stato latente e il matching del flusso delle azioni; (2) un vincolo gaussiano isotropo sullo stato latente. Un unico modello può così eseguire direttamente la policy e pianificare in modo efficiente, valutando le azioni campionate mediante le dinamiche latenti apprese. Rispetto alle precedenti JEPA per la pianificazione visiva, LeWAM porta il tasso di successo dal 28,6% all’89,7% nei compiti caratterizzati da numerosi contatti e dal 10,9% al 46,2% nei compiti con un lungo orizzonte temporale, ed esegue la pianificazione 32,3 volte più velocemente. Nei compiti di clonazione del comportamento, sia in simulazione sia su robot fisici, LeWAM è inoltre competitivo rispetto ai metodi consolidati. Le analisi di interpretabilità mostrano che LeWAM apprende rappresentazioni pertinenti al controllo e distribuzioni delle azioni che guidano efficacemente la pianificazione e il controllo in fase di test.