Authors: Sateesh Kumar, Shivin Dass, Georgios Pavlakos, Roberto Martín-Martín

🌐 Project Page | 📄 Paper (arXiv)

1. Create the environment

   conda create -n collage python=3.10
   conda activate collage

2. Install the repository

   git clone https://github.com/yourusername/collage_retrieval.git
   cd collage_retrieval_code/retrieval
   pip install -e .
   cd collage_retrieval_code/robomimic_policy
   pip install -e .

1. Download the LIBERO dataset

   python data/download_libero.py 

2. Download embeddings

We provide precomputed embeddings for sub-trajectory based features (DINO, Optical Flow, PointNet) used in this work for the LIBERO dataset. Download them from here.

3. Run retrieval

   python retrieval/retrieval_libero.py --task_name <TASK_NAME>

You may need to change the paths in retrieval_libero.py and libero_hdf5_config.py to point to your local directories. The sub-trajectroy retrieval process needs to be done seperately for DINO, Optical Flow, and PointNet features. After running the above commands for all features, you should be ready for weight estimation and policy training.

python retrieval/lang_retrieval_openai.py --task_name <TASK_NAME>

4. Estimate weights For estimating weights, first train policies for each feature, using:

   python robomimic_policy/robomimic/scripts/config_gen/bc_xfmr_gen_validate.py --name <experiment-name> --train_ds_path <path to your train dataset>

   python robomimic_policy/robomimic/scripts/train.py --config <path to config file>

   This logs the policy log likelihoods to the wandb file. You can use those as scores in the weight estimation script to obtain the normalized weights for each feature:

   python robomimic_policy/robomimic/scripts/get_weights.py --task_name <TASK_NAME>

5. Training final COLLAGE policy

   Generate, merged hdf5 file with all features:

   python data/concatenate_hdf5_files.py --files <path to DINO.hdf5> <path to FLOW.hdf5> <path to POINTNET.hdf5> ... --names <DINO> <FLOW> <POINTNET> ... --task_name <TASK_NAME>

   Finally, to train the COLLAGE policy:

   python robomimic_policy/robomimic/scripts/config_gen/bc_xfmr_gen.py --name <experiment-name> --train_ds_path <path to your train dataset> 

   python robomimic_policy/robomimic/scripts/train_4_feats.py --name <experiment-name> --co_dataset {train_path} --dino_weight {dino_weight} --flow_weight {flow_weight}  --pnet_weight {pointnet_weight}

@inproceedings{kumar2025collage,
  title={{COLLAGE}: Adaptive Fusion-based Retrieval for Augmented Policy Learning},
  author={Kumar, Sateesh and Dass, Shivin and Pavlakos, Georgios and Mart{\'\i}n-Mart{\'\i}n, Roberto},
  conference={CoRL},
  year={2025}
}

This codebase borrows concepts and code from the STRAP repositories (WEIRDLabUW/STRAP). The policy training code was original developed by robocasa and later adapted by authors of STRAP and us.

Dino features were obtained using DINOv2

Optical flow features were obtained using GMFlow.

PointNet features were obtained using Pointnet_Pointnet2_pytorch.

For language based retrieval, we use OpenAI's text-embedding-ada-002 model. Please ensure you have access to the OpenAI API and set up your API key in retrieval/utils/constants.py.