This repository provides the official PyTorch implementation for Contrastive Learning of Structured World Models (C-SWMs). It addresses the challenge of learning structured world models from raw sensory data, enabling the discovery of object-based representations without direct supervision. The target audience includes researchers and engineers in reinforcement learning and representation learning.

How It Works

C-SWMs employ a contrastive learning approach to model compositional environments. Each state embedding is structured as a set of object representations and their relations, managed by a graph neural network. This design facilitates the discovery of objects from pixel observations and allows for independent manipulation of objects by an agent.

Quick Start & Requirements

• Installation: Requires Python 3.6/3.7, PyTorch 1.2, OpenAI Gym 0.12.0, atari-py 0.1.4, scikit-image 0.15.0, and Matplotlib 3.0.2.
• Dataset Generation: Scripts are provided for generating datasets for 2D Shapes, 3D Cubes, Atari Pong, Space Invaders, and a 3-Body Gravitational Physics simulation.
• Training & Evaluation: Example commands for training and evaluating models on these datasets are included.
• Documentation: arXiv preprint

Highlighted Details

• Learns interpretable object-based representations.
• Outperforms pixel reconstruction models in structured environments.
• Evaluated on compositional environments, Atari games, and physics simulations.

Maintenance & Community

• The project is associated with authors from the University of Amsterdam.
• No explicit community channels (e.g., Discord, Slack) are mentioned.

Licensing & Compatibility

• The repository does not explicitly state a license. However, the authors are affiliated with a university, suggesting potential academic use. Commercial use or linking with closed-source projects would require clarification.

Limitations & Caveats

The project specifies older versions of Python and key libraries (PyTorch 1.2, Gym 0.12.0), which may pose compatibility challenges with current development environments. The lack of an explicit license could be a barrier for commercial adoption.