🤖 machine learning

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

本文介绍了 KAGE-Bench，这是一个基于 KAGE-Env 平台构建的高速 JAX 原生基准测试，它系统地隔离并评估了特定视觉分布偏移对强化学习智能体的影响，揭示了背景和光度变化往往会导致灾难性失败，而智能体外观的变化则较为温和。

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

发布于 2026-07-02

📖 1 分钟阅读☕ 轻松阅读

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

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在教一个机器人玩像《超级马力欧》这样的视频游戏。你通过向它展示数千小时的游戏过程来训练它：背景是蓝天，地面是绿草，角色是一个红色的水管工。机器人学会了跳过水管并躲避敌人。

现在，想象你把同一个游戏关卡交给了这个机器人。游戏的机制（水管、重力、敌人）完全相同。但天空变成了漆黑一片，地面变成了霓虹粉色，而水管工换上了一套小丑服。

在很多情况下，机器人会立即崩溃。它不再懂得如何玩这个游戏了，尽管“规则”并未改变。它太过于依赖游戏的“外观”，而不是学习游戏的“逻辑”。

这篇论文介绍了一个名为 KAGE-Bench 的新工具，专门研究为什么会发生这种情况以及如何解决它。以下是使用简单类比对该论文进行的拆解：

1. 问题所在：“分心的学生”

目前的 AI 智能体就像是那些只会死记硬背特定考试答案，而不是学习学科知识的学生。如果你改变了试卷上的字体或墨水的颜色，即使题目本身没变，学生也会陷入恐慌并考试不及格。

以前的测试试图衡量这一点，但它们非常混乱。它们会同时改变字体、纸张颜色以及题目的难度。这让人无法判断学生不及格是因为字体问题，还是因为数学题变难了。

2. 解决方案：“受控实验室”（KAGE-Env）

作者构建了一个新的视频游戏环境，称为 KAGE-Env。你可以把它看作是一个超级快速的数字实验室。

“旋钮”类比： 想象一个拥有 93 个不同旋钮的控制面板。每个旋钮都控制着一个特定的视觉元素：背景颜色、亮度、玩家角色的形状，或者是漂浮的干扰物。
神奇之处： 研究人员可以一次只转动一个旋钮。他们可以在保持游戏物理特性、奖励和规则完全不变的情况下，将背景从黑色变为白色。
速度： 这个实验室运行在一种名为 JAX 的特殊技术上，速度极快。它在单个图形卡上可以实现 每秒 3300 万步游戏步骤。为了让你有个直观的概念，这就像是在同时运行一百万个平行宇宙的游戏。这使得他们能够在过去测试几个变化所需的时间内，测试成千上万种视觉变化。

3. 基准测试：“已知轴向”测试（KAGE-Bench）

利用这个快速实验室，他们创建了一个标准化的测试，称为 KAGE-Bench。他们并没有向 AI 投掷随机的变化，而是创建了 34 个特定的“挑战”。

每个挑战都隔离了一种视觉变化（即一个“轴”）。例如：

背景测试： 在黑色背景下训练，在充满噪点和静态干扰的背景下进行测试。
滤镜测试： 使用正常颜色进行训练，在测试时使用会让一切变绿的“色调偏移”。
干扰物测试： 在清晰的屏幕上训练，在有看起来完全像玩家的漂浮形状干扰下进行测试。

4. 他们的发现：“脆弱”的 AI

他们将一种标准的 AI（称为 PPO-CNN）针对这些挑战进行了测试，并发现了一些令人惊讶的事情：

有些变化是致命的： 如果你改变了背景或添加了光照滤镜（比如让屏幕看起来像是在水下），AI 的表现会崩塌。它的胜率会从 90% 骤降到几乎为 0%。
有些变化是无害的： 出人意料的是，改变玩家角色的外观（比如把小丑皮肤换成骷髅皮肤）并没有对 AI 造成太大影响。它仍然可以很好地玩游戏。
“向前移动”的陷阱： 有时候，即使 AI 未能完成关卡，它也会继续向前移动（行走）。如果你只观察“它走了多远”，你会认为它表现得还不错。但如果你观察“它是否完成了关卡？”，你会发现它实际上已经彻底失败了。这表明简单的评分指标可能会掩盖严重的问题。

