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

2026-06-12

AI Summary

This work addresses the challenges of low sample efficiency and insufficient decision focus in visual reinforcement learning by proposing an object-centric planning framework. Operating within a slot-structured latent space derived from a frozen object encoder, the method integrates a Transformer-based world model with Monte Carlo Tree Search (MCTS), augmented by an action-slot fusion mechanism to accurately predict object-level state transitions. Furthermore, an object-causal attention mechanism dynamically guides the policy and value networks to attend to task-relevant entities. By explicitly incorporating object-level inductive biases, the approach achieves substantial performance gains over both object-centric and monolithic baselines across eight benchmark tasks in Object-Centric Visual RL, ManiSkill, RoboSuite, and VizDoom, demonstrating notably higher average normalized scores especially during early training stages.

Abstract

We introduce COMET (Causal Object-centric Model for Efficient Tree search), a model-based reinforcement learning algorithm that performs Monte Carlo Tree Search in a slot-structured latent space. COMET pairs a frozen unsupervised object-centric encoder with a transformer-based world model, in which actions are bound to objects through a novel action-slot fusion mechanism that is used in slot transition prediction. Policy and value heads use object-causal attention, modulating token interactions by learned per-slot relevance scores so that decision-making concentrates on task-relevant entities. COMET adds an explicit object-level inductive bias to MuZero-style latent planning. Across eight visually and dynamically diverse tasks from the Object-Centric Visual RL benchmark, ManiSkill, Robosuite, and VizDoom, COMET achieves a higher mean normalized score during the early stages of training compared to object-centric and monolithic baselines.

Problem

• causal
• object-centric
• Monte Carlo Tree Search
• model-based reinforcement learning
• latent planning

Innovation

• object-centric
• causal attention
• slot-based representation
• Monte Carlo Tree Search
• action-slot fusion