MoFlow-IMLE

The official PyTorch implementation of CVPR'25 paper named "MoFlow: One-Step Flow Matching for Human Trajectory Forecasting via Implicit Maximum Likelihood Estimation Distillation"

📢 News

• [2025/02] 🎉 Paper accepted to CVPR 2025!
• [2025/03] 🚀 Code repository is now public
• [2025/03] 📊 Preprocessed datasets are available in the ./data/[datasets] folder
• [2025/03] 🔥 Model checkpoints and samples from MoFlow of NBA are up on Hugging Face
• [2025/03] 🔥 Model checkpoints of MoFlow on ETH-UCY and SDD datasets are up on Hugging Face

📝 Overview

Human trajectory forecasting is a challenging task that involves two key objectives:

1. Predicting future trajectories with high precision.
2. Generating diverse future movements that capture the inherent uncertainty in human decision-making.

We address this challenge by introducing MoFlow, a novel Motion prediction conditional Flow matching model that generates K-shot future trajectories for all agents in a given scene. Our key contributions include:

• A novel flow matching loss function that ensures both accuracy and diversity in trajectory predictions
• An innovative distillation method for flow models using Implicit Maximum Likelihood Estimation (IMLE)
• A one-step flow matching approach that significantly reduces computational complexity

Our MoFlow can generate diverse trajectories that are physically and socially plausible. Moreover, our one-step student model IMLE is 100 times faster than the teacher flow model during sampling!

🛠️ Environment Setup

# Create a new conda environment
conda create -n moflow python=3.11
conda activate moflow

# Install PyTorch (adjust version based on your CUDA version)
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

# Install other dependencies
pip install -r requirements.txt

📊 Datasets

This project supports three major trajectory datasets:

1. NBA SportVU Dataset

• Description: High-resolution basketball player trajectories from NBA games
• Contents: Detailed player positions and movement data
• Location: Preprocessed data available in data/nba/
• Source: Compatible with datasets used in LED and GroupNet
• We select the first 32500 scenes from the training data and 12500 scenes from the test data. The original datasets can be found in this Google Drive

2. ETH-UCY Dataset

• Description: Pedestrian trajectories across 5 diverse scenes (ETH, Hotel, Univ, Zara1, and Zara2)

• Location: Preprocessed data available in data/eth_ucy/

• Sources:

  • Based on the original ETH-UCY dataset used in SocialGAN
  • Dataset files courtesy of EigenTrajectory

• Note: This version differs from the one used in LED. You can specify --data_source to select the version of your interest.

3. Stanford Drone Dataset (SDD)

• Description: Diverse trajectory data from a university campus environment
• Contents: Multi-agent trajectories including pedestrians, cyclists, and vehicles
• Location: Preprocessed data available in data/sdd/
• Versions:
  • Primary version sourced from TUTR
  • Additional experiments conducted using NSP version

🗂️ Project Structure

• README.md: Project documentation and overview.
• cfg/: Configuration files for different datasets:
  • eth_ucy/, nba/, sdd/: Each contains cor_fm.yml and imle.yml for specific model configurations.
• data/: Dataset-related files and scripts:
  • Data loaders: dataloader_eth_ucy.py, dataloader_nba.py, dataloader_sdd.py.
  • eth_ucy/: Extensive subdirectories (LED, original, raw, etc.) with .npy and .pkl files for train/test splits (e.g., eth_data_train.npy, zara1_test.pkl) and a download script (download_eth_ucy_dataset.sh).
  • nba/: Contains nba_train.npy and nba_test.npy.
  • sdd/: Includes sdd_train.pkl, sdd_test.pkl, and NSP variants (sdd_nsp_train.pkl).
  • store_pickle_eth_files.py: Utility for processing ETH dataset files.
• eval_*.py: Evaluation scripts for ETH (eval_eth.py), NBA (eval_nba.py), and SDD (eval_sdd.py).
• fm_*.py: Flow-matching scripts for ETH (fm_eth.py), NBA (fm_nba.py), and SDD (fm_sdd.py).
• imle_*.py: IMLE scripts for ETH (imle_eth.py), NBA (imle_nba.py), and SDD (imle_sdd.py).
• images/: Visual assets, including moflow-architecture.png and moflow-imle.png.
• models/: Model definitions and utilities:
  • Core files: backbone.py, backbone_eth_ucy.py, flow_matching.py, imle.py.
  • Submodules: context_encoder/ (e.g., eth_encoder.py), motion_decoder/ (e.g., mtr_decoder.py), utils/ (e.g., common_layers.py).
• requirements.txt: Project dependencies.
• trainer/: Training logic in denoising_model_trainers.py and imle_trainers.py.
• utils/: Helper scripts like config.py, normalization.py, and utils.py.

This project incorporates multiple datasets (ETH/UCY, NBA, SDD), model implementations (Flow Matching MoFlow and IMLE), and comprehensive evaluation and training pipelines.

