diff --git a/src/primaite/notebooks/Training-an-SB3-Agent.ipynb b/src/primaite/notebooks/Training-an-SB3-Agent.ipynb
index 8d6789ee..8e18c5c1 100644
--- a/src/primaite/notebooks/Training-an-SB3-Agent.ipynb
+++ b/src/primaite/notebooks/Training-an-SB3-Agent.ipynb
@@ -48,9 +48,9 @@
     "from stable_baselines3 import PPO\n",
     "\n",
     "EPISODE_LEN = 128\n",
-    "NO_STEPS = EPISODE_LEN * 10\n",
-    "BATCH_SIZE = EPISODE_LEN * 10\n",
-    "TOTAL_TIMESTEPS = 5e3 * EPISODE_LEN\n",
+    "NUM_EPISODES = 10\n",
+    "NO_STEPS = EPISODE_LEN * NUM_EPISODES\n",
+    "BATCH_SIZE = 32\n",
     "LEARNING_RATE = 3e-4"
    ]
   },
@@ -69,7 +69,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model.learn(total_timesteps=TOTAL_TIMESTEPS)\n"
+    "model.learn(total_timesteps=NO_STEPS)\n"
    ]
   },
   {