AXIOM：用可扩展的以对象为中心的模型在几分钟内学会玩游戏

摘要 (Abstract)

当前的深度强化学习（DRL）方法在各种领域取得了最先进的性能，但与人类学习相比，其数据效率低下，而人类学习则利用了关于物体及其相互作用的核心先验知识。主动推断（Active Inference）为整合感官信息与先验知识提供了一个有原则的框架，用以学习世界模型并量化其自身信念和预测的不确定性。然而，主动推断模型通常是为具有特定知识的单一任务精心设计的，因此缺乏DRL方法典型的领域灵活性。为了弥合这一差距，我们提出了一种新颖的架构，它集成了一套关于以对象为中心的动力学和相互作用的最小但富有表现力的核心先验，以在低数据情境下加速学习。由此产生的方法，我们称之为AXIOM，将贝叶斯方法通常的数据效率和可解释性与通常与DRL相关的跨任务泛化能力相结合。AXIOM将场景表示为对象的组合，其动力学被建模为分段线性轨迹，捕捉稀疏的对象-对象相互作用。生成模型的结构通过从单个事件中增长和学习混合模型在线扩展，并通过贝叶斯模型缩减进行周期性优化以引导泛化。AXIOM仅用10,000个交互步骤就能掌握各种游戏，其参数数量远少于DRL，并且没有基于梯度的优化的计算开销。这一成果表明，将结构化先验（如对象持久性和分段平滑运动）融入贝叶斯代理，是实现快速、通用和可解释人工智能的关键一步。通过动态调整模型复杂性以适应环境，AXIOM不仅在样本效率上超越了现有技术，还提供了一个透明的决策过程，其内部状态（如对象位置、类型和动力学模式）可以直接检查和理解。这种可解释性对于在需要信任和可验证性的关键应用中部署AI系统至关重要。我们的工作为构建能够在没有大量先验工程的情况下快速构建新环境结构模型的自适应代理铺平了道路。

引言：当AI学会“看见”世界

大家好，我是这篇论文的作者之一。今天，我想和大家聊聊一个我们团队非常兴奋的工作——AXIOM。在人工智能的星辰大海中，我们一直在追寻一个圣杯：如何让机器像人类一样，快速、高效地学习和适应新环境？

想象一下，你第一次玩一款新的电子游戏。你可能只需要几分钟，甚至几十秒，就能掌握基本规则：哪个是“我”，哪个是“敌人”，我能做什么，什么东西能吃，什么东西要躲。但对于今天最顶尖的深度强化学习（DRL）算法，比如那些在Atari游戏或围棋上击败人类的AI，它们往往需要数百万甚至数十亿次的“试错”才能达到类似的水平。这背后的差距，正是我们研究的出发点。

我们人类之所以学习得快，是因为我们的大脑中预装了一套强大的“物理引擎”和“直觉软件”。我们天生就知道，世界是由一个个独立的“物体”组成的；物体会持续存在，不会凭空消失；它们的运动通常是平滑连续的，除非受到外力作用。这些根深蒂固的认知，我们称之为“核心先验”（core priors）。它们是我们理解和与世界互动的基石。

而传统的DRL，就像一个从零开始学习的婴儿，面对一堆像素数据，它不知道什么是“球”，什么是“墙”，只能通过海量数据，慢慢地、痛苦地发现像素之间的统计规律。这个过程不仅效率低下，而且学到的模型往往像一个“黑箱”，我们很难理解它到底在想什么。

所以，我们问自己：能不能给AI也装上这样一套“核心先验”？让它从一开始就以“对象”的视角来理解世界？这就是AXIOM（Active eXpanding Inference with Object-centric Models）的核心思想。我们希望构建一个智能体，它不仅能玩游戏，更能“理解”游戏。它看到的世界不再是像素的海洋，而是一个由玩家、敌人、金币、障碍物等独立对象构成的动态舞台。

