用于 Monte Carlo Tree Search 规划的因果对象中心模型

# 面向蒙特卡洛树搜索规划的因果物体中心模型  
来源：https://arxiv.org/html/2606.14418  

Rodion Vakhitov  
MIRAI  
Moscow, Russia  
vakhitov\.r@miriai\.org  

&Leonid Ugadiarov  
CogAILab & MIRAI  
Moscow, Russia  

Alexey Skrynnik  
CogAILab & MIRAI  
Moscow, Russia  

&Aleksandr Panov  
CogAILab & MIRAI  
Moscow, Russia  

###### 摘要  
我们提出 COMET（用于高效树搜索的因果物体中心模型），这是一种基于模型的强化学习算法，在槽结构化的潜在空间中进行蒙特卡洛树搜索。COMET 将冻结的无监督物体中心编码器与基于Transformer的世界模型相结合，在该模型中，动作通过一种新颖的动作-槽融合机制绑定到物体上，该机制用于槽转移预测。策略和价值头使用物体因果注意力，通过学习到的每个槽的相关性分数调节令牌交互，使决策集中在任务相关实体上。COMET 为 MuZero 风格的潜在规划引入了显式的物体层级归纳偏置。在来自物体中心视觉强化学习基准、ManiSkill、Robosuite 和 VizDoom 的八个视觉和动态多样性任务中，COMET 在训练早期相较于物体中心和单体基线取得了更高的平均归一化分数。  

