mlsanityData to 7 Oct 2026

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

arXiv 2601.14232 ↗First posted 20 Jan 2026cs.LG cs.AI cs.CVcited by 0

Not on the map. Why?

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

Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 39 of its 49 references that are in the map. First posted on arXiv 20 Jan 2026.

Builds on

39 of 49 references.

References · 39

  1. 2017Proximal Policy Optimization Algorithmscited by 14,441
  2. 2020Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixelscited by 573
  3. 2019Leveraging Procedural Generation to Benchmark Reinforcement Learningcited by 478
  4. 2018Gibson Env: Real-World Perception for Embodied Agentscited by 477
  5. 2020Reinforcement Learning with Augmented Datacited by 452
  6. 2021Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulationcited by 375
  7. 2018Quantifying Generalization in Reinforcement Learningcited by 345
  8. 2024ManiSkill3: GPU Parallelized Robotics Simulation and Rendering for Generalizable Embodied AIcited by 216
  9. 2020Generalization in Reinforcement Learning by Soft Data Augmentationcited by 132
  10. 2020Phasic Policy Gradientcited by 121
  11. 2021Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentationcited by 101
  12. 2021The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixelscited by 97
27 more
  1. 2020Improving Generalization in Reinforcement Learning with Mixture Regularizationcited by 81
  2. 2021Decoupling Value and Policy for Generalization in Reinforcement Learningcited by 81
  3. 2024Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learningcited by 80
  4. 2020Automatic Data Augmentation for Generalization in Deep Reinforcement Learningcited by 71
  5. 2025MuJoCo Playgroundcited by 57
  6. 2019Obstacle Tower: A Generalization Challenge in Vision, Control, and Planningcited by 54
  7. 2023XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAXcited by 54
  8. 2023Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAXcited by 52
  9. 2025Memory, Benchmark & Robots: A Benchmark for Solving Complex Tasks with Reinforcement Learningcited by 34
  10. 2025Don't Blind Your VLA: Aligning Visual Representations for OOD Generalizationcited by 34
  11. 2022Stabilizing Off-Policy Deep Reinforcement Learning from Pixelscited by 32
  12. 2022Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learningcited by 32
  13. 2023RL-ViGen: A Reinforcement Learning Benchmark for Visual Generalizationcited by 27
  14. 2021Cross-Trajectory Representation Learning for Zero-Shot Generalization in RLcited by 20
  15. 2023Explore to Generalize in Zero-Shot RLcited by 20
  16. 2023Object-Centric Learning with Slot Mixture Modulecited by 11
  17. 2021WarpDrive: Extremely Fast End-to-End Deep Multi-Agent Reinforcement Learning on a GPUcited by 10
  18. 2022Local Feature Swapping for Generalization in Reinforcement Learningcited by 9
  19. 2025Object-Centric Latent Action Learningcited by 9
  20. 2025Augmented Reality for RObots (ARRO): Pointing Visuomotor Policies Towards Visual Robustnesscited by 7
  21. 2024Improving Generalization on the ProcGen Benchmark with Simple Architectural Changes and Scalecited by 6
  22. 2024DMC-VB: A Benchmark for Representation Learning for Control with Visual Distractorscited by 5
  23. 2025Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAXcited by 5
  24. 2022Bootstrap State Representation using Style Transfer for Better Generalization in Deep Reinforcement Learningcited by 4
  25. 2024Make the Pertinent Salient: Task-Relevant Reconstruction for Visual Control with Distractionscited by 4
  26. 2024Symbolic Disentangled Representations for Imagescited by 2
  27. 2025PERELMAN: Pipeline for scientific literature meta-analysis. Technical reportcited by 1

Built on by

Not cited yet.