bikesnbees/torch_tem

Implementation of the Tolman Eichenbaum Machine in pytorch

Torch TEM

Table of Contents

• About the Project
• Getting Started
  • Installation
  • Model Training
  • Model Analysis
• Contact
• Acknowledgements

About The Project

This is an implementation of the Tolman-Eichenbaum Machine in pytorch, written from scratch by following the Supplementary Material of the original paper. It is extensively annotated and tries to follow the notation and terminology from the publication as closely as possible.

Getting Started

You need to install python >= 3.6.0 and pytorch >= 1.6.0.

Installation

Clone the repo

git clone https://github.com/jbakermans/torch_tem.git

Model Training

With the repo as working directory, train a model by running

python run.py

Model parameters are specified in parameters.py.

Model Analysis

After training a model, analyse a model and plot analysis results by running

python test.py

You will need to specify the correct model run in test.py.

Contact

Jacob Bakermans - jacob.bakermans [at] gmail.com

Project Link: https://github.com/jbakermans/torch_tem

Acknowledgements

Many thanks to James Whittington for advice and assistance throughout.