PaperlayerView exact revision on arXiv

Source-linked AI summary

Contrastive Learning of Structured World Models

Thomas Kipf, Elise van der Pol, Max Welling

arXiv:1911.12247v2stat.MLcs.AIcs.LG

TL;DR

Learning structured world models from raw sensory data remains difficult, particularly when object discovery is unsupervised and pixel reconstruction can misallocate modeling capacity. C-SWMs use contrastive learning over object-factorized state representations with graph-neural-network transition models. Across structured environments, Atari games, and multi-object physics, they learn interpretable abstractions, predict transitions many steps ahead, generalize compositionally, and discover objects without supervision.

  • Problem

    Learning structured descriptions of objects and relations from raw visual data is challenging, and pixel-space objectives can overlook relevant features or emphasize irrelevant ones.

  • Method

    C-SWMs use object-factorized latent state representations, graph neural networks for transitions, and an object-level contrastive loss trained on real versus corrupted experiences.

  • Results

    C-SWMs learn interpretable object-level abstractions, predict transitions many steps into the future, generalize to novel configurations, and identify objects without supervision.

  • Takeaways & Limitations

    Structured representations provide an inductive bias for generalization while retaining an unstructured model as a possible special case.

  • Takeaways & Limitations

    The formulation assumes deterministic environments and the Markov property, so it does not model stochastic transitions or observations and lacks memory for information beyond the current state and action.

Abstract

from arXiv · show

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.

1 INTRODUCTION

C-SWMs address the challenge of learning object- and relation-based world descriptions from raw observations without relying on pixel reconstruction. They use contrastive learning and graph-based transition models, and experiments show interpretable abstractions, long-horizon prediction, generalization, and unsupervised object discovery.

  • Motivation: Pixel-space reconstruction can ignore small but relevant features or spend capacity on visually rich yet irrelevant features.Examples include bullets in Atari games and static backgrounds.
  • Approach: Contrastive learning distinguishes real state-action-state experiences from corrupted or fake triples, adapting an approach associated with relational graph embeddings.The contrastive formulation is motivated by scoring true facts against corrupted facts in graph embedding methods.
  • Approach: C-SWMs learn one abstract state variable per object and model environment transitions with a graph neural network.The model operates on latent abstract representations extracted from observations.
  • Contribution: C-SWMs introduce an object-level contrastive loss for unsupervised learning of object-based representations.The loss connects contrastive learning of state abstractions with relational graph embeddings.
  • Results: Experiments show interpretable object-level abstractions, accurate multi-step transition prediction, combinatorial generalization, and unsupervised object discovery.The evaluation uses a ranking-based strategy across the paper's experimental settings.

2 STRUCTURED WORLD MODELS

Structured World Models learn object-factorized state representations and action-conditioned transitions from offline experience, using relational interactions among object slots. C-SWM combines visual object extraction, shared object encoding, GNN message passing, and object-level contrastive learning.

  • The framework learns abstract state representations that retain information needed to predict follow-up states after actions, using an encoder and abstract transition model.
  • Contrastive Learning: The model adapts TransE-style energy learning by replacing relation embeddings with an action- and source-state-conditioned transition function.
  • Object-Oriented State Factorization: C-SWM factorizes states and actions across object slots, sharing transition-model parameters to encourage generalization and facilitate object discovery.
  • Encoder and Object Extractor: The encoder separates a CNN object extractor from a shared MLP object encoder, with each feature map serving as an object mask and producing an object representation.
  • Relational Transition Model: The relational transition model is a GNN that represents pairwise object interactions and is invariant to the ordering of object representations.
  • Relational Transition Model: The GNN uses shared node and edge update functions for message passing over a fully connected scene graph, whose computational complexity is O(K^2).The authors note that nearest-neighbor connectivity could reduce this to linear complexity, but leave that extension for future work.

3 RELATED WORK

Related work studies structured environment models, contrastive representation learning, and state representation learning. C-SWM is positioned against pixel-based losses and autoencoder- or adversarial-based approaches.

  • Structured Models of Environments: Structured environment models use graph neural networks to model object or agent dynamics and their interactions, often improving predictive accuracy.
  • Structured Models of Environments: Several prior methods learn structured models directly from pixels but rely on pixel-based loss functions, unlike C-SWM.
  • Contrastive Learning: Contrastive learning constructs related positive examples and contrasting negative examples, with applications including graph and word representation learning.
  • State Representation Learning: State representation learning commonly uses autoencoders or adversarial learning, while some methods learn representations without decoding back into pixel space.

4 EXPERIMENTS

