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

📄 arXiv: 2601.14232v1 📥 PDF

作者: Egor Cherepanov, Daniil Zelezetsky, Alexey K. Kovalev, Aleksandr I. Panov

分类: cs.LG, cs.AI, cs.CV

发布日期: 2026-01-20

备注: 38 pages, 44 figures, 3 tables

🔗 代码/项目: PROJECT_PAGE

💡 一句话要点

提出KAGE-Bench，用于快速评估强化学习中已知轴视觉泛化能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 强化学习 视觉泛化 分布偏移 基准测试 JAX PPO CNN

📋 核心要点

  1. 现有强化学习基准测试难以分离视觉偏移的来源，阻碍了对视觉泛化能力的系统分析。
  2. 论文提出KAGE-Env和KAGE-Bench，通过可控的视觉轴分解观察过程，从而隔离和评估视觉泛化能力。
  3. 实验表明，不同视觉轴的偏移对性能影响不同，背景和光度偏移影响较大，而智能体外观偏移影响较小。

📝 摘要（中文）

基于像素的强化学习智能体在纯视觉分布偏移下通常会失效，即使潜在动力学和奖励保持不变。然而，现有的基准测试混淆了多种偏移源，阻碍了系统分析。我们引入了KAGE-Env，一个JAX原生的2D平台游戏环境，它将观察过程分解为独立可控的视觉轴，同时保持底层控制问题不变。通过这种构造，改变一个视觉轴仅通过像素策略的诱导状态条件动作分布来影响性能，为视觉泛化提供了一个清晰的抽象。在此环境的基础上，我们定义了KAGE-Bench，一个包含六个已知轴套件的基准测试，包括34个训练-评估配置对，用于隔离单个视觉偏移。使用标准的PPO-CNN基线，我们观察到强烈的轴依赖性失败，其中背景和光度偏移经常导致成功崩溃，而智能体外观偏移相对良性。一些偏移保留了前进运动，但破坏了任务完成，表明仅凭回报可能会掩盖泛化失败。最后，完全矢量化的JAX实现能够在单个GPU上实现高达3300万环境步/秒的速度，从而能够快速且可重复地扫描视觉因素。

🔬 方法详解

问题定义：基于像素的强化学习智能体在面对视觉分布偏移时表现不佳，即使潜在的动力学和奖励函数没有改变。现有的基准测试通常将多种视觉偏移因素混合在一起，难以对智能体的视觉泛化能力进行细致的分析和诊断。因此，需要一个能够解耦不同视觉因素影响的基准测试环境，以便更好地理解和改进智能体的视觉泛化能力。

核心思路：论文的核心思路是将观察过程分解为多个独立可控的视觉轴。通过控制每个视觉轴的变化，可以单独评估其对智能体性能的影响。这种分解使得研究人员能够更清晰地理解哪些视觉因素对泛化能力影响最大，并针对性地改进算法。

技术框架：KAGE-Bench建立在KAGE-Env之上，KAGE-Env是一个JAX原生的2D平台游戏环境。该环境将观察过程分解为多个独立可控的视觉轴，例如背景、光照、智能体外观等。KAGE-Bench包含六个已知轴套件，每个套件包含多个训练-评估配置对，用于隔离单个视觉偏移。研究人员可以使用这些配置对来评估智能体在不同视觉偏移下的泛化能力。

关键创新：KAGE-Bench的关键创新在于其对观察过程的分解和可控性。通过将观察过程分解为多个独立可控的视觉轴，KAGE-Bench能够提供一个更清晰和可控的实验环境，用于研究智能体的视觉泛化能力。与现有的基准测试相比，KAGE-Bench能够更好地隔离和评估单个视觉因素的影响。

关键设计：KAGE-Env使用JAX实现，具有高度的并行性和可扩展性。这使得研究人员能够在单个GPU上快速地进行大量的实验。KAGE-Bench使用标准的PPO-CNN作为基线算法，并提供了详细的实验结果和分析。论文还讨论了不同视觉轴偏移对智能体行为的影响，例如，一些偏移会保留前进运动，但破坏任务完成。

📊 实验亮点

实验结果表明，不同视觉轴的偏移对智能体性能的影响差异显著。背景和光度偏移通常会导致性能大幅下降，而智能体外观偏移的影响相对较小。使用JAX实现使得KAGE-Bench在单个GPU上能够达到3300万环境步/秒的速度，极大地提高了实验效率。

🎯 应用场景

该研究成果可应用于提升机器人和智能体在复杂视觉环境中的鲁棒性和泛化能力。例如，在自动驾驶领域，可以利用KAGE-Bench评估和改进自动驾驶系统在不同光照、天气和场景下的性能。在游戏AI领域，可以用于训练更具适应性的游戏智能体。

📄 摘要（原文）

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