Sintesi Tecnica: TaskSense

Problematica

I modelli del mondo per il controllo visivo si basano tipicamente sulla ricostruzione di osservazioni complete per apprendere stati latenti compatti. Sebbene efficace in benchmark standardizzati con sfondi statici, questo approccio incoraggia implicitamente le rappresentazioni latenti a preservare informazioni attraverso l'intero input visivo, indipendentemente dalla loro rilevanza per il compito di controllo. In ambienti in cui il contenuto rilevante per il compito occupa solo una piccola frazione dell'osservazione o dove sono presenti distrazioni e rumore di fondo, questo obiettivo di ricostruzione consuma una preziosa capacità rappresentativa su contenuti visivi irrilevanti per il compito. Questa discrepanza tra la ricostruzione visiva e gli obiettivi di controllo crea un bias nelle dinamiche latenti, diluendo i segnali di apprendimento per le caratteristiche rilevanti per il controllo e degradando severamente le prestazioni a valle in presenza di distrazioni visive. Al contrario, fare affidamento esclusivamente sulla supervisione del compito (ad esempio, la predizione della ricompensa) fornisce spesso un segnale troppo debole per apprendere rappresentazioni visive ricche. Le soluzioni esistenti che eliminano completamente la ricostruzione sacrificano le capacità generative necessarie per la pianificazione basata sull'immaginazione, mentre i metodi che applicano l'attenzione o il masking dopo l'estrazione delle feature non riescono a controllare quale informazione entri nel modello delle dinamiche latenti fin dall'inizio.

Metodologia

Gli autori propongono TaskSense, un framework di modellazione del mondo incentrato sul compito che impone la rilevanza per il compito prima della codifica latente. TaskSense aumenta l'architettura DreamerV3 con un meccanismo di attenzione spaziale stocastica differenziabile condizionato dallo stato di credenza latente precedente.

Componenti Core

1. Attenzione Spaziale Stocastica Condizionata dal Latente:

   • Invece di predire l'attenzione per ogni pixel, il modello definisce l'attenzione su un reticolo grossolano $\Lambda$ che copre l'osservazione.
   • Un MLP leggero predice i logit di attenzione basandosi sullo stato latente precedente $s_{t-1}$ (specificamente lo stato nascosto $h_{t-1}$ e il latente stocastico $z_{t-1}$), permettendo al modello di decidere dove concentrarsi in base alla sua evoluzione della comprensione del compito piuttosto che sull'aspetto visivo istantaneo.
   • Il campo di attenzione è modellato come una collezione di variabili indipendenti stocastiche utilizzando una distribuzione Binary Concrete (Relaxed Bernoulli). Ciò consente di prestare attenzione a più regioni o di ignorarle simultaneamente, evitando i vincoli competitivi dell'attenzione softmax.
   • Durante l'addestramento vengono campionati maschere differenziabili; durante l'inferenza viene utilizzata la media deterministica.

2. Codifica Consapevole dell'Attenzione (Attention-Aware Encoding):

   • La maschera di attenzione campionata viene sovracampionata (upsampled) e utilizzata per modulare l'osservazione tramite moltiplicazione elemento per elemento, attenuando le regioni a bassa attenzione.
   • Per risolvere l'ambiguità tra pixel a bassa intensità e pixel ad alta intensità soppressi, la mappa di attenzione sovracampionata viene concatenata come canale di input aggiuntivo all'encoder.

3. Obiettivo di Dinamica Inversa Ausiliario:

   • Per evitare che il regolarizzatore di sparsità faccia collassare l'attenzione verso una soluzione banale (ignorando tutte le osservazioni), viene introdotto un obiettivo ausiliario.
   • Il modello predice l'azione $a_{t-1}$ eseguita tra due passi temporali consecutivi utilizzando la differenza temporale delle feature codificate ($f_t - f_{t-1}$).
   • Questo obiettivo ancora il meccanismo di attenzione alle regioni rilevanti per il controllo, incoraggiando la ritenzione di caratteristiche visive che sono predittive del comportamento dell'agente.

4. Ricostruzione Condizionata dall'Attenzione:

   • Invece di ricostruire l'osservazione completa, il decoder ricostruisce solo le regioni con attenzione ($o^{att}_t$).
   • Il decoder è condizionato sulla mappa di attenzione campionata (tramite stop-gradient) per garantire che sappia quali regioni sono state trattenute, impedendo che lo stato latente venga penalizzato per l'informazione mancante che è stata intenzionalmente filtrata.

