Learning Abstract World Models with a Group-Structured Latent Space

Riepilogo Tecnico: Apprendimento di Modelli del Mondo Astratti con uno Spazio Latente Strutturato per Gruppi

Enunciato del Problema

L'apprendimento di modelli astratti significativi dei Processi Decisionali di Markov (MDP) è essenziale per migliorare la generalizzazione da dati limitati nell'apprendimento per rinforzo (RL). Sebbene i metodi esistenti di apprendimento del modello del mondo riescano a incorporare dati ad alta dimensionalità in spazi astratti a dimensionalità inferiore per migliorare l'efficienza del campione, gli approcci all'avanguardia richiedono spesso grandi dataset di addestramento per raggiungere prestazioni sovrumane. Una limitazione chiave è la mancanza di integrazione di strutture geometriche note (simmetrie) intrinseche nella dinamica dell'ambiente. Senza tali prior, gli agenti potrebbero faticare a generalizzare a scenari non visti o richiedere dati eccessivi per apprendere simmetrie semplici, come le invarianze rotazionali o traslazionali.

Metodologia

Gli autori propongono un framework che impone prior geometriche sulla varietà di rappresentazione a bassa dimensionalità di un modello di transizione appreso. La metodologia centrale coinvolge:

1. Spazio Latente Strutturato per Gruppi: Invece di trattare lo spazio degli stati astratti $Z$ come uno spazio vettoriale euclideo standard, gli autori lo modellano come uno spazio quoziente (ad esempio, $\mathbb{R}/k\mathbb{Z}$) o un prodotto di tali spazi. Questa scelta riflette i gruppi di simmetria sottostanti dell'ambiente (ad esempio, gruppi ciclici per le rotazioni).
2. Azioni di Gruppo come Dinamiche di Transizione: Il modello di transizione $\tau: Z \times A \to Z$ è vincolato a operare tramite un'azione di gruppo. Nello specifico, il prossimo stato latente è previsto come lo stato corrente combinato con un vettore di transizione appreso tramite un'operazione di gruppo $\oplus$:
   $$\hat{z}_{t+1} = z_t \oplus \Delta(z_t, a; \theta_{trans})$$
   Per spazi quoziente come $\mathbb{R}/k\mathbb{Z}$, questa operazione è implementata tramite aritmetica modulare, codificando naturalmente le condizioni al contorno periodiche.
3. Rappresentazioni Disaccoppiate: Per gestire ambienti contenenti sia caratteristiche simmetriche (strutturate) che non simmetriche (non strutturate), il framework promuove il disaccoppiamento. Viene applicata una regolarizzazione di sparsità al vettore di transizione $\Delta$, forzando le caratteristiche che dovrebbero rimanere invarianti sotto una specifica azione ad essere zero. Ciò è ottenuto mascherando coordinate specifiche del vettore di transizione basandosi su una mappatura $\sigma(a)$ che definisce quali dimensioni latenti sono influenzate dall'azione $a$.
4. Obiettivo di Apprendimento: Il modello è addestrato utilizzando un obiettivo di apprendimento contrastivo auto-supervisionato (perdita InfoNCE) per allineare le transizioni previste con la verità fondamentale, combinato con una perdita di previsione della ricompensa e una perdita di regolarizzazione del volume per mantenere lo spazio astratto compatto e interpretabile. A differenza di approcci precedenti che impongono l'equivarianza attraverso architetture di reti neurali con pesi accoppiati (che possono essere computazionalmente costose), questo metodo impone la struttura direttamente all'interno della geometria dello spazio latente.

Contributi Chiave

• Prior Geometrici nello Spazio Latente: Il documento dimostra come incorporare strutture simmetriche note (ad esempio, simmetrie rotazionali e traslazionali) selezionando spazi latenti appropriati (spazi quoziente) e azioni di gruppo corrispondenti, piuttosto che modificare l'architettura della rete neurale.
• Gestione di Simmetrie Miste: Il framework apprende con successo rappresentazioni che combinano caratteristiche strutturate (simmetriche) e non strutturate, permettendo al modello di scalare ad ambienti complessi come le viste in prima persona 3D dove non tutte le dinamiche sono simmetriche.
• Migliorata Generalizzazione: L'approccio porta a rappresentazioni più semplici e disaccoppiate che generalizzano meglio alle transizioni non viste e richiedono meno campioni di addestramento rispetto alle baseline completamente non strutturate.
• Prestazioni RL a valle: I modelli astratti appresi servono come prior efficaci per compiti RL a valle, migliorando significativamente l'efficienza del campione in regimi a basso dato.

Risultati Sperimentali

Gli autori hanno valutato il loro approccio su diversi ambienti:

• MDP Giocattolo (Passaggio e Toro): In un semplice mondo a griglia ciclico e in un MDP a toro (prodotto di due gruppi ciclici), il modello strutturato per gruppi ha previsto con precisione transizioni non viste dove le baseline non strutturate hanno fallito. Il modello ha recuperato la topologia corretta (ad esempio, mappando lo spazio latente su un toro in $\mathbb{R}^3$) utilizzando solo 2 dimensioni latenti.
• MiniGrid (Vista dall'Alto): Il metodo ha disaccoppiato con successo l'orientamento (modellato come un gruppo ciclico) dalla posizione (modellata come spazio euclideo), risultando in una rappresentazione di spazio prodotto $\mathbb{R}/2\pi\mathbb{Z} \times \mathbb{R}^2$.
• VizDoom (Vista in Prima Persona 3D): In un ambiente ad alta dimensionalità con input RGB, il modello ha mappato dati visivi complessi in uno spazio latente tridimensionale. Ha catturato correttamente la natura ciclica della rotazione e la natura traslazionale del movimento, superando le baseline nelle metriche di generalizzazione (Hits@k, Mean Reciprocal Rank) e nei compiti RL a valle.
• Metriche Quantitative: In tutti gli ambienti, il modello "AWM + Prior Geometrici" ha superato significativamente le baseline non strutturate (AWM senza prior) e altri metodi (PRAE, Matrice di Rotazione) nella generalizzazione a transizioni non viste, in particolare in regimi a basso dato. Nei compiti RL a valle (Toro, MiniGrid, VizDoom), gli agenti che utilizzano il prior geometrico hanno ottenuto ricompense cumulative più elevate con meno passaggi di addestramento.

Significato e Affermazioni

Il documento afferma che sfruttando prior geometrici attraverso la scelta dello spazio latente e delle azioni di gruppo, è possibile apprendere modelli del mondo più efficienti dal punto di vista del campione e generalizzabili senza il sovraccarico computazionale delle reti neurali equivarianti. Gli autori sottolineano che il loro metodo permette la combinazione flessibile di caratteristiche simmetriche e non simmetriche, rendendolo applicabile a scenari complessi e reali come i giochi in prima persona 3D. Il lavoro suggerisce che rispettare la simmetria sottostante dell'ambiente nella rappresentazione astratta porta a "rappresentazioni più semplici e disaccoppiate", che a loro volta facilitano una migliore previsione e pianificazione. Gli autori forniscono il codice completo per garantire la riproducibilità.