MABe_2022_TVAE: a Trajectory Variational Autoencoder baseline for the 2022 Multi-Agent Behavior challenge

This repository contains jupyter notebooks that implement a trajectory variational autoencoder (tVAE) baseline model for embedding the mouse and fly trajectory datasets for the 2022 Multi-Agent Behavior Challenge. Performance of these notebooks is as follows:

Where the "task F1" values are F1 scores on specific sample evaluation tasks. Note, while this baseline is outperforming PCA for the mice, it actually does significantly worse than PCA for the flies! (Why? Who knows!)

How to use these notebooks

Clone this repository and open train_mouse_tvae.ipynb or train_fly_tvae.ipynb in a jupyter notebook session. Follow notebook instructions to make your own submission, then play with the model architecture and parameters inside the tvae directory to see if you are able to improve performance!

Note, you'll need to download the mouse and/or fly datasets into the provided mouse_data and fly_data directories to use this code.