Riepilogo Tecnico: KAGE-Bench

Problematica

Gli agenti di reinforcement learning (RL) addestrati su osservazioni di pixel ad alta dimensionalità sono notoriamente fragili di fronte a spostamenti di distribuzione puramente visivi, anche quando la dinamica latente, i kernel di transizione e le funzioni di ricompensa rimangono invariati. Sebbene gli agenti possano convergere e performare bene in distribuzione (in-distribution), spesso falliscono catastroficamente sotto cambiamenti di aspetto, illuminazione o rumore visivo. Un ostacolo primario alla diagnosi di questi fallimenti è la mancanza di benchmark di valutazione in grado di isolare specifiche fonti di spostamento visivo. I benchmark esistenti (ad esempio, Procgen, Distracting Control Suite) spesso intrecciano molteplici fonti di variazione — come l'aspetto dello sfondo, la geometria, la dinamica e i distrattori — rendendo impossibile attribuire i gap di performance a specifici fattori visivi. Inoltre, molti ambienti di RL basati su pixel sono computazionalmente costosi, ostacolando studi di ablazione su larga scala e test sistematici di ipotesi.

Metodologia

Gli autori introducono KAGE-Env e KAGE-Bench per affrontare queste limitazioni attraverso una filosofia di design a "assi noti" (known-axis).

KAGE-Env: Una Piattaforma Platformer Fattorizzata Visivamente

KAGE-Env è un ambiente platformer 2D nativo in JAX, progettato per fattorizzare il processo di osservazione in assi visivi indipendentamente controllabili, mantenendo fissa la dinamica di controllo latente.

• Architettura: L'ambiente è implementato interamente in JAX con compilazione end-to-end JIT ed esecuzione vettorizzata (utilizzando vmap e lax.scan). Ciò consente una massiccia parallelizzazione, raggiungendo fino a 33 milioni di step di ambiente al secondo su una singola GPU con $2^{16}$ ambienti paralleli.
• Fattorizzazione: Il kernel di osservazione $O_\xi(\cdot|s)$ è parametrizzato da una configurazione visiva $\xi$. Fondamentalmente, il kernel di transizione $P$ e la funzione di ricompensa $r$ sono invarianti rispetto a $\xi$. Ciò assicura che qualsiasi differenza di performance tra l'addestramento e la valutazione derivi esclusivamente da come una policy basata su pixel $\pi(a|o)$ risponde a diverse renderizzazioni dello stesso stato latente $s$.
• Configurazione: L'ambiente è controllato tramite un singolo file YAML, che espone 93 parametri esplicitamente controllabili suddivisi in categorie quali l'immagine/colore dello sfondo, l'aspetto dell'agente/sprite, i personaggi non giocanti (NPC), i distrattori, i filtri fotometrici (luminosità, tonalità, rumore), gli effetti di illuminazione e il layout.

KAGE-Bench: Il Protocollo di Benchmark

Costruito su KAGE-Env, KAGE-Bench definisce un protocollo di valutazione standardizzato composto da sei suite a "assi noti" e 34 coppie di configurazione addestramento-valutazione.

• Isolamento: In ogni coppia, la configurazione di addestramento ($\xi_{train}$) e la configurazione di valutazione ($\xi_{eval}$) differiscono solo in un sottoinsieme designato di parametri di rendering (un asse specifico), mentre tutti gli altri parametri e la dinamica sottostante rimangono identici.
• Fondamento Teorico: Il paper formalizza che la generalizzazione visiva in queste condizioni è equivalente a uno spostamento nella policy di stato indotta $\pi_\xi$. Marginalizzando il kernel di osservazione, la policy basata su pixel $\pi(a|o)$ induce una distribuzione di azione condizionata allo stato $\pi_\xi(a|s)$. Il benchmark dimostra che i gap di performance sono esattamente le differenze tra queste policy indotte che agiscono sullo stesso MDP latente.
• Metriche: Oltre al classico ritorno episodico (episodic return), il benchmark riporta metriche a livello di traiettoria: distanza percorsa, progresso normalizzato verso l'obiettivo e successo del task binario. Questa distinzione è vitale, poiché il ritorno da solo può mascherare i fallimenti in cui un agente compie un progresso in avanti ma non completa il compito.

