KAGE-Bench: 강화학습을 위한 고속의 알려진 축(Known-Axis) 기반 시각적 일반화 평가

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

픽셀 기반 강화학습 에이전트는 잠재 역학(latent dynamics)과 보상이 변하지 않는 상황에서도 순수 시각적 분포 변화(distribution shift)가 발생하면 종종 실패하지만, 기존 벤치마크들은 여러 변화 요인을 복잡하게 얽어놓아 체계적인 분석을 어렵게 합니다. 본 논문에서는 근본적인 제어 문제는 고정한 채 관측 과정을 독립적으로 제어 가능한 시각적 축으로 분해하는 JAX 네이티브 2D 플랫폼 환경인 KAGE-Env를 소개합니다. 구조적으로 볼 때, 시각적 축의 변화는 오직 픽셀 정책(pixel policy)에 의해 유도된 상태 조건부 행동 분포를 통해서만 성능에 영향을 미치므로, 시각적 일반화를 위한 명확한 추상화를 제공합니다. 이 환경을 기반으로 우리는 개별적인 시각적 변화를 분리하여 평가하는 34개의 학습-평가 구성 쌍을 포함한 6개의 알려진 축 제품군(suite)으로 이루어진 벤치마크인 KAGE-Bench를 제안합니다. 표준 PPO-CNN 베이스라인을 사용한 실험에서 축에 따라 뚜렷한 실패 양상이 관찰되었습니다. 배경 및 광도 변화는 종종 성공률을 급격히 떨어뜨리는 반면, 에이전트의 외형 변화는 비교적 영향이 적었습니다. 또한 일부 변화에서는 전진 움직임은 유지되지만 과제 완수는 실패하는 경우가 있어, 단순한 보상(return) 지표만으로는 일반화 실패를 파악하기 어려울 수 있음을 보여주었습니다. 마지막으로, 완전히 벡터화된 JAX 구현은 단일 GPU에서 초당 최대 3,300만 환경 스텝을 처리할 수 있어, 시각적 요인에 대한 빠르고 재현 가능한 실험(sweep)을 가능하게 합니다. 코드: https://avanturist322.github.io/KAGEBench/.

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