diff --git a/examples/agent_examples/whittington_2020_example.ipynb b/examples/agent_examples/whittington_2020_example.ipynb
index a672847..6c8ae43 100644
--- a/examples/agent_examples/whittington_2020_example.ipynb
+++ b/examples/agent_examples/whittington_2020_example.ipynb
@@ -24,10 +24,8 @@
     "import importlib\n",
     "import os\n",
     "import pickle\n",
-    "import pandas as pd\n",
-    "import matplotlib.pyplot as plt\n",
-    "import numpy as np\n",
     "\n",
+    "import matplotlib.pyplot as plt\n",
     "import pandas as pd\n",
     "\n",
     "from neuralplayground.comparison import GridScorer\n",
diff --git a/examples/arena_examples/arena_examples.ipynb b/examples/arena_examples/arena_examples.ipynb
index 3f7693f..f8fd5c0 100644
--- a/examples/arena_examples/arena_examples.ipynb
+++ b/examples/arena_examples/arena_examples.ipynb
@@ -528,8 +528,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from neuralplayground.agents.whittington_2020 import Whittington2020\n",
-    "import neuralplayground.agents.whittington_2020_extras.whittington_2020_parameters as parameters"
+    "import neuralplayground.agents.whittington_2020_extras.whittington_2020_parameters as parameters\n",
+    "from neuralplayground.agents.whittington_2020 import Whittington2020"
    ]
   },
   {
diff --git a/neuralplayground/arenas/discritized_objects.py b/neuralplayground/arenas/discritized_objects.py
index 60b91ba..20ea52b 100644
--- a/neuralplayground/arenas/discritized_objects.py
+++ b/neuralplayground/arenas/discritized_objects.py
@@ -2,7 +2,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
-from matplotlib.colors import LogNorm
 
 from neuralplayground.arenas.arena_core import Environment
 from neuralplayground.plotting.plot_utils import make_plot_trajectories