README.md

Understanding Physical Dynamics with Counterfactual World Modeling

Rahul Venkatesh*¹ · Honglin Chen*^1* · Kevin Feigelis*¹ · Daniel M. Bear¹ · Khaled Jedoui¹ · Klemen Kotar¹ · Felix Binder² · Wanhee Lee¹ · Sherry Liu¹ · Kevin A. Smith³ · Judith E. Fan¹ · Daniel L. K. Yamins¹

(* equal contribution)

¹Stanford    ²UCSD    ³MIT

This work presents the Counterfactual World Modeling (CWM) framework. CWM is capable of counterfactual prediction and extraction of vision structures useful for understanding physical dynamics.

📣 News

• 2024-06-01: Release project page and codes

🔨 Installation

git clone https://github.com/rahulvenkk/cwm_release.git
pip install -e .

✨ Usage

To download and use a pre-trianed model run the following

from cwm.model.model_factory import model_factory
model = model_factory.load_model('vitbase_8x8patch_3frames_1tube')

This will automatically initialize the appropriate model class and download the specified weights to your $CACHE directory.

🔄 Pre-training

To train the model run the following script

./scripts/pretrain/3frame_patch8x8_mr0.90_gpu.sh