Obiettivo di Addestramento

La funzione di perdita totale combina le perdite standard delle dinamiche latenti di DreamerV3 ($L_{dyn}$) con i nuovi componenti:
$$L_{WM} = L_{dyn} + L_{rec} + \lambda_{sparse}L_{sparse} + \lambda_{inv}L_{inv}$$
Dove $L_{rec}$ è la perdita di ricostruzione sulle osservazioni filtrate, $L_{sparse}$ incoraggia mappe di attenzione compatte e $L_{inv}$ è la perdita di dinamica inversa.

Contributi Chiave

1. Framework TaskSense: Un framework di modellazione del mondo basato sulla ricostruzione che impone la rilevanza per il compito prima della codifica latente attraverso un'attenzione spaziale stocastica condizionata dal latente.
2. Meccanismo di Apprendimento Auto-Supervisionato: Un approccio innovativo che ancora l'attenzione stocastica alle regioni rilevanti per il controllo combinando la supervisione della dinamica inversa con la ricostruzione sulle osservazioni filtrate.
3. Validazione Empirica: Dimostrazione che TaskSense mantiene prestazioni competitive sui benchmark standard pur superando significativamente i baseline sotto distrazioni visive.

Risultati Sperimentali

Gli autori hanno valutato TaskSense su DeepMind Control Suite (DMC) e Distracting Control Suite (DCS), che presenta sfondi di immagini naturali dinamiche.

• Prestazioni Standard: Sul DMC standard (sfondi statici), TaskSense eguaglia le prestazioni del baseline DreamerV3 nella maggior parte dei compiti, dimostrando che ricostruire solo un sottoinsieme dell'osservazione è sufficiente per un efficace pianificazione latente.
• Robustezza alle Distrazioni: Sul Distracting Control Suite, TaskSense supera costantemente DreamerV3. Ad esempio, sul compito Walker Run, TaskSense ha ottenuto un ritorno medio di 531.2 rispetto ai 286.2 di DreamerV3. Ciò conferma che filtrare il contenuto visivo irrilevante prima della codifica migliora la robustezza.
• Analisi Qualitativa: Le visualizzazioni delle mappe di attenzione apprese sul compito Cheetah Run mostrano che TaskSense si concentra costantemente sulle regioni rilevanti per il controllo (ad esempio, le gambe del cheetah) sopprimendo il torso e gli sfondi distraenti.
• Studi di Ablazione:
  • Rimuovere l'obiettivo della dinamica inversa porta a un'attenzione diffusa che non riesce a localizzare le regioni rilevanti per il controllo, causando un significativo degrado delle prestazioni.
  • Sostituire l'attenzione stocastica con l'attenzione deterministica risulta in un calo significativo delle prestazioni, evidenziando l'importanza di esporre il modello a diverse osservazioni parziali durante l'addestramento.
  • Rimuovere la ricostruzione condizionata dall'attenzione causa un degrado minore ma percepibile, indicando la necessità di allineare l'obiettivo di ricostruzione con l'input filtrato.

Significato e Rivendicazioni

L'articolo sostiene che TaskSense affronti un disallineamento fondamentale nei modelli del mondo attuali: la tendenza a preservare tutte le informazioni visive invece di concentrarsi su ciò che conta per il controllo. Spostando il focus sul filtraggio a livello di osservazione piuttosto che sul raffinamento post-hoc della rappresentazione, TaskSense consente ai modelli del mondo basati sulla ricostruzione di allocare efficientemente la capacità rappresentativa verso le informazioni rilevanti per il compito.

Gli autori affermano che questo approccio consente ai modelli del mondo di mantenere i benefici generativi della ricostruzione (essenziali per la pianificazione basata sull'immaginazione) pur ottenendo robustezza contro le distrazioni visive che tipicamente degradano i modelli standard. Il lavoro suggerisce che imporre la rilevanza per il compito all'interfaccia dell'osservazione sia una direzione promettente per migliorare i modelli del mondo visivi e l'apprendimento delle rappresentazioni nel reinforcement learning basato sul modello.

L'articolo riconosce un limite: assume che le caratteristiche visive predittive delle azioni immediate (catturate dalla dinamica inversa) siano sinonimo di quelle rilevanti per il successo del compito a lungo termine. In scenari in cui queste divergono, il segnale di supervisione attuale potrebbe concentrare l'attenzione su caratteristiche sub-ottimali, una direzione identificata per il lavoro futuro.