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