9 April 2025
Solving mazes is a classic problem in computer science and artificial
intelligence, and humans have been constructing mazes for thousands of
years. Although finding the shortest path through a maze is a solved
problem, this very fact makes it an excellent testbed for studying how
machine learning algorithms solve problems and represent spatial
information. We introduce maze-dataset, a user-friendly
Python library for generating, processing, and visualizing datasets of
mazes. This library supports a variety of maze generation algorithms
providing mazes with or without loops, mazes that are connected or not,
and many other variations. These generation algorithms can be configured
with various parameters, and the resulting mazes can be filtered to
satisfy desired properties. Also provided are tools for converting mazes
to and from various formats suitable for a variety of neural network
architectures, such as rasterized images, tokenized text sequences, and
various visualizations. As well as providing a simple interface for
generating, storing, and loading these datasets,
maze-dataset is extensively tested, type hinted,
benchmarked, and documented.
While maze generation itself is straightforward, the architectural challenge comes from building a system supporting many algorithms with configurable parameters, property filtering, and representation transformation. This library aims to greatly streamline the process of generating and working with datasets of mazes that can be described as subgraphs of an n × n lattice with boolean connections and, optionally, start and end points that are nodes in the graph. Furthermore, we place emphasis on a wide variety of possible text output formats aimed at evaluating the spatial reasoning capabilities of Large Language Models (LLMs) and other text-based transformer models.
For interpretability and behavioral research, algorithmic tasks offer benefits by allowing systematic data generation and task decomposition, as well as simplifying the process of circuit discovery (Räuker et al., 2023). Although mazes are well suited for these investigations, we found that existing maze generation packages (Cobbe et al., 2019; Ehsan, 2022; Harries et al., n.d.; Németh, 2019; Schwarzschild, Borgnia, Gupta, Bansal, et al., 2021) lack support for transforming between multiple representations and provide limited control over the maze generation process.
A multitude of public and open-source software packages exist for generating mazes (Ehsan, 2022; Németh, 2019; Schwarzschild, Borgnia, Gupta, Bansal, et al., 2021). However, nearly all of these packages produce mazes represented as rasterized images or other visual formats rather than the underlying graph structure, and this makes it difficult to work with these datasets.
Most prior works provide mazes in visual or raster formats, and we provide a variety of similar output formats:
RasterizedMazeDataset,
utilizing as_pixels(),
which can exactly mimic the outputs provided in
easy-to-hard-data (Schwarzschild, Borgnia, Gupta, Bansal, et al.,
2021) and can be configured to be similar to the outputs of
Németh (2019)as_ascii()
provides a format similar to (Oppenheim,
2018; Singla,
2023)MazePlot
provides a feature‑rich plotting utility with support for multiple
paths, heatmaps over positions, and more. This is similar to the outputs
of (Alance AB, 2019;
Ehsan, 2022; Guo et al., 2011;
Nag, 2020)The text format provided by SolvedMaze(...).as_tokens()
is similar to that of (Liu & Wu, 2023) but with many more
options, detailed in section:
.
Preserving metadata about the generation algorithm with the dataset itself is essential for studying the effects of distributional shifts. Our package efficiently stores the dataset along with its metadata in a single human-readable file (M. Ivanitskiy, n.d.). As far as we are aware, no existing packages do this reliably.
Storing mazes as images or adjacency matrices is not only difficult to work with, but also inefficient. We use a highly efficient method detailed in section: .
Our package is easily installable with source code freely available. It is extensively tested, type hinted, benchmarked, and documented. Many other maze generation packages lack this level of rigor and scope, and some (Ayaz et al., 2008) appear to simply no longer be accessible.
We direct readers to our examples,
docs,
and notebooks
for more information. Our package can be installed from PyPi via
pip install maze-dataset, or directly from the git
repository (Michael I. Ivanitskiy et al.,
2023a).
Datasets of mazes are created from a MazeDatasetConfig
configuration object, which allows specifying the number of mazes, their
size, the generation algorithm, and various parameters for the
generation algorithm. Datasets can also be filtered after generation to
satisfy certain properties. Custom filters can be specified, and some
filters are included in MazeDatasetFilters.
Internally, mazes are SolvedMaze
objects, which have path information and a tensor optimized for storing
sub-graphs of a lattice. These objects can be converted to and from
several formats, shown in [fig:output-fmts], to maximize
their utility in different contexts.
In previous work, maze tasks have been used with Recurrent
Convolutional Neural Network (RCNN) derived architectures (Schwarzschild, Borgnia, Gupta, Huang, et al.,
2021). To facilitate the use of our package in this context,
we replicate the format of (Schwarzschild, Borgnia, Gupta, Bansal, et al.,
2021) and provide the RasterizedMazeDataset
class which returns rasterized pairs of (input, target) mazes as shown
in [fig:e2h-raster].
Autoregressive transformer models can be quite sensitive to the exact format of input data, and may even use delimiter tokens to perform reasoning steps (Pfau et al., 2024; Spies et al., 2024). To facilitate systematic investigation of the effects of different representations of data on text model performance, we provide a variety of text output formats, with an example given in [fig:token-regions]. We utilize Finite State Transducers (Gallant, 2015) for efficiently storing valid tokenizers.
We benchmarks for generation time across various configurations in [tab:benchmarks] and [fig:benchmarks]. Experiments were performed on a standard GitHub runner without parallelism. Additionally, maze generation under certain constraints may not always be successful, and for this we provide a way to estimate the success rate of a given configuration, described in [fig:sre].
Using an adjacency matrix for storing mazes would be memory
inefficient by failing to exploit the highly sparse structure, while
using an adjacency list could lead to a poor lookup time. This package
utilizes a simple, efficient representation of mazes as subgraphs of a
finite lattice, detailed in [fig:maze-impl], which we call a LatticeMaze.
Our package is implemented in Python(Rossum, 1995), and makes use of the extensive scientific computing ecosystem, including NumPy (Harris et al., 2020) for array manipulation, plotting tools Waskom (2021), Jupyter notebooks (Kluyver et al., 2016), and PySR (Cranmer, 2023) for symbolic regression.
This package was originally built for the needs of the (Michael I.
Ivanitskiy et al., 2023b) project, which aims to investigate
spatial planning and world models in autoregressive transformer models
trained on mazes (Michael Igorevich Ivanitskiy, Spies, et al.,
2023; Michael Igorevich Ivanitskiy, Shah, et al.,
2023; Spies et al., 2024). It was extended for
work on understanding the mechanisms by which recurrent convolutional
and implicit networks (Fung et al., 2022) solve mazes given a
rasterized view (Knutson
et al., 2024), which required matching the pixel-padded and
endpoint constrained output format of (Schwarzschild, Borgnia, Gupta, Bansal, et al.,
2021). Ongoing work using maze-dataset aims to
investigate the effects of varying the tokenization format on the
performance of pretrained LLMs on spatial reasoning.
This package has also been utilized in work by other groups:
By (Nolte et al., 2024) to compare the effectiveness of transformers trained with the MLM-𝒰 (Kitouni et al., 2024) multistep prediction objective against standard autoregressive training for multi-step planning on our maze task.
By (Wang et al., 2024) and (Chen et al., 2024) to study the effectiveness of imperative learning.
By (Zhang et al., 2025) to introduce a novel framework for reasoning diffusion models.
By (Dao & Vu, 2025) to improve spatial reasoning in LLMs with GRPO.