Dyn-O: Building Structured World Models with Object-Centric Representations

TL;DR

Dyn-O introduces an advanced object-centric world model capable of handling complex, cluttered environments directly from pixels, outperforming previous models like DreamerV3 in prediction accuracy and enabling detailed feature manipulation.

Contribution

The paper presents Dyn-O, a novel structured world model that improves object-centric representations and dynamics modeling in complex visual environments, enhancing generalization and manipulation capabilities.

Findings

Outperforms DreamerV3 in Procgen prediction tasks

Handles complex, cluttered scenes from pixel data

Enables detailed manipulation of object features

Abstract

World models aim to capture the dynamics of the environment, enabling agents to predict and plan for future states. In most scenarios of interest, the dynamics are highly centered on interactions among objects within the environment. This motivates the development of world models that operate on object-centric rather than monolithic representations, with the goal of more effectively capturing environment dynamics and enhancing compositional generalization. However, the development of object-centric world models has largely been explored in environments with limited visual complexity (such as basic geometries). It remains underexplored whether such models can generalize to more complex settings with diverse textures and cluttered scenes. In this paper, we fill this gap by introducing Dyn-O, an enhanced structured world model built upon object-centric representations. Compared to prior work in object-centric representations, Dyn-O improves in both learning representations and modeling dynamics. On the challenging Procgen games, we find that our method can learn object-centric world models directly from pixel observations, outperforming DreamerV3 in rollout prediction accuracy. Furthermore, by decoupling object-centric features into dynamics-agnostic and dynamics-aware components, we enable finer-grained manipulation of these features and generate more diverse imagined trajectories.