KAGE-Bench: Schnelle Bewertung der visuellen Generalisierung mit bekannter Achse für Reinforcement Learning

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

January 20, 2026

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

cs.AI

papers.abstract

Pixelbasierte Reinforcement-Learning-Agenten versagen häufig unter rein visuellen Distributionsverschiebungen, selbst wenn die latenten Dynamiken und Belohnungen unverändert bleiben. Bestehende Benchmarks vermengen jedoch häufig mehrere Verschiebungsquellen und erschweren eine systematische Analyse. Wir stellen KAGE-Env vor, eine in JAX native 2D-Plattformer-Umgebung, die den Beobachtungsprozess in unabhängig steuerbare visuelle Achsen faktorisiert, während das zugrundeliegende Steuerungsproblem konstant bleibt. Durch diesen Aufbau beeinflusst die Variation einer visuellen Achse die Leistung ausschließlich über die induzierte zustandsbedingte Aktionsverteilung einer Pixel-Policy, was eine saubere Abstraktion für visuelle Generalisierung bietet. Aufbauend auf dieser Umgebung definieren wir KAGE-Bench, einen Benchmark mit sechs bekannten Achsen-Suites, die 34 Trainings-Evaluations-Konfigurationspaare umfassen und individuelle visuelle Verschiebungen isolieren. Unter Verwendung einer standardmäßigen PPO-CNN-Baseline beobachten wir starke achsenabhängige Fehler: Hintergrund- und photometrische Verschiebungen lassen den Erfolg häufig zusammenbrechen, während Verschiebungen in der Agenten-Erscheinung vergleichsweise harmlos sind. Einige Verschiebungen erhalten die Vorwärtsbewegung aufrecht, unterbrechen aber die Aufgabenabschlussfähigkeit, was zeigt, dass der Return allein Generalisierungsfehler verschleiern kann. Schließlich ermöglicht die vollständig vektorisierte JAX-Implementierung bis zu 33 Millionen Umgebungsschritte pro Sekunde auf einer einzelnen GPU, was schnelle und reproduzierbare Sweeps über visuelle Faktoren ermöglicht. Code: 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/.