KAGE-Bench: Быстрая оценка визуального обобщения с известной осью для обучения с подкреплением

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

January 20, 2026

Авторы: Egor Cherepanov, Daniil Zelezetsky, Alexey K. Kovalev, Aleksandr I. Panov

cs.AI

Аннотация

Пиксельные агенты обучения с подкреплением часто демонстрируют сбои при чисто визуальном сдвиге распределения, даже когда латентная динамика и функция вознаграждения остаются неизменными. Однако существующие бенчмарки смешивают несколько источников сдвига, что затрудняет систематический анализ. Мы представляем KAGE-Env — платформер на JAX, который факторизует процесс наблюдения в независимо управляемые визуальные оси при фиксированной базовой задаче управления. По построению, изменение визуальной оси влияет на производительность только через индуцированное условно-стационарное распределение действий пиксельной политики, обеспечивая чёткую абстракцию для исследования визуальной обобщающей способности. На основе этой среды мы определяем KAGE-Bench — бенчмарк из шести наборов с известными осями, включающий 34 пары конфигураций "обучение-оценка", которые изолируют отдельные визуальные сдвиги. Используя стандартный базовый алгоритм PPO-CNN, мы наблюдаем выраженные сбои, зависящие от оси: сдвиги фона и фотометрические параметры часто приводят к полному провалу, тогда как изменения внешнего вида агента оказываются сравнительно безопасными. Некоторые сдвиги сохраняют поступательное движение, но нарушают выполнение задачи, демонстрируя, что одна лишь совокупная награда может маскировать failures обобщения. Наконец, полностью векторизованная реализация на JAX позволяет достигать до 33 миллионов шагов среды в секунду на одном GPU, обеспечивая быстрое и воспроизводимое исследование визуальных факторов. Код: 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/.