Performance Asymmetry in Model-Based Reinforcement Learning

1. Il Problema: Asimmetria delle Prestazioni (Performance Asymmetry)

Gli autori identificano un problema critico e finora nascosto nel campo dell'Apprendimento per Rinforzo basato su Modello (MBRL) sul benchmark Atari100k. Sebbene gli agenti MBRL abbiano recentemente superato la media delle prestazioni umane (Human-Normalized Score, HNS > 1.0), le metriche aggregate convenzionali (come la media aritmetica) mascherano una grave asimmetria delle prestazioni.

• Definizione: Gli agenti MBRL eccellono in modo sproporzionato in alcuni compiti (chiamati Agent-Optimal), spesso quelli con punteggi umani bassi, ma falliscono drasticamente in altri compiti (chiamati Human-Optimal), dove gli esseri umani ottengono punteggi elevati.
• Il Fenomeno: L'agente SOTA attuale (DIAMOND, basato su pixel e diffusion) ottiene un HNS medio complessivo di 1.46, ma è il peggiore tra tutti i baseline sui compiti Human-Optimal. Esiste un divario di prestazioni di 21 volte tra i due sottoinsiemi di compiti.
• Causa Radice: L'asimmetria è aggravata dall'uso di input pixel-based e dalla natura dei compiti Human-Optimal, che presentano spazi di azione più complessi e un alto numero di compiti di tipo "Shooter" (dove l'agente deve usare un'arma a distanza). Questi compiti soffrono della "maledizione della dimensionalità" quando elaborati direttamente dagli input visivi.

2. Metodologia e Soluzione Proposta

Per affrontare questo problema, il paper introduce un nuovo framework di valutazione e un nuovo modello del mondo.

A. Nuova Metrica di Valutazione: Sym-HNS

Gli autori criticano la media aritmetica per la sua invarianza ai valori estremi e per il fatto che i grandi valori dominano i piccoli.

• Partitioning: Dividono i 26 giochi di Atari100k in due sottoinsiemi bilanciati di 13 compiti ciascuno: Agent-Optimal (HNS medio > 0.75) e Human-Optimal (HNS medio < 0.75).
• Sym-HNS: Introducono una nuova metrica aggregata basata sulla media armonica dei punteggi medi sui due sottoinsiemi:
  $$\text{Sym-HNS} = \frac{2 \cdot \mu_{AO} \cdot \mu_{HO}}{\mu_{AO} + \mu_{HO}}$$
  Questa metrica penalizza le prestazioni sbilanciate, costringendo i modelli a essere competitivi su entrambi i tipi di compiti per ottenere un punteggio alto.

B. Il Modello: JEDI (Joint Embedding DIffusion)

Per risolvere l'asimmetria, gli autori propongono JEDI, un modello del mondo latent-based (basato su spazio latente) che integra la diffusione.

• Architettura: JEDI è un modello end-to-end che apprende un encoder per comprimere gli input pixel in uno spazio latente e un modello di dinamica basato sulla diffusione (Diffusion Dynamics Model) in quello spazio latente.
• Training con JEPA: A differenza dei metodi precedenti che usano la ricostituzione (reconstruction loss) per addestrare l'encoder, JEDI utilizza un approccio ispirato a JEPA (Joint Embedding Predictive Architecture). L'encoder viene addestrato direttamente dai gradienti della funzione di perdita di denoising del modello di diffusione latente, senza bisogno di un decoder di ricostituzione o di estrazione di stati nascosti.
• Vantaggi Tecnici:
  • Efficienza: La compressione latente riduce drasticamente la complessità computazionale rispetto agli input pixel.
  • Stabilità: Viene utilizzata una funzione di clamp differenziabile (tanh) per stabilizzare il processo di denoising nello spazio latente.
  • Compatibilità: È il primo lavoro che dimostra come una funzione di perdita di diffusione "off-the-shelf" sia direttamente compatibile con il paradigma JEPA.

3. Risultati Chiave

I risultati sperimentali su Atari100k (con 5 semi diversi) confermano l'efficacia di JEDI:

1. Riduzione dell'Asimmetria: JEDI inverte la tendenza negativa osservata negli agenti pixel-based. Ottiene il punteggio SOTA su Sym-HNS, dimostrando un bilanciamento superiore tra i compiti.
2. Prestazioni su Compiti Difficili:
   • Human-Optimal: JEDI supera significativamente i baseline sui compiti Human-Optimal, raggiungendo per la prima volta prestazioni sovrumane sul gioco BankHeist.
   • Visually Bottlenecked: Contrariamente alla credenza comune che i modelli latenti falliscano nei compiti visivamente complessi, JEDI ottiene il SOTA su Breakout (un compito Agent-Optimal visivamente denso) e rimane competitivo su Assault.
3. Efficienza Computazionale:
   • Velocità: JEDI addestra 2 volte più velocemente e ha una velocità di inferenza 3 volte superiore rispetto a DIAMOND.
   • Memoria: Utilizza il 43% in meno di memoria GPU grazie alla compressione latente (12x).
4. Robustezza: JEDI dimostra una maggiore robustezza in ambienti stocastici (con frame skipping casuale) rispetto agli agenti latenti non basati su diffusione.

4. Contributi Principali

1. Scoperta dell'Asimmetria: Identificazione e analisi del problema di "overfitting sui compiti" negli agenti MBRL, dove le medie aggregate nascondono fallimenti critici su compiti complessi.
2. Nuovo Framework di Valutazione: Introduzione della partizione Agent/Human-Optimal e della metrica Sym-HNS per una valutazione più equa e bilanciata.
3. Analisi Causale: Dimostrazione che l'asimmetria è causata dalla maledizione della dimensionalità negli input pixel per compiti complessi (spazi di azione grandi, dinamica "Shooter") e che la diffusione nei pixel aiuta solo i compiti visivamente semplici ma non risolve il problema della complessità logica.
4. JEDI: Sviluppo di un nuovo modello del mondo end-to-end che combina spazio latente e diffusione, addestrato con un framework stile JEPA, ottenendo prestazioni SOTA su tutte le metriche chiave (Sym-HNS, Human-Optimal, Breakout) con maggiore efficienza.

5. Significato e Impatto

Questo lavoro è fondamentale perché sposta il paradigma di valutazione nel MBRL da una semplice ricerca della "media più alta" a un'analisi più sfumata del bilanciamento delle prestazioni. Dimostra che:

• Gli agenti MBRL attuali non stanno necessariamente migliorando in modo generale, ma stanno semplicemente specializzandosi su un sottoinsieme di compiti.
• L'uso di modelli latenti, se combinato correttamente con tecniche avanzate come la diffusione e l'apprendimento self-supervised (JEPA), può superare i limiti degli input pixel sia in termini di efficienza che di capacità di ragionamento su compiti complessi.
• Apre nuove direzioni di ricerca per comprendere le caratteristiche fondamentali dei compiti in cui gli umani eccellono rispetto a quelli in cui gli agenti eccellono, suggerendo che la vera intelligenza artificiale richiede una capacità di generalizzazione bilanciata su tutti i tipi di scenari.