🚀 Usage

<img src="images/moflow-architecture.png" alt="MoFlow-IMLE architecture" width="50%">

Our teacher MoFlow model and student IMLE model share almost the same architecture. However, the student IMLE model does not require time conditioning, eliminating the neural network component that processes it. We can load the weights from our pre-trained teacher model via --load_pretrained to accelerate the training further.

🎯 Training and Evaluation Pipeline

For each dataset, we need to train our MoFlow teacher model first:

1. Teacher Model Training

### NBA dataset
python fm_nba.py --exp <exp_name> --tied_noise --fm_in_scaling --checkpt_freq 5 --batch_size 192 --init_lr 1e-3

### ETH dataset
python3 fm_eth.py --exp <exp_name> --rotate --rotate_time_frame 6 --subset eth --tied_noise --fm_in_scaling --checkpt_freq 1 --batch_size 32 --init_lr 1e-4 
python3 fm_eth.py --exp <exp_name> --rotate --rotate_time_frame 6 --subset hotel --tied_noise --fm_in_scaling --checkpt_freq 1 --batch_size 48 --init_lr 1e-4 
python3 fm_eth.py --exp <exp_name> --rotate -rotate_time_frame 6 --subset univ --tied_noise --fm_in_scaling --checkpt_freq 1 --batch_size 48 --init_lr 1e-4 
python3 fm_eth.py --exp <exp_name> --rotate --rotate_time_frame 6 --subset zara1 --tied_noise --fm_in_scaling --checkpt_freq 1 --batch_size 32 --init_lr 1e-4 
python3 fm_eth.py --exp <exp_name> --rotate --rotate_time_frame 6 --subset zara2 --tied_noise --fm_in_scaling --checkpt_freq 1 --batch_size 32 --init_lr 1e-4 

### SDD dataset
python fm_sdd.py --exp <exp_name> --rotate --rotate_time_frame 6 --tied_noise --fm_in_scaling --checkpt_freq 1 --batch_size 48 --init_lr 1e-4 --perturb_ctx 0.03

2. Teacher Model Sampling

### NBA dataset
python3 eval_nba.py --ckpt_path <path_to_nba_teacher_checkpoint> \
--batch_size 1000 --sampling_steps 100 --solver lin_poly --lin_poly_p 5 --lin_poly_long_step 1000 --save_samples --eval_on_train

### ETH dataset
python3 eval_eth.py \
--ckpt_path <path_to_eth_teacher_checkpoint> \
--subset eth --rotate --rotate_time_frame 6 \
--batch_size 1000 --sampling_steps 100 --solver lin_poly --lin_poly_p 5 --lin_poly_long_step 1000 --save_samples --eval_on_train

python3 eval_eth.py \
--ckpt_path <path_to_hotel_teacher_checkpoint> \
--subset hotel --rotate --rotate_time_frame 6 \
--batch_size 1000 --sampling_steps 100 --solver lin_poly --lin_poly_p 5 --lin_poly_long_step 1000 --save_samples --eval_on_train

python3 eval_eth.py \
--ckpt_path <path_to_univ_teacher_checkpoint> \
--subset univ --rotate --rotate_time_frame 6 \
--batch_size 1000 --sampling_steps 100 --solver lin_poly --lin_poly_p 5 --lin_poly_long_step 1000 --save_samples --eval_on_train

python3 eval_eth.py \
--ckpt_path <path_to_zara1_teacher_checkpoint> \
--subset zara1 --rotate --rotate_time_frame 6 \
--batch_size 1000 --sampling_steps 100 --solver lin_poly --lin_poly_p 5 --lin_poly_long_step 1000 --save_samples --eval_on_train

python3 eval_eth.py \
--ckpt_path <path_to_zara2_teacher_checkpoint> \
--subset zara2 --rotate --rotate_time_frame 6 \
--batch_size 1000 --sampling_steps 100 --solver lin_poly --lin_poly_p 5 --lin_poly_long_step 1000 --save_samples --eval_on_train

### SDD dataset
python eval_sdd.py --ckpt_path <path_to_sdd_teacher_checkpoint> \
--rotate --rotate_time_frame 6 --batch_size 1000 --sampling_steps 100 --solver lin_poly --lin_poly_p 5 --lin_poly_long_step 1000 --save_samples --eval_on_train

Let's train our IMLE student model now. First, we need to move the pickle files containing the teacher MoFlow samples from [results_dir]/samples/ to ./data/[datasets]/imle/. Completing this step is essential to the successful training of the student model.

3. Student Model Training

### NBA dataset
python3 imle_nba.py --exp <exp_name> \
--checkpt_freq 1 --epochs 50 --batch_size 48 --init_lr 1e-3 --num_to_gen 20 \
--load_pretrained --ckpt_path <path_to_nba_teacher_checkpoint>

### E