SOLD: Slot Object-Centric Latent Dynamics Models for Relational Manipulation Learning from Pixels

TL;DR

SOLD introduces an object-centric model-based RL framework that learns latent dynamics directly from pixel inputs by decomposing scenes into slot-based representations. It combines a SAVi-based object encoder/decoder, an autoregressive object-centric dynamics model, and a Slot Aggregation Transformer to predict rewards, values, and actions from structured latent states. The approach yields interpretable attention over task-relevant objects and demonstrates strong performance, outperforming DreamerV3 and TD-MPC2 on relational manipulation benchmarks, while generalizing to non-object-centric tasks. This work advances sample efficiency and interpretability in visual robotic manipulation by leveraging object-centric representations for both world modeling and behavior learning.

Abstract

Learning a latent dynamics model provides a task-agnostic representation of an agent's understanding of its environment. Leveraging this knowledge for model-based reinforcement learning (RL) holds the potential to improve sample efficiency over model-free methods by learning from imagined rollouts. Furthermore, because the latent space serves as input to behavior models, the informative representations learned by the world model facilitate efficient learning of desired skills. Most existing methods rely on holistic representations of the environment's state. In contrast, humans reason about objects and their interactions, predicting how actions will affect specific parts of their surroundings. Inspired by this, we propose Slot-Attention for Object-centric Latent Dynamics (SOLD), a novel model-based RL algorithm that learns object-centric dynamics models in an unsupervised manner from pixel inputs. We demonstrate that the structured latent space not only improves model interpretability but also provides a valuable input space for behavior models to reason over. Our results show that SOLD outperforms DreamerV3 and TD-MPC2 - state-of-the-art model-based RL algorithms - across a range of benchmark robotic environments that require relational reasoning and manipulation capabilities. Videos are available at https://slot-latent-dynamics.github.io/.

Figures (17)

• Figure 1: Suite of visual environments requiring relational reasoning and low-level manipulation to be solved.
• Figure 2: $\text{SOLD}$ is trained by concurrently making the world model consistent with replayed experiences and learning behaviors through latent imagination.
• Figure 3: Open-loop predictions of our object-centric dynamics model. Starting from a single context frame, our model predicts the next 50 frames by propagating slot representations forward without access to any intermediate images.
• Figure 4: Achieved returns over the training duration across the eight benchmark environments. The dashed vertical line represents the offset for our method to account for the samples used during pre-training.
• Figure 5: $\text{SOLD}$ discovers objects relevant for task completion in an unsupervised manner over long horizons. We depict the normalized attention of the [out] token of the actor over the object tokens using Attention Rollout abnar2020. The full slot history is shown in Figure \ref{['fig:full_sequence']}.