第一块基石：用“插槽”分割世界 (Slot Mixture Model)

AXIOM的第一步，就是要学会“看见”物体。我们采用了一种名为“槽混合模型”（Slot Mixture Model, sMM）的技术。你可以把它想象成一个自动化的Photoshop图层工具。

当一张游戏画面（一堆像素）输入时，sMM会自动创建几个“插槽”（slots），每个插槽就像一个透明的图层，专门负责捕捉画面中的一个独立对象。比如，一个插槽可能抓住了玩家控制的黄色小方块，另一个插C槽抓住了移动的蓝色敌人，还有一个插槽抓住了静止的绿色砖块。

每个插槽都会学习它所负责对象的连续属性，比如位置 \((p_x, p_y)\)、颜色 \((r, g, b)\) 和形状（比如宽度和高度 \((e_x, e_y)\)）。更酷的是，这个过程是动态的。如果画面中突然出现了一个新对象，AXIOM会自动“扩容”，增加一个新的插槽来跟踪它。如果一个对象消失了，对应的插槽就会被标记为“未使用”，等待被重新激活。这种灵活性，使得AXIOM能够适应对象数量动态变化的环境。

第二块基石：给物体贴上“身份标签” (Identity Mixture Model)

仅仅把物体分开还不够，我们还需要知道“这是什么”。这就是“身份混合模型”（Identity Mixture Model, iMM）的工作。iMM像一个聪明的分类员，它会检查每个插槽捕捉到的对象属性（主要是颜色和形状），然后给它贴上一个“身份标签”。

例如，所有“黄色、方形”的对象可能都被归为“玩家”这一类；所有“蓝色、圆形”的对象被归为“敌人”；所有“红色、小圆形”的被归为“子弹”。通过这种方式，AXIOM学会了识别不同类型的对象，即使同一类型的对象在屏幕上有很多个实例。这对于泛化至关重要。当模型学会了“敌人”的行为模式后，无论出现多少个新的敌人，只要它们的外观相似，模型都能立刻认出并做出正确反应。

第三块基石：预测物体的“下一步” (Transition & Recurrent Models)

理解了“是什么”之后，下一步就是预测“会发生什么”。这部分是AXIOM的“大脑核心”，由两个紧密协作的模型完成：过渡混合模型（Transition Mixture Model, tMM）和循环混合模型（Recurrent Mixture Model, rMM）。

tMM：运动模式的“字典”

tMM 的作用是学习各种基本的运动模式。你可以把它想象成一本“动作字典”。比如，字典里可能有这样的词条：

• 模式1： 自由落体（y方向速度增加）
• 模式2： 向左匀速移动（x方向速度为负常数）
• 模式3： 碰撞后反弹（速度方向反转）
• 模式4： 静止不动（速度为零）

这些运动模式都是简单的线性动态系统，比如 \( x_t = D \cdot x_{t-1} + b \)，其中 \(x_t\) 是物体在时间 \(t\) 的状态（位置、速度等），\(D\) 和 \(b\) 是描述该运动模式的矩阵和向量。tMM会在线学习和扩展这本字典，当遇到一种新的运动方式时，它就会创建一个新的词条。

rMM：智能的“调度员”

光有字典还不够，还需要一个聪明的调度员来决定在特定情境下该用哪个词条。这就是rMM的工作。rMM会观察全局信息，特别是物体之间的相互作用。

比如，rMM会注意到：“当‘玩家’这个对象和‘墙壁’这个对象的距离小于某个阈值时，并且玩家正在向墙壁移动，那么下一步应该激活tMM中的‘碰撞后反弹’模式。” 或者，“当‘子弹’和‘敌人’的位置重叠时，会产生+1的奖励，并且它们俩都会消失。”

