Flow Equivariant World Models

Anonymous Authors

Welcome to the codebase for the paper Flow Equivariant World Models: Memory for Partially Observed Dynamic Environments! This repository contains the code and commands for training and running inference on FloWM models in the 2D (MNIST World) and 3D (Dynamic Block World) environments. It also contains the code for running the Diffusion Forcing Transformer (DFoT) and Diffusion State Space Model (DFoT-SSM) baselines. Please follow the steps below to run experiments.

Contents:

• Setup
  • Environment
  • Dataset
  • Model Checkpoints
    • Checkpoint map
• Inference Commands
• Training Commands
• Code Walkthrough

For this anonymous version, we do not have the datasets and checkpoints available. However, they are all ready and will be available upon official release. If requested by a reviewer, we can figure out a way to anonymously host them.

Setup

Environment

This command will create a conda environment named flowm. If you are using cuda version 13, then please also attach the flag --cuda13. The other default is --cuda12.4. If the flash attention install breaks, see here for more help:

bash ./setup_flowm_env.sh

If you would like to run the dfot-ssm baseline, then please run the following command instead: This script is more fragile and you may need to look into the issues for causal_conv1d here and for mamba ssm here.

bash ./setup_flowm_env.sh --download-ssm
conda activate flowm

To track runs on wandb, make a copy of configurations/my_secrets_template.yaml, rename it my_secrets.yaml, and then fill in your wandb api key there. You will also need to fill in entity and project in config.yaml.

Dataset

Dataset download will be available upon official release.

After downloading, the datasets will appear in the ./data/ folder. When running the code, metadata will be automatically generated. If you only downloaded a subset of the data, then please delete the corresponding metadata for the split you are now downloading (which will be generated as empty). For instance, if you only downloaded the tex blockworld validation split and now want to download the tex blockworld training split, delete ./data/blockworld/metadata/tex_training.pt before running the training code.

Each configuration has training and validation splits. The blockworld dataset has configs dynamic (main results in paper), static, and tex. The mnist_world dataset has configs dynamic_po (main results in paper), static_po, dynamic_fo, and dynamic_fo_no_sm, where po means partially observed, fo means fully observed, and no_sm means no self motion. Change the configurations to the script accordingly to download the splits of need. Only the validation splits are necessary to replicate the paper's results.

Model Checkpoints

Model checkpoint download will be available upon official release.

Checkpoint map

To keep track of downloaded checkpoints, we provide a default checkpoint map at configurations/ckpt_map/default.yaml. This creates a pointer to downloaded checkpoints when the code actually runs.

Inference Commands

Please see the wiki for full commands! For example, you can run inference on the Textured Blockworld Validation set for 280 frames using FloWM with the following command:

python -m main shortcode=exp/blockworld/flowm/infer/metrics_140/tex_70ctx +name=infer_blockworld_flowm_tex_70ctx_140 algorithm=flowm_video dataset=blockworld ckpt_map=default

Training Commands

Please see the wiki for full commands! For example, you can train FloWM on the Textured Blockworld Training set with the following command:

python -m main shortcode=exp/blockworld/flowm/train/tex +name=blockworld_flowm_10m_tex_50c90p algorithm=flowm_video dataset=blockworld ckpt_map=default

Code Walkthrough

Please see our wiki for more information on the code itself. This wiki includes information on the configurations, FloWM and baseline algorithm implementations, VAE information, adding a new dataset, and more.