Dreamweaver: Learning Compositional World Models from Pixels

Author: Junyeob Baek, Yi-Fu Wu, Gautam Singh, Sungjin Ahn

Published to ICLR 2025

[paper] [project] [openreview]

Overview

Hightlighted Demo

Installation

git clone https://github.com/ahn-ml/dreamweaver-release.git
cd dreamweaver-release
pip install -r requirements.txt

Dataset

Download the required datasets using the following commands:

pip install gdown

# Download moving-sprites dataset
gdown "https://drive.google.com/uc?export=download&confirm=t&id=19qc4M_tlSTp1nOuuIK8L5gPEzi0FhKHb"

# Downalod moving-clevr dataset
gdown "https://drive.google.com/uc?export=download&confirm=t&id=1_-DO8Vs3dABCHmhBzNi3awKRppZwYc8W"

# Downalod moving-clevrtex dataset
gdown "https://drive.google.com/uc?export=download&confirm=t&id=1ZqT_SfkPAzNxNY00Zm9UGao7BnJVm9gr"

# Downalod dancing-sprites dataset
gdown "https://drive.google.com/uc?export=download&confirm=t&id=1PME2UU0_EdZZxo6ZQQvT-BHs-GlJ4KW7"

# Downalod dancing-clevr dataset
gdown "https://drive.google.com/uc?export=download&confirm=t&id=1wGFbJDXKJrFOo14FKdYmyrNv7xNrQh-k"

# Downalod moving-sprites-ood dataset
gdown "https://drive.google.com/uc?export=download&confirm=t&id=1qDiP1Q79ErBJehCEgP6HcONe07qPRhJ9"

# Downalod dancing-clevr-ood dataset
gdown "https://drive.google.com/uc?export=download&confirm=t&id=1Z6jfsnIYzY9OEXW2QW42StMiJHsm-pIf"

Make sure the datasets are placed in the src/dataset/data directory before running the training script.

How to Train

bash train.sh # run this first for sanity-check and change the hyperparameters appropreately

Citation

@inproceedings{
        baek2025dreamweaver,
        title={Dreamweaver: Learning Compositional World Models from Pixels},
        author={Junyeob Baek and Yi-Fu Wu and Gautam Singh and Sungjin Ahn},
        booktitle={The Thirteenth International Conference on Learning Representations},
        year={2025},
        url={https://openreview.net/forum?id=e5mTvjXG9u}
        }

Acknowledgements

Some modules in this code are referenced by singhgautam/sysbinder. We thank the authors for their great works!