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

초록

픽셀 기반 강화 학습 에이전트는 잠재 역학과 보상이 변경되지 않은 경우에도 순수 시각적 분포 변화 하에서 종종 실패하지만, 기존 벤치마크는 여러 변화 원인을 복합적으로 다루어 체계적 분석을 어렵게 합니다. 본 연구에서는 기본 제어 문제는 고정한 채 관찰 과정을 독립적으로 제어 가능한 시각적 축으로 분해하는 JAX 기반 2D 플랫폼 게임 환경인 KAGE-Env를 소개합니다. 설계 상 시각적 축을 변경하는 것은 픽셀 정책의 상태-조건부 행동 분포를 통해 성능에만 영향을 미쳐 시각적 일반화를 위한 명확한 추상화를 제공합니다. 이 환경을 기반으로 개별 시각적 변화를 분리한 34개의 학습-평가 구성 쌍으로 이루어진 6개의 알려진 축 세트로 구성된 벤치마크인 KAGE-Bench를 정의합니다. 표준 PPO-CNN 기준 모델을 사용하여 배경 및 광도 측정 변화가 성공률을 급격히 저하시키는 등 축에 따라 현저한 실패 패턴을 관찰한 반면, 에이전트 외형 변화는 상대적으로 영향을 적게 미쳤습니다. 여러 변화는 과제 완수를 방해하면서 전진 운동은 유지하여, 단순한 보상 합계만으로는 일반화 실패를 파악하기 어렵다는 점을 보여줍니다. 마지막으로, 완전히 벡터화된 JAX 구현을 통해 단일 GPU에서 초당 최대 3,300만 환경 단계를 처리할 수 있어 시각적 요소에 대한 빠르고 재현 가능한 탐색이 가능합니다. 코드: 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/.