Skip to content

Repository files navigation

Whole-Body Conditioned Egocentric Video Prediction, NeurIPS 2025
Official PyTorch Implementation

This repo contains the official PyTorch implementation of Whole-Body Conditioned Egocentric Video Prediction- the Conditional Diffusion Transformer (CDiT) model training and evaluation code. See the project page for additional results.

Whole-Body Conditioned Egocentric Video Prediction
Yutong Bai, Danny Tran, Amir Bar, Yann LeCun, Trevor Darrell, Jitendra Malik
UC Berkeley (BAIR), AI at Meta, New York University

Setup

First, download and set up the repo:

git clone https://github.com/DannyTran123/peva
cd peva

Data

To download and preprocess data, please follow the instructions at this github repo here for data preparation.

Requirements:

Option 1: Using uv (Recommended)

This project now supports the uv package manager for faster and more reliable dependency management:

# Quick setup with the provided script
uv sync

Additional Dependencies

Some packages may need to be installed from GitHub:

# If using uv
uv add "dreamsim @ git+https://github.com/ssundaram/dreamsim.git"
uv add "cdfvd @ git+https://github.com/universome/fvd.git"
uv add "aitviewer @ git+https://github.com/eth-ait/aitviewer.git"

Option 2: Using conda/mamba (Legacy)

mamba create -n peva python=3.10
mamba activate peva
pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu126
mamba install ffmpeg
pip3 install decord einops evo transformers diffusers tqdm timm notebook dreamsim torcheval lpips ipywidgets

Training

Using uv with torchrun (Recommended):

export NUM_NODES=8
export HOST_NODE_ADDR=<HOST_ADDR>
export CURR_NODE_RANK=<NODE_RANK>

uv run torchrun \
  --nnodes=${NUM_NODES} \
  --nproc-per-node=8 \
  --node-rank=${CURR_NODE_RANK} \
  --rdzv-backend=c10d \
  --rdzv-endpoint=${HOST_NODE_ADDR}:29500 \
  train.py --config config/peva_cdit_xl.yaml --ckpt-every 2000 --eval-every 10000 --bfloat16 1 --epochs 300

Using conda/mamba (Legacy):

export NUM_NODES=8
export HOST_NODE_ADDR=<HOST_ADDR>
export CURR_NODE_RANK=<NODE_RANK>

torchrun \
  --nnodes=${NUM_NODES} \
  --nproc-per-node=8 \
  --node-rank=${CURR_NODE_RANK} \
  --rdzv-backend=c10d \
  --rdzv-endpoint=${HOST_NODE_ADDR}:29500 \
  train.py --config config/peva_cdit_xl.yaml --ckpt-every 2000 --eval-every 10000 --bfloat16 1 --epochs 300

Or using submitit and slurm (8 machines of 8 gpus):

python submitit_train_cw.py --nodes 8 --partition learn --qos explore --config ./config/peva_cdit_xl.yaml --ckpt-every 2000 --eval-every 10000 --bfloat16 1 --epochs 300

Or locally on one GPU for debug:

python train.py --config config/peva_cdit_xl.yaml --ckpt-every 2000 --eval-every 10000 --bfloat16 0 --epochs 300

Note: torch compile can lead to ~40% faster training speed. However, it might lead to instabilities and inconsistent behvaior across different pytorch versions. Use carefuly. Specify using torch compile in the config file.

Evaluation

directory to save evaluation results: export RESULTS_FOLDER=/path/to/res_folder/

Evaluate on single time step prediction

1. Prepare ground truth frames for evaluation (one-time)

python isolated_peva_infer.py \
    --exp config/peva_cdit_xl.yaml \
    --datasets nymeria \
    --batch_size 96 \
    --num_workers 12 \
    --eval_type time \
    --output_dir ${RESULTS_FOLDER} \
    --gt 1

2. Predict future state given action

python isolated_peva_infer.py \
    --exp config/peva_cdit_xl.yaml \
    --ckp 0100000 \
    --datasets nymeria \
    --batch_size 64 \
    --num_workers 12 \
    --eval_type time \
    --output_dir ${RESULTS_FOLDER} \
    --compile 1

