KAGE-Bench: Evaluación Rápida de Generalización Visual con Ejes Conocidos para Aprendizaje por Refuerzo

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

Resumen

Los agentes de aprendizaje por refuerzo basados en píxeles a menudo fallan ante cambios puramente visuales en la distribución, incluso cuando la dinámica latente y las recompensas permanecen inalteradas. Sin embargo, los puntos de referencia existentes entrelazan múltiples fuentes de variación y dificultan un análisis sistemático. Presentamos KAGE-Env, una plataforma 2D nativa de JAX que factoriza el proceso de observación en ejes visuales controlables de forma independiente, manteniendo fijo el problema de control subyacente. Por diseño, la variación de un eje visual afecta al rendimiento únicamente a través de la distribución de acciones condicionada al estado inducida en una política de píxeles, proporcionando una abstracción limpia para la generalización visual. Sobre este entorno, definimos KAGE-Bench, un benchmark de seis suites de ejes conocidos que comprende 34 pares de configuraciones de entrenamiento-evaluación que aíslan cambios visuales individuales. Utilizando un baseline estándar PPO-CNN, observamos fallos importantes dependientes del eje: los cambios en el fondo y las variaciones fotométricas suelen colapsar el éxito, mientras que las alteraciones en la apariencia del agente son comparativamente benignas. Varios cambios preservan el movimiento hacia adelante pero interrumpen la finalización de la tarea, demostrando que la recompensa por sí sola puede ocultar fallos de generalización. Finalmente, la implementación completamente vectorizada en JAX permite hasta 33 millones de pasos de entorno por segundo en una sola GPU, posibilitando barridos rápidos y reproducibles sobre factores visuales. 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/.