Maze Transformer
This repo is built to facilitate the training and analysis of autoregressive transformers on maze-solving tasks.
![]() |
![]() |
![]() |
Installation
pip install git+ssh://git@github.com/aisc-understanding-search/maze-transformer.git
Note: if you want to install the library in colab, follow the steps in this Colab notebook.
You can run all tests via
make test(this will take a few minutes)
Usage
Most of the functionality is demonstrated in the ipython notebooks in the notebooks/ folder.
demo_dataset.ipynbhow to easily create a dataset of mazes, utilities for filtering the generates mazes via properties, and basic visualizationtrain_model.ipynbconfiguration setup and training a basic modeleval_model.ipynbloading a trained model, and computing various metrics on its performance on a datasetplot_attention.ipynbvarious attention visualization utilitiesdemo_latticemaze.ipynbinternals of theLatticeMazeandSolvedMazeobjects, and advanced visualizationtrain_model_hallway.ipynbtraining a model on a customized dataset
Development
Prerequisites
-
Install Poetry
-
Install Python 3.10
- It's a good idea to use pyenv to manage python versions
- If using pyenv, you'll need to update your Poetry config for it to use the pyenv Python version:
poetry config virtualenvs.prefer-active-python true
-
Install dev dependencies
poetry config virtualenvs.in-project true poetry install --with dev -
Run unit, integration, and notebook tests
make test -
(Optional) If you want to work with the jupyter notebooks in VSCode
- create a jupyter kernel with
poetry run ipython kernel install --user --name=maze-transformer - Restart VSCode
- In VSCode, select the python interpreter located in
maze-transformer/.venv/binas your juptyer kernel
- create a jupyter kernel with
Testing & Static analysis
make help will print all available commands.
-
all tests via
make test- unit tests via
make unit - integration tests via
make integration - notebook tests via
make test_notebooks
- unit tests via
-
formatter (black, pycln, and isort) via
make format- formatter in check-only mode via
make check-format
- formatter in check-only mode via



