quick maze task

.gitignore

@@ -17,4 +17,5 @@ data/*
 examples/*
 !examples/01_mnist.ipynb
 !examples/02_inference.ipynb
+!examples/03_mazes.ipynb
 checkpoints

examples/03_mazes.ipynb

@@ -0,0 +1,515 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "04a72c0e",
+   "metadata": {},
+   "source": [
+    "# The Continuous Thought Machine – Tutorial 03: Mazes [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/SakanaAI/continuous-thought-machines/blob/main/examples/01_mnist.ipynb) [![arXiv](https://img.shields.io/badge/arXiv-2505.05522-b31b1b.svg)](https://arxiv.org/abs/2505.05522)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b05cf27b",
+   "metadata": {},
+   "source": [
+    "### Maze Solving\n",
+    "\n",
+    "In Section 4 of the [technical report](https://arxiv.org/pdf/2505.05522), we showcase how a CTM can be used to solve 2D mazes. Typically in the literature, the task of solving a maze is described as a form of binary classification: by ensuring the output space matches the dimensions of the input space, a model can classify, for each pixel in the input image, if the pixel belongs to the path though the maze. While this approach has seen success, it exludes the need to think in a more natural fashion. We seek to design a more challening task, where a more-human like solution is required.\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bc2f769e",
+   "metadata": {},
+   "source": [
+    "### Task Details\n",
+    "\n",
+    "Instead of classifying if each pixel in the maze is or is not along the solution path, we instead constrain the output space, such that the model must output a plan: an entire trajectoru of actions corresponding to the steps an agent must take to go from the start position to goal position. Specifically, for each internal tick of the CTM, the model produces a sequence of actions. We can then use a cross-entropy loss function which compares these actions to the ground-truth trajectory required to solve the maze."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ec34151e",
+   "metadata": {},
+   "source": [
+    "### Task Details\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c257dbd3",
+   "metadata": {},
+   "source": [
+    "In addition to install some dependencies, we also clone the CTM repo (assuming this tutorial is being ran in Colab), so that we can access the base CTM model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "c1ccfdcf",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Requirement already satisfied: gdown in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (5.2.0)\n",
+      "Requirement already satisfied: beautifulsoup4 in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (from gdown) (4.13.4)\n",
+      "Requirement already satisfied: filelock in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (from gdown) (3.17.0)\n",
+      "Requirement already satisfied: requests[socks] in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (from gdown) (2.32.3)\n",
+      "Requirement already satisfied: tqdm in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (from gdown) (4.67.1)\n",
+      "Requirement already satisfied: soupsieve>1.2 in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (from beautifulsoup4->gdown) (2.7)\n",
+      "Requirement already satisfied: typing-extensions>=4.0.0 in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (from beautifulsoup4->gdown) (4.13.2)\n",
+      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (from requests[socks]->gdown) (3.4.1)\n",
+      "Requirement already satisfied: idna<4,>=2.5 in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (from requests[socks]->gdown) (3.10)\n",
+      "Requirement already satisfied: urllib3<3,>=1.21.1 in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (from requests[socks]->gdown) (2.3.0)\n",
+      "Requirement already satisfied: certifi>=2017.4.17 in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (from requests[socks]->gdown) (2025.4.26)\n",
+      "Requirement already satisfied: PySocks!=1.5.7,>=1.5.6 in /home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages (from requests[socks]->gdown) (1.7.1)\n",
+      "^C\n",
+      "\u001b[31mERROR: Operation cancelled by user\u001b[0m\u001b[31m\n",
+      "\u001b[0mfatal: destination path 'continuous-thought-machines' already exists and is not an empty directory.\n"
+     ]
+    }
+   ],
+   "source": [
+    "!pip install gdown\n",
+    "!pip install mediapy\n",
+    "!git clone https://github.com/SakanaAI/continuous-thought-machines.git\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6397bf56",
+   "metadata": {},
+   "source": [
+    "Download the training and test data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "634a61d8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!gdown \"https://drive.google.com/uc?id=1Z8FFnZ7pZcu7DfoSyfy-ghWa08lgYl1V\"\n",
+    "!unzip \"small-mazes.zip\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1ab57a96",
+   "metadata": {},
+   "source": [
+    "Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "24ffe416",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/ciaran_sakana_ai/.conda/envs/atm/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "import sys\n",
+    "sys.path.append(\"./continuous-thought-machines\")\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from tqdm.auto import tqdm\n",
+    "import os\n",
+    "\n",
+    "# From CTM repo\n",
+    "from models.ctm import ContinuousThoughtMachine as CTM\n",
+    "from data.custom_datasets import MazeImageFolder\n",
+    "from tasks.mazes.plotting import make_maze_gif\n",
+    "from tasks.image_classification.plotting import plot_neural_dynamics"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4407a4a8",
+   "metadata": {},
+   "source": [
+    "Prepare the data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "cc28def0",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Loading mazes:   1%|▏         | 134/9000 [00:00<00:15, 583.49it/s]"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Loading mazes: 100%|██████████| 9000/9000 [00:15<00:00, 572.59it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Solving all mazes...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Loading mazes: 100%|██████████| 1000/1000 [00:01<00:00, 643.73it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Solving all mazes...\n"
+     ]
+    }
+   ],
+   "source": [
+    "data_root = './small-mazes'\n",
+    "train_data = MazeImageFolder(root=f'{data_root}/train/', which_set='train', maze_route_length=50)\n",
+    "test_data = MazeImageFolder(root=f'{data_root}/test/', which_set='test', maze_route_length=50)\n",
+    "\n",
+    "trainloader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True, num_workers=0, drop_last=True)\n",
+    "testloader = torch.utils.data.DataLoader(test_data, batch_size=32, shuffle=True, num_workers=1, drop_last=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0e164024",
+   "metadata": {},
+   "source": [
+    "Next, let's define the loss for the maze task.\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "63e75f71",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def maze_loss(predictions, certainties, targets, cirriculum_lookahead=5, use_most_certain=True):\n",
+    "    \"\"\"\n",
+    "    Computes the maze loss with auto-extending cirriculum.\n",
+    "\n",
+    "    Predictions are of shape: (B, route_length, class, internal_ticks),\n",
+    "        where classes are in [0,1,2,3,4] for [Up, Down, Left, Right, Wait]\n",
+    "    Certainties are of shape: (B, 2, internal_ticks), \n",
+    "        where the inside dimension (2) is [normalised_entropy, 1-normalised_entropy]\n",
+    "    Targets are of shape: [B, route_length]\n",
+    "\n",
+    "    cirriculum_lookahead: how far to look ahead in the auto-cirriculum\n",
+    "\n",
+    "    use_most_certain will select either the most certain point or the final point. For baselines,\n",
+    "        the final point proved the only usable option. \n",
+    "    \n",
+    "    \"\"\"\n",
+    "    # Predictions reshaped to: [B*route_length, 5, internal_ticks]\n",
+    "    predictions_reshaped = predictions.flatten(0,1)\n",
+    "    # Targets reshaped to: [B*route_length, internal_ticks]\n",
+    "    targets_reshaped = torch.repeat_interleave(targets.unsqueeze(-1), \n",
+    "                                               predictions.size(-1), -1).flatten(0,1).long()\n",
+    "    \n",
+    "    # Losses are of shape [B, route_length, internal_ticks]\n",
+    "    losses = nn.CrossEntropyLoss(reduction='none')(predictions_reshaped, targets_reshaped)\n",
+    "    losses = losses.reshape(predictions[:,:,0].shape)\n",
+    "    \n",
+    "    # Below is the code for auto-cirriculum\n",
+    "    # Find where correct, and make sure to always push +5 beyond that\n",
+    "    iscorrects = (predictions.argmax(2) == targets.unsqueeze(-1)).cumsum(1)\n",
+    "    correct_mask = (iscorrects == torch.arange(1, iscorrects.size(1)+1, device=iscorrects.device).reshape(1, -1, 1))\n",
+    "    correct_mask[:,0,:] = 1\n",
+    "    upto_where = correct_mask.cumsum(1).argmax(1).max(-1)[0]+cirriculum_lookahead\n",
+    "    loss_mask = torch.zeros_like(losses)\n",
+    "    for bi in range(predictions.size(0)):\n",
+    "        loss_mask[bi, :upto_where[bi]] = 1\n",
+    "\n",
+    "    # Reduce losses along route dimension\n",
+    "    # Will now be of shape [B, internal_ticks]\n",
+    "    losses = (losses * loss_mask).sum(1)/(loss_mask.sum(1))\n",
+    "\n",
+    "    loss_index_1 = losses.argmin(dim=1)\n",
+    "    loss_index_2 = certainties[:,1].argmax(-1)\n",
+    "    if not use_most_certain:\n",
+    "        loss_index_2[:] = -1\n",
+    "    \n",
+    "    batch_indexer = torch.arange(predictions.size(0), device=predictions.device)\n",
+    "    loss_minimum_ce = losses[batch_indexer, loss_index_1]\n",
+    "    loss_selected = losses[batch_indexer, loss_index_2]\n",
+    "\n",
+    "    loss = ((loss_minimum_ce + loss_selected)/2).mean()\n",
+    "    return loss, loss_index_2, upto_where.detach().cpu().numpy()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "89cb8dd7",
+   "metadata": {},
+   "source": [
+    "We define a helper function to update the training progress bar with key metrics:\n",
+    "\n",
+    "Displayed Metrics:\n",
+    "\n",
+    "- Train Loss & Test Loss: Standard loss values during training\n",
+    "- Train Acc (Step): The average accuracy for individual directional predictions across all steps in the maze trajectories\n",
+    "- Statistics on the internal ticks:\n",
+    "    - Average internal tick where the model has highest certainty\n",
+    "    - Standard deviation of these certainty peaks\n",
+    "    - Range (min ↔ max) of internal ticks where peak certainty occurs\n",
+    "    - The position upto which the loss is taken, on average\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "fb57caee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_pbar_desc(train_loss, train_accuracy_finegrained, test_loss, optimizer, where_most_certain, upto_where):\n",
+    "    \"\"\"A helper function to create a description for the tqdm progress bar\"\"\"\n",
+    "    pbar_desc = f'Train Loss={train_loss if isinstance(train_loss, float) else train_loss.item():0.3f}. Train Acc(step)={train_accuracy_finegrained:0.3f}. Test Loss={test_loss if isinstance(test_loss, float) else test_loss.item():0.3f}. LR={optimizer.param_groups[-1][\"lr\"]:0.6f}.'\n",
+    "    pbar_desc += f' Where_certain={where_most_certain.float().mean().item():0.2f}+-{where_most_certain.float().std().item():0.2f} ({where_most_certain.min().item():d}<->{where_most_certain.max().item():d}). Upto: {sum(upto_where) / len(upto_where):0.2f}.'\n",
+    "    return pbar_desc"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "02de7c62",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def train(model, trainloader, testloader, device='cpu', training_iterations=10000, test_every=1000, lr=1e-4, log_dir='./logs'):\n",
+    "\n",
+    "    os.makedirs(log_dir, exist_ok=True)\n",
+    "    \n",
+    "    model.train()\n",
+    "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)\n",
+    "    iterator = iter(trainloader)\n",
+    "    \n",
+    "    train_losses = []\n",
+    "    test_losses = []\n",
+    "    \n",
+    "    with tqdm(total=training_iterations) as pbar:\n",
+    "        for stepi in range(training_iterations):\n",
+    "\n",
+    "            try:\n",
+    "                inputs, targets = next(iterator)\n",
+    "            except StopIteration:\n",
+    "                iterator = iter(trainloader)\n",
+    "                inputs, targets = next(iterator)\n",
+    "    \n",
+    "            inputs, targets = inputs.to(device), targets.to(device)\n",
+    "            \n",
+    "            optimizer.zero_grad()\n",
+    "            \n",
+    "            predictions_raw, certainties, _ = model(inputs)\n",
+    "\n",
+    "            # Reshape: (B, SeqLength, 5, Ticks)\n",
+    "            predictions = predictions_raw.reshape(predictions_raw.size(0), -1, 5, predictions_raw.size(-1))\n",
+    "            \n",
+    "            # Compute loss\n",
+    "            train_loss, where_most_certain, upto_where = maze_loss(predictions, certainties, targets, use_most_certain=True)\n",
+    "\n",
+    "            train_accuracy_finegrained = (predictions.argmax(2)[torch.arange(predictions.size(0), device=predictions.device), :, where_most_certain] == targets).float().mean().item()\n",
+    "\n",
+    "            train_loss.backward()\n",
+    "            optimizer.step()\n",
+    "        \n",
+    "            train_losses.append(train_loss)\n",
+    "\n",
+    "            if stepi % test_every == 0 or stepi == 0:\n",
+    "                model.eval()\n",
+    "                with torch.no_grad():\n",
+    "                    test_loss_per_batch = []\n",
+    "                    for inputs, targets in testloader:\n",
+    "                        inputs = inputs.to(device)\n",
+    "                        targets = targets.to(device)\n",
+    "                        \n",
+    "                        predictions_raw, certainties, _ = model(inputs)\n",
+    "                        predictions = predictions_raw.reshape(predictions_raw.size(0), -1, 5, predictions_raw.size(-1))\n",
+    "                        \n",
+    "                        test_loss, _, _ = maze_loss(predictions, certainties, targets, use_most_certain=True)\n",
+    "                        test_loss_per_batch.append(test_loss.item())\n",
+    "\n",
+    "                    test_loss = sum(test_loss_per_batch) / len(test_loss_per_batch)\n",
+    "                    test_losses.append(test_loss)\n",
+    "\n",
+    "                    create_maze_gif_visualization(model, testloader, device, log_dir)\n",
+    "                model.train()\n",
+    "\n",
+    "            pbar_desc = make_pbar_desc(train_loss, train_accuracy_finegrained, test_loss, optimizer, where_most_certain, upto_where)\n",
+    "            pbar.set_description(pbar_desc)\n",
+    "            pbar.update(1)\n",
+    "                    \n",
+    "    return train_losses, test_losses\n",
+    "\n",
+    "def create_maze_gif_visualization(model, testloader, device, log_dir):\n",
+    "    \"\"\"\n",
+    "    Create GIF visualization of maze solving with attention tracking\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    model.eval()\n",
+    "    with torch.no_grad():\n",
+    "        # Get a test batch\n",
+    "        inputs_viz, targets_viz = next(iter(testloader))\n",
+    "        inputs_viz = inputs_viz.to(device)\n",
+    "        targets_viz = targets_viz.to(device)\n",
+    "        \n",
+    "        batch_index_to_viz = 0\n",
+    "        \n",
+    "        predictions_raw, certainties, _, pre_activations, post_activations, attention_tracking = model(inputs_viz, track=True)\n",
+    "        \n",
+    "        # Reshape predictions\n",
+    "        predictions = predictions_raw.reshape(predictions_raw.size(0), -1, 5, predictions_raw.size(-1))\n",
+    "        \n",
+    "        # Reshape attention tracking for visualization\n",
+    "        att_shape = (model.kv_features.shape[2], model.kv_features.shape[3])\n",
+    "        attention_tracking = attention_tracking.reshape(attention_tracking.shape[0], attention_tracking.shape[1], -1, att_shape[0], att_shape[1])\n",
+    "\n",
+    "        plot_neural_dynamics(post_activations, 100, log_dir, axis_snap=True)\n",
+    "        \n",
+    "        # Create maze GIF with attention visualization\n",
+    "        maze_input = (inputs_viz[batch_index_to_viz].detach().cpu().numpy() + 1) / 2\n",
+    "        maze_predictions = predictions[batch_index_to_viz].detach().cpu().numpy()\n",
+    "        maze_targets = targets_viz[batch_index_to_viz].detach().cpu().numpy()\n",
+    "        maze_attention = attention_tracking[:, batch_index_to_viz] if attention_tracking.ndim > 2 else attention_tracking\n",
+    "\n",
+    "        # Generate the maze GIF\n",
+    "        make_maze_gif(\n",
+    "            maze_input,\n",
+    "            maze_predictions,\n",
+    "            maze_targets,\n",
+    "            maze_attention,\n",
+    "            log_dir\n",
+    "        )\n",
+    "        \n",
+    "        predictions_raw, certainties, _ = model(inputs_viz)\n",
+    "        predictions = predictions_raw.reshape(predictions_raw.size(0), -1, 5, predictions_raw.size(-1))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2c180995",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using neuron select type: random-pairing\n",
+      "Synch representation size action: 32\n",
+      "Synch representation size out: 32\n",
+      "Model parameters: 3,304,684\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|          | 0/100000 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Processing frames for maze plotting: 100%|██████████| 50/50 [00:05<00:00,  8.52it/s]\n",
+      "Train Loss=1.458. Train Acc(step)=0.197. Test Loss=1.613. LR=0.000100. Where_certain=31.31+-11.91 (13<->49). Upto: 6.44.:   1%|          | 565/100000 [03:16<8:28:46,  3.26it/s] "
+     ]
+    }
+   ],
+   "source": [
+    "# Set device\n",
+    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+    "\n",
+    "# Define the model\n",
+    "model = CTM(\n",
+    "    iterations=50,          # Number of thinking steps\n",
+    "    d_model=512,            # Model dimension\n",
+    "    d_input=128,             # Input dimension\n",
+    "    heads=8,                # Attention heads\n",
+    "    n_synch_out=32,         # Output synchronization neurons\n",
+    "    n_synch_action=32,      # Action synchronization neurons\n",
+    "    synapse_depth=8,        # Synapse network depth\n",
+    "    memory_length=15,       # Memory length\n",
+    "    deep_nlms=True,         # Use deep memory\n",
+    "    memory_hidden_dims=16,  # Memory hidden dimensions\n",
+    "    backbone_type='resnet34-2',  # Feature extractor\n",
+    "    out_dims=50 * 5,        # Output dimensions (route_length * 5 directions)\n",
+    "    prediction_reshaper=[50, 5],  # Reshape to [route_length, directions]\n",
+    "    dropout=0.1,\n",
+    "    do_layernorm_nlm=False,\n",
+    "    positional_embedding_type='none'\n",
+    ").to(device)\n",
+    "\n",
+    "# Initialize model parameters with dummy forward pass\n",
+    "sample_batch = next(iter(trainloader))\n",
+    "dummy_input = sample_batch[0][:1].to(device)\n",
+    "with torch.no_grad():\n",
+    "    _ = model(dummy_input)\n",
+    "\n",
+    "print(f'Model parameters: {sum(p.numel() for p in model.parameters()):,}')\n",
+    "\n",
+    "# Train the model\n",
+    "log_dir = './maze_training_logs'\n",
+    "train_losses, test_losses = train(\n",
+    "    model=model,\n",
+    "    trainloader=trainloader,\n",
+    "    testloader=testloader,\n",
+    "    device=device,\n",
+    "    training_iterations=100000,\n",
+    "    lr=1e-4,\n",
+    "    log_dir=log_dir\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "atm",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

tasks/mazes/plotting.py

@@ -88,7 +88,7 @@ def draw_path(x, route, valid_only=False, gt=False, cmap=None):
 
     return x
 
-def make_maze_gif(inputs, predictions, targets, attention_tracking, save_location):
+def make_maze_gif(inputs, predictions, targets, attention_tracking, save_location, verbose=True):
     """
     Expect inputs, predictions, targets as numpy arrays
     """
@@ -130,7 +130,7 @@ def make_maze_gif(inputs, predictions, targets, attention_tracking, save_locatio
     cmap_viridis = plt.get_cmap('viridis')
     step_linspace = np.linspace(0, 1, predictions.shape[-1])  # For sampling colours
     with tqdm(total=predictions.shape[-1], initial=0, leave=True, position=1, dynamic_ncols=True) as pbar:
-        pbar.set_description('Processing frames for maze plotting')
+        if verbose: pbar.set_description('Processing frames for maze plotting')
         for stepi in np.arange(0, predictions.shape[-1], 1):
             fig, axes = plt.subplot_mosaic(mosaic, figsize=aspect_ratio)
             for ax in axes.values():