C-SWMs are evaluated across structured grid worlds, Atari games, and a multi-object physics simulation using latent-space ranking metrics and qualitative analyses. They learn object-based representations and transitions, with stronger generalization in structured settings but limitations on difficult Atari scenes and deterministic assumptions.

  • Experimental setup: C-SWMs are evaluated on interacting-object grid worlds, Atari Pong and Space Invaders, and a 3-body physics simulation using random-policy experience.Observations are color images for grid worlds and concatenated frames for Atari and physics environments.
  • Evaluation: Latent-space ranking metrics compare predicted next-state representations with encoded true observations and reference states, avoiding pixel-reconstruction or downstream-task metrics.The evaluation reports Hits at Rank 1 and Mean Reciprocal Rank.
  • Qualitative results: In grid worlds, C-SWM discovers object-specific filters and position-related latent coordinates without direct supervision, while its transition graph captures grid structure and blocked-action effects.The learned transition model uses latent representations rather than visual inputs to determine whether actions have effects.
  • Qualitative results: In 3-body physics, C-SWM learns object-specific encoders and latent representations tracking location and velocity, with transitions that generalize to unseen instances.The physics comparison reports strong short-term predictions, with C-SWM slightly ahead in the 10-step setting.
  • Ablations and comparisons: Reconstruction-based baselines generally generalize less well to unseen scenes, while removing relational interactions or state factorization degrades future-state prediction and generalization.The effect is especially strong for VAE-based World Models because their prior further constrains latent representations.
  • Limitations: Atari results have high variance: both World Model and C-SWM struggle with perfect long-term prediction, and the best object-slot count differs between Space Invaders and Atari Pong.The authors suggest selecting K on a validation set when it is unknown beforehand.
  • Limitations: The feed-forward CNN extractor cannot disambiguate multiple instances of identical objects, while the deterministic Markov formulation excludes stochastic worlds and memory-dependent transitions.The paper identifies iterative disambiguation, probabilistic extensions, and memory mechanisms as future directions.

5 CONCLUSIONS

The conclusions argue that structured world models provide stronger inductive biases for generalization without eliminating unstructured solutions. C-SWMs learn object-oriented models that generalize better to unseen situations, motivating future use in planning and reinforcement learning.

  • Conclusions: Structured world models strengthen inductive biases for generalization without necessarily constraining model generality, since unstructured representations remain possible as a special case.In the 3-body setting, identical representations across object slots can effectively ignore pairwise interactions.
  • Conclusions: C-SWMs learn object-oriented models that generalize better to unseen situations, likely because structure permits a transition model of significantly lower complexity.The paper proposes model-based planning and reinforcement learning as future applications for counterfactual reasoning about actions and novel interactions.

A.1 OBJECT-SPECIFIC REPRESENTATIONS

C-SWM learns object-specific latent representations whose transition structure is regular across visually different objects, unlike a pixel-loss baseline. Latent-space ranking and pixel-space comparisons show strong long-horizon prediction, while Atari representations are less interpretable.

  • Baseline comparison: The pixel-loss SWM baseline loses regular latent structure, making transitions harder to generalize to unseen environment instances.This comparison uses the same 3D cubes visualization as the C-SWM model.
  • C-SWM representations: C-SWM produces nearly identical abstract transition graphs for visually different objects in the 3D cubes environment.The graph is shown separately for each object slot, with edge colors denoting action types.
  • Cross-environment visualizations: C-SWM object representations are qualitatively visualized for 3-body physics and unseen Atari Pong and Space Invaders instances.The 3-body visualizations use two random seeds, while the Atari examples show object filters and slot-specific transition graphs.
  • Limitations: Atari latent object representations are generally less interpretable because objects are non-exchangeable, actions have indirect effects, and visual features are shared.These properties differ from the more exchangeable objects and localized action effects in the structured environments.
  • Pixel-space comparison: Ranking-based latent evaluation captures qualitative predictive behavior that pixel-space MSE does not, despite C-SWM’s near-perfect qualitative predictions.The World Model baseline can preserve several object positions over many steps, while its errors compound over time; the authors therefore encourage latent ranking evaluations.
  • Hinge-loss comparison: A hinge applied only to the negative energy term significantly improves results over the full hinge variant.The authors attribute this likely to the full hinge being minimized by increasing embedding norms.

A.4 MULTIPLE FEATURE MAPS PER OBJECT SLOT

