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

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/

When we train a RL agent to play a video game by letting it learn directly from the pixels on screen, it can become very good at the game. But these agents are surprisingly fragile: change something purely cosmetic - the background scenery, the brightness of the image, or the look of the character - and the agent often falls apart, even though the game itself plays exactly the same way. We don't yet understand this well, partly because the test environments people use tend to change many things at once, making it impossible to tell which kind of visual change actually trips the agent up. To study this cleanly, we built KAGE-Env, a simple 2D platformer game in which we can switch each visual feature on or off independently - the background, the lighting, the character's appearance, and so on - while keeping the actual rules and goals of the game completely unchanged. On top of it we created KAGE-Bench, a set of carefully designed tests that isolate one visual change at a time, so we can measure exactly which kinds of changes hurt performance. Testing a standard RL agent, we found the failures depend strongly on what changes: swapping the background or altering colors and brightness often causes the agent to fail completely, while changing how the character looks barely matters. We also found that the agent sometimes keeps moving forward and looks like it's doing fine by one common measure, yet never actually finishes the level - a warning that the usual scorecards can hide real failures. Finally, our game runs extremely fast (tens of millions of steps per second on a single graphics card), making it cheap and easy for other researchers to run these experiments and reproduce our results. Website: https://avanturist322.github.io/KAGEBench/