技術要約：計画のための因果的オブジェクト中心モデルとモンテカルロ木探索

問題提起

モデルベース強化学習（MBRL）は、環境のダイナミクスを学習して想像上の経験を生成することで、現実世界での相互作用の必要性を減らし、大きな成功を収めてきた。しかし、高次元で非定常な、複数の相互作用するオブジェクトから構成される環境において、正確なワールドモデルを学習することは依然として困難である。

既存の視覚的RLアプローチは、通常、入力画像を単一の包括的な表現へと変換する畳み込みニューラルネットワーク（CNN）エンコーダに依存している。これらのモノリシックな表現は、意思決定に不可欠なオブジェクトレベルの構造や相互作用を捉えられないことが多い。複雑なシーンでは、動的な背景、無関係な実体、あるいは小さくてもタスクに関連するオブジェクトの存在が、エージェントの性能を著しく低下させる可能性がある。人間は世界を個別のエンティティ（実体）の集合として認識し、それによって効率的な推論を可能にしているが、標準的なRLエージェントは、明示的な教師あり学習やインスタンスセグメンテーションマスクなしでは、これらのエンティティを分離することに苦慮する。

手法：COMET

著者らは、スロット構造化された潜在空間内でモンテカルロ木探索（MCTS）を実行するモデルベース強化学習アルゴリズムであるCOMET（Causal Object-centric Model for Efficient Tree search）を提案する。COMETは、主に以下の3つのコンポーネントを統合している。

1. 固定されたオブジェクト中心エンコーダ

COMETは、固定された教師なしオブジェクト中心エンコーダ（例：SLATE、DINOSAUR、またはSlot Contrast）を利用して、視覚的な観測を一連のオブジェクト中心の潜在表現、すなわち「スロット」 ($\bar{s}_t = \{s^1_t, \dots, s^n_t\}$) へとマッピングする。

• スロットの初期化: 静止画モデルにおけるスロット・アテンションの置換不変性と確率的な性質に対処するため、著者らは時刻 $t+1$ のスロット表現を時刻 $t$ のスロットを用いて初期化する。これにより、エピソード内におけるオブジェクト割り当ての時系列的な一貫性が確保される。
• 固定（Freezing）: エンコーダはランダムなポリシーを通じて収集されたデータを用いて事前学習され、RLフェーズ中は安定したオブジェクトレベルの特徴を提供するために固定された状態で保持される。

2. アクションとスロットを融合させたTransformerベースのワールドモデル

COMETの核となるのは、これらのスロット上で動作するTransformerベースのワールドモデルである。標準的なMuZeroスタイルのアプローチがモノリシックな状態ベクトルを処理するのに対し、COMETはオブジェクトスロットのシーケンスを処理する。

• アクションとオブジェクトの結合（Action-Object Binding）: アクションをオブジェクトスロットと融合させる新しいメカニズムを採用している。すべてのオブジェクト潜在変数を単一のベクトルに圧縮して次状態を予測する（これは情報のボトルネックを生む）のではなく、アクションの埋め込みを各スロットに独立して結合する。これらのスロット条件付きアクション埋め込みは、共有MLPプロジェクターを通過した後、Transformerバックボーンを介して、次状態のスロット ($\hat{s}_{t+1}$) と報酬を予測する。
• アーキテクチャ: モデルは、nanoGPTベースのTransformerバックボーンを用いたUniZeroスタイルの学習パイプラインに従う。モデルは、現在のスロット集合と実行されたアクションに基づいて、将来のスロットと報酬を予測する。

3. 方策および価値のためのオブジェクト因果アテンション

方策（Policy）および価値（Value）ヘッドは、**オブジェクト因果アテンション（object-causal attention）**を備えた修正されたTransformerアーキテクチャを利用する。

