Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Object-Centric World Models Meet Monte Carlo Tree Search

Rodion Vakhitov, Leonid Ugadiarov, Aleksandr Panov

TL;DR

The paper addresses learning efficient, structured world models for object-rich environments by introducing ObjectZero, which uses a object-centric encoder to produce a fixed set of slots $\bar{s}_t=\{s^1_t,\dots,s^K_t\}$ and a graph neural network to predict dynamics, rewards, and policy, coupled with Monte Carlo Tree Search planning. It demonstrates that integrating a frozen SLATE/DINOSAUR-based slot extractor with a GNN world model and MCTS planning yields faster convergence and competitive performance across Object Reaching and Block Lifting tasks compared to strong baselines like EZ-V2 and DreamerV3. The key contributions include (i) a fully object-centric MBRL architecture, (ii) a graph-based dynamics/reward/value/policy model operating on object slots, and (iii) empirical evidence that object-centric planning can match or exceed performance of monolithic approaches in complex, interactive environments. This work advances the practical viability of object-centric representations for scalable, sample-efficient RL in dynamic, object-rich domains, with potential impact on robotics and embodied AI.

Abstract

In this paper, we introduce ObjectZero, a novel reinforcement learning (RL) algorithm that leverages the power of object-level representations to model dynamic environments more effectively. Unlike traditional approaches that process the world as a single undifferentiated input, our method employs Graph Neural Networks (GNNs) to capture intricate interactions among multiple objects. These objects, which can be manipulated and interact with each other, serve as the foundation for our model's understanding of the environment. We trained the algorithm in a complex setting teeming with diverse, interactive objects, demonstrating its ability to effectively learn and predict object dynamics. Our results highlight that a structured world model operating on object-centric representations can be successfully integrated into a model-based RL algorithm utilizing Monte Carlo Tree Search as a planning module.

Object-Centric World Models Meet Monte Carlo Tree Search

TL;DR

The paper addresses learning efficient, structured world models for object-rich environments by introducing ObjectZero, which uses a object-centric encoder to produce a fixed set of slots and a graph neural network to predict dynamics, rewards, and policy, coupled with Monte Carlo Tree Search planning. It demonstrates that integrating a frozen SLATE/DINOSAUR-based slot extractor with a GNN world model and MCTS planning yields faster convergence and competitive performance across Object Reaching and Block Lifting tasks compared to strong baselines like EZ-V2 and DreamerV3. The key contributions include (i) a fully object-centric MBRL architecture, (ii) a graph-based dynamics/reward/value/policy model operating on object slots, and (iii) empirical evidence that object-centric planning can match or exceed performance of monolithic approaches in complex, interactive environments. This work advances the practical viability of object-centric representations for scalable, sample-efficient RL in dynamic, object-rich domains, with potential impact on robotics and embodied AI.

Abstract

In this paper, we introduce ObjectZero, a novel reinforcement learning (RL) algorithm that leverages the power of object-level representations to model dynamic environments more effectively. Unlike traditional approaches that process the world as a single undifferentiated input, our method employs Graph Neural Networks (GNNs) to capture intricate interactions among multiple objects. These objects, which can be manipulated and interact with each other, serve as the foundation for our model's understanding of the environment. We trained the algorithm in a complex setting teeming with diverse, interactive objects, demonstrating its ability to effectively learn and predict object dynamics. Our results highlight that a structured world model operating on object-centric representations can be successfully integrated into a model-based RL algorithm utilizing Monte Carlo Tree Search as a planning module.
Paper Structure (20 sections, 5 equations, 4 figures, 3 tables)

This paper contains 20 sections, 5 equations, 4 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Object-Centric Representation Learning
  4. Object-Centric Representations and RL
  5. Background
  6. Reinforcement Learning
  7. EfficientZeroV2
  8. ObjectZero
  9. Slots Extractor
  10. Dynamics Model
  11. Policy model, Reward model and Value model
  12. Environments
  13. Causal World
  14. Robosuite
  15. Experiments
  16. ...and 5 more sections

Figures (4)

  • Figure 1: ObjectZero training overview. The slots extractor encodes observations into object-centric representations. Separate GNN models predict next-state slots, rewards, policy and value outputs. Gumbel search generates target policy $\pi_t$ and value $z_t$.
  • Figure 2: Success rate and return averaged over 30 episodes and three seeds for ObjectZero, EZ-V2, DreamerV3, ROCA, OCRL, OC-CA and OC-SA in the Object Reaching Task in Causal World and the Block Lifting Task in Robosuite. Shaded areas indicate standard deviation
  • Figure 3: Examples of observations and attention maps produced by the SLATE model in the Causal World Object Reaching task.
  • Figure 4: Examples of observations and attention maps produced by the DINOSAUR model in the Robosuite Block Lifting task.