AXIOM: Learning to Play Games in Minutes withExpanding Object-Centric Models
AXIOM:借助扩展对象中心模型实现分钟级游戏策略学习
https://arxiv.org/pdf/2505.24784
https://github.com/VersesTech/axiom
摘要
当前的深度强化学习(DRL)在多个领域实现了最先进的性能,但在数据效率方面仍不如人类学习,因为人类能够利用关于物体及其相互作用的核心先验知识。主动推理(active inference)提供了一个有原则的框架,将感知信息与先验知识结合,用于学习世界模型,并量化自身信念和预测的不确定性。然而,主动推理模型通常为特定任务手工设计,缺乏DRL所具备的跨领域灵活性。为此,我们提出了一种新架构AXIOM,它整合了关于以物体为中心的动力学和交互的最简但具有表达能力的核心先验,以加速在低数据场景下的学习。AXIOM结合了贝叶斯方法的数据效率和可解释性,以及DRL的跨任务泛化能力。该模型将场景表示为物体的组合,其动力学建模为分段线性轨迹,捕捉稀疏的物体间交互。生成模型结构通过从单个事件中增长并学习混合模型,并通过贝叶斯模型简化定期优化,从而实现泛化。AXIOM在仅10,000次交互步骤内掌握多种游戏,参数数量远少于DRL方法,且无需基于梯度优化的计算开销。
1 引言
强化学习(RL)作为一种灵活的框架,在复杂任务中取得了显著成功。然而,现有方法存在多个缺陷:需要大量训练数据、依赖大容量经验回放缓冲区,并专注于最大化累积奖励而缺乏结构化探索。这与人类学习形成对比——人类依靠核心先验快速泛化到新任务。核心先验代表了塑造感知与学习的基本组织原则(超先验),是构建更复杂知识结构的基础。例如,这些先验使人类能直观理解“物体在无外力时沿平滑轨迹运动”,并有助于因果推理、掌握行为与结果之间的关系。将视觉场景分解为物体表示已被证明在样本效率、泛化性和鲁棒性方面具有潜力。这些问题天然适合贝叶斯智能体架构,如主动推理(active inference),它为整合先验知识提供了理论基础,支持持续适应而不遗忘旧知识。已有研究指出这种方法与人类认知过程高度一致,其中信念随新证据不断更新。尽管有这些理论优势,主动推理的应用多局限于小规模任务,依赖精心设计的先验,难以达到DRL在多样领域的广泛适用性。
为弥合这一差距,我们提出了一种新的主动推理架构,融合了最小但具表达性的关于物体及交互的核心先验。具体来说,我们提出了AXIOM(基于对象中心模型的主动扩展推理),其包含三个关键组件:(1) 高斯混合模型,将视觉输入解析为以物体为中心的表示,并自动扩展以容纳新物体;(2) 转移混合模型,发现运动原型(如下落、滑动、弹跳);(3) 多物体潜在特征上的稀疏关系混合模型,学习由物体状态、动作、奖励和动态模式共同驱动的因果相关交互。AXIOM的学习算法具有三重效率:第一,它采用变分贝叶斯更新逐帧顺序学习,无需经验回放缓冲或梯度计算,支持对数据分布变化的在线适应;第二,其混合结构可通过添加新成分解释新数据、合并冗余成分降低模型复杂度,实现快速结构学习;第三,通过维护参数后验分布,AXIOM可以在策略选择中引入信息寻求目标,从而进行考虑不确定性的探索。
为验证模型效果,我们引入Gameworld 10k基准测试,一套专为评估智能体在10,000次交互内玩不同像素游戏效率的新环境。许多现有RL基准(如Arcade Learning Environment或MuJoCo)强调长视野信用分配、复杂物理或视觉复杂性,往往掩盖了快速学习与泛化的真正挑战。因此,Gameworld 10k中的每个游戏均遵循类似模式:视觉场景中包含多个物体、一个可控得分玩家物体,以及遵循连续轨迹且交互稀疏的其他物体。我们设计了10款视觉元素简化的游戏(使用不同形状大小的单色精灵图),以聚焦于动态建模与控制机制的研究,而非过度复杂的物体分割模型。Gameworld环境还允许精确控制游戏特征与动力学,便于测试系统如何适应对游戏因果或视觉结构的稀疏干预(如物体形状与颜色)。在该基准上,我们的智能体在低数据环境下(10,000步)表现优于主流强化学习模型,且不依赖任何基于梯度的优化方法。虽然我们尚未将AXIOM部署到RL文献中常见的复杂控制任务规模,但我们的成果标志着朝向构建具备紧凑、可解释世界模型并在多领域实现快速决策智能体的重要进展。我们的主要贡献如下:
提出AXIOM:一种全新的以物体为中心的主动推理智能体,具备在线学习、可解释、样本高效、适应性强且计算成本低的特点。
为展示AXIOM的有效性,我们引入了一个可修改的新基准套件,专门针对具有物体与稀疏交互的环境中样本高效学习问题。
我们展示了这种无需梯度的方法在样本效率与绝对性能上均可超越最先进的深度学习方法,且在线学习方案对环境扰动具有鲁棒性。
2 方法
2.1 模型的增长与剪枝
快速结构学习。秉承快速结构学习 [23] 的理念,AXIOM 使用一种在线增长启发式方法动态扩展其全部四个混合模块(sMM、iMM、tMM、rMM):依次处理每一个新数据点,判断它是由现有成分最好地解释,还是应当创建一个新的成分,并随后更新所选成分的参数。
贝叶斯模型缩减(BMR)。每 ΔTBMR=500 帧,我们最多采样 2000 个已使用的 rMM 成分,使用祖先采样从模型生成的数据中计算它们之间的互期望对数似然,并贪婪地测试合并候选。若某一合并操作能减少在剩余变量采样数据条件下,关于奖励和下一 tMM 开关状态的多项分布的期望自由能,则接受该合并;否则回滚。BMR 使得 AXIOM 能够从单一事件中泛化动力学规律,例如通过合并多个单事件聚类,学习到当球撞击屏幕底部时会获得负奖励(见第 3 节,图 4a)。
2.2 规划
AXIOM 使用主动推理(active inference)进行规划 [33];它在不同的策略(动作序列)条件下展开对未来轨迹的预测,然后使用期望自由能(expected free energy)对这些策略进行推理,所选择的策略 是使期望自由能最小化的那个策略:
每时间步的期望效用在规划时通过已学习的模型和 slot 潜变量进行评估,并在整个规划视野中随时间步累加。期望信息增益(公式(10)右边第二项)是根据 rMM 的后验 Dirichlet 计数计算得出的,用于衡量采取当前考虑中的策略时,能在多大程度上获得关于 rMM 开关状态的信息。有关规划的更多细节见附录 A.11。
3 结果
为了评估 AXIOM,我们在 Gameworld 环境中将其与两种最先进的基于样本高效、像素输入的深度强化学习基线方法进行比较:BBF 和 DreamerV3。
基准测试环境
Gameworld 环境的设计目标是让人类学习者在几分钟内就能解决,从而确保学习过程不依赖于脆弱的探索机制或复杂的信用分配。该套件包括 10 个多样化的游戏,借助大型语言模型生成,灵感来自 ALE 和经典电子游戏,同时保持了轻量级和结构化的设计。
Gameworld 环境可在 https://github.com/VersesTech/gameworld 获取。图 2 展示了其中包含的游戏的多样性与视觉简洁性。为评估鲁棒性,Gameworld 10k 支持受控干预,如物体颜色或形状的变化,以测试智能体在表面领域变化下的泛化能力。
基线方法
BBF [34] 建立在 SR-SPR [35] 的基础上,代表了目前最样本高效的无模型方法之一。我们针对 Gameworld 10k 套件对其预处理进行了调整,将帧跳过替换为对连续两帧的最大池化;其余所有已发表的超参数保持不变。
其次,DreamerV3 [36] 是一种基于世界模型的智能体,在仅使用像素输入的游戏和控制任务中表现优异;我们使用已发布的设置,但将训练比设为 1024,批量大小为 16(有效训练比为 64:1)。
我们选择这两个基线是因为它们代表了从原始像素中进行样本高效学习的最先进水平。请注意,对于 BBF 和 DreamerV3,我们将图像帧分别缩放至 84×84 和 96×96 像素(遵循已发布实现),而 AXIOM 则直接操作 Gameworld 的完整分辨率 210×160 帧。
奖励表现
图 3 显示了在 Gameworld 10k 套件上,从第 0 步到第 10000 步的每步奖励的 1000 步移动平均值(均值 ± 1 标准差,共 10 个种子)。表 1 显示了 AXIOM、BBF 和 DreamerV3 在 10k 交互步结束时所达到的累积奖励。
AXIOM 在每一个 Gameworld 环境中都达到了高于或与 BBF 和 DreamerV3 相当的平均累积奖励。值得注意的是,AXIOM 不仅在多个游戏中达到了更高的峰值得分,而且收敛速度明显更快,通常在前 5000 步就已获得大部分最终奖励,而 BBF 和 DreamerV3 几乎需要完整的 10000 步才能达到类似效果。
对于那些在 10k 步时 BBF 和 Dreamer 表现接近随机的游戏,我们确认其性能最终确实有所提升,排除了这些游戏本身对这些架构来说本质上过于困难的可能性(详见附录 E.1)。
综合来看,这表明 AXIOM 的以物体为中心的世界模型,结合其快速的在线结构学习和推理算法,可以显著减少实现高性能所需的交互次数。
固定交互距离可以带来更高的累积奖励,因为智能体无需花费动作去学习该距离,但这样做要求为每个游戏单独调优这一交互距离。这说明了如何将关于特定领域的额外知识融入像 AXIOM 这样的贝叶斯模型中,以进一步提高样本效率。
引入公式 (10) 中的信息增益项可以使智能体在某些游戏中更快地获得奖励(例如 Bounce),但在其他游戏(例如 Gold)中反而会导致平均奖励增长更慢,因为它鼓励访问信息丰富但奖励为负的状态。
BMR 对需要空间泛化的游戏(如 Gold 和 Hunt)至关重要,但在 Cross 游戏中却会损害性能,因为早期合并聚类会削弱信息增益项并抑制探索。有关详细讨论请参见附录 E.2。
计算成本 表 2 比较了在单块 A100 GPU 上测得的模型大小以及每步训练时间(模型更新与规划)。
尽管由于使用了大量基于模型的 rollout,AXIOM 存在一定的规划开销,但其模型更新远比 BBF 更高效,在每样本的墙上时间(wall-clock time)方面具有更优的权衡。
AXIOM 的以物体为中心的扩展模型会根据环境复杂度收敛到一个足够的复杂度,而 BBF 和 DreamerV3 的模型大小则是固定的(且大得多)。
可解释性 与传统的深度强化学习方法不同,AXIOM 拥有一个结构化的、以物体为中心的模型,其潜在变量和参数可以用人类可读的方式直接解释(例如:形状、颜色、位置)。AXIOM 的转移混合模型还将复杂的轨迹分解为更简单的线性子序列。
图 4a 展示了在 Impact 游戏中,由 AXIOM 想象出的轨迹以及基于奖励条件的 rMM 聚类结果。图 4a 中间面板展示了潜在空间中的想象轨迹,可以直接从对应物体的颜色和位置来解读。
由于循环混合模型(rMM)将开关状态建立在各种与游戏和物体相关的特征之上,我们可以将这些开关变量设定为依赖于不同的游戏特征,并通过可视化展示 rMM 所学到的关联关系(例如,奖励与空间位置之间的联系)。
图 4a 最右侧面板展示了与奖励(绿色)和惩罚(红色)相关的 rMM 聚类在空间中的分布。这些聚类的分布解释了 AXIOM 对“在空间中的哪些位置会遇到奖励”的信念,例如当玩家未接到球时预期会在屏幕底部出现惩罚(图 4a 最右下面板的红色聚类)。
图 4b 显示了在训练过程中活跃的 rMM 成分数量急剧下降。通过主动合并聚类以最小化简化模型所对应的期望自由能,贝叶斯模型缩减(BMR)在保持或提升性能的同时提高了计算效率(见表 1)。这种合并后的成分使得模型能够超越训练数据进行插值,从而增强泛化能力。
这种自动简化揭示了实现最优性能所需的最简动力学集合,使 AXIOM 的决策过程变得透明且稳健。
图 4c 表明,随着训练的推进,每步的信息增益逐渐减少,而期望效用则逐步上升,反映出世界模型变得可靠后,智能体从探索阶段转向了利用阶段。
干扰鲁棒性 最后,我们在系统性地对游戏机制施加干扰的情况下测试了 AXIOM 的表现。在这里,我们在第 5000 步时对每个物体的颜色或形状进行了扰动。
图 4d 显示,AXIOM 对形状扰动具有很强的鲁棒性,因为它仍能通过 iMM 正确推断出物体类型。而在面对颜色扰动时,AXIOM 会新增身份类型并重新学习其动力学特性,导致性能略有下降,随后逐渐恢复。
由于 AXIOM 的世界模型具有良好的可解释性结构,我们可以预先为其注入关于可能颜色扰动的知识,在 iMM 推理步骤中仅使用形状信息,并在之后根据形状重新映射被扰动的 slot,从而恢复性能。
有关更多细节,请参见附录 E.3。
局限性与未来工作。我们的工作受到这样一个事实的限制:核心先验本身是由人工设计的,而不是从数据中自主发现的。未来的工作将聚焦于开发能够从数据中自动推断这类核心先验的方法,这将使我们的方法能够应用于更复杂的领域,如 Atari 或 Minecraft [36],在这些领域中,潜在的生成过程不那么透明,但仍受类似的因果原理所支配。我们相信,这一方向代表了构建自适应智能体的关键一步,这样的智能体能够在无需显式设计领域知识的情况下,快速构建新环境的结构化模型。
A 完整模型细节
AXIOM 的世界模型是一个隐马尔可夫模型(hidden Markov model),其具有以物体为中心的潜在状态空间。该模型本身包含两个主要组成部分:
- 以物体为中心、类似 slot-attention 的似然模型
- 递归切换状态空间模型(recurrent switching state space model)。
这一递归切换状态空间模型被应用于由似然模型识别出的每一个物体或“slot”,并使用分段线性轨迹来建模每个物体的动力学行为。
不同于大多数其他潜在状态空间模型(包括其他以物体为中心的模型),AXIOM 的一个显著特点是其可适应的复杂度——它通过在线增长和剪枝机制(见算法1 和 算法2)来迭代扩展与缩减模型,使其结构能够匹配其所交互世界的复杂度。这包括自动推断场景中物体的数量,以及描述所有物体运动所需的动态模式数量。
这种方法受到最近提出的快速结构学习方法[23] 的启发,后者旨在从零开始自动学习数据集的层次化生成模型。
符号说明前言
大写加粗符号表示矩阵或向量值随机变量的集合,小写加粗符号表示多元变量。
A.1 生成模型
该模型将感知与动力学因式分解为两个独立的生成模块:
A.2 Slot Mixture Model (sMM)
A.3 运动与存在潜变量
这个二值门控变量随后调节与身份模型(iMM)、转移混合模型(tMM)和递归混合模型(rMM)相关的各种似然函数的输入精度,从而有效地对这些模型在未跟踪或不存在的 slot 上的学习过程进行“掩码”处理。
最终效果是:那些被推断为正在运动且存在的 slot 保持完整的精度,而其他任何组合状态都会将该 slot 的输入协方差压缩至 0,从而在参数学习中去除其充分统计量的影响。
A.4 相互作用变量
A.7 转移混合模型
A.8 循环混合模型
循环混合模型(Recurrent Mixture Model,简称 rMM)用于直接从当前槽级别的特征中推断转移模型的开关状态。这种将开关状态依赖于连续特征的建模方式,与循环切换线性动态系统(recurrent Switching Linear Dynamical System,rSLDS)[19] 中所使用的结构相同。然而,与 rSLDS 不同的是:rSLDS 通常使用一种判别式映射(例如 softmax 或 stick-breaking 参数化形式)来从连续状态中推断开关状态,而 rMM 则通过一个在混合连续-离散槽状态上的混合模型[30] 来生成式地恢复这种依赖关系。
通过这种方式,“选择”用于对 tMM 进行条件建模的开关状态,实际上是通过对具有特定条件关系(在此背景下,是一种联合混合似然关系)的其他隐变量和观测变量进行推理而自然产生的。具体来说,rMM 将连续变量和离散变量的分布建模为由另一个每槽隐含分配变量 驱动的混合模型。rMM 定义了一个关于连续和离散槽特定信息元组的混合似然。
多元正态分布成分的参数配备了NIW 先验(正态-逆-威沙特先验),而离散的分类似然(Categorical likelihood)的参数则配备了狄利克雷先验。与 AXIOM 的所有其他模块一样,我们为的混合权重配备了一个截断的 stick-breaking 先验,其中最后一个第 M个伪计数参数用于调节添加新 rMM 成分的倾向性。
另外,我们还使用了门控变量 来过滤用于动态学习的槽:通过放大那些未被推断为正在移动且存在的槽所对应的协方差,从而降低其在模型中的影响。
固定距离变体(Fixed distance variant)
Aviate(飞行)
在这个环境中,玩家控制一只小鸟,需要穿越一系列垂直的管道。小鸟在重力作用下下落,玩家可以通过执行“拍打”动作让小鸟跳跃。目标是引导小鸟穿过管道之间的狭窄水平空隙,避免与管道结构的任何部分或屏幕的顶部/底部边缘发生碰撞。一旦与管道碰撞,或飞出屏幕顶部或底部,将获得负奖励,并结束游戏。
Bounce(反弹)
这个环境模拟了一个简化版的经典游戏 Pong,玩家控制一个球拍,与由 AI 控制的对手对打一个球。玩家有三个离散动作:向上移动球拍、向下移动球拍,或保持不动,这会影响球接触时的垂直轨迹。目标是将球击过对手的球拍以得分(+1 奖励),同时防止对手击中自己(-1 奖励)。游戏是回合制的,一方得分后游戏重置。
Cross(穿越)
受经典 Atari 游戏 Freeway 启发,该环境要求玩家扮演一个黄色方块,穿越多车道的道路而不被汽车撞到。玩家有三个离散动作:向上移动、向下移动或保持原地不动,用于在八个不同的车道之间进行垂直移动。各种颜色和速度的汽车持续从左向右行驶并在屏幕边缘循环出现。目标是到达屏幕顶部以获得正奖励;若与汽车碰撞,则玩家会被重置回屏幕底部并受到负奖励。
Driver(驾驶)
这是一个基于车道的驾驶模拟环境,玩家从俯视视角控制一辆汽车,在多车道道路上行驶。玩家可以选择三个离散动作:停留在当前位置、向左移动或向右移动,从而实现变道。目标是尽可能行驶更远的距离,避免与其他出现在车道上并以不同速度向下行驶的车辆发生碰撞。与另一辆车碰撞会带来负奖励并结束游戏。
Explode(爆炸)
这个游戏灵感来自街机经典游戏 Kaboom!,玩家控制屏幕底部的一个水平水桶,任务是接住从移动轰炸机上掉落的炸弹。玩家可以执行三个离散动作:保持静止、向左移动或向右移动,以便精确地在水平方向上拦截下落的投射物。轰炸机会持续在屏幕顶部来回移动,并周期性释放炸弹,这些炸弹在下落过程中会加速。成功用水桶接到炸弹可获得正奖励,而让炸弹掉出屏幕则会导致负奖励。
Fruits(水果)在这个游戏中,玩家扮演一个角色,必须在躲避危险岩石的同时收集下落的水果。玩家可以执行三个离散动作之一:向左移动、向右移动或保持不动,控制屏幕底部的水平移动。各种颜色的水果从顶部掉落,被玩家的“隐形篮子”接住时会获得正奖励。相反,以深灰色矩形表示的岩石如果被接住,则会结束游戏并带来负奖励。
Gold(黄金)在这个游戏中,玩家控制一个由黄色方块表示的角色,以俯视视角在一个草地上移动,目标是收集金币并避开狗。玩家可选择五个离散动作之一:静止不动、向上、向右、向下或向左移动,从而实现灵活的屏幕导航。金币是静态的可收集物品,接触时给予正奖励;而狗则在屏幕上动态移动,作为障碍物存在,若发生碰撞则结束游戏并带来负奖励。
Hunt(狩猎)这个游戏采用类似网格的多车道环境,玩家以俯视视角控制角色进行导航。玩家有四个离散动作可供选择:向左、向右、向上或向下移动,实现在游戏区域内的二维自由移动。屏幕上不断出现沿车道水平移动的物品和障碍物。玩家的目标是收集有益物品以获得正奖励,同时巧妙规避有害障碍物,避免受到负奖励的影响,鼓励策略性路径规划。
Impact(冲击)该环境模拟经典街机游戏《打砖块》(Breakout),玩家控制屏幕底部的一个水平球拍,用来反弹球以击碎砖墙。玩家有三个离散动作:将球拍向左移动、向右移动或保持静止。目标是通过用球击中砖块来将其全部清除,每击碎一块砖获得一次正奖励。如果球漏过球拍,玩家将获得负奖励且游戏重置。当所有砖块都被清除后游戏结束。
Jump(跳跃)这是一个横向卷轴式的无尽奔跑游戏,玩家控制一个角色持续向前奔跑,并面对各种障碍物。玩家有两个离散动作:不采取任何操作,或发起跳跃以避开不同类型的障碍物。与障碍物碰撞将导致负奖励,并立即重置游戏。
E 补充结果与消融实验E.1 在 100K 步长下的基线性能
将实际运行时间延长至 100,000 次交互步数后,模型基础方法(model-based)与无模型方法(model-free)之间的对比更加明显。
在Hunt游戏中,DreamerV3 在整个训练过程中几乎没有任何显著进展,其表现始终接近于随机策略的基准,而BBF则持续提升,并最终达到了与我们的以对象为中心的智能体在仅 10,000 步内所达到的平均回合回报相当的水平。
在Gold游戏中,两个基线方法在 100,000 步内确实实现了学习,但它们的最终性能仍低于我们的智能体在短短 10,000 步内所达到的表现水平(见图6)。
打开网易新闻 查看精彩图片E.2 消融实验无信息增益(No information gain)
当禁用信息增益时,我们得到了图7中的紫色曲线。总体来看,从表面上看,信息增益对大多数游戏的影响似乎不大。然而,这其实是可以预期的:如图4c所示,在Explode游戏中,信息增益仅在前几百步中驱动性能提升,之后便由期望效用主导。
从累积奖励的角度来看,信息增益实际上在许多因玩家与物体互动而产生负奖励的游戏中损害了表现。这是因为这些交互事件在初期会被预测为信息量丰富,从而鼓励智能体多次尝试经历它们。这一点在Cross游戏中尤为明显:在经历了最初几次碰撞后,没有信息增益的消融模型立即决定完全不再尝试穿越道路。
图8可视化了rMM生成的聚类结果,清楚地展示了在Cross游戏中,没有信息增益会扼杀探索行为。因此,我们认为信息增益在困难探索任务中将发挥更重要的作用,这是一个值得未来研究的方向。
打开网易新闻 查看精彩图片无贝叶斯模型简化(No Bayesian Model Reduction, BMR)
图7中的橙色曲线显示了禁用贝叶斯模型简化(BMR)的影响。BMR 对 Gold 和 Hunt 这两个允许玩家在二维区域内自由移动的游戏具有关键影响。在这些情况下,BMR 能够通过合并相似聚类来实现对动态和物体交互的空间泛化能力。
打开网易新闻 查看精彩图片唯一的例外再次是 Cross 游戏:在这里,禁用 BMR 反而产生了表现最好的智能体。这一现象再次与信息增益的相互作用有关。由于 BMR 会合并相似的聚类,未发生碰撞的向上移动会被归入一个经常访问的单一聚类。这使得该聚类从信息增益的角度来看变得不那么“有信息量”,导致智能体更倾向于先去与不同的车辆发生碰撞。
然而,当禁用 BMR 时,每个空间位置都会拥有自己的聚类,智能体会被吸引去访问那些较少观察到的位置,例如屏幕顶部。如果我们绘制图8c中最终的 rMM 聚类结果,也可以定性地看到这种差异。
这引出了一个问题:在学习过程中何时进行 BMR 最合适?显然,BMR 对于将已观察到的事件泛化到新情境至关重要,但如果在学习早期阶段就启用它,反而可能阻碍学习过程。这种相互作用机制仍需进一步研究。
规划展开次数与采样数量(Planning rollouts and samples)
在规划阶段,我们在每个时间步都进行 rollout(轨迹展开)采样,因此存在一个明显的权衡关系:即策略数量与每次 rollout 的采样数量之间的平衡——既要考虑计算时间成本(见图5),又要保证所找到计划的质量。
我们进行了网格搜索,调整 rollout 数量 [64, 128, 256, 512] 和每条 rollout 的采样次数 [1, 3, 5],并评估了三个随机种子下的表现。结果显示在图9中,表明性能差异并不显著,但平均而言,更多的 rollout 和每条 rollout 中抽取多个样本略优于单一样本的情况。
因此,在我们的主要评估中,我们使用了512 条策略,每条策略采样3 个样本。但图5和图9的结果表明,当计算时间受限时,将策略数量减少到 128 或 64 是一种可行的方式,可以在不牺牲性能的前提下提高效率。
打开网易新闻 查看精彩图片E.3 扰动实验(Perturbations)
Gameworld 10k 基准测试的一个优势在于它能够在不同环境中施加统一的扰动,从而让我们可以量化不同模型对视觉特征变化的鲁棒性。
在当前实验中,我们引入了两种类型的扰动:
评估模型鲁棒性的方法
我们在训练中途(即第5,000步)应用每种扰动,并绘制出 Axiom、Dreamer 和 BBF 在各游戏中平均奖励的变化曲线(见图11)。
我们的结果还表明:
重映射槽身份扰动(Remapped Slot Identity Perturbations)
在图11中以紫色线表示的这项扰动实验,旨在展示 AXIOM 世界模型的“白盒”性质及其可解释性。
打开网易新闻 查看精彩图片在这个实验中,我们首先像之前一样执行标准的“颜色扰动”,然后我们将关于“物体颜色不可靠”的知识编码进 AXIOM 的世界模型中。
具体来说:
在实践中,这意味着颜色发生变化的对象不会被快速分配新的身份,从而使得相同的基于身份的动力学(即 rMM 中的聚类)仍然可以用于预测和解释这些对象的行为,即使它们的颜色发生了变化。
这也解释了为什么在某些游戏中,使用这种“颜色重映射”技巧后,扰动几乎没有任何负面影响,尤其是在那些可以通过形状轻松推断对象身份的游戏(如 Explode)中。
图12 展示了有和没有“重映射技巧”的 iMM 身份槽;图11d 则展示了所有游戏的性能影响。
打开网易新闻 查看精彩图片对于某些游戏中具有相同形状的对象(例如 Hunt 中的奖励物和障碍物,或 Fruits 中的水果和岩石),这种重映射技巧无效,因为仅凭形状信息不足以推断对象类型并据此对动力学建模。
在这种情况下,可以考虑使用更多特征来推断对象身份,比如位置或动力学特征。然而,扩展我们的模型以整合这些信息以进一步提升鲁棒性,将是我们未来的工作方向。
因为 AXIOM 学到的潜在对象特征可以直接解释为帧中物体的颜色,
我们可以在提取对象身份的推理步骤中(即 iMM 的推理步骤),去除与颜色相关的潜在维度的影响,
并仅使用形状信息来进行对象类型推断。
- BBF 和 Dreamer 对形状变化也具有鲁棒性
对于颜色扰动,Dreamer(与 Axiom 类似)有时会出现短暂的性能下降(例如在 Explode 游戏中),但随后也会恢复。
相比之下,BBF 对这两种扰动都几乎没有受到影响。我们推测这种鲁棒性源于扰动是在 BBF 尚未收敛时就施加的,因此视觉特征的变化对其学习过程影响较小。
在形状扰动下,Axiom 表现出较强的鲁棒性。我们认为这归功于其身份模型(iMM),该模型即使在物体外观发生改变的情况下,仍能成功地将新形状映射到已有的对象身份上。
在颜色扰动下,Axiom 的表现通常会下降——这是因为身份模型最初会将扰动后的精灵视为新的对象——但随后它迅速恢复,因为它会将这些新身份重新分配给之前学习的动力学模式。
- 颜色扰动
(Color perturbation):改变所有精灵(sprite)和背景的颜色(见图10b);
- 形状扰动
(Shape perturbation):将基本图形从正方形变为圆形和三角形(见图10c)。
打开网易新闻 查看精彩图片
深度强化学习的早期突破,例如使用深度 Q 网络玩 Atari 游戏 [39],是无模型的(model-free),并且需要在数百万帧图像上进行训练才能达到人类水平的表现。为此,近年来的研究开始转向基于模型的强化学习(model-based reinforcement learning),通过学习环境的“世界模型”来减少所需的环境交互次数,从而实现更好的泛化能力 [40, 41]。
一个著名的例子是 Dreamer 系列模型,它依赖于连续和离散状态空间的混合结构来建模环境动力学 [36, 42, 43]。这类世界模型模拟了人类认知的某些方面,例如对物理现象的直觉理解和物体追踪 [5, 7]。为此,人们也在这一类架构中引入先验知识,使得特定类型的世界结构可以被更快、更好地学习。例如,在物体层面上建模交互已被证明可以在多个任务中提升样本效率、泛化能力和鲁棒性 [9–12]。
近年来,随着 IODINE [44] 和 Slot Attention [45] 等模型的提出,物体分割领域也取得了显著进展。这些模型利用自注意力机制的优势,强制槽位隐变量之间竞争解释图像像素的能力。Slot Attention 中使用的自注意力形式与用于拟合高斯混合模型的 E 步和 M 步密切相关 [46, 47],这也启发了我们:AXIOM 通过槽混合模型(sMM)的推理与学习从图像中分割出物体。
后续在这些开创性工作基础上的改进包括:
- Latent Slot Diffusion
:使用扩散模型进一步提升了原始方法;
- SlotSSM [48]
:不仅将物体分解作为图像分割的归纳偏置,还用于视频预测。
最近一些提出以对象为中心、基于模型的方法包括:
- FOCUS
:验证了此类方法在机器人操作任务中低数据量下的泛化优势 [49];
- OC-STORM 和 SSWM
:利用以对象为中心的信息来预测环境的动力学和奖励 [14, 50];
- SPARTAN
:提出使用大型 Transformer 架构来识别稀疏的局部因果模型,以准确预测未来物体状态 [13]。
与 OC-STORM 不同的是,后者使用预训练视觉基础模型和分割掩码提取对象特征,而 AXIOM 则无需对象级别的监督即可在线识别并分割物体(尽管目前仅测试了单色多边形等简单对象)。此外,AXIOM 还能在线增长和修剪其以对象为中心的状态空间,并像 OC-STORM 一样使用其世界模型生成的轨迹进行规划。
贝叶斯推理(Bayesian Inference)
我们的模型中的推理、学习和规划源自主动推理框架(active inference framework),该框架允许我们将贝叶斯原理与强化学习相结合,通过最小化预期自由能(expected free energy)来平衡奖励最大化与信息增益 [15, 16]。
为了学习环境的结构,我们借鉴了快速结构学习方法[24]:首先向模型中添加混合成分 [51],然后使用贝叶斯模型简化(Bayesian model reduction)对其进行剪枝 [21, 22, 24]。
我们在时间混合建模方面的做法,与近期关于结构学习高斯混合模型的研究有概念上的相似之处,这些研究在强化学习背景下自适应地确定感知和转移建模所需的成分数量 [52]。
AXIOM 的模型与最初的快速结构学习方法 [23] 的一个重要区别在于:
AXIOM 使用了更结构化的先验(表现为 sMM 的以对象为中心的因子化和 tMM 的分段线性结构);
并且使用了连续混合模型似然函数,而非纯粹离散的似然。
我们所使用的转移混合模型(tMM)是一种截断的无限切换线性动态系统(SLDS)[29, 53, 54]。特别是,我们依赖于一种称为循环 SLDS(recurrent SLDS)的最新形式 [19],它引入了开关状态对连续状态的依赖关系,从而解决了标准 SLDS 的两个关键限制:状态无关的转移和上下文盲的动力学建模。
我们的创新之处在于如何处理 rSLDS 中的循环连接:我们采用了一种生成式模型(generative model)而非判别式模型来建模开关状态。这使得我们可以更加灵活地根据各种信息源(包括连续和离散信息)对开关状态进行条件建模,并且开关状态的依赖关系在连续特征上是二次的(quadratic),从而克服了原始 rSLDS 中使用 softmax 似然时对线性可分性的内在假设 [19, 55]。
原文链接:https://arxiv.org/pdf/2505.24784
热门跟贴