Codebase for the AAAI-2025 paper "Data Augmentation for Instruction Following Policies via Trajectory Segmentation"
- Anaconda3/Miniconda3
conda create -n play_segmentation python=3.11
conda activate play_segmentation
pip install -e .
Overview of the different components of Play Segmentation
- First get the necessary Datasets BabyAI/CALVIN.
- Train a Labelling Model on the annotated datasets.
- Train Segmentation Models.
- Train a policy.
To get the CALVIN dataset follow the instructions here.
To generate the segmented as well as unsegmented trajectories in the BabyAI environment follow the instructions here.
Note: The labelling model also needs to be trained to obtain pretrained embeddings of videos for the segmentation methods TriDet and Play Segmentation.
To train the labelling model follow the instructions here.
To apply the trained segmentation models to extract labelled segments from play trajectories, follow the instructions here.
To train the different segmentation models follow the instructions here. At the moment the following segmentation models are implemented:
- Play Segmentation (Ours)
- UnLoc (https://arxiv.org/abs/2308.11062)
- TriDet (https://arxiv.org/abs/2303.07347) / Code
To train an imitation learning policy for BabyAI and a policy via MCIL for CALVIN see here.
If you use this Codebase, please cite:
@article{hopner2025data,
title={Data Augmentation for Instruction Following Policies via Trajectory Segmentation},
author={H{\"o}pner, Niklas and Tiddi, Ilaria and van Hoof, Herke},
journal={arXiv preprint arXiv:2503.01871},
year={2025}
}