From 49be3b36391fa95a61205a387445a7b5fc92c561 Mon Sep 17 00:00:00 2001
From: Charlie Crane <charlie.crane@t-s.qinetiq.com>
Date: Thu, 18 Apr 2024 10:23:26 +0100
Subject: [PATCH] #2453 - Updates to Training-an-SB3-agent.ipynb following
 review

---
 src/primaite/notebooks/Training-an-SB3-Agent.ipynb | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

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"
    ]
   },
   {