Contributi Chiave

1. KAGE-Env: Un ambiente di RL ad alto throughput, nativo in JAX, con 93 parametri visivi controllabili, capace di simulazione vettorizzata fino a 33M di step/secondo su una singola GPU.
2. KAGE-Bench: Un protocollo di benchmark che isola gli spostamenti della distribuzione visiva per costruzione, caratterizzato da sei suite (Agente, Sfondo, Distrattori, Effetti, Filtri, Layout) e 34 coppie specifiche di addestramento-valutazione.
3. Diagnosi Empirica: Una valutazione sistematica della generalizzazione visiva utilizzando un baseline standard PPO-CNN, quantificando come diversi assi visivi impattano la performance e identificando classi specifiche di spostamenti che inducono un degrado severo.

Risultati

Utilizzando un baseline standard PPO-CNN addestrato da singoli frame RGB, gli autori hanno osservato forti fallimenti dipendenti dall'asse:

• Degrado Severo: Filtri (es. spostamento di tonalità, saturazione, luminosità) ed Effetti (es. intensità dell'illuminazione, conteggio) hanno causato i collassi di performance più significativi. I tassi di successo sono scesi rispettivamente dell'86,8% e dell'80,5%, spesso mentre l'agente manteneva ancora un certo movimento in avanti.
• Spostamenti dello Sfondo: I cambiamenti nell'immagine o nel colore dello sfondo hanno portato a un calo del 53,3% nel tasso di successo e di un 30,5% nella distanza percorsa, indicando che gli spostamenti dello sfondo compromettono sia la locomozione che il completamento del task.
• Spostamenti Benigni: Gli spostamenti dell'aspetto dell'agente (es. cambiare lo sprite da uno scheletro a un clown) sono stati comparativamente benigni, causando solo un calo del 21,1% nel tasso di successo.
• Il Gap "Movimento vs Completamento": Diversi spostamenti (in particolare Distrattori e Layout) hanno preservato il movimento in avanti (piccoli gap di distanza di circa il 3-4%) ma hanno causato enormi cali nel completamento del task (gap del Tasso di Successo del 30-60%). Ciò evidenzia come le metriche basate sul ritorno da sole possano mascherare i fallimenti di generalizzazione dove un agente si muove ma non risolve il task.
• Baseline di Robustezza: Gli autori hanno testato diversi metodi di robustezza (RAD-PPO con cutout-color, crop-resize, CURL, DrQ, SVEA). Sebbene RAD-PPO con cutout-color abbia mostrato alcuni miglioramenti sulle configurazioni difficili, i metodi che si affidano ad aumentazioni geometriche (crop-resize) non sono riusciti a convergere, probabilmente perché hanno interrotto le posizioni spaziali assolute necessarie per la precisione del timing dei salti e l'evitamento degli ostacoli in questo contesto platformer.

Significato e Rivendicazioni

Il paper sostiene che KAGE-Bench fornisca un framework veloce, riproducibile e diagnostico per valutare il RL basato su pixel sotto variazione visiva controllata.

• Attribuzione: Fattorizzando il processo di osservazione, il benchmark consente l'attribuzione precisa dei fallimenti di generalizzazione a specifici fattori visivi, rimuovendo gli effetti confondenti di dinamiche intrecciate o cambiamenti nella struttura del task presenti in altri benchmark.
• Scalabilità: L'implementazione nativa in JAX affronta il collo di bottiglia computazionale della valutazione della generalizzazione visiva, rendendo fattibili ampie scansioni (sweeps) sui parametri visivi.
• Utilità Diagnostica: I risultati dimostrano che la difficoltà della generalizzazione visiva non è uniforme; varia sostanzialmente tra gli assi. Il benchmark funge da strumento per identificare quali specifici fattori visivi (es. spostamenti fotometrici rispetto all'aspetto dell'agente) siano più critici affinché una determinata architettura di policy apprenda le invarianti.

Gli autori posizionano KAGE-Bench non come un proxy completo per tutte le sfide del mondo reale del RL, ma come uno strumento diagnostico controllato per supportare un'analisi più sistematica della robustezza visiva e per guidare lo sviluppo di algoritmi capaci di gestire specifici spostamenti visivi su assi noti.