Contrastive Learning of Structured World Models

By Thomas Kipf, Elise van der Pol, Max Welling. 2019

341 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Representation & Self-Supervised Learning, Evaluation & Benchmarks

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Representation & Self-Supervised Learning

In Evaluation & Benchmarks

Citations in the corpus

Object-Centric Learning with Slot Attention (2020, methodology): “[46] Thomas Kipf, Elise van der Pol, and Max Welling.”
On the Binding Problem in Artificial Neural Networks (2020, background): “…for learning abstract relations using spiking neural networks (Doumas et al., 2008), and shares similarities to more recent contrastive learning objectives that require a neural network to infer relations from a dataset of positive and negative pairings (Kipf et al., 2020; Hadsell et al., 2006).”
Systematic Evaluation of Causal Discovery in Visual Model Based Reinforcement Learning (2021, methodology): “We also evaluate the learned models on ranking metrics in the latent space as well as reconstruction-based metrics in the observation space [Kipf et al., 2019].”
Contrastive Representation Learning: A Framework and Review (2020, methodology): “C-SWM [60] State + Action Next state Random state CNN and GCN Addition Pair loss InfoGraph [97] Global graph Local graph Random local graph GIN Readout summation JSD MI estimator”
BlockGAN: Learning 3D Object-aware Scene Representations from Unlabelled Images (2020, background): “Future work can adopt more powerful relational learning models [25, 45, 51] to learn more complex object interactions such as inter-object shadowing or reﬂections.”
Learning Dynamic Belief Graphs to Generalize on Text-Based Games (2020, background): “[23] propose contrastive methods to learn latent structured world models (C-SWMs) as state representations for vision-based environments.”
Neural Production Systems (2021, methodology): “For learning action conditioned world-models, we follow the same experimental setup as (Kipf et al., 2019).”
Space-Time Correspondence as a Contrastive Random Walk (2020, methodology): “Approaches include dynamic programming, particle filtering, various graph-based combinatorial optimization, and more recently, graph neural networks [119, 87, 8, 83, 15, 117, 49, 48, 54, 11, 13, 52].”
Graph Self-Supervised Learning: A Survey (2021, background): “CSWM [157] introduces a same-scale node-level contrastive learning strategy to train structured world models that simulate multi-object systems.”
Conditional Contrastive Learning: Removing Undesirable Information in Self-Supervised Representations (2021, background): “One major way of self-supervised learning is contrastive learning [5, 18, 31], which tries to construct pairs of related data (termed positive pairs) and pairs of unrelated samples (termed negative pairs) and learn a model that scores the positive and negative pairs differently [21].”
Conditional Object-Centric Learning from Video (2021, background): “…2019; Engelcke et al., 2020; Crawford & Pineau, 2019; Lin et al., 2020; Locatello et al., 2020; Du et al., 2021a) or videos (van Steenkiste et al., 2018; Kosiorek et al., 2018; Stelzner et al., 2019; Kipf et al., 2020; Crawford & Pineau, 2020; Creswell et al., 2021) without explicit supervision.”
Entity Abstraction in Visual Model-Based Reinforcement Learning (2019, result): “Comparison with similar work: The closest three works to OP3 are the Transporter [3], COBRA [5], and C-SWMs [52].”
Offline Reinforcement Learning from Images with Latent Space Models (2020, background): “…space (Levine et al., 2016; Finn et al., 2016; Watter et al., 2015; Banijamali et al., 2018; Zhang et al., 2019; Hafner et al., 2019, 2020; Ha and Schmidhuber, 2018; Kipf et al., 2019; Chen et al., 2020) and can either learn a policy within the model or deploy shooting-based planning methods.”
Unsupervised Object Keypoint Learning using Local Spatial Predictability (2020, methodology): “Graph neural networks (GNNs) have been previously used for intuitive physics models (Battaglia et al., 2016; Chang et al., 2017; van Steenkiste et al., 2018; Janner et al., 2019; Veerapaneni et al., 2019; Kipf et al., 2020; Stanić et al., 2021).”
Hierarchical Planning for Long-Horizon Manipulation with Geometric and Symbolic Scene Graphs (2020, methodology): “Many works have demonstrated that factorized representations of objects and entities improve a model’s ability to reason about novel scene and new tasks [2, 7, 18, 21].”
Visually Grounded Compound PCFGs (2020, background): “In the vision and machine learning community, unsupervised induction of structured image representations (aka scene graphs or world models) has been receiving increasing attention (Eslami et al., 2016; Burgess et al., 2019; Kipf et al., 2020).”
Physion: Evaluating Physical Prediction from Vision in Humans and Machines (2021, methodology): “Not all vision models were equally far off, though: among those whose encoders and dynamics simulators were fully unsupervised, SVG , a model with only convolutional latent states, performed nearly at chance levels; OP3 , an object-centric model trained by rendering pixel-level future predictions (b”
CLINE: Contrastive Learning with Semantic Negative Examples for Natural Language Understanding (2021, neutral): “…(Velickovic et al., 2019), visual representations (van den Oord et al., 2018; He et al., 2020; Chen et al., 2020), response representations (Lin et al., 2020b; Su et al., 2020), text representations (Iter et al., 2020; Ding et al., 2021) and structured world models (Kipf et al., 2020).”