3. Report metrics compared to GT (LPIPS, DreamSim, FID)

python isolated_peva_eval.py \
    --datasets nymeria \
    --gt_dir ${RESULTS_FOLDER}/gt \
    --exp_dir ${RESULTS_FOLDER}/peva_cdit_xl \
    --eval_types time

Results are saved in ${RESULTS_FOLDER}/peva_cdit_xl/nymeria

Evaluate on following ground truth trajectories

python isolated_peva_infer.py \
    --exp config/peva_cdit_xl.yaml \
    --datasets nymeria \
    --batch_size 96 \
    --num_workers 12 \
    --eval_type rollout \
    --output_dir ${RESULTS_FOLDER} \
    --gt 1 \
    --rollout_fps_values 1,4

2. Simulate a GT trajectory using PEVA

python isolated_peva_infer.py \
    --exp config/peva_cdit_xl.yaml \
    --ckp 0100000 \
    --datasets <dataset_name> \
    --batch_size 64 \
    --num_workers 12 \
    --eval_type rollout \
    --output_dir ${RESULTS_FOLDER} \
    --rollout_fps_values 1,4

Specify the atomic action using the --atomic_action_type flag.

Available options:

  • Locomotion

    • move_forward
    • move_rotate_left
    • move_rotate_right
  • Left Hand

    • move_l_hand_up
    • move_l_hand_down
    • move_l_hand_left
    • move_l_hand_right
  • Right Hand

    • move_r_hand_up
    • move_r_hand_down
    • move_r_hand_left
    • move_r_hand_right
  • Pose Estimation

    • pose_estimation
python isolated_peva_infer.py \
    --exp config/peva_cdit_xl.yaml \
    --ckp 0100000 \
    --datasets <dataset_name> \
    --batch_size 64 \
    --num_workers 12 \
    --eval_type rollout \
    --output_dir ${RESULTS_FOLDER} \
    --rollout_fps_values 1,4 \
    --atomic_action_type move_forward

3. Report metrics compared to GT trajectories (LPIPS, DreamSim, FID)

    python isolated_peva_eval.py \
        --datasets nymeria \
        --gt_dir ${RESULTS_FOLDER}/gt \
        --exp_dir ${RESULTS_FOLDER}/peva_cdit_xl \
        --rollout_fps_values 1,4 \
        --eval_types rollout

Results are saved in ${RESULTS_FOLDER}/peva_cdit_xl/<dataset_name>

For the different atomic actions:

    python isolated_peva_eval.py \
        --datasets nymeria \
        --gt_dir ${RESULTS_FOLDER}/gt \
        --exp_dir ${RESULTS_FOLDER}/peva_cdit_xl_move_forward \
        --rollout_fps_values 1,4 \
        --num_sec_eval 2 \
        --eval_types rollout

Results are saved in ${RESULTS_FOLDER}/peva_cdit_xl_move_forward/<dataset_name>

BibTeX

@article{bai2025whole,
  title={Whole-Body Conditioned Egocentric Video Prediction},
  author={Bai, Yutong and Tran, Danny and Bar, Amir and LeCun, Yann and Darrell, Trevor and Malik, Jitendra},
  journal={arXiv preprint arXiv:2506.21552},
  year={2025}
}

Acknowledgments

The authors thank Rithwik Nukala for his help in annotating atomic actions. We thank Katerina Fragkiadaki, Philipp Krähenbühl, Bharath Hariharan, Guanya Shi, Shubham Tsunami and Deva Ramanan for the useful suggestions and feedbacks for improving the paper; Jianbo Shi for the discussion regarding control theory; Yilun Du for the support on Diffusion Forcing; Brent Yi for his help in human motion related works and Alexei Efros for the discussion and debates regarding world models. This work is partially supported by the ONR MURI N00014-21-1-2801.

License

The code and model weights are licensed under Creative Commons Attribution-NonCommercial 4.0 International. See LICENSE.txt for details.

About

Official code for the NeurIPS 2025 paper "Whole-Body Conditioned Egocentric Video Prediction".

Resources

Code of conduct

Contributing

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages