flowm

KO | EN

gitlite — search

#python #nodejs #docker #typescript #mongodb #reactjs #onedionys #package-manager #tea #javascript #ai #java

★ 54 Open GitHub ↗

Full model implementation for Flow Equivariant World Models (ICML 2026), world models with memory for dynamic scenes

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<h1 align="center"> Flow Equivariant World Models (ICML 2026)</h1>

 
 <a href="https://hansenlillemark.com/">Hansen Jin Lillemark*1, 2</a>
 ·
 <a href="https://huskydoge.github.io/">Benhao Huang*3</a>
 ·
 <a href="https://fnzhan.com/">Fangneng Zhan1</a>
 ·
 <a href="https://yilundu.github.io/">Yilun Du1</a>
 ·
 <a href="https://akandykeller.github.io/">T. Anderson Keller1</a>
 ·
 
 *Equal contribution 1Harvard 2UC San Diego 3Carnegie Mellon University
 
 <h3 align="center"><a href="https://arxiv.org/abs/2601.01075">Paper</a> | <a href="https://flowequivariantworldmodels.github.io/">Website</a> | <a href="https://huggingface.co/flowm123">Models & Datasets</a></h3>

Welcome to the codebase for the paper [Flow Equivariant World Models: Memory for Partially Observed Dynamic Environments](https://flowequivariantworldmodels.github.io/), accepted at ICML 2026! 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), Diffusion State Space Model (DFoT-SSM), and DreamerV3-based RSSM baselines. Please follow the steps below to run experiments.

Contents:
- [Environment Setup](#environment)
- [Dataset Download](#dataset)
- [Model Checkpoint Download](#model-checkpoints)
- [Inference commands (replicating results)](#inference-commands)
- [Training commands](#training-commands)
- [Code Walkthrough](#code-walkthrough)
- [Citation](#citation)

# 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](https://github.com/Dao-AILab/flash-attention/issues/1708#issuecomment-2995556229):

```
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](https://github.com/state-spaces/mamba/issues/55#issuecomment-1858638484) and [for mamba ssm here](https://github.com/state-spaces/mamba/issues/719).
```
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

To download the datasets used in the paper, please use the script `./download_datasets.sh`. This script specifies the dataset, split, and configuration. To download all splits and all data, use the following command:

```bash
bash ./download_datasets.sh --download-workers 16 --extract-workers 16
```

See the file for more options. To download just the `blockworld tex validation` split, run the following:

```bash
bash ./download_datasets.sh --dataset blockworld --configs tex --splits validation --download-workers 16 --extract-workers 16
```

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.

For the code to generate the 3D Dynamic Blockworld dataset, please see this repo: [Link](https://github.com/hlillemark/Miniworld).

## Model Checkpoints
If intending to run from scratch, you may skip this part (except the VAE if training DFoT or DFoT-SSM from scratch). We provide checkpoints for the primary result splits.

The following script will download all checkpoints for all models for all datasets (FloWM, DFoT, DFoT-SSM, VAE), requiring ~15GB of space:

```bash
bash ./download_ckpts.sh
```

### 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. Downloading from the above script should work automatically, but look into this file if you find any errors.

# Inference Commands

Please see [the wiki](https://github.com/hlillemark/flowm/wiki/Inference-and-Reproducing-Results) for full commands! For example, you can run inference on the Textured Blockworld Validation set for 280 frames using FloWM with the following command:

```bash
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](https://github.com/hlillemark/flowm/wiki/Training) for full commands! For example, you can train FloWM on the Textured Blockworld Training set with the following command:

```bash
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](https://github.com/hlillemark/flowm/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.

# Citation

If you find our work interesting or helpful, please cite it using the following BibTex:

```bibtex
@inproceedings{
lillemark2026flow,
title={Flow Equivariant World Models: Structured Memory for Dynamic Environments},
author={Hansen Lillemark and Benhao Huang and Fangneng Zhan and Yilun Du and T. Anderson Keller},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://arxiv.org/abs/2601.01075}
}
```

# Acknowledgement

This repo is forked from [Boyuan Chen](https://boyuan.space/)'s research template [repo](https://github.com/buoyancy99/research-template).

flowm

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README