🤖 machine learning

JEDI: Joint Embedding Diffusion World Model for Online Model-Based Reinforcement Learning

JEDI introduce il primo modello di mondo a diffusione latente end-to-end online che unifica l'apprendimento di rappresentazioni predittive in stile JEPA con obiettivi di diffusione per ottenere un'efficienza superiore e prestazioni competitive nell'apprendimento per rinforzo basato su modelli, rispetto sia agli approcci basati su pixel sia a quelli latenti addestrati separatamente.

Autori originali: Jing Yu Lim, Rushi Shah, Zarif Ikram, Samson Yu, Haozhe Ma, Tze-Yun Leong, Dianbo Liu

Pubblicato 2026-05-14

📖 4 min di lettura☕ Lettura da pausa caffè

CC BY 4.0

Autori originali: Jing Yu Lim, Rushi Shah, Zarif Ikram, Samson Yu, Haozhe Ma, Tze-Yun Leong, Dianbo Liu

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di dover insegnare a un robot a giocare ai videogiochi. Per farlo in modo efficiente, il robot ha bisogno di un "modello del mondo" una simulazione mentale di come funziona il gioco, così da poter esercitarsi nella sua mente prima di giocare effettivamente.

Per lungo tempo, il modo migliore per costruire questa simulazione mentale era far sì che il robot cercasse di ricostruire lo schermo del gioco pixel per pixel, come un pittore che cerca di copiare esattamente una fotografia. Questo è preciso ma molto lento e richiede una quantità enorme di memoria informatica (come cercare di trasportare un'intera biblioteca nello zaino).

Recentemente, una nuova tecnica chiamata Diffusione è diventata popolare. Pensa alla diffusione come a uno scultore che inizia con un blocco di argilla rumorosa e piena di interferenze e scolpisce lentamente il rumore per rivelare una statua chiara. Questo è ottimo per comprendere scene complesse, ma farlo pixel per pixel è ancora troppo pesante per un robot da eseguire in tempo reale.

Un altro approccio ha cercato di comprimere il gioco in una piccola "sintesi" astratta (uno spazio latente) prima, ma queste sintesi erano spesso addestrate separatamente e non imparavano le regole del gioco abbastanza bene da sole.

Ecco JEDI (Joint Embedding Diffusion).

Gli autori di questo articolo hanno creato un nuovo metodo che combina il meglio di entrambi i mondi. Ecco come funziona JEDI, usando semplici analogie:

1. La "Fotografia Mentale" contro il "Dipinto"

Vecchio Metodo (Diffusione a Pixel): Immagina il robot che cerca di memorizzare ogni singolo granello di sabbia su una spiaggia per comprendere l'oceano. È preciso, ma ci vuole un'eternità e consuma tutte le sue risorse cerebrali.
Il Metodo di JEDI: Invece di memorizzare la sabbia, JEDI insegna al robot a scattare una fotografia mentale dell'essenza dell'oceano (le onde, il colore, il movimento). Non si preoccupa dei singoli granelli di sabbia. Comprime lo schermo del gioco in una piccola e efficiente "sintesi" che cattura solo ciò che conta per vincere.

2. L'Addestramento tramite "Gioco di Indovinelli"

Come impara il robot a scattare queste fotografie?

Il Vecchio Metodo: Il robot veniva spesso addestrato mostrandogli un'immagine e chiedendogli di rip dipingerla esattamente. Se mancava un dettaglio, riceveva una penalità. È come se uno studente venisse valutato sulla calligrafia piuttosto che sulla comprensione della storia.
Il Metodo JEDI: JEDI utilizza un gioco di indovinelli predittivo.
1. Il robot vede lo stato attuale del gioco.
2. Gli viene chiesto di indovinare come sarà la prossima "fotografia mentale".
3. Per renderlo più difficile (e più intelligente), il computer prende la reale prossima fotografia, le aggiunge un po' di "rumore statico" (come sfocarla) e chiede al robot di rimuovere il rumore per riportarla alla chiarezza.
4. Crucialmente, il robot impara questo mentre gioca. Non ha bisogno di un insegnante separato che gli mostri come dipingere; impara le regole del gioco provando a prevedere il futuro.

3. Perché è una Grande Novità (I Risultati)

L'articolo afferma che JEDI è un importante aggiornamento per tre motivi principali:

È più leggero: Poiché JEDI lavora con piccole "fotografie mentali" invece di interi fotogrammi video, utilizza il 43% in meno di memoria informatica. È come passare dal portare uno zaino pesante pieno di libri al portare un unico quaderno intelligente.
È più veloce: Il robot può pensare (campionare) 3 volte più velocemente e addestrarsi 2,5 volte più velocemente rispetto ai precedenti migliori metodi basati sui pixel.
Gioca in modo diverso: Questa è la parte più sorprendente. L'articolo ha scoperto che JEDI non gioca solo più velocemente; gioca in modo diverso.
- Su giochi che erano già facili per altri robot, JEDI era buono ma non sempre il migliore in assoluto.
- Tuttavia, sui giochi più difficili e caotici (come gli sparatutto con molti bersagli in movimento), JEDI ha brillato. Sembra gestire il caos meglio perché la sua "fotografia mentale" si concentra sulla strategia piuttosto che distrarsi con il rumore visivo.

La Conclusione

L'articolo sostiene che non è necessario essere un pittore perfetto (ricostruendo ogni pixel) per essere un grande stratega. Insegnando al robot a prevedere il "succo" del futuro usando un gioco di rimozione del rumore (diffusione), si ottiene un sistema che è più veloce, più economico da eseguire e sorprendentemente migliore nella gestione di situazioni difficili e caotiche.

Gli autori chiamano questo metodo JEDI e affermano che è la prima volta che questo specifico metodo di "indovinello predittivo" è stato combinato con successo con la "rimozione del rumore" (diffusione) per insegnare a un robot a giocare ai videogiochi online, senza bisogno di insegnanti pre-addestrati.

Riepilogo Tecnico: JEDI: Modello del Mondo a Diffusione di Incorporamento Congiunto per l'Apprendimento per Rinforzo Basato su Modello Online

Enunciato del Problema

L'Apprendimento per Rinforzo Basato su Modello Online (MBRL) si affida a modelli del mondo appresi per prevedere gli esiti futuri e migliorare l'efficienza del campionamento. Sebbene i modelli di diffusione siano emersi come backbones potenti per la modellazione del mondo in domini visivamente complessi, gli approcci attuali affrontano una tensione significativa:

Diffusione a Pixel: Modelli come DIAMOND operano direttamente nello spazio dei pixel. Sebbene efficaci, sono computazionalmente costosi in termini di memoria (VRAM), tempo di campionamento e durata dell'addestramento.
Diffusione Latente: Spostare la diffusione nello spazio latente migliora l'efficienza ma spesso si basa su latenti addestrati separatamente (ad esempio, tramite obiettivi di ricostruzione) piuttosto che su obiettivi di modello del mondo end-to-end. Questi approcci, come Horizon Imagination (HI), non hanno ancora eguagliato le prestazioni dei modelli di diffusione basati sui pixel e mancano dell'apprendimento rappresentativo predittivo end-to-end che ha guidato i progressi moderni nel MBRL.

La domanda centrale affrontata è se un modello del mondo a diffusione possa apprendere la propria rappresentazione latente end-to-end dall'obiettivo di denoising utilizzando un framework di Architettura di Incorporamento Predittivo Congiunto (JEPA), combinando così l'efficienza della diffusione latente con le prestazioni dell'apprendimento predittivo end-to-end.

Metodologia: JEDI

Gli autori introducono JEDI (Joint Embedding Diffusion), il primo modello del mondo a diffusione latente online end-to-end costruito su un framework JEPA.

Architettura Principale

JEDI apprende simultaneamente un codificatore di spazio latente ( $E_\phi$ ) e un modello di dinamica. Il sistema è composto da:

Codificatore: Comprime le osservazioni $x_t$ in latenti puliti $z^0_t$ .
Dinamica a Diffusione Latente: Prevede il prossimo stato latente attraverso un denoising condizionato. Invece di una previsione a singolo passo, utilizza una catena di diffusione inversa condizionata dal latente corrente $z^0_t$ e dalle azioni passate per prevedere il latente futuro pulito $z^0_{t+1}$ .
Teste di Ricompensa e Terminazione: Prevedono ricompense e segnali di terminazione direttamente dallo stato latente.

Meccanismo di Addestramento

A differenza dei tradizionali modelli di diffusione latente che si affidano a funzioni di perdita di ricostruzione (ad esempio, VAE) per addestrare il codificatore, JEDI addestra il codificatore end-to-end utilizzando la funzione di perdita di denoising della diffusione.

Obiettivo di Incorporamento Congiunto: Il codificatore viene aggiornato tramite gradienti derivanti dall'obiettivo di denoising condizionato. Il target latente futuro è trattato come un'"etichetta" derivata dall'ambiente, e il modello impara a denoisare da una versione rumorosa di questo latente futuro fino allo stato pulito, condizionato dallo stato corrente.
Stop-Gradient: Per prevenire il collasso della rappresentazione e stabilizzare l'addestramento, viene applicato uno stop-gradient al target latente futuro durante gli aggiornamenti del codificatore, seguendo i principi di JEPA e TDMPC2.
Tecniche di Stabilizzazione: Il modello impiega un clamping tanh differenziabile sulle attivazioni latenti e un meccanismo di commutazione casuale tra l'output del denoisatore e l'output del codificatore per la condizione del passo successivo, per bilanciare coerenza e supervisione.

Motivazione Teorica

Il documento fornisce una giustificazione teorica che collega JEDI al principio del Collo di Bottiglia dell'Informazione.

JEPA come Collo di Bottiglia: L'obiettivo standard JEPA è dimostrato inducere un collo di bottiglia dell'informazione variazionale, premiando rappresentazioni che mantengono le informazioni utili per prevedere il futuro mentre comprimono i dettagli irrilevanti.
Decomposizione Diffusione-JEPA: Gli autori dimostrano che la diffusione latente condizionata ammette una decomposizione simile del collo di bottiglia dell'informazione variazionale. La funzione di perdita di denoising condizionata agisce come termine predittivo all'interno di questa struttura a collo di bottiglia. Ciò suggerisce che il denoising della diffusione fornisce un segnale di apprendimento predittivo end-to-end fondato su principi per il codificatore latente, piuttosto che servire meramente come surrogato generativo.

Contributi Chiave

Primo Modello del Mondo a Diffusione Latente End-to-End: JEDI è il primo modello MBRL online ad addestrare il proprio codificatore latente end-to-end utilizzando un obiettivo di denoising della diffusione all'interno di un framework JEPA, eliminando la necessità di pre-addestramento separato o funzioni di perdita di ricostruzione.
Connessione Teorica: Il documento stabilisce un legame matematico tra gli obiettivi di diffusione condizionata e il collo di bottiglia dell'informazione variazionale, spiegando perché il denoising può supportare efficacemente l'apprendimento rappresentativo predittivo.
Guadagni di Efficienza: Spostandosi dallo spazio dei pixel a uno spazio latente appreso compatto ( $16 \times 8 \times 8$ ) e addestrando end-to-end, JEDI riduce significativamente il sovraccarico computazionale rispetto alla diffusione basata sui pixel.

Risultati Sperimentali

JEDI è stato valutato sul benchmark Atari100k (26 compiti, 5 semi) e sull'ambiente Craftium (presa di decisioni incarnata 3D).

Prestazioni

Atari100k: JEDI ottiene prestazioni competitive, eguagliando o superando basi solide (DIAMOND, DreamerV3, STORM) sulle metrica della Media Interquartile (IQM) e del Gap di Ottimalità. Ottiene risultati all'avanguardia (SOTA) in 7 giochi e prestazioni da secondo classificato nel numero di compiti sovrumani.
Craftium: JEDI supera Horizon Imagination (HI) in tutti e quattro i compiti testati, dimostrando che latenti predittivi end-to-end possono supportare prestazioni elevate in ambienti 3D senza addestramento latente basato sulla ricostruzione.
Stocasticità: Negli esperimenti con frame-skipping stocastico, JEDI supera costantemente DreamerV3, suggerendo che i modelli di diffusione potrebbero essere più robusti all'alta incertezza aleatoria.

Efficienza

Rispetto alla baseline basata sui pixel DIAMOND, JEDI dimostra significativi miglioramenti di efficienza:

Utilizzo VRAM: Utilizza il 43% in meno di VRAM.
Velocità di Campionamento: Raggiunge un campionamento del modello del mondo >3 volte più veloce.
Velocità di Addestramento: Si addestra >2,5 volte più velocemente (38 ore contro 98 ore su A100 per Atari100k).
Dimensione del Modello: Mantiene un conteggio di parametri simile (13,5M) a DIAMOND ma supera significativamente HI (97M parametri) in termini di efficienza, mantenendo prestazioni competitive.

Spostamento del Profilo di Prestazioni

Una scoperta distinta è che JEDI esibisce un profilo di prestazioni a livello di compito diverso rispetto a DIAMOND. Mentre DIAMOND eccelle nei regimi ad alto HNS (Human Normalized Score), JEDI mostra miglioramenti marcati nei regimi a basso HNS (giochi difficili). Questo è catturato dalla metrica Har-HNS, dove JEDI supera DIAMOND, indicando che fallisce meno severamente nei compiti più difficili. Gli autori attribuiscono ciò all'interfaccia latente end-to-end che riduce la dimensionalità dell'input per l'attore-critico, permettendo una migliore ottimizzazione in giochi complessi con spazi di azione elevati.

Significato e Affermazioni

Il documento afferma che JEDI dimostra che la diffusione latente predittiva è un'alternativa valida ed efficiente sia alla diffusione nello spazio dei pixel che ai modelli latenti addestrati separatamente per il MBRL online.

Oltre il Calcolo: I risultati suggeriscono che i latenti predittivi end-to-end cambiano il comportamento del modello del mondo stesso, non solo il costo computazionale. Lo spostamento nei profili di prestazioni indica che la rappresentazione appresa altera fondamentalmente il modo in cui l'agente interagisce con ambienti complessi.
Approccio Fondato su Principi: Radicando il metodo nella teoria del collo di bottiglia dell'informazione, il documento posiziona la diffusione latente predittiva come una direzione fondata su principi per la scalabilità del MBRL online.
Limitazioni: Gli autori notano modestamente che le prove sono attualmente limitate a Atari100k e Craftium, e il metodo non è stato ancora testato su domini come DMControl o Procgen. Riconoscono inoltre che le prestazioni dipendono da scelte progettuali specifiche (ad esempio, clamping, commutazione casuale) e che la derivazione teorica motiva le proprietà della rappresentazione piuttosto che garantire l'ottimizzazione dell'intero obiettivo variazionale.