DAWM: Diffusion Action World Models for Offline Reinforcement Learning via Action-Inferred Transitions

Riepilogo Tecnico: DAWM – Modelli del Mondo Azione-Diffusione per l'Apprendimento per Rinforzo Offline tramite Transizioni Inferite dall'Azione

1. Enunciato del Problema

I modelli del mondo basati sulla diffusione hanno dimostrato forti capacità nella sintesi di traiettorie realistiche a lungo orizzonte per l'apprendimento per rinforzo (RL) offline. Tuttavia, i metodi esistenti affrontano un problema critico di compatibilità con gli algoritmi standard di RL offline basati sul valore (ad es. TD3BC, IQL) che si affidano all'apprendimento della differenza temporale (TD) a un passo.

• Limitazione dei Modelli Senza Azione: Metodi come il Modello del Mondo Diffusione (DWM) generano stati futuri e ricompense, ma non generano esplicitamente le azioni. Di conseguenza, non possono costruire transizioni complete $(s, a, r, s')$, rendendoli incompatibili con l'apprendimento TD a un passo, che richiede coerenza tra stati, azioni e ricompense per una regolarizzazione efficace del critico.
• Limitazione della Modellazione Congiunta: Gli approcci che tentano di modellare congiuntamente stati, ricompense e azioni spesso soffrono di un aumento della complessità e dell'instabilità dell'addestramento, portando a una riduzione delle prestazioni pratiche.
• Limitazione dei Metodi Basati sulla Pianificazione: I pianificatori basati sulla diffusione (ad es. Diffuser, Decision Diffuser) che recuperano le azioni tramite Modelli di Dinamica Inversa (IDM) o modellazione congiunta si affidano tipicamente al campionamento iterativo e alla selezione delle ricompense post-hoc. Ciò comporta elevati costi di inferenza e bypassa l'efficienza dell'apprendimento TD a un passo.

La sfida principale è generare dati sintetici che supportino un apprendimento TD a un passo stabile ed efficiente, senza incorrere nell'instabilità dell'addestramento della modellazione congiunta o nei costi di inferenza della pianificazione iterativa.

2. Metodologia: DAWM

Gli autori propongono DAWM (Diffusion Action World Model), un framework modulare che disaccoppia la generazione delle traiettorie dall'inferenza delle azioni per produrre transizioni complete e etichettate adatte all'RL offline. Il framework consta di due fasi principali:

Fase 1: Generazione Condizionata di Traiettorie Diffusione

• Un modello di diffusione $p_\theta$ viene addestrato per generare sequenze future di stati e ricompense condizionati allo stato corrente $s_t$, all'azione $a_t$ e a un ritorno atteso ($R_{tg}$) target.
• Il modello definisce la distribuzione condizionata $p_\theta(\hat{s}_{t+1:t+H+1}, \hat{r}_{t:t+H} \mid s_t, a_t, R_{tg})$.
• L'obiettivo di addestramento minimizza l'errore atteso tra il rumore previsto e quello reale utilizzando un obiettivo di denoising con punteggio e guida senza classificatore.
• Crucialmente, questo modello genera traiettorie senza azione (solo stati e ricompense), semplificando il compito generativo evitando le complessità della modellazione congiunta delle azioni all'interno del processo di diffusione.

Fase 2: Completamento dell'Azione tramite Modello di Dinamica Inversa (IDM)

• Per convertire le traiettorie senza azione in transizioni complete, viene impiegato un Modello di Dinamica Inversa (IDM), $f_\phi$, addestrato separatamente.
• L'IDM ricostruisce le azioni mancanti modellando la distribuzione $f_\phi(a_t \mid s_{t-m:t-1}, s_t, s_{t+1})$ come una distribuzione Gaussiana diagonale.
• L'IDM viene addestrato su dati offline reali per massimizzare la verosimiglianza delle azioni osservate date le transizioni di stato.
• All'inferenza, l'IDM viene applicato alle traiettorie generate dalla diffusione per inferire le azioni, producendo tuple complete $(s, a, r, s')$.

Addestramento della Politica Offline

• Le transizioni complete sintetizzate vengono utilizzate per addestrare agenti RL offline utilizzando algoritmi standard basati sul valore, specificamente TD3BC (che aggiunge un regolarizzatore di clonazione comportamentale a TD3) e IQL (che esegue una regressione supervisionata pesata dall'vantaggio).
• Le funzioni Q vengono aggiornate utilizzando target di Bellman calcolati dalle transizioni ricostruite, abilitando aggiornamenti TD a un passo efficienti.

3. Contributi Chiave

1. Framework Modulare: Gli autori propongono un framework efficiente che combina la generazione di stati-ricompense basata sulla diffusione con un IDM addestrato separatamente per l'inferenza delle azioni. Questo design evita l'instabilità dell'addestramento della modellazione congiunta producendo al contempo transizioni complete.
2. Compatibilità con l'Apprendimento TD a Un Passo: Generando transizioni complete $(s, a, r, s')$, DAWM abilita l'uso di algoritmi conservativi di RL offline (TD3BC, IQL) che si affidano all'apprendimento TD a un passo, offrendo un'alternativa computazionalmente efficiente ai pianificatori di diffusione iterativi.
3. Prestazioni Empiriche: Il documento dimostra che gli algoritmi conservativi di RL offline traggono beneficio significativo dall'addestramento su queste traiettorie aumentate, superando costantemente i precedenti baseline basati sulla diffusione (DWM, DD) su molteplici compiti di locomozione D4RL.
4. Fattibilità dei Dati Sintetici: I risultati indicano che gli agenti addestrati su dati sintetici generati da DAWM possono raggiungere prestazioni paragonabili o superiori rispetto agli agenti addestrati su dataset offline reali, evidenziando il potenziale delle traiettorie sintetiche come alternativa pratica.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su 9 benchmark standard di locomozione D4RL (Hopper, Walker2d, HalfCheetah) con qualità dei dati variabile (Medium, Medium-Replay, Medium-Expert).

• Confronto con DWM: DAWM ha superato costantemente il baseline DWM senza azione. Con TD3BC, DAWM ha ottenuto un miglioramento medio del 9,3%, e con IQL del 9,5%. I guadagni sono stati particolarmente pronunciati in compiti di controllo complessi come Walker2d.
• Confronto con Metodi Basati sulla Pianificazione (DD): DAWM ha superato il Decision Diffuser (DD) in termini di ritorno normalizzato. Inoltre, DAWM ha dimostrato un'efficienza significativamente superiore, raggiungendo un'inferenza circa 4,0 volte più veloce rispetto a DD, poiché evita il campionamento iterativo durante l'inferenza.
• Confronto con Dati Reali: Gli agenti addestrati su dati generati da DAWM hanno ottenuto ritorni normalizzati paragonabili agli agenti addestrati su dati reali utilizzando TD3BC/IQL standard, e in alcuni casi (ad es. Hopper-MR con TD3BC) li hanno superati. Mentre SSORL (un baseline che utilizza IDM su dati reali) ha funzionato bene, gli autori notano che SSORL beneficia del filtraggio per le traiettorie più lunghe, un fattore non presente nella generazione a lunghezza fissa di DAWM.
• Studi di Ablazione:
  • Volume dei Dati (DAWM-T): Ridurre il numero di coppie di condizionamento di un fattore 8 (DAWM-T) ha comportato solo modesti cali di prestazioni, suggerendo che il meccanismo di completamento dell'azione è un fattore più critico del mero volume dei dati.
  • Orizzonte di Generazione: DAWM ha mantenuto prestazioni robuste su diversi orizzonti di generazione ($H=1, 3, 7$), indicando che il framework di apprendimento TD a un passo è stabile indipendentemente dalla lunghezza del rollout.

5. Significato e Affermazioni

Il documento afferma che DAWM colma il divario tra la modellazione generativa efficiente e l'apprendimento delle politiche basato su TD. Il suo significato principale risiede in:

• Abilitazione dell'Apprendimento TD a Un Passo: Fornisce un meccanismo per utilizzare dati generati dalla diffusione per algoritmi standard e stabili di RL offline che erano precedentemente incompatibili con i modelli del mondo senza azione.
• Efficienza Computazionale: Offre una soluzione pratica per l'RL generativo offline su larga scala riducendo i costi di inferenza rispetto ai metodi di pianificazione iterativa.
• Alternativa Pratica: I risultati suggeriscono che le traiettorie sintetiche generate tramite questo approccio modulare possono servire come un'alternativa praticabile e di alta qualità ai dataset offline reali, potenzialmente mitigando i problemi di scarsità di dati nelle applicazioni reali dove l'esplorazione è costosa o pericolosa.

Gli autori concludono che, sebbene il metodo mostri promesse, il lavoro futuro indagherà la sua estensione al controllo visivo, alla robotica ad alta dimensionalità e alla robustezza in condizioni di dinamiche imperfette.