🤖 AI

Causal Object-Centric Models for Planning with Monte Carlo Tree Search

本文介绍了 COMET，这是一种通过结合冻结的对象中心编码器与具有动作-槽位融合（action-slot fusion）和对象因果注意力（object-causal attention）机制的 Transformer 架构世界模型的基于模型的强化学习算法，该算法增强了蒙特卡洛树搜索规划，使其在多种视觉和动态任务中的早期阶段表现优于现有基准模型。

原作者： Rodion Vakhitov, Leonid Ugadiarov, Alexey Skrynnik, Aleksandr Panov

发布于 2026-06-15

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Rodion Vakhitov, Leonid Ugadiarov, Alexey Skrynnik, Aleksandr Panov

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

想象一下，你正在试图教一个机器人玩一款复杂的视频游戏或解决一个谜题。最大的挑战不仅仅是告诉机器人“该做什么”，而是帮助它弄清楚在混乱的场景中“什么才是重要的”。

这篇论文介绍了一种名为 COMET（基于因果对象中心化模型的高效树搜索）的新型人工智能系统。你可以把 COMET 想象成一个机器人，它看到的不仅仅是一张由像素组成的巨大、模糊的色块，而是将世界看作是由不同的角色和道具组成的集合，就像一位导演在观察一个有着演员和家具的舞台。

以下是它的工作原理，通过简单的概念进行拆解：

1. “槽位”（Slot）系统：整理混乱

大多数 AI 系统观察图像时，会尝试一次性理解整个画面。COMET 则不同。它将图像分解为多个“槽位”。

类比： 想象一个忙碌的厨房。普通的 AI 看到的是一堆锅碗瓢盆和食材混杂在一起的巨大混乱。而 COMET 就像一位副厨，能立即将所有东西分类到不同的碗里：一个碗装洋葱，一个碗装刀具，一个碗装炉灶。
工作原理： COMET 使用一个特殊的“冻结”（预训练且不可更改）工具，将视觉世界分解为这些独立的物体“槽位”。它不需要被专门教导如何这样做，它会自动完成。

2. “脑内电影”（世界模型）

一旦 COMET 将物体分拣到各个槽位中，它就需要规划下一步行动。它通过运行一段“心理模拟”来实现这一点。

类比： 在你真正伸手去拿杯子之前，你可能会想象：“如果我抓住把手，杯子就会被提起来。”COMET 也会针对其“槽位”中的每一个物体进行这样的思考。它会在脑海中运行一部关于接下来会发生什么的“心理电影”。
转折点： 在许多 AI 系统中，动作（如“抓取”）只是发送给整个世界的一个单一指令。COMET 使用了一个巧妙的技巧，叫做动作-槽位融合（Action-Slot Fusion）。它将“抓取”指令专门附加到“杯子”这个槽位上，同时忽略“炉灶”这个槽位。这就像是给舞台上的特定演员下达特定指令，而不是对着整个观众席大喊大叫。

3. “聚焦过滤器”（因果注意力）

这是 COMET 最聪明的部分。在一个复杂的场景中，并非每个物体都重要。如果你正在玩一个需要推动红色方块的游戏，那么蓝色方块和背景中的树木就无关紧要。

类比： 想象你在一个拥挤的房间里试图听一个人说话。普通的 AI 会试图同时听取所有人的声音，这会让它感到困惑。COMET 拥有一个“聚焦过滤器”。它会为房间里的每个人分配一个“相关性得分”。它调高了你需要倾听的那个人（红色方块）的音量，并将其他人（树木、蓝色方块）的音量降为零。
结果： 当 COMET 做出决策时，它只关注那些真正影响结果的物体。这使得它学习得更快、更聪明。

4. “树搜索”（规划）

为了决定该做什么，COMET 使用了一种名为**蒙特卡洛树搜索（MCTS）**的方法。

类比： 想象一位棋手在预判。他们会想：“如果我走这一步，对手可能会走那一步。如果他们那样做，我可以走这一步……”他们在脑海中构建了一棵可能的路径之树。
COMET 的版本： COMET 不是用模糊的图像来构建这棵树，而是使用其清晰的“物体槽位”。它在脑海中模拟数千种未来的场景，但由于它只追踪重要的物体，因此它比其他系统能更高效地完成这项工作。

他们的研究发现了什么？

研究人员在八项不同的任务上测试了 COMET，范围从简单的 2D 谜题到复杂的 3D 机械臂运动以及视频游戏场景（例如抵御怪物进攻防线）。

结果： COMET 的学习速度比其他系统更快，尤其是在训练的早期阶段。
为什么这很重要： 它证明了通过教导 AI 将世界视为独立的、相互作用的物体（而不是单张图像），并让它只关注真正重要的部分，AI 可以更高效地学习解决问题。

简而言之： COMET 是一个通过先将世界分类为不同的碎片，然后运行仅关注对当前任务有意义的碎片的心理模拟，从而学会玩游戏的机器人。这使得它比那些试图同时处理整个世界的机器人，成为了一个更高效的学习者。

技术摘要：用于蒙特卡洛树搜索的因果面向对象模型

问题陈述

基于模型的强化学习（MBRL）通过学习环境动力学来生成想象的经验，从而减少了对现实世界交互的需求，并在该领域取得了显著成功。然而，在由多个相互作用物体组成的高维、非平稳环境中，学习准确的世界模型仍然具有挑战性。

