Automatic training for the dynamic foraging task

Overview

Here is a diagram of the automatic training system we've been devloping for the dynamic foraging task (those blue arrows on top of our existing foraging behavior pipeline)

This repo is the red circle in the diagram. It will be running on Code Ocean and do the following things (very similar to mTrack):

1. Define the training curriculum (mTrack's "Regimen")
2. Retrieve data for each mouse from the behavioral master table (df_behavior on S3 that stores session-wise metrics)
3. Evaluate daily performance based on the curriculum and make decisions of the next training stage (stored in df_manager)
4. Push the decisions back to S3 by uploading df_manager, from which our python GUI can access and automatically set the training parameters on the next day.

Key elements

• Task schema defines schema for training parameters
• Curriculum schema defines schema for the curriculum, especially the evaluate_transitions method
• Auto train manager fetches data from df_behavior and updates df_manager (or "tables" on any other database)
• Curriculum manager manages all available pre-generated curriculums (on any S3 bucket).

Usage

1. Design the curriculum, i.e., all training stages and transition rules.
   • Here is an example curriculum for the dynamic foraging task.
   • Here is automatically generated json file
   • Here are automatically rendered diagrams for stage transitions rules and parameters (click the images to try the hover feature 😊)

rules	parameters

2. Create AutoTrainManager and connect it to the behavior database df_behavior.

3. Feed all necessary metrics to Auto train manager and let it run.

   • Here is an open-loop simulation with our old mice
     

4. To add other tasks, users should add their own task and curriculum schemas.

Demo

• Demo notebook for the curriculum schema.
• Demo notebook for a full automation workflow (auto train manager)
• Demo notebook for the curriculum manager

Notes

System upgrade checklist

• Upgrade all training rigs
• Upgrade CO capsule (terminate all running workers and start a new one)
• Upgrade streamlit app

Compared with SLIMS/mTrack

See this thread

Instructions for adding new curriculums

Since graphviz in Code Ocean has some unsolved bug, it is recommended to install the library locally in a conda environment to create new curriculums.

1. Clone this repo to your local computer
2. Create conda environment by

   conda create -n autotrain python=3.8
   conda activate autotrain

3. Install the library in editable mode

   pip install -e .

4. Install graphviz via conda (which installs necessary dot .bin files for you)

   conda install python-graphviz

5. Set up AWS credential. (ask Han for this step)
6. Test the installation by running demo script code\aind_auto_train\curriculums\dummy_task.py
7. Create you own curriculum based on existing curriculums in code\aind_auto_train\curriculums\
8. Run your new .py file and check the generated diagrams under the default folder {your windows user folder}/capsule/scratch/saved_curriculums/
9. Run the script code\aind_auto_train\curriculum_manager.py to upload new curriculum files to AWS S3 bucket s3://aind-behavior-data/foraging_auto_training/saved_curriculums/
10. Check the new curriculum on the Streamlit app.