🤖 machine learning

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

This paper introduces KAGE-Bench, a high-speed JAX-native benchmark built on the KAGE-Env platform that systematically isolates and evaluates the impact of specific visual distribution shifts on reinforcement learning agents, revealing that background and photometric changes often cause catastrophic failure while agent-appearance shifts are more benign.

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

Published 2026-07-02

📖 5 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to play a video game like Super Mario. You train the robot by showing it thousands of hours of gameplay where the background is a blue sky, the ground is green grass, and the character is a red plumber. The robot learns to jump over pipes and avoid enemies.

Now, imagine you hand that same robot a new level. The game mechanics (the pipes, the gravity, the enemies) are exactly the same. But the sky is now pitch black, the ground is neon pink, and the plumber is wearing a clown suit.

In many cases, the robot will crash immediately. It doesn't know how to play the game anymore, even though the "rules" haven't changed. It got too used to the look of the game rather than the logic of it.

This paper introduces a new tool called KAGE-Bench to study exactly why this happens and how to fix it. Here is a breakdown of the paper using simple analogies:

1. The Problem: The "Distracted Student"

Current AI agents are like students who memorize the answers to a specific test rather than learning the subject. If you change the font on the test paper or the color of the ink, the student panics and fails, even if the questions are the same.

Previous tests tried to measure this, but they were messy. They would change the font, the paper color, and the difficulty of the questions all at once. It was impossible to tell if the student failed because of the font or because the math got harder.

2. The Solution: The "Controlled Lab" (KAGE-Env)

The authors built a new video game environment called KAGE-Env. Think of this as a super-fast, digital laboratory.

The "Knob" Analogy: Imagine a control panel with 93 different knobs. Each knob controls one specific visual thing: the background color, the brightness, the shape of the player character, or the presence of floating distractors.
The Magic: The researchers can turn only one knob at a time. They can change the background from black to white while keeping the game physics, the rewards, and the rules exactly the same.
The Speed: This lab runs on a special technology (JAX) that is incredibly fast. It can run 33 million game steps per second on a single graphics card. To put that in perspective, it's like running a million parallel universes of this game simultaneously. This allows them to test thousands of visual changes in the time it used to take to test just a few.

3. The Benchmark: The "Known-Axis" Test (KAGE-Bench)

Using this fast lab, they created a standardized test called KAGE-Bench. Instead of throwing random changes at the AI, they created 34 specific "challenges."

Each challenge isolates one visual change (an "axis"). For example:

The Background Test: Train on a black background, test on a noisy, static-filled background.
The Filter Test: Train with normal colors, test with a "hue shift" that turns everything green.
The Distractor Test: Train with a clear screen, test with floating shapes that look exactly like the player.

4. What They Discovered: The "Brittle" AI

They tested a standard AI (called PPO-CNN) against these challenges and found some surprising things:

Some changes are deadly: If you change the background or add lighting filters (like making the screen look like it's underwater), the AI's performance crashes. It goes from winning 90% of the time to winning almost 0%.
Some changes are harmless: Surprisingly, changing the appearance of the player character (like swapping a clown skin for a skeleton skin) didn't hurt the AI much. It could still play the game fine.
The "Moving Forward" Trap: Sometimes, the AI kept moving forward (walking) even when it failed to finish the level. If you only looked at "how far it walked," you'd think it was doing okay. But if you looked at "did it finish the level?", it was actually failing completely. This shows that simple scores can hide big problems.

5. Why This Matters

The paper argues that to build robots or self-driving cars that work in the real world, we need to know exactly which visual changes break our AI.

Old way: "Our robot failed in the rain. Maybe it's the rain? Maybe it's the mud? Maybe it's the light?" (Too many variables).
KAGE-Bench way: "We know our robot fails specifically when the lighting changes to 'low contrast,' but it handles 'rain' just fine." (Precise diagnosis).

Summary

The authors built a super-fast, digital video game lab where they can tweak the visuals like a sound engineer adjusting equalizer knobs. They used this to prove that AI is currently very fragile: it can handle a new character skin but will crash if the background changes color. Their goal is to give researchers a clear, fast way to find these weaknesses so they can build AI that is truly robust, not just lucky.

Technical Summary: KAGE-Bench

Problem Statement

Reinforcement learning (RL) agents trained on high-dimensional pixel observations are notoriously brittle when faced with purely visual distribution shifts, even when the underlying latent dynamics, transition kernels, and reward functions remain unchanged. While agents may converge and perform well in-distribution, they often fail catastrophically under changes in appearance, lighting, or visual noise. A primary obstacle to diagnosing these failures is the lack of evaluation benchmarks that can isolate specific sources of visual shift. Existing benchmarks (e.g., Procgen, Distracting Control Suite) often entangle multiple sources of variation—such as background appearance, geometry, dynamics, and distractors—making it impossible to attribute performance gaps to specific visual factors. Furthermore, many pixel-based RL environments are computationally expensive, hindering large-scale ablation studies and systematic hypothesis testing.

Methodology

The authors introduce KAGE-Env and KAGE-Bench to address these limitations through a "known-axis" design philosophy.

KAGE-Env: A Factorized Visual Platformer

KAGE-Env is a JAX-native 2D platformer environment designed to factorize the observation process into independently controllable visual axes while keeping the latent control problem fixed.

Architecture: The environment is implemented entirely in JAX with end-to-end JIT compilation and vectorized execution (using vmap and lax.scan). This allows for massive parallelism, achieving up to 33 million environment steps per second on a single GPU with $2^{16}$ parallel environments.
Factorization: The observation kernel $O_\xi(\cdot|s)$ is parameterized by a visual configuration $\xi$ . Crucially, the transition kernel $P$ and reward function $r$ are invariant to $\xi$ . This ensures that any performance difference between training and evaluation arises solely from how a pixel-based policy $\pi(a|o)$ responds to different renderings of the same latent state $s$ .
Configuration: The environment is controlled via a single YAML file, exposing 93 explicitly controllable parameters across categories such as background imagery/color, agent appearance/sprites, non-player characters (NPCs), distractors, photometric filters (brightness, hue, noise), lighting effects, and layout.

KAGE-Bench: The Benchmark Protocol

Built upon KAGE-Env, KAGE-Bench defines a standardized evaluation protocol comprising six "known-axis" suites and 34 train-evaluation configuration pairs.

Isolation: In each pair, the training configuration ( $\xi_{train}$ ) and evaluation configuration ( $\xi_{eval}$ ) differ only in a designated subset of rendering parameters (one specific axis), while all other parameters and the underlying dynamics remain identical.
Theoretical Foundation: The paper formalizes that visual generalization under these conditions is equivalent to a shift in the induced state policy $\pi_\xi$ . By marginalizing the observation kernel, the pixel policy $\pi(a|o)$ induces a state-conditional action distribution $\pi_\xi(a|s)$ . The benchmark demonstrates that performance gaps are exactly differences between these induced policies acting on the same latent MDP.
Metrics: Beyond standard episodic return, the benchmark reports trajectory-level metrics: distance traveled, normalized progress toward the goal, and binary task success. This distinction is vital, as return alone can obscure failures where an agent makes forward progress but fails to complete the task.

Key Contributions

KAGE-Env: A high-throughput, JAX-native RL environment with 93 controllable visual parameters, capable of vectorized simulation at up to 33M steps/second on a single GPU.
KAGE-Bench: A benchmark protocol that isolates visual distribution shifts by construction, featuring six suites (Agent, Background, Distractors, Effects, Filters, Layout) and 34 specific train-evaluation pairs.
Empirical Diagnosis: A systematic evaluation of visual generalization using a standard PPO-CNN baseline, quantifying how different visual axes impact performance and identifying specific classes of shifts that induce severe degradation.

Results

Using a standard PPO-CNN baseline trained from single RGB frames, the authors observed strong axis-dependent failures:

Severe Degradation: Filters (e.g., hue shift, saturation, brightness) and Effects (e.g., lighting intensity, count) caused the most significant performance collapses. Success rates dropped by 86.8% and 80.5% respectively, often while the agent maintained some forward motion.
Background Shifts: Changes in background imagery or color led to a 53.3% drop in success rate and a 30.5% drop in distance traveled, indicating that background shifts impair both locomotion and task completion.
Benign Shifts: Agent appearance shifts (e.g., changing the sprite from a skeleton to a clown) were comparatively benign, causing only a 21.1% drop in success rate.
The "Motion vs. Completion" Gap: Several shifts (notably Distractors and Layout) preserved forward motion (small distance gaps of ~3-4%) but caused massive drops in task completion (Success Rate gaps of 30-60%). This highlights that return-based metrics alone can mask generalization failures where an agent moves but fails to solve the task.
Robustness Baselines: The authors tested several robustness methods (RAD-PPO with cutout-color, crop-resize, CURL, DrQ, SVEA). While RAD-PPO with cutout-color showed some improvement on hard configurations, methods relying on geometric augmentations (crop-resize) failed to converge, likely because they disrupted the absolute spatial positions required for precise jump timing and obstacle avoidance in this platformer setting.

Significance and Claims

The paper claims that KAGE-Bench provides a fast, reproducible, and diagnostic framework for evaluating pixel-based RL under controlled visual variation.

Attribution: By factorizing the observation process, the benchmark enables precise attribution of generalization failures to specific visual factors, removing the confounding effects of entangled dynamics or task structure changes found in other benchmarks.
Scalability: The JAX-native implementation addresses the computational bottleneck of visual generalization evaluation, making exhaustive sweeps over visual parameters feasible.
Diagnostic Utility: The results demonstrate that visual generalization difficulty is not uniform; it varies substantially across axes. The benchmark serves as a tool to identify which specific visual factors (e.g., photometric shifts vs. agent appearance) are most critical for a given policy architecture to learn invariances against.

The authors position KAGE-Bench not as a comprehensive proxy for all real-world RL challenges, but as a controlled diagnostic tool to support more systematic analysis of visual robustness and to guide the development of algorithms capable of handling specific, known-axis visual shifts.