KAGE-Bench: Avaliação Rápida de Generalização Visual com Eixos Conhecidos para Aprendizagem por Reforço

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

January 20, 2026

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

cs.AI

Resumo

Os agentes de aprendizagem por reforço baseados em pixels frequentemente falham sob mudanças de distribuição puramente visuais, mesmo quando a dinâmica latente e as recompensas permanecem inalteradas. No entanto, os *benchmarks* existentes envolvem múltiplas fontes de mudança e dificultam uma análise sistemática. Apresentamos o KAGE-Env, uma plataforma 2D nativa em JAX que fatoriza o processo de observação em eixos visuais controláveis de forma independente, mantendo o problema de controlo subjacente fixo. Por construção, a variação de um eixo visual afeta o desempenho apenas através da distribuição de ações condicionada ao estado induzida por uma política de pixels, fornecendo uma abstração clara para a generalização visual. Com base neste ambiente, definimos o KAGE-Bench, um *benchmark* de seis *suites* de eixos conhecidos, compreendendo 34 pares de configurações de treino-avaliação que isolam mudanças visuais individuais. Utilizando uma linha de base padrão PPO-CNN, observamos falhas fortemente dependentes do eixo: mudanças no fundo e fotométricas frequentemente colapsam o sucesso, enquanto mudanças na aparência do agente são comparativamente benignas. Várias mudanças preservam o movimento para a frente, mas impedem a conclusão da tarefa, mostrando que a recompensa por si só pode ocultar falhas de generalização. Por fim, a implementação totalmente vetorizada em JAX permite até 33 milhões de passos de ambiente por segundo numa única GPU, possibilitando varreduras rápidas e reproduzíveis sobre fatores visuais. Código: 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/.