rMM通过学习这种“If-Then”规则，将独立的物体动力学（tMM）与它们之间的稀疏交互联系起来，从而构建了一个完整的、可预测的世界模型。它不仅考虑了物体自身的状态，还考虑了物体的类型（来自iMM）、玩家的动作（a）、以及可能产生的奖励（r）。

学习与进化：动态生长与贝叶斯剪枝

AXIOM最独特的特性之一是它的“成长性”。与那些大小固定、参数庞大的神经网络不同，AXIOM的模型是从一个极简的初始状态开始，随着与环境的交互，动态地“生长”和“修剪”自己。

在线生长（Online Growing）

当AXIOM遇到一个全新的、现有模型无法解释的现象时——比如一种新的物体、一种新的运动模式、或者一种新的交互规则——它不会强行用旧知识去拟合，而是会“长”出一个新的组件来专门处理这个新情况。这就像我们学习新知识，不是去修改已有的常识，而是在大脑中开辟一个新的区域来存储它。这个过程是逐帧发生的，无需梯度下降，也无需巨大的回放缓冲区（Replay Buffer），因此极为高效。

贝叶斯模型缩减（Bayesian Model Reduction, BMR）

然而，无限制的生长会导致模型变得冗余和臃肿。比如，智能体可能在屏幕左下角学到了“撞墙会扣分”，又在右上角学到了同样的事情，从而创建了两个独立的规则。为了解决这个问题，AXIOM会周期性地进行“自省”，这个过程我们称之为贝叶斯模型缩减（BMR）。

BMR会检查所有已有的规则（即rMM中的组件），如果发现有两个或多个规则实际上描述的是同一类事件（比如“任何地方撞墙都会扣分”），它就会将它们合并成一个更通用的规则。这个合并的原则是看是否能让整个模型的“证据”（Bayesian model evidence）最大化，通俗地说，就是找到对数据最简洁、最普适的解释。通过BMR，AXIOM能够从具体的、零散的经验中提炼出普适的规律，实现真正的泛化。

决策：在“功利”与“好奇心”之间权衡

拥有了一个强大的世界模型后，AXIOM如何决策呢？这里我们引入了“主动推断”（Active Inference）的框架。简单来说，智能体在选择下一步行动时，会考虑两个因素：

1. 预期效用（Utility）： 这个行动能给我带来多少奖励？这很直接，就是最大化得分。
2. 信息增益（Information Gain）： 这个行动能让我对世界了解多少？也就是满足“好奇心”。

在主动推断中，智能体会选择一个能最小化“预期自由能”（Expected Free Energy）的行动序列。这个自由能 \(G(\pi)\) 可以漂亮地分解为上述两项：

在游戏初期，世界模型还不完善，充满了不确定性。此时，“信息增益”项会占据主导，驱使智能体去探索未知的区域和交互，即使这可能会暂时导致扣分。这就像一个孩子，会故意去触摸各种东西，看看会发生什么。随着学习的深入，模型变得越来越准确，不确定性降低，“预期效用”项的权重就会增加，智能体的行为会变得更加“功利”，专注于执行能获得最高分的操作。

结论：迈向更通用、更高效的AI

通过将以对象为中心的核心先验与动态的、自组织的贝叶斯学习框架相结合，AXIOM在我们的Gameworld 10k基准测试中，仅用1万次交互（对DRL来说是极低的样本量）就掌握了10种不同的像素游戏，其性能超越了SOTA的DRL模型如DreamerV3和BBF。更重要的是，AXIOM的模型是完全可解释的——我们可以随时“打开”它的大脑，查看它识别出了哪些对象，学习到了哪些动力学规则，以及它对世界不确定性的估计。

我们相信，AXIOM代表了构建通用人工智能的一条有前途的道路。未来的工作将致力于让模型能自动发现这些核心先验，而不是由我们手动设计，从而让它能够应对更复杂的、真实世界的挑战。我们的最终目标，是创造出能够像我们一样，带着深刻的理解和好奇心去探索、学习和创造的智能体。