现有的视觉强化学习方法通常依赖于卷积神经网络（CNN）编码器，该编码器产生单个整体性的输入图像表示。这些单体化表示往往无法捕捉到对于决策至关重要的物体级结构和交互。在复杂的场景中，动态背景、无关实体或微小但与任务相关的物体都可能显著降低智能体的性能。虽然人类将世界感知为由离散实体组成的，从而实现高效推理，但标准的强化学习智能体在没有显式监督或实例分割掩码的情况下，很难隔离这些实体。

方法论：COMET

作者引入了 COMET（用于高效树搜索的因果面向对象模型），这是一种在槽位结构化（slot-structured）潜空间内执行蒙特卡洛树搜索（MCTS）的基于模型的强化学习算法。COMET 集成了三个主要组件：

1. 冻结的面向对象编码器

COMET 利用一个冻结的、无监督的面向对象编码器（例如 SLATE、DINOSAUR 或 Slot Contrast），将视觉观测映射为一组面向对象的潜在表示，或称为“槽位”（slots）（ $\bar{s}_t = \{s^1_t, \dots, s^n_t\}$ ）。

槽位初始化： 为了解决静态图像模型中槽位注意力（slot attention）的置换不变性和随机性问题，作者使用来自时间 $t$ 的槽位来初始化时间 $t+1$ 的槽位表示。这确保了在一个回合（episode）内物体分配的临时一致性。
冻结： 该编码器通过随机策略收集的数据进行预训练，并在强化学习阶段保持冻结，以提供稳定的物体级特征。

2. 基于 Transformer 的世界模型与动作-槽位融合

COMET 的核心是一个基于 Transformer 的世界模型，它在这些槽位上运行。与处理单体化状态向量的标准 MuZero 式方法不同，COMET 处理的是一系列物体槽位。

动作-物体绑定： 一种新颖的机制将动作与物体槽位进行融合。不同于将所有物体潜在表示压缩成单个向量以预测下一状态（这会产生瓶颈），COMET 将动作嵌入独立地与每个槽位进行拼接。这些经过槽位调节的动作嵌入被传递给一个共享的 MLP 投影器，然后通过 Transformer 主干网络来预测下一状态槽位（ $\hat{s}_{t+1}$ ）和奖励。
架构： 该模型遵循类似 UniZero 的训练流程，使用基于 nanoGPT 的 Transformer 主干网络。它根据当前的槽位集合和所采取的动作来预测未来的槽位和奖励。

3. 用于策略和价值的物体因果注意力

策略和价值头利用了配备了**物体因果注意力（object-causal attention）**的改进型 Transformer 架构。

因果得分： 对于每个潜在槽位 $h^k_t$ ，网络会估计一个因果得分 $\alpha^k_t \in [0, 1]$ ，代表对应的物体对当前决策具有相关性的概率。
注意力调制： 这些得分用于调制缩放点积注意力。一个学习到的权重矩阵 $W_t$ 将预定义的因果图结构提升到标记（token）层面，使得注意力机制能够专注于任务相关的实体（因果物体），同时抑制无关物体（非因果物体）。这使得智能体能够将其规划重点放在当前与环境交互的特定子集物体上。

核心贡献

COMET 算法： 引入了一种基于 MCTS 的面向对象 MBRL 算法，该算法将冻结的物体级表示与用于在结构化潜空间中进行规划的 Transformer 世界模型相结合。
动作-槽位融合机制： 一种在统一的 Transformer 主干网络内将动作与物体中心槽位进行绑定的新颖方法。这实现了面向物体的世界建模和策略/价值预测，而不会将物体信息坍缩为单个向量。
物体因果注意力： 通过集成学习到的因果得分来调制标记交互，从而使策略在决策过程中能够动态地专注于任务相关的实体。
实证评估： 在包括 Object-Centric Visual RL 基准、ManiSkill、Robosuite 和 VizDoom 在内的八个视觉和动态多样化的任务中进行了全面的评估。

实验结果

COMET 与以下基线进行了评估：

面向对象基线： OCRL（带有物体编码器的无模型 PPO）、OCDreamer（面向对象的 RSSM）以及 SOLD（无监督基于槽位的动力学）。
单体化基线： UniZero（标准的 MuZero 式潜空间规划）。

性能发现：

样本效率： 在多样化的任务套件中，与强大的单体化 MCTS 基线方法和其他面向对象的 RL 基线相比，COMET 在训练早期阶段实现了更高的平均归一化分数（Mean Normalized Score）。
任务特性：
- 在视觉简单的任务中（例如 Object Goal, Object Comparison），COMET 利用高质量的物体表示和因果注意力来准确识别相关物体，从而实现更快的收敛。
- 在连续控制任务中（例如 Block Lifting, Cube Pushing），COMET 展示了稳定的性能，尽管它并不在所有特定指标上都优于所有基线（例如，它在 Block Lifting 上并未超越竞争对手）。
泛化能力： 该方法在视觉复杂度变化的各种环境中表现出鲁棒性，从单色背景上的 2D 图形到复杂的 3D 机器人操作以及第一人称射击场景。

重要性与主张

本文认为，显式建模单个物体对决策的重要性可以提高在面向物体环境中的策略学习。通过在 MuZero 式的潜空间规划中引入显式的物体级归纳偏置，COMET 使模型能够对单个实体及其交互进行推理，而不是将场景视为一个整体性的斑块。

作者声称，COMET 有效地解决了在复杂多物体环境中表示观测值的挑战。通过解耦物体表示并使用因果注意力来过滤相关性，该算法实现了更高的样本效率和更集中的决策。这项工作表明，物体中心化的归纳偏置是提高视觉强化学习中基于模型的规划能力的有效路径，特别是在无法获得实例分割或需要无监督表示学习的情况下。