Additional feature maps per object slot do not clearly improve Space Invaders performance, while robustness tests show that environmental randomness reduces predictive performance.

  • Multiple feature maps: Multiple feature maps per object slot provide no clear advantage on Space Invaders.The comparison keeps other model settings unchanged and evaluates ranking scores over multi-step latent prediction.
  • Environment variants: Adding a no-op action has little effect on C-SWM, whereas adding a randomly moving object reduces predictive performance.The results are reported as ranking scores over four runs on hold-out environment instances.
  • Environment variants: Because C-SWM is fully deterministic, random object motion introduces a challenge that could potentially be addressed by explicitly modeling environmental stochasticity.The random-object variant gives one object an independently sampled action at every turn.
  • Stability: Discovered object representations can vary between runs and may become less suitable for forward prediction or generalization on some datasets.The authors associate initialization dependence with Atari environments and the action-free 3-body physics simulation.

A.7 TRAINING TIME

C-SWM training time is comparable to the World Model baseline across the reported environments, with the Atari models training in under 20 minutes.

  • Training time: C-SWM and the World Model baseline have comparable overall training times.Both train for approximately 1 hour on 2D shapes, 2 hours on 3D cubes, and 30 minutes on 3-body physics using one GTX1080Ti GPU.
  • Training time: The Atari Pong and Space Invaders models typically train in less than 20 minutes.The PAIG baseline is a notable exception, requiring approximately 6 hours on a TitanX Pascal GPU under recommended settings.

B DATASETS

The evaluation spans controlled grid worlds, visually challenging 3D blocks, restricted Atari environments, and an action-free gravitational three-body simulation.

  • Data generation: Training experiences are generated from random object placements by uniformly sampling an object and object-specific action at each time step.The standard grid-world setup gives one object an action per step.
  • 2D Shapes: The 2D Shapes environment is a 5 × 5 grid world containing five uniquely colored and shaped objects at random positions.One selected object can move one grid position in a cardinal direction at each time step.
  • 3D Blocks: The 3D Blocks environment preserves 2D Shapes dynamics while adding perspective changes and partial occlusions through different rendering.It tests robustness to visual variation rather than changing the underlying object-motion rules.
  • Atari environments: Atari Pong is restricted to the first interaction between the ball and player-controlled paddle, discarding the initial 58 random interactions.This restriction removes early interactions where the opponent behaves independently of player action.
  • Atari environments: Space Invaders uses a deterministic Gym variant and discards the first 50 random interactions before populating the experience buffer.The environment is processed and restricted similarly to Pong.
  • 3-body physics: The 3-body physics environment evolves through classical gravitational dynamics without actions.The model receives concatenated consecutive frames to provide implicit velocity information.

C EVALUATION METRICS

The evaluation uses rank-based metrics to assess how closely predicted state representations match encoded true observations. These include a k-nearest-neighbor accuracy score and mean reciprocal rank.

  • The k-nearest-neighbor score is 1 when the prediction falls within the k nearest reference representations around the encoded true observation, and 0 otherwise.The neighborhood includes the true node itself, so the score tests whether the prediction’s rank is at most k.
  • Mean reciprocal rank is defined as the average inverse rank of the predicted representation.

D ARCHITECTURE AND HYPERPARAMETERS

The models use convolutional object extraction, object-wise multilayer encoders, and graph-based transition modeling, with contrastive and reconstruction-based baselines for comparison. Architectures and training settings vary across environments and include ablations that remove relational structure, factorization, or contrastive learning.

  • Object extractor: Object extractors use environment-specific CNNs whose final feature maps produce five object slots and sigmoid object masks.The 3D cubes extractor uses four convolutional layers, while the other environments use simpler or two-layer CNN variants.
  • Object encoder: Flattened object features are encoded by an MLP with two 512-unit hidden layers, ReLU activations, and LayerNorm before the second hidden-layer activation.The resulting object representations have environment-dependent input dimensionalities, including 2500 dimensions for 3D cubes and 1000 for Atari and physics environments.
  • Transition model: The GNN transition model uses node and edge MLPs sharing the object encoder’s two-layer, 512-unit, LayerNorm-and-ReLU architecture.
  • Baselines: The VAE-based World Model baseline first trains a 32-dimensional autoencoder or VAE, then freezes it and trains a transition model on latent representations.The second stage uses mean-squared error; for the VAE, transitions use predicted latent means.
  • Baselines: Binary cross entropy was generally more stable to optimize and produced better results than mean squared error for decoder-based baselines.
  • Ablations: Ablations remove GNN edge updates, scene factorization, or contrastive learning to test the roles of relational transitions, object factorization, and the contrastive objective.
Loading 1911.12247v2…