KAGE-Bench: Snelle evaluatie van visuele generalisatie met bekende assen voor reinforcement learning

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

January 20, 2026

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

cs.AI

Samenvatting

Pixelgebaseerde reinforcement learning-agenten falen vaak onder puur visuele distributieverschuivingen, zelfs wanneer de latente dynamiek en beloningen ongewijzigd blijven, maar bestaande benchmarks verstrengelen meerdere bronnen van verschuiving en belemmeren systematische analyse. Wij introduceren KAGE-Env, een inheems JAX 2D-platformspel dat het observatieproces factoriseert in onafhankelijk bestuurbare visuele assen, terwijl het onderliggende controleprobleem ongewijzigd blijft. Door constructie beïnvloedt het variëren van een visuele as de prestaties uitsluitend via de geïnduceerde staat-voorwaardelijke actieverdeling van een pixelbeleid, wat een zuivere abstractie biedt voor visuele generalisatie. Voortbouwend op deze omgeving definiëren we KAGE-Bench, een benchmark van zes bekende-assen-suites bestaande uit 34 train-evaluatie configuratieparen die individuele visuele verschuivingen isoleren. Met een standaard PPO-CNN-basislijn observeren we sterke as-afhankelijke mislukkingen, waarbij achtergrond- en fotometrische verschuivingen de succesratio vaak doen instorten, terwijl verschuivingen in agent-uiterlijk relatief goedaardig zijn. Verschillende verschuivingen behouden voorwaartse beweging maar breken taakvoltooiing, wat aantoont dat beloning alleen generalisatiefouten kan verhullen. Ten slotte maakt de volledig gevectoriseerde JAX-implementatie tot 33M omgevingsstappen per seconde mogelijk op een enkele GPU, wat snelle en reproduceerbare sweeps over visuele factoren mogelijk maakt. 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/.