5. 为什么这很重要

论文指出，为了构建能在现实世界中运行的机器人或自动驾驶汽车，我们需要准确知道究竟是哪些视觉变化破坏了我们的 AI。

旧方法： “我们的机器人在雨中失败了。也许是因为雨？也许是因为泥泞？也许是因为光线？”（变量太多）。
KAGE-Bench 方法： “我们知道我们的机器人特别是在光照变为‘低对比度’时会失败，但它能很好地应对‘雨天’。”（精准诊断）。

总结

作者构建了一个超级快速的数字视频游戏实验室，他们可以像音响工程师调整均衡器旋钮一样微调视觉效果。他们利用这个实验室证明了目前的 AI 是非常脆弱的：它可以应对更换角色皮肤，但如果背景颜色改变，它就会崩溃。他们的目标是为研究人员提供一种清晰、快速的方法来寻找这些弱点，从而构建出真正具有鲁棒性（稳健性）的 AI，而不仅仅是靠运气表现良好的 AI。

技术摘要：KAGE-Bench

问题陈述

在基于高维像素观测的强化学习（RL）训练中，智能体在面对纯粹的视觉分布偏移时表现得极其脆弱，即使潜在的隐状态动力学、转移核（transition kernels）和奖励函数保持不变。虽然智能体可能在分布内收敛并表现良好，但在面对外观、光照或视觉噪声的变化时，往往会发生灾难性的失败。诊断这些失败的主要障碍在于缺乏能够隔离特定视觉偏移来源的评估基准。现有的基准测试（例如 Procgen、Distracting Control Suite）通常将多种来源的变动（如背景外观、几何结构、动力学和干扰物）纠缠在一起，使得无法将性能差距归因于特定的视觉因素。此外，许多基于像素的 RL 环境计算成本高昂，阻碍了大规模消融研究和系统的假设检验。

方法论

作者引入了 KAGE-Env 和 KAGE-Bench，通过“已知轴”（known-axis）的设计理念来解决这些局限性。

KAGE-Env：一个因子化的视觉平台跳跃游戏

KAGE-Env 是一个 JAX 原生的 2D 平台跳跃环境，旨在将观测过程分解为可独立控制的视觉轴，同时保持潜在控制问题固定不变。

架构： 该环境完全使用 JAX 实现，支持端到端 JIT 编译和向量化执行（使用 vmap 和 lax.scan）。这实现了大规模并行化，在单个 GPU 上通过 $2^{16}$ 个并行环境，最高可实现每秒 3300 万个环境步。
因子化： 观测核 $O_\xi(\cdot|s)$ 由视觉配置 $\xi$ 参数化。至关重要的是，转移核 $P$ 和奖励函数 $r$ 对 $\xi$ 是不变的。这确保了训练与评估之间的任何性能差异，仅源于像素策略 $\pi(a|o)$ 对同一潜在状态 $s$ 的不同渲染方式的响应。
配置： 环境通过单个 YAML 文件进行控制，暴露了 93 个显式可控参数，涵盖背景图像/颜色、智能体外观/精灵（sprites）、非玩家角色（NPCs）、干扰物、光度滤波器（亮度、色调、噪声）、光照效果和布局等类别。

KAGE-Bench：基准测试协议

基于 KAGE-Env，KAGE-Bench 定义了一个标准化的评估协议，包含六个“已知轴”套件和 34 个“训练-评估”配置对。

