lpwm

<p align="center"> <img src="https://img.shields.io/badge/conference-ICLR%202026-orange" /> <img src="https://img.shields.io/badge/pytorch-%E2%89%A52.6-blue" /> <img src="https://img.shields.io/badge/license-MIT-green" /> <img src="https://img.shields.io/badge/arXiv-2603.04553-b31b1b" /> </p>

Official PyTorch implementation of the paper "Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modeling".

<h1 align="center"> <br> Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modeling <br> </h1> <h2 align="center"> ICLR 2026 Oral <br> </h2> <h3 align="center"> <a href="https://taldatech.github.io">Tal Daniel</a> • <a href="https://carl-qi.github.io/">Carl Qi</a> • <a href="https://danhrmti.github.io/">Dan Haramati</a> • <a href="https://lambda.ai/research">Amir Zadeh</a> • <a href="https://lambda.ai/research">Chuan Li</a> • <a href="https://avivt.github.io/avivt/">Aviv Tamar</a> • <a href="https://www.cs.cmu.edu/~dpathak/">Deepak Pathak</a> • <a href="https://r-pad.github.io/">David Held</a> </h3> <h3 align="center">Official repository of Deep Latent Particles v3 (DLPv3) & LPWM</h3> <h4 align="center"><a href="https://arxiv.org/abs/2603.04553">Arxiv</a> • <a href="https://taldatech.github.io/lpwm-web">Project Website</a> • <a href="https://youtu.be/aZeaCyXJjYI">Video</a> • <a href="https://openreview.net/forum?id=lTaPtGiUUc">OpenReview</a></h4> <h4 align="center"> <a href="https://colab.research.google.com/github/taldatech/lpwm"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a> </h4> <p align="center"> <img src="https://github.com/taldatech/lpwm-web/blob/main/assets/sketchy_particlegrid.gif?raw=true" height="150"><br> <img src="https://github.com/taldatech/lpwm-web/blob/main/assets/langtable_476_lpwm_gen_comp.gif?raw=true" height="130"> <img src="https://github.com/taldatech/lpwm-web/blob/main/assets/bridge_313_lang_comp.gif?raw=true" height="130"><br> </p>

Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modeling

Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modeling<br> Tal Daniel , Carl Qi, Dan Haramati, Amir Zadeh, Chuan Li, Aviv Tamar, Deepak Pathak, David Held<br>

Abstract: We introduce Latent Particle World Model (LPWM), a self-supervised object-centric world model scaled to real-world multi-object datasets and applicable in decision-making. LPWM autonomously discovers keypoints, bounding boxes, and object masks directly from video data, enabling it to learn rich scene decompositions without supervision. Our architecture is trained end-to-end purely from videos and supports flexible conditioning on actions, language, and image goals. LPWM models stochastic particle dynamics via a novel latent action module and achieves state-of-the-art results on diverse real-world and synthetic datasets. Beyond stochastic video modeling, LPWM is readily applicable to decision-making, including goal-conditioned imitation learning, as we demonstrate in the paper.

Citation

Daniel, T., Qi, C., Haramati, D., Zadeh, A., Li, C., Tamar, A., Pathak, D., & Held, D. (2026). Latent particle world models: Self-supervised object-centric stochastic dynamics modeling. In The Fourteenth International Conference on Learning Representations (ICLR 2026). https://openreview.net/forum?id=lTaPtGiUUc

@inproceedings{
daniel2026latent,
title={Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modeling},
author={Tal Daniel and Carl Qi and Dan Haramati and Amir Zadeh and Chuan Li and Aviv Tamar and Deepak Pathak and David Held},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=lTaPtGiUUc}
}

Quickstart

# Install environment
conda env create -f environment.yml
conda activate dlp

# Train LPWM on Sketchy
python train_lpwm.py --dataset sketchy

# Generate videos with a pretrained model
python generate_lpwm_video_prediction.py --help

• lpwm
• Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modeling
  • Citation
  • Prerequisites
  • Model Zoo - Pretrained Models
  • Datasets
  • Conditioning Types, Multi-View and Examples
  • Attribute Variable Names
  • DLPv3 and LPWM - Training
  • Training Logs, Progress Bar and Saved Images/Videos
  • LPWM - Evaluation
    • PSNR, SSIM, LPIPS
    • FVD
  • LPWM - Video Prediction and Generation with Pre-trained Models
  • Example Usage, Documentation and Notebooks
    • DLPv3 and LPWM Walkthrough Tutorial Notebook
  • Repository Organization

Prerequisites

We provide an environment.yml file which installs the required packages in a conda environment named torch. Alternatively, you can use pip to install requirements.txt.

• Create the environment with: conda env create -f environment.yml.

| Library | Version | Notes | |-------------------|--------------|-------------------------------------------------| | Python | > = 3.9 | - | | torch | > = 2.6.0 | - | | torchvision | > = 0.21 | - | | matplotlib | > = 3.10.0 | - | | numpy | > = 1.24.3 | - | | h5py | > = 3.13.0 | Some datasets (e.g., Balls) use H5/HDF | | py-opencv | > = 4.11 | For plotting | | tqdm | > = 4.67.0 | - | | scipy | > = 1.15 | - | | scikit-image | > = 0.25.2 | Required to generate the "Shapes" dataset | | ffmpeg | = 4.2.2 | Required to generate video files | | accelerate | > = 1.5.0 | For multi-GPU training | | imageio | > = 2.6.1 | For creating video GIFs | | piqa | > = 1.3.1 | For image evaluation metrics: LPIPS, SSIM, PSNR | | einops | > = 0.81 | - | | huggingface-hub | > = 0.29 | Downloading checkpoint and datasets | | notebook | > = 6.5.4 | To run Jupyter Notebooks |

For a manual installation guide, see docs/installation.md.

Model Zoo - Pretrained Models

• We provide pre-trained checkpoints for datasets used in the paper.
• All model checkpoints should be placed inside the /checkpoints directory.

The following table lists the available pre-trained checkpoints and where to download them.

| Model Type | Dataset | Link | |---------------|-------------------------|--------------------------------------------------------------------------------------| | LPWM | Sketchy (128x128) | MEGA.nz | | LPWM-Action | Sketchy (128x128) | MEGA.nz | | LPWM | BAIR (128x128) | MEGA.nz | | LPWM-Language | LangaugeTable (128x128) | MEGA.nz | | LPWM-Language | Bridge (128x128) | MEGA.nz |

Datasets

| Dataset | Notes | Link | |----------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------| | Sketchy | Original dataset from Deepmind. We use a subset and provide the pre-processed data on HF. | HF | | Bridge | Pre-processed videos from the BRIDGE dataset with T5-