Contrastive Learning of Structured World Models

TL;DR

This paper introduces C-SWMs, a contrastive learning framework for structured world models that discover objects and relations from raw data, outperforming reconstruction-based models in complex environments.

Contribution

The paper presents a novel contrastive learning approach for structured world models that learn object and relation representations without supervision.

Findings

C-SWMs outperform reconstruction-based models in structured environments.

C-SWMs learn interpretable object-based representations.

Effective in environments with multiple interacting objects.

Abstract

A structured understanding of our world in terms of objects, relations, and hierarchies is an important component of human cognition. Learning such a structured world model from raw sensory data remains a challenge. As a step towards this goal, we introduce Contrastively-trained Structured World Models (C-SWMs). C-SWMs utilize a contrastive approach for representation learning in environments with compositional structure. We structure each state embedding as a set of object representations and their relations, modeled by a graph neural network. This allows objects to be discovered from raw pixel observations without direct supervision as part of the learning process. We evaluate C-SWMs on compositional environments involving multiple interacting objects that can be manipulated independently by an agent, simple Atari games, and a multi-object physics simulation. Our experiments demonstrate that C-SWMs can overcome limitations of models based on pixel reconstruction and outperform typical representatives of this model class in highly structured environments, while learning interpretable object-based representations.