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

背景

在强化学习（Reinforcement Learning, RL）领域，尤其是基于模型的强化学习（Model-based RL），如何高效地处理高维视觉输入并做出长期规划，一直是一个核心挑战。传统的端到端方法往往将环境视为一个黑盒，直接学习从像素到动作的映射，这导致样本效率低下且缺乏可解释性。

近年来，以 MuZero 为代表的算法通过在潜在空间（latent space）中进行蒙特卡洛树搜索（Monte Carlo Tree Search, MCTS）来绕过显式的环境模型，显著提升了在复杂任务中的表现。然而，MuZero 及其变体通常处理的是“单体”（monolithic）的潜在表示，即整个场景被编码为一个单一的向量或特征图。这种表示方式忽略了现实世界中物体（objects）的独立性和结构性，导致模型难以捕捉物体间的因果关系，且在面对动态变化的场景时泛化能力受限。

与此同时，“以对象为中心”（Object-Centric）的表示学习旨在将视觉输入分解为独立的对象槽（slots），每个槽对应一个具体的实体。这种方法虽然提高了可解释性和模块化，但如何将这些离散的、结构化的对象表示有效地整合进基于树的规划算法中，仍是一个未充分探索的领域。

核心内容

本文提出了 COMET（Causal Object-centric Model for Efficient Tree search），这是一种基于模型的强化学习算法，旨在通过在结构化潜在空间中进行蒙特卡洛树搜索，实现高效且因果驱动的规划。COMET 的核心创新在于将“以对象为中心”的表示学习与 MuZero 风格的潜在规划相结合，并引入了显式的对象级归纳偏置（inductive bias）。

1. 架构设计：冻结编码器与 Transformer 世界模型

COMET 的架构由两个主要部分组成：

• 冻结的无监督对象中心编码器（Frozen Unsupervised Object-Centric Encoder）： 该编码器负责将原始视觉输入（如图像）分解为一系列结构化的“槽”（slots）。每个槽代表场景中的一个独立对象或实体。在 COMET 中，这个编码器是冻结的（即训练过程中参数不更新），它提供了一个稳定、解耦的对象表示基础。
• 基于 Transformer 的世界模型（Transformer-based World Model）： 这是一个动态模型，用于预测下一个时间步的潜在状态。它接收当前的对象槽表示和动作，并输出下一个时间步的对象槽表示。

2. 动作-槽融合机制（Action-Slot Fusion）

为了实现动作对对象的影响建模，COMET 提出了一种新颖的动作-槽融合机制（action-slot fusion mechanism）。在该机制中，动作不再仅仅作用于全局潜在状态，而是被绑定到特定的对象槽上。这种绑定方式使得世界模型能够更精确地预测动作如何改变特定对象的状态，从而在槽转换预测（slot transition prediction）中捕捉到更细粒度的动态变化。

3. 对象因果注意力（Object-Causal Attention）

在策略头（Policy Head）和价值头（Value Head）中，COMET 使用了对象因果注意力（object-causal attention）。这是一种特殊的注意力机制，它通过学习每个槽的相关性得分（per-slot relevance scores）来调节 token 之间的交互。

• 动态聚焦： 这种机制允许模型在决策过程中，根据当前任务的需求，动态地关注与任务最相关的实体（entities）。
• 因果性： 通过限制注意力模式，模型能够更清晰地捕捉对象之间的因果关系，避免无关信息的干扰，从而提高决策的准确性和效率。

4. 在 MCTS 中的应用

COMET 将上述组件整合到蒙特卡洛树搜索（MCTS）框架中。在搜索过程中，COMET 在结构化的对象潜在空间中进行模拟和回溯。由于引入了显式的对象级归纳偏置，COMET 能够更有效地探索状态空间，特别是在那些涉及多个独立交互物体的场景中。

5. 实验结果

作者在多个视觉上动态多样的基准测试中评估了 COMET，包括：

• Object-Centric Visual RL benchmark
• ManiSkill
• Robosuite
• VizDoom

实验结果显示，与基于对象的中心（object-centric）和单体（monolithic）基线方法相比，COMET 在训练早期阶段取得了更高的平均归一化得分（mean normalized score）。这表明 COMET 在样本效率和初始学习阶段具有显著优势。

关键要点

• COMET 算法： 提出了一种名为 COMET 的基于模型的强化学习算法，结合了对象中心表示和蒙特卡洛树搜索。
• 结构化潜在空间： COMET 在“槽”（slot）结构化的潜在空间中进行 MCTS，而非传统的单体向量空间。
• 冻结编码器： 使用冻结的无监督对象中心编码器，确保对象表示的稳定性和解耦性。
• 动作-槽融合： 引入新颖的动作-槽融合机制，将动作绑定到特定对象，以精确预测对象状态的转换。
• 对象因果注意力： 策略和价值头采用对象因果注意力机制，通过学习每槽相关性得分，使决策聚焦于任务相关的实体。
• 显式归纳偏置： 为 MuZero 风格的潜在规划添加了显式的对象级归纳偏置，增强了模型对物体间因果关系的建模能力。
• 性能优势： 在 ManiSkill、Robosuite 和 VizDoom 等八个多样化任务中，COMET 在训练早期阶段的表现优于现有的对象中心和单体基线方法。

意义与影响

COMET 的提出标志着基于模型的强化学习在结构化表示和因果推理方面的重要进展。其意义主要体现在以下几个方面：

1. 提升样本效率： 通过引入对象中心的表示和因果注意力，COMET 能够更快速地学习环境动态，特别是在训练初期。这对于需要大量交互数据的现实世界应用（如机器人控制）至关重要。
2. 增强可解释性： 将潜在状态分解为独立的对象槽，使得模型的决策过程更加透明。研究人员可以追踪特定对象对最终决策的贡献，从而更好地理解模型的内部机制。
3. 改进泛化能力： 显式的对象级归纳偏置有助于模型更好地泛化到未见过的场景和物体组合中，因为模型学习的是对象之间的通用交互规则，而非特定的像素模式。
4. 桥接表示学习与规划： COMET 成功地将无监督的对象中心表示学习与有监督的树搜索规划相结合，为未来研究如何在复杂视觉环境中进行高效、可解释的决策提供了新的范式。

总之，COMET 通过结合对象中心表示、因果注意力和蒙特卡洛树搜索，为解决高维视觉强化学习中的样本效率和可解释性问题提供了有力的工具。随着对象中心表示学习技术的进一步成熟，此类方法有望在机器人学、游戏 AI 和自动驾驶等领域发挥更大的作用。