KAGE-Bench : Évaluation Rapide de la Généralisation Visuelle à Axe Connu pour l'Apprentissage par Renforcement

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

Les agents d'apprentissage par renforcement basés sur pixels échouent souvent face à un simple changement de distribution visuelle, même lorsque la dynamique latente et les récompenses restent inchangées. Cependant, les benchmarks existants entremêlent plusieurs sources de variation et entravent une analyse systématique. Nous présentons KAGE-Env, une plateforme 2D native JAX qui factorise le processus d'observation selon des axes visuels indépendamment contrôlables, tout en maintenant inchangé le problème de contrôle sous-jacent. Par construction, la variation d'un axe visuel n'affecte les performances qu'à travers la distribution d'actions conditionnelle à l'état induite par une politique pixel, offrant une abstraction épurée pour l'étude de la généralisation visuelle. Sur cette base, nous définissons KAGE-Bench, un benchmark de six suites à axes connus comprenant 34 paires de configurations entraînement-évaluation qui isolent des changements visuels individuels. En utilisant un modèle de référence PPO-CNN standard, nous observons des échecs prononcés dépendants de l'axe : les changements d'arrière-plan et photométriques annulent souvent toute réussite, tandis que les modifications d'apparence de l'agent sont relativement bénignes. Certaines variations préservent le mouvement vers l'avant tout en empêchant l'accomplissement de la tâche, montrant que le seul retour cumulé peut masquer des échecs de généralisation. Enfin, l'implémentation entièrement vectorisée dans JAX permet d'atteindre jusqu'à 33 millions d'étapes environnementales par seconde sur un seul GPU, autorisant des explorations rapides et reproductibles des facteurs visuels. 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/.