This work is a component of the Yale Social Robotics Lab robot recycling project led by PhD student Debsmita Ghose: Project Page, Public Project Repo, Paper.
Our problem is framed as a recycling robot learning to remove correct recyclables from a recycling stream based on a few examples provided by a human. The model in this repo leverages the contrastive learning framework SimSiam, which learns meaningful representations of images from unlabeled datasets. Using SimSiam on its own does not take full advantage of the small set of human examples during training. Therefore, we add a human supervised head to pull the human-selected images close together in the latent space, creating a representation that is tailored to the human's needs. The final project opted for the Contrastive Clustering learning framework rather than SimSiam since it provided better results.
The model was implemented in PyTorch. The train_runner and avg_eval_runner scripts are provided for training and evaluating the model on new data. Instructions can be found within the scripts. The model options and data paths can be set in the file config/config.yaml. For usage on new data, the CropsDataSet in data/crops.py class should be modified to load in data from the desired dataset.