KAGE-Bench: Valutazione Rapida della Generalizzazione Visiva ad Asse Noto per l'Apprendimento per Rinforzo

KAGE-Bench: Fast Known-Axis Visual Generalization Evaluation for Reinforcement Learning

January 20, 2026

Autori: Egor Cherepanov, Daniil Zelezetsky, Alexey K. Kovalev, Aleksandr I. Panov

cs.AI

Abstract

Gli agenti di apprendimento per rinforzo basati su pixel spesso falliscono in presenza di un puro shift distributivo visivo, anche quando le dinamiche latenti e le ricompense rimangono invariate. Tuttavia, i benchmark esistenti mescolano molteplici fonti di variazione, ostacolando un'analisi sistematica. Introduciamo KAGE-Env, una piattaforma 2D nativa in JAX che scompone il processo di osservazione in assi visivi controllabili indipendentemente, mantenendo fisso il problema di controllo sottostante. Per costruzione, la variazione di un asse visivo influisce sulle prestazioni solo attraverso la distribuzione delle azioni condizionata allo stato indotta in una politica pixel, fornendo un'astrazione pulita per la generalizzazione visiva. Basandoci su questo ambiente, definiamo KAGE-Bench, un benchmark composto da sei suite ad asse noto che include 34 coppie di configurazioni addestrazione-valutazione, progettato per isolare singoli shift visivi. Utilizzando una baseline standard PPO-CNN, osserviamo forti fallimenti dipendenti dall'asse: gli shift del fotometrico e dello sfondo causano spesso un collasso completo del successo, mentre gli shift dell'aspetto dell'agente sono relativamente benigni. Diversi shift preservano il movimento in avanti ma compromettono il completamento del compito, dimostrando che la ricompensa totale da sola può mascherare fallimenti di generalizzazione. Infine, l'implementazione completamente vettorizzata in JAX consente fino a 33 milioni di step ambientali al secondo su una singola GPU, permettendo analisi rapide e riproducibili sui fattori visivi. Codice: https://avanturist322.github.io/KAGEBench/.

English

Pixel-based reinforcement learning agents often fail under purely visual distribution shift even when latent dynamics and rewards are unchanged, but existing benchmarks entangle multiple sources of shift and hinder systematic analysis. We introduce KAGE-Env, a JAX-native 2D platformer that factorizes the observation process into independently controllable visual axes while keeping the underlying control problem fixed. By construction, varying a visual axis affects performance only through the induced state-conditional action distribution of a pixel policy, providing a clean abstraction for visual generalization. Building on this environment, we define KAGE-Bench, a benchmark of six known-axis suites comprising 34 train-evaluation configuration pairs that isolate individual visual shifts. Using a standard PPO-CNN baseline, we observe strong axis-dependent failures, with background and photometric shifts often collapsing success, while agent-appearance shifts are comparatively benign. Several shifts preserve forward motion while breaking task completion, showing that return alone can obscure generalization failures. Finally, the fully vectorized JAX implementation enables up to 33M environment steps per second on a single GPU, enabling fast and reproducible sweeps over visual factors. Code: https://avanturist322.github.io/KAGEBench/.