From bb88d43b90e0eec249c7af6815dabe9718b8a987 Mon Sep 17 00:00:00 2001
From: Charlie Crane <charlie.crane@t-s.qinetiq.com>
Date: Thu, 18 Apr 2024 13:50:47 +0100
Subject: [PATCH] #2453 - Updating  with some explanation to improve the
 readability of the notebook

---
 .../notebooks/Training-an-SB3-Agent.ipynb      | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/src/primaite/notebooks/Training-an-SB3-Agent.ipynb b/src/primaite/notebooks/Training-an-SB3-Agent.ipynb
index 8e18c5c1..59fd46c4 100644
--- a/src/primaite/notebooks/Training-an-SB3-Agent.ipynb
+++ b/src/primaite/notebooks/Training-an-SB3-Agent.ipynb
@@ -1,5 +1,16 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Training an SB3 Agent\n",
+    "\n",
+    "This notebook will demonstrate how to use primaite to create and train a PPO agent.\n",
+    "\n",
+    "#### First, import `PrimaiteGymEnv` and read our config file"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -69,7 +80,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model.learn(total_timesteps=NO_STEPS)\n"
+    "model.learn(total_timesteps=NO_STEPS)\n",
+    "model.save(\"PrimAITE-PPO-example-agent\")"
    ]
   },
   {
@@ -78,7 +90,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model.save(\"PrimAITE-v3.0.0b7-PPO\")"
+    "model.save(\"PrimAITE-PPO-example-agent\")"
    ]
   },
   {
@@ -88,7 +100,7 @@
    "outputs": [],
    "source": [
     "eval_model = PPO(\"MlpPolicy\", gym)\n",
-    "eval_model = PPO.load(\"PrimAITE-v3.0.0b7-PPO\", gym)"
+    "eval_model = PPO.load(\"PrimAITE-PPO-example-agent\", gym)"
    ]
   },
   {