Sintesi Tecnica: Modelli Causali Centrati sugli Oggetti per la Pianificazione con Ricerca ad Albero Monte Carlo

Definizione del Problema

L'apprendimento per rinforzo basato su modelli (MBRL) ha ottenuto un successo significativo apprendendo la dinamica dell'ambiente per generare esperienze immaginate, riducendo così la necessità di interazioni nel mondo reale. Tuttavia, apprendere modelli del mondo accurati rimane una sfida in ambienti ad alta dimensionalità e non stazionari composti da molteplici oggetti interagenti.

Gli approoli esistenti di RL visivo si affidano tipicamente a encoder a Reti Neurali Convoluzionali (CNN) che producono una singola rappresentazione olistica di un'immagine di input. Queste rappresentazioni monolitiche spesso non riescono a catturare la struttura e le interazioni a livello di oggetto, che sono critiche per il processo decisionale. In scene complesse, la presenza di sfondi dinamici, entità irrilevanti o oggetti piccoli ma rilevanti per il compito può degradare significativamente le prestazioni dell'agente. Mentre gli esseri umani percepiscono il mondo come composto da entità discrete, permettendo un ragionamento efficiente, gli agenti RL standard faticano a isolare queste entità senza supervisione esplicita o maschere di segmentazione delle istanze.

Metodologia: COMET

Gli autori introducono COMET (Causal Object-centric Model for Efficient Tree search), un algoritmo di apprendimento per rinforzo basato su modelli che esegue la Ricerca ad Albero Monte Carlo (MCTS) all'interno di uno spazio latente strutturato a slot. COMET integra tre componenti primarie:

1. Encoder Centrato sugli Oggetti Congelato

COMET utilizza un encoder centrato sugli oggetti non supervisionato e congelato (ad es., SLATE, DINOSAUR o Slot Contrast) per mappare le osservazioni visive in un insieme di rappresentazioni latenti centrate sugli oggetti, o "slot" ($\bar{s}_t = \{s^1_t, \dots, s^n_t\}$).

• Inizializzazione degli Slot: Per affrontare l'invarianza alla permutazione e la natura stocastica dell'attenzione agli slot nei modelli di immagini statiche, gli autori inizializzano le rappresentazioni degli slot al tempo $t+1$ utilizzando gli slot del tempo $t$. Ciò garantisce la coerenza temporale delle assegnazioni degli oggetti all'interno di un episodio.
• Congelamento: L'encoder è preaddestrato su dati raccolti tramite una politica casuale e rimane congelato durante la fase di RL per fornire caratteristiche stabili a livello di oggetto.

2. Modello del Mondo Basato su Transformer con Fusione Azione-Slot

Il cuore di COMET è un modello del mondo basato su transformer che opera su questi slot. A differenza degli approcci standard stile MuZero che elaborano un vettore di stato monolitico, COMET elabora una sequenza di slot di oggetti.

• Binding Azione-Oggetto: Un meccanismo innovativo fonde le azioni con gli slot degli oggetti. Invece di comprimere tutti i latenti degli oggetti in un singolo vettore per predire lo stato successivo (il che creerebbe un collo di bottiglia), l'embedding dell'azione viene concatenato indipendentemente con ciascun oggetto. Questi embedding di slot condizionati dall'azione vengono passati attraverso un proiettore MLP condiviso e poi attraverso un backbone transformer per predire i successivi slot dello stato ($\hat{s}_{t+1}$) e le ricompense.
• Architettura: Il modello segue una pipeline di addestramento stile UniZero utilizzando un backbone transformer basato su nanoGPT. Predice i futuri slot e le ricompense basandosi sull'attuale set di slot e sull'azione intrapresa.

3. Attenzione Causale per Oggetto per Policy e Valore

Le teste della policy e del valore utilizzano un'architettura transformer modificata dotata di attenzione causale per oggetto.

