[Paper Review] Contrastive Learning of Structured World Models

Thomas Kipf, Elise van der Pol+1 more|arXiv (Cornell University)|Nov 27, 2019

Domain Adaptation and Few-Shot Learning70 references68 citations

TL;DR

C-SWMs learn object-based latent representations and an action-conditioned, graph-structured transition model via contrastive learning, enabling unsupervised object discovery and accurate multi-object dynamics prediction in structured environments.

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.

Motivation & Objective

• Motivate learning a structured, object-centric world model to improve generalization and counterfactual reasoning.
• Develop an unsupervised method to discover objects from pixels without direct supervision.
• Propose a contrastive, object-level loss to train object representations and transitions.
• Leverage a graph neural network to model relations and interactions between objects.
• Demonstrate that structured representations improve long-horizon state prediction and generalization.

Proposed method

• Encode observations into a set of object-centric latent representations using a two-part encoder: a CNN-based object extractor and an MLP-based object encoder.
• Model object interactions with a graph neural network that predicts latent state updates via a translational transition (z_t + T(z_t, a_t) ≈ z_{t+1}).
• Train with an object-level contrastive hinge loss that distinguishes true state-action-state triples from corrupted negatives (based on TransE-style energy).
• Adopt an object-factorized latent space Z = Z_1 × … × Z_K and corresponding actions A = A_1 × … × A_K to capture compositional structure and enable parameter sharing.
• Evaluate using ranking metrics (Hits@1, MRR) in latent space for multi-step ahead prediction across diverse environments.

Experimental results

Research questions

• RQ1 Can C-SWMs discover objects from raw pixel observations without supervision?
• RQ2 Do object-centric latent representations and GNN-based transitions enable accurate multi-step state prediction and combinatorial generalization?
• RQ3 Does contrastive learning improve latent representations and predictive accuracy compared to reconstruction-based baselines?
• RQ4 How does object-factorization affect generalization to unseen environment configurations?

Key findings

• C-SWMs learn interpretable object-level representations and accurate transition predictions, outperforming reconstruction-based baselines in highly structured environments.
• In grid-world and physics tasks, C-SWMs achieve near-perfect latent-space prediction for short and mid horizons (1, 5, 10 steps) with high H@1 and MRR, especially when using object-factorized representations and GNN transitions.
• Object discovery emerges without supervision, with per-object latent coordinates closely aligning with true object positions (up to a random linear transform).
• Contrastive loss substantially improves generalization to unseen configurations over pixel-reconstruction losses, particularly in multi-object settings and when using VAE-based decoders.
• Increasing the number of object slots (K) requires validation-based tuning for Atari tasks; iterative/object-centric encoding could further improve robustness.