KAGE-Bench：面向强化学习的快速已知轴视觉泛化评估基准

提出可控视觉轴解耦的2D平台环境，系统评估RL智能体的视觉泛化能力

前置知识

视觉泛化（Visual Generalization）

指基于像素观测训练的强化学习策略在观测过程发生变化（如背景、光照、外观改变）时能否保持性能的能力。核心挑战在于策略可能过度依赖视觉表面特征（如纹理、颜色）而非学习任务相关的不变表示，导致在新的视觉环境下性能崩溃。这种泛化失败与监督学习中的分布偏移类似，但在RL中更复杂，因为智能体需要在环境中交互学习。

这是本文研究的核心问题，理解视觉泛化的定义和挑战是理解KAGE-Bench动机和贡献的基础

部分可观马尔可夫决策过程（POMDP）

POMDP是标准MDP的扩展，智能体无法直接观测到环境的真实状态，只能获得部分观测。在视觉RL中，像素观测就是部分观测，因为原始图像不包含完整的物理状态信息。POMDP由状态空间S、动作空间A、转移核P、奖励函数r、观测空间Ω、观测核O等组成，其中观测核O决定了如何将隐藏状态映射到观测。

KAGE-Bench的形式化框架建立在POMDP之上，理解观测核O的概念对于理解'已知轴'设计至关重要

诱导状态策略（Induced State Policy）

给定视觉配置ξ、观测核Oξ和像素策略π，诱导状态策略πξ定义为对观测变量边际化后的状态-动作条件分布：πξ(a|s) = ∫π(a|o)Oξ(do|s)。这个概念表明，即使像素策略π固定，改变视觉配置ξ也会改变从隐状态到动作的有效映射，因为策略是在不同的渲染结果上评估的。

这是KAGE-Bench理论分析的核心工具，将视觉泛化问题归约为策略偏移问题，使得性能差异可归因于观测核的变化

JAX向量化执行

JAX是Google开发的高性能数值计算库，支持自动微分、JIT编译和硬件加速。通过vmap（向量化映射）和lax.scan等原语，可以将环境模拟并行化到GPU/TPU上，实现大规模批量仿真。相比传统的CPU-bound模拟器，这种设计可将吞吐量提升数个数量级。

KAGE-Env的高性能实现依赖于JAX的向量化能力，这使得大规模视觉参数扫描成为可能

研究动机

现有视觉泛化基准存在两个根本性问题。首先，训练集和评估集之间的视觉变化往往纠缠了多个来源——背景外观、几何结构、动态行为和干扰物同时改变，导致无法将性能下降归因于特定的视觉因素。例如，Procgen依赖程序生成，训练-测试差异通常反映了外观和场景组成的混合变化；Obstacle Tower和LevDoom等3D环境中多个因素共同变化，使得失败归因变得极其困难。其次，许多基于像素的RL环境计算成本高昂，限制了大规模消融实验和假设验证。这种纠缠的评估方式阻碍了对视觉泛化失败机制的系统性分析，也使得新方法的有效性难以被准确评估。

本文的目标是本文的目标是构建一个视觉泛化基准，其中分布偏移的来源通过设计被隔离——每个训练-评估配置对只在一个指定的视觉轴上发生变化，而潜在的控制问题（状态空间、转移动力学、奖励函数）保持完全不变。具体目标包括：（1）开发一个JAX原生的2D平台环境，支持93个可独立控制的视觉参数；（2）构建6个已知轴套件，包含34个训练-评估配置对；（3）实现高达每秒3300万环境步的吞吐量，支持单GPU上的大规模参数扫描。

与已有工作不同的是，KAGE-Bench的独特切入角度在于'已知轴'设计原则。与现有基准不同，本文不是简单地引入随机的视觉变化，而是将观测过程分解为独立可控的视觉轴（背景、角色外观、干扰物、滤镜、光照效果、布局），并构造只在单一轴上变化的训练-评估对。这种设计使得任何性能差异都可以明确归因于观测过程沿指定轴的变化，而非任务结构、动力学或奖励的改变。形式化地说，视觉泛化问题被归约到诱导状态策略的偏移问题，为失败诊断提供了理论基础。

核心方法

KAGE-Bench的核心思路可以用'控制变量法'来类比：就像科学实验中每次只改变一个自变量来观察其影响，KAGE-Bench每次只改变一个视觉轴来评估其对策略性能的影响。整个系统分为两层：底层是KAGE-Env，一个完全用JAX实现的2D平台游戏环境，支持93个可配置参数；上层是KAGE-Bench协议，定义如何选择和配对环境配置来构建已知轴评估套件。技术路线包括：（1）设计一个因子化的观测过程，将视觉渲染分解为独立可控的组件；（2）实现端到端的JIT编译和vmap向量化，支持大规模并行仿真；（3）构建训练-评估配置对，确保潜在控制问题完全一致。

KAGE-Bench最本质的创新是将视觉泛化问题形式化为'诱导状态策略偏移'问题。关键洞察是：对于一个固定的像素策略π，改变视觉配置ξ并不会改变底层的MDM（状态空间、转移核、奖励函数），而是改变了观测核Oξ，进而改变了从隐状态到动作的有效映射πξ。形式上，执行π在视觉POMDP Mξ中的性能等价于执行诱导策略πξ在潜在MDM M中的性能：J(π; Mξ) = J(πξ; M)。这个等价性表明，视觉泛化差距完全可以通过比较不同诱导状态策略在同一潜在MDM中的表现来解释，为故障诊断提供了清晰的归因框架。

方法步骤详情

KAGE-Bench的工作流程包括以下步骤：（1）环境配置：通过YAML文件指定视觉配置ξ，包括背景（128张图片）、角色外观（27个精灵皮肤或9种几何形状×21种颜色）、干扰物、滤镜（亮度、对比度、饱和度、色相等）、光照效果（点光源的强度、衰减、数量、颜色）和布局参数；（2）观测渲染：观测核Oξ将隐状态s映射为128×128的RGB图像，不同的ξ产生不同的渲染结果；（3）策略训练：使用标准PPO-CNN在训练配置ξ_train下训练策略；（4）评估：在评估配置ξ_eval下评估同一策略，记录距离、进度、成功率和回报；（5）泛化差距计算：∆(π) = J(π; D_train) - J(π; D_eval)，以及轨迹级指标（距离、进度、成功率）的差异。

技术新颖性

与现有基准相比，KAGE-Bench的技术新颖性体现在三个方面。首先，已知轴设计使得视觉变化可被精确控制和归因，而Procgen、RL-ViGen等基准的视觉变化是纠缠的。其次，完全向量化的JAX实现支持单GPU上216个并行环境，达到每秒3300万步的吞吐量，比传统CPU模拟器快数个数量级，使得大规模参数扫描变得可行。第三，形式化的理论分析（Theorem 4.2和Corollary 4.3）将视觉泛化归约为诱导状态策略偏移问题，为基准设计和结果解释提供了数学基础，这是现有基准所缺乏的。

**Figure 1**: KAGE-Env中的代表性观测，展示可控的已知轴视觉变化

实验结果

实验使用标准PPO-CNN基线（来自CleanRL库）在6个已知轴套件的34个配置对上进行评估，每个配置运行10个随机种子。核心发现包括：（1）泛化难度强烈依赖于视觉轴——按成功率退化排序，滤镜（∆SR=86.8%）和光照效果（80.5%）导致最严重失败，其次是布局（62.8%）和背景（53.3%），干扰物（30.9%）和角色外观（21.1%）相对温和；（2）背景变化同时损害运动和完成——平均而言，距离和进度下降30.5%，成功率从0.90降至0.42（∆SR=53.3%）；（3）光照和滤镜扰动主要破坏任务完成——距离退化相对温和（∆Dist=11.7%和20.8%），但成功率崩溃（0.83→0.11和0.82→0.16），表明运动和塑形奖励可以持续，但任务完成失败；（4）小的运动差距可能掩盖大的完成差距——干扰物和布局显示小的距离/进度差距（约3-4%）但显著的成功率下降（30.9%和62.8%）。具体配置层面，黑色→噪声背景的成功率差距达98.9%，色相偏移180°达98.8%，4个光源达95.5%，7个与智能体相同的干扰物达92.0%。

查看结构化数据

任务	指标	本文	基线	提升
背景泛化	成功率差距∆SR	53.3%下降	N/A	黑色→噪声背景：98.9%下降
滤镜泛化	成功率差距∆SR	86.8%下降	N/A	色相偏移180°：98.8%下降
光照效果泛化	成功率差距∆SR	80.5%下降	N/A	4个光源：95.5%下降
干扰物泛化	成功率差距∆SR	30.9%下降	N/A	7个相同干扰物：92.0%下降
角色外观泛化	成功率差距∆SR	21.1%下降	N/A	最温和的轴
布局泛化	成功率差距∆SR	62.8%下降	N/A	青色→红色布局

局限与改进

KAGE-Bench存在以下局限性：（1）任务简单性——KAGE-Env是一个简单的2D平台游戏，可能无法代表更复杂的视觉RL任务（如3D导航、精细操作），泛化行为的模式可能不适用于其他领域；（2）基线单一性——实验仅使用PPO-CNN基线，缺乏对其他算法（如数据增强、表示学习方法）的系统评估，无法验证基准是否能区分不同方法的泛化能力；（3）视觉轴正交性假设——虽然设计目标是让视觉轴独立变化，但实际中可能存在交互效应（如背景和光照同时变化），论文未充分探索轴间交互；（4）评估指标的局限性——虽然引入了距离、进度等轨迹级指标，但这些指标仍基于简化的水平位置，可能无法捕捉更细粒度的行为差异；（5）硬件依赖——JAX实现需要GPU支持，在CPU-only环境下的性能优势可能不明显。

独立分析的弱点

KAGE-Bench的主要弱点包括：（1）任务多样性不足——单一的2D平台任务可能无法充分代表视觉RL的多样性，建议扩展到更多任务类型（如导航、操作、驾驶）；（2）缺乏方法对比——论文仅报告PPO-CNN基线结果，缺乏对现有视觉泛化方法（如RAD、DrQ、数据增强、对比学习）的系统评估，使得基准的诊断价值未被充分展示；（3）轴间交互未探索——实际部署中多个视觉因素可能同时变化，建议增加多轴组合的评估套件；（4）缺乏难度递进——虽然有一些简单的'健全性检查'配置对，但缺乏系统性的难度递进设计，建议引入从易到难的梯度评估；（5）评估协议的统计严谨性——最大值统计（maximum-over-training）可能高估泛化能力，建议同时报告平均值和置信区间。

未来方向

基于KAGE-Bench的成果，未来研究可从以下方向展开：（1）系统评估现有视觉泛化方法——使用KAGE-Bench作为标准测试平台，对数据增强（RAD、DrQ）、表示学习（CURL、SPR）、正则化等方法进行轴特定的消融分析；（2）开发轴感知的泛化方法——针对不同视觉轴的特点设计专门的鲁棒性技术，如针对光照变化的亮度不变表示、针对背景变化的前景-背景分离；（3）扩展任务空间——将已知轴设计原则应用到更复杂的环境（如3D导航、机器人操作）；（4）探索视觉轴间的交互效应——研究多轴同时变化时的泛化行为，以及是否存在'轴间协同'或'轴间干扰'现象；（5）理论分析深化——进一步研究诱导状态策略偏移的性质，如哪些视觉变化对哪些类型的任务影响最大。

复现评估

KAGE-Bench在复现性方面具有显著优势：（1）完全开源——代码、环境、配置文件和训练脚本均在GitHub公开（https://avanturist322.github.io/KAGEBench/）；（2）硬件要求灵活——支持多种GPU（H100、A100、V100、T4）和CPU-only环境（Apple M3 Pro），单GPU即可运行大规模实验；（3）配置标准化——所有环境参数通过YAML文件指定，支持精确复现；（4）计算成本可控——使用JAX向量化实现，单GPU上可达到每秒3300万步的吞吐量，使得大规模参数扫描在合理时间内完成；（5）依赖简单——主要依赖JAX和CleanRL，环境配置标准化。总体而言，复现难度较低，适合研究团队快速验证和扩展。

KAGE-Bench：面向强化学习的快速已知轴视觉泛化评估基准 KAGE-Bench: Fast Known-Axis Visual Generalization Evaluation for Reinforcement Learning