- Check
compute_optical_flow.shfor computing optical flow for your dataset. - Check
pretrain.shfor pretraining optical flow VAE. - Check
retrieve.shfor retrieving relevant data. - Check
train.shfor flow-guided training with target and retrieved data.
This codebase is based on BridgeData V2.
If you use this code and/or FlowRetrieval in your work, please cite the paper with:
@misc{lin2024flowretrieval,
title={FlowRetrieval: Flow-Guided Data Retrieval for Few-Shot Imitation Learning},
author={Li-Heng Lin and Yuchen Cui and Amber Xie and Tianyu Hua and Dorsa Sadigh},
year={2024},
eprint={2408.16944},
archivePrefix={arXiv},
primaryClass={cs.RO},
url={https://arxiv.org/abs/2408.16944},
}