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
First, download and set up the repo:
git clone https://github.com/DannyTran123/peva
cd pevaTo download and preprocess data, please follow the instructions at this github repo here for data preparation.
This project now supports the uv package manager for faster and more reliable dependency management:
# Quick setup with the provided script
uv syncSome 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"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 ipywidgetsexport 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 300export 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 300Or 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 300Or locally on one GPU for debug:
python train.py --config config/peva_cdit_xl.yaml --ckpt-every 2000 --eval-every 10000 --bfloat16 0 --epochs 300Note: 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.
directory to save evaluation results:
export RESULTS_FOLDER=/path/to/res_folder/
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 1python 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 1python isolated_peva_eval.py \
--datasets nymeria \
--gt_dir ${RESULTS_FOLDER}/gt \
--exp_dir ${RESULTS_FOLDER}/peva_cdit_xl \
--eval_types timeResults are saved in ${RESULTS_FOLDER}/peva_cdit_xl/nymeria
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,4python 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,4Specify the atomic action using the --atomic_action_type flag.
Available options:
-
Locomotion
move_forwardmove_rotate_leftmove_rotate_right
-
Left Hand
move_l_hand_upmove_l_hand_downmove_l_hand_leftmove_l_hand_right
-
Right Hand
move_r_hand_upmove_r_hand_downmove_r_hand_leftmove_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 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 rolloutResults 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 rolloutResults are saved in ${RESULTS_FOLDER}/peva_cdit_xl_move_forward/<dataset_name>
@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}
}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.
The code and model weights are licensed under Creative Commons Attribution-NonCommercial 4.0 International. See LICENSE.txt for details.