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

arXiv: https://arxiv.org/abs/2601.14232 · HTML: https://mlsanity.com/paper/2601.14232/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in RL for robot locomotion and manipulation, one of 15 directions in Deep reinforcement learning methods.
- Builds on 39 of its 49 references that are in the map.
- First posted on arXiv 20 Jan 2026.

## Placed in

- Deep reinforcement learning methods / [RL for robot locomotion and manipulation](https://mlsanity.com/deeprl/rl-for-robot-locomotion-and-manipulation/)

## Builds on (most cited first, 30 of 39)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [Gibson Env: Real-World Perception for Embodied Agents](https://mlsanity.com/paper/1808.10654/) (2018)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Quantifying Generalization in Reinforcement Learning](https://mlsanity.com/paper/1812.02341/) (2018)
- [ManiSkill3: GPU Parallelized Robotics Simulation and Rendering for Generalizable Embodied AI](https://mlsanity.com/paper/2410.00425/) (2024)
- [Generalization in Reinforcement Learning by Soft Data Augmentation](https://mlsanity.com/paper/2011.13389/) (2020)
- [Phasic Policy Gradient](https://mlsanity.com/paper/2009.04416/) (2020)
- [Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation](https://mlsanity.com/paper/2107.00644/) (2021)
- [The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixels](https://mlsanity.com/paper/2101.02722/) (2021)
- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [Decoupling Value and Policy for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2102.10330/) (2021)
- [Improving Generalization in Reinforcement Learning with Mixture Regularization](https://mlsanity.com/paper/2010.10814/) (2020)
- [MuJoCo Playground](https://mlsanity.com/paper/2502.08844/) (2025)
- [Automatic Data Augmentation for Generalization in Deep Reinforcement Learning](https://mlsanity.com/paper/2006.12862/) (2020)
- [Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2306.09884/) (2023)
- [XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2312.12044/) (2023)
- [Obstacle Tower: A Generalization Challenge in Vision, Control, and Planning](https://mlsanity.com/paper/1902.01378/) (2019)
- [Memory, Benchmark & Robots: A Benchmark for Solving Complex Tasks with Reinforcement Learning](https://mlsanity.com/paper/2502.10550/) (2025)
- [Don't Blind Your VLA: Aligning Visual Representations for OOD Generalization](https://mlsanity.com/paper/2510.25616/) (2025)
- [Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning](https://mlsanity.com/paper/2209.09203/) (2022)
- [Stabilizing Off-Policy Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2207.00986/) (2022)
- [RL-ViGen: A Reinforcement Learning Benchmark for Visual Generalization](https://mlsanity.com/paper/2307.10224/) (2023)
- [Cross-Trajectory Representation Learning for Zero-Shot Generalization in RL](https://mlsanity.com/paper/2106.02193/) (2021)
- [Explore to Generalize in Zero-Shot RL](https://mlsanity.com/paper/2306.03072/) (2023)
- [Object-Centric Learning with Slot Mixture Module](https://mlsanity.com/paper/2311.04640/) (2023)
- [WarpDrive: Extremely Fast End-to-End Deep Multi-Agent Reinforcement Learning on a GPU](https://mlsanity.com/paper/2108.13976/) (2021)
- [Local Feature Swapping for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2204.06355/) (2022)

## Abstract (arXiv)

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/.