• Score di Causalità: Per ogni slot latente $h^k_t$, la rete stima uno score di causalità $\alpha^k_t \in [0, 1]$, che rappresenta la probabilità che l'oggetto corrispondente sia rilevante per la decisione corrente.
• Modulazione dell'Attenzione: Questi score vengono utilizzati per modulare l'attenzione dot-product scalata. Una matrice di peso appresa $W_t$ eleva una struttura di grafo causale predefinita al livello di token, consentendo al meccanismo di attenzione di concentrarsi sulle entità rilevanti per il compito (oggetti causali) mentre sopprime quelle irrilevanti (oggetti non causali). Ciò consente all'agente di focalizzare la pianificazione sul sottoinsieme specifico di oggetti che attualmente interagiscono con l'ambiente.

Contributi Chiave

1. Algoritmo COMET: L'introduzione di un algoritmo MBRL basato su MCTS che combina rappresentazioni a livello di oggetto congelate con un modello del mondo basato su transformer per la pianificazione in uno spazio latente strutturato.
2. Meccanismo di Fusione Azione-Slot: Un metodo innovativo per legare le azioni agli slot centrati sugli oggetti all'interno di un unico backbone transformer. Ciò consente la modellazione del mondo e la predizione di policy/valore centrate sugli oggetti senza collassare l'informazione dell'oggetto in un singolo vettore.
3. Attenzione Causale per Oggetto: L'integrazione di score di causalità appresi per modulare le interazioni tra token, consentendo alla policy di concentrarsi dinamicamente sulle entità rilevanti per il compito durante il processo decisionale.
4. Valutazione Empirica: Valutazione completa attraverso otto compiti diversificati per visione e dinamica, inclusi l'Object-Centric Visual RL benchmark, ManiSkill, Robosuite e VizDoom.

Risultati Sperimentali

COMET è stato valutato rispetto a diversi baseline:

• Baseline Centrate sugli Oggetti: OCRL (PPO con encoder di oggetti), OCDreamer (RSSM centrato sugli oggetti) e SOLD (dinamica basata su slot non supervisionata).
• Baseline Monolitiche: UniZero (pianificazione latente standard stile MuZero).

Risultati sulle Prestazioni:

• Efficienza di Campionamento: Attraverso la suite di compiti diversificata, COMET raggiunge un Mean Normalized Score più elevato nelle prime fasi dell'addestramento rispetto sia ai forti metodi MCTS monolitici che ad altri baseline di RL centrati sugli oggetti.
• Specificità dei Compiti:
  • Nei compiti visivamente semplici (ad es., Object Goal, Object Comparison), COMET sfrutta rappresentazioni di alta qualità degli oggetti e l'attenzione causale per identificare accuratamente gli oggetti rilevanti, portando a una convergenza più rapida.
  • Nei compiti di controllo continuo (ad es., Block Lifting, Cube Pushing), COMET dimostra prestazioni costanti, sebbene non superi universalmente tutti i baseline su ogni specifica metrica (ad es., non ha superato i competitor in Block Lifting).
• Generalizzazione: Il metodo mostra robustezza in ambienti con diversa complessità visiva, dai forme 2D su sfondi monocromatici a complessi scenari di manipolazione robotica 3D e shooter in prima persona.

Significato e Rivendicazioni

Il paper sostiene che modellare esplicitamente l'importanza dei singoli oggetti per il processo decisionale migliori l'apprendimento della policy in ambienti orientati agli oggetti. Introducendo un bias induttivo esplicito a livello di oggetto in una pianificazione latente stile MuZero, COMET consente al modello di ragionare su singole entità e le loro interazioni piuttosto che trattare la scena come una massa olistica.

Gli autori affermano che COMET affronta efficacemente la sfida di rappresentare le osservazioni in ambienti complessi con molteplici oggetti. Decoupling le rappresentazioni degli oggetti e utilizzando l'attenzione causale per filtrare la rilevanza, l'algoritmo ottiene una maggiore efficienza di campionamento e un processo decisionale più mirato. Il lavoro suggerisce che i bias induttivi centrati sugli oggetti siano una via valida ed efficace per migliorare la pianificazione basata su modelli nella RL visiva, in particolare quando non è disponibile la segmentazione delle istanze o quando è richiesto l'apprendimento di rappresentazioni non supervisionate.