• 因果スコア: 各潜在スロット $h^k_t$ について、ネットワークは、対応するオブジェクトが現在の意思決定に関連している確率を表す因果スコア $\alpha^k_t \in [0, 1]$ を推定する。
• アテンションの変調: これらのスコアは、スケールド・ドットプロダクト・アテンションの変調に使用される。学習された重み行列 $W_t$ が、定義済みの因果グラフ構造をトークンレベルへと持ち上げ、アテンション機構が、環境と相互作用している特定のオブジェクト（因果的オブジェクト）に集中し、無関係なオブジェクト（非因果的オブジェクト）を抑制することを可能にする。これにより、エージェントは現在環境と相互作用している特定のオブジェクトのサブセットに対して、計画を集中させることができる。

主な貢献

1. COMETアルゴリズム: 固定されたオブジェクトレベルの表現と、構造化された潜在空間での計画のためのTransformerベースのワールドモデルを組み合わせた、MCTSベースのオブジェクト中心MBRLの導入。
2. アクション・スロット融合メカニズム: 統一されたTransformerバックボーン内で、アクションをオブジェクト中心のスロットに結合する新しい手法。これにより、オブジェクト情報を単一のベクトルに崩壊させることなく、オブジェクト中心のワールドモデリングと方策/価値予測が可能になる。
3. オブジェクト因果アテンション: トークン間の相互作用を変調するために学習された因果スコアを統合し、意思決定中にタスクに関連するエンティティへ動的に焦点を合わせることを可能にする。
4. 実証的評価: Object-Centric Visual RLベンチマーク、ManiSkill、Robosuite、VizDoomを含む、視覚的および動的に多様な8つのタスクにわたる包括的な評価。

実験結果

COMETは以下のベースラインと比較して評価された：

• オブジェクト中心ベースライン: OCRL（オブジェクトエンコーダを用いたモデルフリーPPO）、OCDreamer（オブジェクト中心RSSM）、およびSOLD（教師なしスロットベースダイナミクス）。
• モノリシックベースライン: UniZero（標準的なMuZeroスタイルの潜在計画）。

パフォーマンスの知見:

• サンプル効率: 多様なタスク群において、COMETは強力なモノリシックMCTSベースの手法および他のオブジェクト中心RLベースラインの両方と比較して、学習の初期段階において高い**平均正規化スコア（Mean Normalized Score）**を達成している。
• タスク固有の特性:
  • 視覚的に単純なタスク（例：Object Goal, Object Comparison）において、COMETは高品質なオブジェクト表現と因果アテンションを活用して関連オブジェクトを正確に特定し、高速な収束を実現している。
  • 連続制御タスク（例：Block Lifting, Cube Pushing）において、COMETは一貫した性能を示すが、すべての特定の指標において一様にベースラインを上回るわけではない（例：Block Liftingでは競合他社を上回らなかった）。
• 汎用性: 本手法は、単色の背景上の2D形状から、複雑な3Dロボット操作や一人称視点のシューティングシナリオに至るまで、変化する視覚的複雑性を持つ環境に対して堅牢性を示している。

意義と主張

本論文は、個々のオブジェクトの重要性を明示的にモデル化することが、オブジェクト指向の環境における方策学習を改善すると断じている。MuZeroスタイルの潜在計画に明示的なオブジェクトレベルの帰納バイアスを導入することで、COMETはシーンを包括的な塊として扱うのではなく、個々のエンティティとその相互作用に基づいて推論することを可能にする。

著者らは、COMETが複雑なマルチオブジェクト環境における観測の表現という課題を効果的に解決していると主張している。オブジェクト表現を分離し、因果アテンションを使用して関連性をフィルタリングすることで、アルゴリズムはより高いサンプル効率と、より焦点の絞られた意思決定を実現している。この研究は、インスタンスセグメンテーションが利用できない場合や、教師なし表現学習が必要な場合において、オブジェクト中心の帰納バイアスが、視覚的RLにおけるモデルベース計画を向上させるための実行可能かつ効果的な経路であることを示唆している。