## 1 引言  
人类可以在行动之前通过心理模拟过去的经历或可能的未来结果来推理行动的后果[36 (https://arxiv.org/html/2606.14418#bib.bib9)]。受这一能力启发，世界模型被引入强化学习（RL）中，作为模拟环境并提高学习效率的一种方式[10 (https://arxiv.org/html/2606.14418#bib.bib10)]。在基于模型的强化学习（MBRL）中，智能体学习环境动态模型，并利用它生成想象经验，从而减少与现实世界交互的需求。MBRL 方法已在广泛的任务中取得了强劲性能。显著例子包括 Dreamer 算法家族[11 (https://arxiv.org/html/2606.14418#bib.bib11),12 (https://arxiv.org/html/2606.14418#bib.bib12),13 (https://arxiv.org/html/2606.14418#bib.bib13)]，它们采用潜在世界模型进行长时域想象；基于模型预测路径积分（MPPI）控制的规划方法[14 (https://arxiv.org/html/2606.14418#bib.bib15),15 (https://arxiv.org/html/2606.14418#bib.bib16)]；以及将蒙特卡洛树搜索 (MCTS)[5 (https://arxiv.org/html/2606.14418#bib.bib58),33 (https://arxiv.org/html/2606.14418#bib.bib60)] 与学习模型相结合的方法[33 (https://arxiv.org/html/2606.14418#bib.bib60)]。尽管取得了这些进展，但在高维、非平稳且由多个相互作用物体组成的环境中，学习准确的世界模型仍然困难。视觉环境面临的挑战之一在于如何有效表示观测。大多数现有方法依赖卷积神经网络（CNN）编码器[21 (https://arxiv.org/html/2606.14418#bib.bib21)]，它从输入图像生成一个单一的整体表示。然而，这种表示可能无法捕捉物体层级结构和相互作用，而后者通常对决策至关重要[32 (https://arxiv.org/html/2606.14418#bib.bib22)]。在复杂场景中，微小但任务相关的物体、动态背景或大量无关实体可能显著降低智能体性能[22 (https://arxiv.org/html/2606.14418#bib.bib61)]。  

参见标题  
图1：COMET 中的物体中心表示。观测 \(O_t\) 被转换为一组物体表示 \(\bar{s}_t\)，并估计其因果得分 \(\bar{\alpha}_t\)。通过聚焦于最相关的物体及其交互，规划可以在树搜索过程中集中于场景中与任务相关的元素。  

相比之下，人类将世界感知为由离散实体（如物体）组成[39 (https://arxiv.org/html/2606.14418#bib.bib23)]，这实现了高效推理和规划。物体中心强化学习将环境表示为一系列物体层级组件，每个组件对应一个单独物体。当存在实例分割掩码时，可以使用 CNN 编码器提取物体表示；或者使用监督分割模型[4 (https://arxiv.org/html/2606.14418#bib.bib24),20 (https://arxiv.org/html/2606.14418#bib.bib25),31 (https://arxiv.org/html/2606.14418#bib.bib26)]，但它们需要标注数据。大量工作转向无监督物体中心表示学习[19 (https://arxiv.org/html/2606.14418#bib.bib27),23 (https://arxiv.org/html/2606.14418#bib.bib28),24 (https://arxiv.org/html/2606.14418#bib.bib35),8 (https://arxiv.org/html/2606.14418#bib.bib37),37 (https://arxiv.org/html/2606.14418#bib.bib29),38 (https://arxiv.org/html/2606.14418#bib.bib38),35 (https://arxiv.org/html/2606.14418#bib.bib30),48 (https://arxiv.org/html/2606.14418#bib.bib31),25 (https://arxiv.org/html/2606.14418#bib.bib54),7 (https://arxiv.org/html/2606.14418#bib.bib32)]，直接从原始图像中发现结构化表示，使其适用于无需外部监督的强化学习。  

维持物体层级世界模型的物体中心 MBRL 方法可以显式表示物体动态和交互，从而实现更聚焦和可解释的决策。许多现实世界和模拟环境本质上是面向物体的：场景由多个物体组成，它们的交互决定奖励。然而，在任何给定时间步，通常只有一小部分物体参与当前决策相关的交互。例如，在机器人操作任务中，机器人在某一时刻通常只与一个物体交互。因此，动作通常只影响少量物体的状态，而剩余物体对当前决策基本无关。受这一观察启发，我们假设显式建模单个物体对决策的重要性可以改进策略学习。为此，我们提出 COMET，一种基于 MCTS 的物体中心 MBRL 算法。在 COMET 中，世界模型为物体中心表示维护解耦的潜在变量。策略和价值模型使用基于Transformer的架构[43 (https://arxiv.org/html/2606.14418#bib.bib59)] 在这些潜在变量之上，并结合物体因果注意力机制。每个网络将物体令牌与用于动作或价值预测的专用目标令牌一起处理，同时注意力通过学习到的因果分数进行调制，以强调任务相关物体。  

总之，我们的主要贡献如下：  

- •我们提出 COMET，一种基于 MCTS 的物体中心 MBRL 算法，它将冻结的物体层级表示与基于Transformer的世界模型相结合，在物体结构化的潜在空间中进行规划。  
- •我们提出一种新颖的动作-物体绑定机制，将动作与物体中心槽融合，在统一的Transformer主干中实现动作与物体之间的学习绑定，从而支持物体中心世界建模以及策略/价值预测。  
- •我们在多样化的面向物体的视觉控制任务上评估 COMET，包括物体中心基准环境和机器人操作任务，并表明它在任务间表现一致，平均而言，与强单体 MCTS 基 MBRL 方法和物体中心 RL 基线相比，实现了更高的样本效率。  

参见标题  
图2：COMET 训练概览。冻结的槽提取器将观测映射为槽，这些槽由Transformer主干处理，产生潜在表示 \(h_t^1, h_t^2, \dots, h_t^n\)。这些潜在表示连同可学习的目标令牌，被送入策略和价值Transformer以预测动作分布或价值。接下来，动作嵌入独立地与每个槽拼接，并通过共享的 MLP 投影器，产生槽条件化的动作嵌入 \(a_t^1, a_t^2, \dots, a_t^n\)。这些由Transformer主干处理得到 \(z_t^1, z_t^2, \dots, z_t^n\)，用于预测下一状态（下一时刻的槽）和奖励。  

## 2 相关工作  
### 2.1 物体中心表示学习  
越来越多的研究聚焦于直接从原始感官输入中学习结构化的物体中心表示，无需人工标注。这些方法不是将场景编码为单个全局向量，而是将观测分解为可以独立处理的实体集合。一个关键机制是 Slot Attention[24 (https://arxiv.org/html/2606.14418#bib.bib35)]，它通过竞争性交叉注意力将固定数量的潜在槽迭代分配给输入的不同部分。后续工作将此想法扩展到序列数据。SAVi[18 (https://arxiv.org/html/2606.14418#bib.bib36)] 和 SAVi++[8 (https://arxiv.org/html/2606.14418#bib.bib37)] 利用运动线索（如光流和深度）引入时间一致性，使槽能够跨帧持久存在。其他方法专注于使用更具表现力的生成模型改进重建质量。SLATE 和 STEVE[38 (https://arxiv.org/html/2606.14418#bib.bib38)] 将离散潜在令牌化（dVAE[42 (https://arxiv.org/html/2606.14418#bib.bib39)]）与基于Transformer的解码器以及基于 Slot Attention 的分组相结合。相比之下，DINOSAUR[35 (https://arxiv.org/html/2606.14418#bib.bib30)] 使用预训练 DINO[1 (https://arxiv.org/html/2606.14418#bib.bib41)] 表示的特征级目标替代像素重建，以学习语义上有意义的物体。更近期的 Slot Contrast[25 (https://arxiv.org/html/2606.14418#bib.bib54)] 通过对比对应物体表示，强制时间步之间槽的对齐，从而在动态场景中实现更鲁棒的跟踪并减少槽模糊性。并非所有物体中心模型都依赖 Slot Attention。Deep Latent Particles (DLP)[6 (https://arxiv.org/html/2606.14418#bib.bib56)] 将图像表示为低维粒子，这些粒子解耦空间位置和外观。另一个不同方向是人工藏本振荡神经元 (AKOrN)[26 (https://arxiv.org/html/2606.14418#bib.bib57)]，它引入振荡神经动力学，其中同步神经元形成对应于物体或部分的连贯群体。  

### 2.2 物体中心强化学习  
近期研究将物体中心表示融入基于模型的强化学习，以更好地捕捉环境的组合结构。COBRA[45 (https://arxiv.org/html/2606.14418#bib.bib44)] 学习一个基于从 MONet[3 (https://arxiv.org/html/2606.14418#bib.bib45)] 获得的潜在槽的转移模型，并将其与内在动机结合以提高数据效率。FOCUS[9 (https://arxiv.org/html/2606.14418#bib.bib46)] 使用编码器-解码器架构，通过学习到的掩码将场景分割为物体特定的潜在变量。OC-STORM[49 (https://arxiv.org/html/2606.14418#bib.bib47)] 采用时空Transformer，在物体中心和像素级表示上进行联合推理以进行动态建模。COBRA 受限于缺乏对物体交互的显式建模，从而限制了其捕捉关系动态的能力，而我们的方法则解决了这一问题。相比之下，FOCUS 和 OC-STORM 依赖标注的分割掩码，限制了它们在完全无监督环境中的适用性。  

与我们的设置更接近的是 STICA[28 (https://arxiv.org/html/2606.14418#bib.bib33)]，它提出了一个物体中心基于模型的 RL 框架，结合了基于槽的表示与基于Transformer的世界模型以及带有物体因果注意力的决策模块。SOLD[27 (https://arxiv.org/html/2606.14418#bib.bib48)] 在无监督情况下直接从像素学习物体中心潜在动态，通过一个动作条件化的基于槽的动态模型和一个用于策略和价值学习的槽聚合Transformer。Object-Centric Dreamer[41 (https://arxiv.org/html/2606.14418#bib.bib55)] (OCDreamer) 通过用物体中心 RSSM 替换 RSSM，并引入 GNN 显式建模预测和控制中的物体交互，扩展了 Dreamer。  

除了基于模型的方法，物体中心表示也用于无模型 RL。OCRL[47 (https://arxiv.org/html/2606.14418#bib.bib50)] 将基于Transformer的物体编码器集成到 PPO[34 (https://arxiv.org/html/2606.14418#bib.bib51)] 中，从而灵活利用不同的物体中心特征。类似地，OC-CA 和 OC-SA[40 (https://arxiv.org/html/2606.14418#bib.bib52)] 使用 Slot Attention 作为特征提取器，并研究其跨环境的泛化能力。  

## 3 COMET  
参见标题  
图3：平均归一化分数 (6 (https://arxiv.org/html/2606.14418#S4.E6)) 相对于归一化步数。归一化参数列于附录A (https://arxiv.org/html/2606.14418#A1)。左图：所有考虑任务（Object Goal、Object Interaction、Object Comparison、Property Comparison、Object Reaching、Block Lifting、Cube Pushing 和 Defend The Line）上除 SOLD 外所有算法的平均归一化分数。右图：与 SOLD 兼容的连续控制任务（Object Reaching、Block Lifting 和 Cube Pushing）的平均归一化分数。  

COMET 是一种基于 MCTS 的物体中心模型强化学习算法，在槽结构化的潜在空间中进行规划。该方法结合了三个组件：一个将视觉观测映射为物体槽的冻结物体中心编码器、一个预测未来槽和奖励的基于Transformer的世界模型，以及配备物体因果注意力的策略/价值头。我们的实现基于 LightZero 框架[29 (https://arxiv.org/html/2606.14418#bib.bib43)]，并遵循 UniZero 训练流程[30 (https://arxiv.org/html/2606.14418#bib.bib34)]。与在单体状态嵌入上操作的 UniZero 不同，COMET 将每个观测表示为一组物体中心槽，并在这些槽上进行动态预测和决策。这种设计为 MuZero 风格的潜在规划引入了显式的物体层级归纳偏置，使模型能够推理单个实体及其交互。与 UniZero 类似，COMET 使用统一的Transformer主干，采用基于 nanoGPT 的架构[17 (https://arxiv.org/html/2606.14418#bib.bib63)]。  

### 3.1 槽提取器  
槽提取器将图像观测 \(o_t\) 映射为一组物体中心潜在表示。具体来说，它生成一个无序向量集合 \(\bar{s}_t = \{s_t^1, \dots, s_t^n\}\)，其中 \(n\) 是一个固定的超参数，通常定义为场景中最大物体数量加一个背景槽。基于静态图像观测训练的槽架构的一个关键挑战是，由于槽注意力的置换不变性和随机性，槽的排序在时间步之间无法保证一致性。为了缓解此问题并确保同一 episode 内物体表示的时间一致性，我们在时间步 \(t+1\) 使用时间步 \(t\) 获得的槽来初始化槽表示。这鼓励槽随时间稳定分配给底层物体。相比之下，在基于视频的物体中心模型中不会出现此问题，因为时间一致性直接在架构中建模（例如，在 Slot Contrast 中，通过……）。