SeFA-Policy: Fast and Accurate Visuomotor Policy Learning with Selective Flow Alignment

[Project page] [Paper]

Rong Xue^1*, Jiageng Mao^1*, Mingtong Zhang¹, Yue Wang¹,

¹University of Southern California,

To reproduce our simulation benchmark results, install our conda environment on a Linux machine with Nvidia GPU. On Ubuntu 20.04 you need to install the following apt packages for mujoco:

$ sudo apt install -y libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf

We provide two ways of installation:

1. Dockerfile under repo root is for docker installation:

   $ docker build -t sefapolicy:latest .

2. We recommend Mambaforge instead of the standard anaconda distribution for faster installation:

   $ mamba env create -f environment.yaml
   $ mamba activate sefapolicy

   but you can use conda as well:

   $ conda env create -f environment.yaml
   $ conda activate sefapolicy

   Then install extra packages:

   (sefapolicy)$ pip install garage --no-deps
   (sefapolicy)$ cd third_party
   (sefapolicy)[third_party]$ cd dexart-release && pip install -e . && cd ..
   (sefapolicy)[third_party]$ cd rrl-dependencies && pip install -e mj_envs/. && pip install -e mjrl/. && cd .. && cd ..
   (sefapolicy)$ pip install natsort

Under the repo root, create data subdirectory:

[sefa_policy]$ mkdir data && cd data

For Push-T, Franka Kitchen and RoboMimic, download the corresponding zip file from https://diffusion-policy.cs.columbia.edu/data/training/.

[data]$ wget https://diffusion-policy.cs.columbia.edu/data/training/pusht.zip
[data]$ wget https://diffusion-policy.cs.columbia.edu/data/training/kitchen.zip
[data]$ wget https://diffusion-policy.cs.columbia.edu/data/training/robomimic_image.zip

Extract training data:

[data]$ unzip pusht.zip && rm -f pusht.zip && unzip kitchen.zip && rm -f kitchen.zip && unzip robomimic_image.zip && rm -f robomimic_image.zip

For the 50 tasks in Meta-World, download mt50 into data folder and extract the zip file:

[data]$ unzip metaworld.zip && rm -f metaworld.zip

For Adroit tasks, follow 3D Diffusion Policy to generate training samples and save them into data/adroit.

Activate conda environment:

[sefa_policy]$ mamba activate sefapolicy

Login to wandb (if you haven't already). If you do not want to use wandb, please skip this step and pass use_wandb=false during training.

(sefapolicy)[sefa_policy]$ wandb login

The configuration files of ours and baselines are provided under folder sefa_policy/config. Here we take Meta-World dataset as an example. A task.env_name is required when training tasks in Meta-World.

This will launch training with default seed 42 on GPU 0:

python train.py --config-name=train_sefa_unet_hybrid_metaworld task.env_name=[env_name]

When generating couplings via Selective Flow Alignment,

1. pass mode=gen_data,
2. set base_ckpt as the designated path to your base checkpoint,
3. name sefa_name as your SeFA data's name.

python train.py --config-name=train_sefa_unet_hybrid_metaworld task.env_name=[env_name] mode=gen_data base_ckpt=[path to your base checkpoint] sefa_name=[your SeFA data's name]

When training SeFA Policy,

1. pass mode=sefa,
2. set base_ckpt as the path where your base checkpoint is saved,
3. pass sefa_name as the same SeFA data's name used when generating couplings via Selective Flow Alignment.

python train.py --config-name=train_sefa_unet_hybrid_metaworld task.env_name=[env_name] mode=sefa base_ckpt=[path to your base checkpoint] sefa_name=[your SeFA data's name]

Run the evaluation script:

(sefapolicy)[sefa_policy]$ python eval.py --checkpoint [path to your checkpoint] --output_dir data/metaworld_eval_output --device cuda:0

This will generate the following directory structure:

(sefapolicy)[sefa_policy]$ tree data/metaworld_eval_output
data/metaworld_eval_output
├── eval_log.json
└── media
    ├── 1fxtno84.mp4
    └── 3p7jqn32.mp4

1 directory, 3 files

eval_log.json contains metrics that is logged to wandb during training:

(sefapolicy)[sefa_policy]$ cat data/metaworld_eval_output/eval_log.json
{
  "test/mean_success_rate": ...,
  "test/sim_max_reward_[seed_0]": ...,
  "test/sim_max_reward_[seed_1]": ...,
...
  "train/sim_video_1": "data/metaworld_eval_output/media/2fo4btlf.mp4"
}

This repository is released under the MIT license. See LICENSE for additional details.

• The Diffusion Policy baseline is adapted from Cheng Chi's implementation.
• The Robomimic tasks and ObservationEncoder are used extensively in this project.
• The Push-T task is adapted from IBC.
• The Adroit tasks are adapted from 3D Diffusion Policy.