隔离： 在每一对配置中，训练配置（ $\xi_{train}$ ）和评估配置（ $\xi_{eval}$ ）仅在指定的渲染参数子集（一个特定的轴）上存在差异，而所有其他参数及底层动力学保持一致。
理论基础： 论文形式化地指出，在这些条件下，视觉泛化等价于诱导状态策略 $\pi_\xi$ 的偏移。通过对观测核进行边缘化处理，像素策略 $\pi(a|o)$ 诱导出一个状态条件动作分布 $\pi_\xi(a|s)$ 。该基准测试证明，性能差距正是这些在相同潜在 MDP 上运行的诱导策略之间的差异。
指标： 除了标准的单回合回报（episodic return）外，该基准测试还报告了轨迹级指标：移动距离、归一化目标进度和二元任务成功率。这种区分至关重要，因为仅靠回报可能会掩盖那些智能体虽有前进动作但未能完成任务的失败情况。

核心贡献

KAGE-Env： 一个高吞吐量的 JAX 原生 RL 环境，具有 93 个可控视觉参数，能够在单个 GPU 上实现高达 33M 步/秒的向量化模拟。
KAGE-Bench： 一个通过构造手段隔离视觉分布偏移的基准测试协议，包含六个套件（智能体、背景、干扰物、效果、滤波器、布局）和 34 个特定的训练-评估对。
经验诊断： 使用标准的 PPO-CNN 基线对视觉泛化进行系统性评估，量化了不同视觉轴如何影响性能，并识别出了会导致严重退化的特定类别的偏移。

结果

使用从单帧 RGB 图像训练的标准 PPO-CNN 基线，作者观察到了强烈的轴相关失效现象：

严重退化： 滤波器（例如色调偏移、饱和度、亮度）和效果（例如光照强度、数量）导致了最显著的性能崩溃。成功率分别下降了 86.8% 和 80.5%，且智能体通常仍能保持一定的向前运动。
背景偏移： 背景图像或颜色的变化导致成功率下降 53.3%，移动距离下降 30.5%，表明背景偏移会同时损害运动能力和任务完成能力。
良性偏移： 智能体外观的偏移（例如将精灵从骷髅变为小丑）相对而言是良性的，仅导致成功率下降 21.1%。
“运动 vs. 完成”差距： 若干偏移（特别是干扰物和布局）保留了向前运动（距离差距仅约 3-4%），但导致了任务完成率的大幅下降（成功率差距达 30-60%）。这凸显了基于回报的指标可能会掩盖泛化失败——即智能体虽然在移动，但未能解决任务。
鲁棒性基线： 作者测试了几种鲁棒性方法（带有 cutout-color 的 RAD-PPO、crop-resize、CURL、DrQ、SVEA）。虽然带有 cutout-color 的 RAD-PPO 在困难配置下显示出一些改进，但依赖几何增强的方法（crop-resize）未能收敛，这可能是因为它们破坏了该平台跳跃游戏中精确跳跃时机和障碍物规避所需的绝对空间位置。

意义与主张

论文声称 KAGE-Bench 提供了一个快速、可复现且具诊断性的框架，用于评估受控视觉变化下的基于像素的 RL。

归因： 通过因子化观测过程，该基准测试能够实现对泛化失败原因的精确归因，消除了其他基准测试中存在的纠缠动力学或任务结构变化的混淆效应。
可扩展性： JAX 原生的实现解决了视觉泛化评估中的计算瓶颈，使得对视觉参数进行详尽的搜索变得可行。
诊断效用： 结果表明，视觉泛化的难度并非均匀分布；它在不同轴之间存在显著差异。该基准测试作为一个工具，可以识别哪些特定的视觉因素（例如光度偏移对比智能体外观）对于给定策略架构学习不变性最为关键。

作者将 KAGE-Bench 定位为不是一个涵盖所有现实世界 RL 挑战的全面代理，而是一个受控的诊断工具，旨在支持对视觉鲁棒性的更系统分析，并指导开发能够处理特定“已知轴”视觉偏移的算法。