Reinforcement_Learning_WRRF_Optimization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "99b5b9d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reinforcement Learning WRRF Optimization\n",
    "# By: Henry Croll, Kaoru Ikuma, and Soumik Sarkar\n",
    "# Copyright © - 2024\n",
    "# Iowa State University Research Foundation, Inc.\n",
    "\n",
    "# This notebook was developed for training and evaluation of reinforcement learning agents applied to water resource recovery facility optimization.\n",
    "# This notebook requires additional files and software to function: \n",
    "# - SUMO22beta from Dynamita with Digital Twin Toolkit\n",
    "# - scheduler.py from Dynamita\n",
    "# - tool.py from Dynamita\n",
    "# - SUMO model .dll file for simulation\n",
    "# - SUMO state .xml file for simulation starting state\n",
    "# - Influent.xlsx file with influent data\n",
    "\n",
    "# Initial Stablebaselines3 agents and test environment based on Deep RL Tutorial by Nicholas Renotte\n",
    "# Create test custom environment based on Tutorial by Nicholas Renotte\n",
    "\n",
    "# Agent tuned for open control of BSM1 WWTF simulation\n",
    "\n",
    "# Updated 12/10/2024"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e6c0e8e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scheduler as ds\n",
    "import tool as dtool\n",
    "import time \n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import shutil\n",
    "import tensorflow\n",
    "\n",
    "import os\n",
    "import random\n",
    "import gym\n",
    "from gym import Env, spaces\n",
    "from stable_baselines3 import PPO, DQN, A2C, TD3\n",
    "from stable_baselines3.common.vec_env import DummyVecEnv\n",
    "from stable_baselines3.common.evaluation import evaluate_policy\n",
    "from stable_baselines3.common.monitor import Monitor\n",
    "from stable_baselines3.common.utils import set_random_seed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ee8f092",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create folder for logs\n",
    "log_path = os.path.join('Test_Training', 'Test_Logs')\n",
    "\n",
    "# define what model and state to use\n",
    "SUMOmodel = \"MODEL_NAME_HERE.dll\"\n",
    "SUMOstate = \"STATE_NAME_HERE\"\n",
    "\n",
    "# Create active xml \n",
    "SUMOstate_active = \"ACTIVE_STATE_NAME_HERE\"\n",
    "path = 'YOUR_PATH_HERE\\\\Dynamita\\\\Sumo22\\\\PythonAPI\\\\YOUR_PATH_HERE\\\\'\n",
    "src = path + SUMOstate + \".xml\"\n",
    "dst = path + SUMOstate_active + \".xml\"\n",
    "\n",
    "shutil.copyfile(src, dst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fd544159",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This callback is called when new data is available. It depends on the Sumo__Datacomm variable. \n",
    "# In this example it happens every hour.     \n",
    "def data_callback(job, data):\n",
    "    # Get the information regarding the running simulation\n",
    "    jobData = ds.sumo.getJobData(job)\n",
    "    \n",
    "    # Sumo__Time is in milliseconds, convert it to days.\n",
    "    data[\"Sumo__Time\"] /= dtool.day\n",
    "    \n",
    "    # print progress to the console\n",
    "    # print(f\"DC #{job} - {data['Sumo__Time']}\")\n",
    "    \n",
    "    # Save the data for the job. We could process it right here \n",
    "    # (print, save it in a file, database, etc).\n",
    "    # But in this example we collect it and export it once the \n",
    "    # simulation is finished.\n",
    "    for item in data:     \n",
    "        jobData[\"results\"][item] = data[item]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2ac89c5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This callback is called when Sumo wants to say something, like a command was successful or not,\n",
    "# an error occured or the simulation is finished.\n",
    "def msg_callback(job, msg):\n",
    "    # Make sure we see the messages\n",
    "    #print(\"MSG #\" + str(job) + \": '\" + msg + \"'\")\n",
    "    \n",
    "    # If the simulation is finished, we tell Sumo to save the state - we'll get another message when the state is saved\n",
    "    if (ds.sumo.isSimFinishedMsg(msg)):\n",
    "        jobData = ds.sumo.getJobData(job)   \n",
    "        ds.sumo.sendCommand(job, f\"save {SUMOstate_active}.xml\")\n",
    "        \n",
    "    # When the state is finished we can clean up the simulation\n",
    "    elif msg.startswith(\"530045\"):\n",
    "        jobData = ds.sumo.getJobData(job)        \n",
    "        ds.sumo.finish(job)\n",
    "        jobData[\"finished\"] = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e81ef56f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def SUMO_step(self, set_air1, set_air2, set_air3, set_air4, set_air5, set_recycle, set_ras):\n",
    "    \n",
    "    # start a Sumo simulation\n",
    "    job = ds.sumo.schedule(\n",
    "                        SUMOmodel, # model dll path extracted from a .sumo file\n",
    "                        \n",
    "                        # Commands to the SumoCore to setup the simulation\n",
    "        commands      = [\n",
    "                        # Load a base state that includes parameters set from the Sumo GUI. \n",
    "                        # If this is omitted every state and parameter will be to it's factory default\n",
    "                         f\"load {SUMOstate_active}.xml\",\n",
    "                         \"maptoic\",\n",
    "            \n",
    "                        # Send new influent parameters\n",
    "            \n",
    "                        f\"set Sumo__Plant__Influent2__param__Q {self.influentdf['Q'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__SU {self.influentdf['SU'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__SB {self.influentdf['SB'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__XU {self.influentdf['XU'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__XB {self.influentdf['XB'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__XOHO {self.influentdf['XOHO'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__XNITO {self.influentdf['XNITO'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__XE {self.influentdf['XE'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__SO2 {self.influentdf['SO2'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__SNOx {self.influentdf['SNOx'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__SNHx {self.influentdf['SNHx'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__SN_B {self.influentdf['SN,B'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__XN_B {self.influentdf['XN,B'][self.influent]};\",\n",
    "                        f\"set Sumo__Plant__Influent2__param__SALK {self.influentdf['SALK'][self.influent]};\",\n",
    "\n",
    "                        # Send new control parameters\n",
    "                         f\"set Sumo__Plant__Sideflowdivider1__param__Qpumped_target {set_recycle};\",\n",
    "                         f\"set Sumo__Plant__CSTR1__param__DOSP {set_air1};\",\n",
    "                         f\"set Sumo__Plant__CSTR2__param__DOSP {set_air2};\",\n",
    "                         f\"set Sumo__Plant__CSTR3__param__DOSP {set_air3};\",\n",
    "                         f\"set Sumo__Plant__CSTR4__param__DOSP {set_air4};\",\n",
    "                         f\"set Sumo__Plant__CSTR5__param__DOSP {set_air5};\",\n",
    "                         f\"set Sumo__Plant__Clarifier1__param__Qsludge_target {set_ras};\",\n",
    "            \n",
    "                        # Set simulation length and how often we want data\n",
    "                         f\"set Sumo__StopTime {15*dtool.minute};\",\n",
    "                         f\"set Sumo__DataComm {15*dtool.minute};\", \n",
    "                                       \n",
    "                        # Choose simulation mode from dynamic/steady\n",
    "                         \"mode dynamic;\",\n",
    "#                         \"loadtsv inf_data.tsv;\"\n",
    "            \n",
    "                        # Start the simulation\n",
    "                        # If you omit this, you can still start the simulation with sumo.sendCommand(job1, \"start\")\n",
    "                        # but we've not seen a usecase where it is needed.\n",
    "                         \"start;\"],\n",
    "                                                 \n",
    "                        # list the variables we're interested in\n",
    "        variables     = self.variables,\n",
    "                        \n",
    "           \n",
    "                        # when a jobData is persistent it doesn't get cleared when the simulation is finished\n",
    "                        # so we can store simulation results in it.\n",
    "        jobData       = {\n",
    "                            ds.sumo.persistent: True,\n",
    "                            \"results\" : { },\n",
    "                            \"finished\": False,\n",
    "                            \"index\" : self.i\n",
    "                        },\n",
    "        );\n",
    "    \n",
    "    # wait until the simulation is finished\n",
    "    jobData = ds.sumo.getJobData(job)\n",
    "    while (jobData[\"finished\"] != True):\n",
    "        time.sleep(0.0001)\n",
    "    \n",
    "        \n",
    "    simdata = ds.sumo.jobData.copy()\n",
    "    \n",
    "    # clean up the job data from memory\n",
    "    ds.sumo.deleteJobData\n",
    "    #ds.sumo.cleanup()\n",
    "    \n",
    "    return simdata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a6c3a4e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define reward function\n",
    "def getreward(self, data):     \n",
    "\n",
    "    ME = 24*0.005*(1000*sum(map(lambda x : x<20, (self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR1__kLaGO2'], \n",
    "                                      self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR2__kLaGO2'])))\n",
    "         + 1333*sum(map(lambda x : x<20, (self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR3__kLaGO2'],\n",
    "                                    self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR4__kLaGO2'],\n",
    "                                    self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR5__kLaGO2']))))\n",
    "\n",
    "    AE = (9.458/(1.8*1000)*(1000*(self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR1__kLaGO2'] \n",
    "                              + self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR2__kLaGO2']) + \n",
    "                            1333*(self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR3__kLaGO2']\n",
    "                              + self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR4__kLaGO2']\n",
    "                              + self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR5__kLaGO2'])))\n",
    "    \n",
    "    PE = (0.004 * self.resultdf.iloc[self.total-1]['Sumo__Plant__Pipe11__Q'] + \n",
    "            0.05 * self.resultdf.iloc[self.total-1]['Sumo__Plant__Sludge1__Q'] + \n",
    "            0.008 * self.resultdf.iloc[self.total-1]['Sumo__Plant__Pipe14__Q'])\n",
    "    \n",
    "    SP = 0 #0.75/1000 * self.resultdf.iloc[self.total-1]['Sumo__Plant__Sludge1__XTSS'] * self.resultdf.iloc[self.total-1]['Sumo__Plant__Sludge1__Q']\n",
    "    \n",
    "    EQ = (2*self.resultdf.iloc[self.total-1]['Sumo__Plant__Effluent1__XTSS']\n",
    "          + self.resultdf.iloc[self.total-1]['Sumo__Plant__Effluent1__TCOD']\n",
    "          + 10*self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR5__SNOx']\n",
    "          + 30*self.resultdf.iloc[self.total-1]['Sumo__Plant__CSTR5__SKN']\n",
    "          + 2*self.resultdf.iloc[self.total-1]['Sumo__Plant__Effluent1__TBOD_5'])*self.resultdf.iloc[self.total-1]['Sumo__Plant__Effluent1__Q']/1000\n",
    "        \n",
    "    reward = 0 - ME - AE - PE - EQ - 5*SP\n",
    "    \n",
    "    return reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a43f9af8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Environment\n",
    "# Write custom environment\n",
    "class SUMOEnv(Env):\n",
    "    def __init__(self):\n",
    "        \n",
    "        ### Define SUMO Parameters ###\n",
    "        \n",
    "        # free all resources of Sumo before we continue\n",
    "        ds.sumo.cleanup()\n",
    "        \n",
    "        # set the number of parallel simulations\n",
    "        ds.sumo.setParallelJobs(1)\n",
    "        \n",
    "        # for debugging:\n",
    "        # ds.sumo.setLogDetails(6)\n",
    "\n",
    "        # Define influent dataframe\n",
    "        self.influentdf = pd.read_excel('Influent.xlsx')\n",
    "\n",
    "        # setup callbacks for Sumo to call\n",
    "        ds.sumo.message_callback = msg_callback\n",
    "        ds.sumo.datacomm_callback = data_callback\n",
    "        \n",
    "        # list the variables we're interested in for Evaluation\n",
    "        self.sumo_time = [\"Sumo__Time\"]      \n",
    "        self.influent = [\"Sumo__Plant__Influent2__Q\",\n",
    "                    \"Sumo__Plant__Influent2__TCOD\",\n",
    "                    \"Sumo__Plant__Influent2__TKN\",\n",
    "                    \"Sumo__Plant__Influent2__TP\"]\n",
    "        self.operation = [\"Sumo__Plant__Pipe11__Q\",\n",
    "                    \"Sumo__Plant__Pipe14__Q\",\n",
    "                    \"Sumo__Plant__Pipe15__Q\",\n",
    "                    \"Sumo__Plant__CSTR1__SO2\",\n",
    "                    \"Sumo__Plant__CSTR2__SO2\",\n",
    "                    \"Sumo__Plant__CSTR3__SO2\",\n",
    "                    \"Sumo__Plant__CSTR4__SO2\",\n",
    "                    \"Sumo__Plant__CSTR5__SO2\",\n",
    "                    \"Sumo__Plant__CSTR1__kLaGO2\",\n",
    "                    \"Sumo__Plant__CSTR2__kLaGO2\",\n",
    "                    \"Sumo__Plant__CSTR3__kLaGO2\",\n",
    "                    \"Sumo__Plant__CSTR4__kLaGO2\",\n",
    "                    \"Sumo__Plant__CSTR5__kLaGO2\",\n",
    "                    \"Sumo__Plant__CSTR1__SNHx\",\n",
    "                    \"Sumo__Plant__CSTR2__SNHx\",\n",
    "                    \"Sumo__Plant__CSTR3__SNHx\",\n",
    "                    \"Sumo__Plant__CSTR4__SNHx\",\n",
    "                    \"Sumo__Plant__CSTR5__SNHx\",\n",
    "                    \"Sumo__Plant__CSTR1__SNOx\",\n",
    "                    \"Sumo__Plant__CSTR2__SNOx\",\n",
    "                    \"Sumo__Plant__CSTR3__SNOx\",\n",
    "                    \"Sumo__Plant__CSTR4__SNOx\",\n",
    "                    \"Sumo__Plant__CSTR5__SNOx\",\n",
    "                    \"Sumo__Plant__CSTR1__XTSS\", \n",
    "                    \"Sumo__Plant__CSTR5__ORP\",\n",
    "                    \"Sumo__Plant__CSTR5__SNOx\",\n",
    "                    \"Sumo__Plant__CSTR5__TKN\",\n",
    "                    \"Sumo__Plant__CSTR5__TCOD\",\n",
    "                    \"Sumo__Plant__CSTR5__TBOD_5\",\n",
    "                    \"Sumo__Plant__CSTR5__SKN\",\n",
    "                    \"Sumo__Plant__CSTR5__TBOD_5\",\n",
    "                    \"Sumo__Plant__SRT1\"]\n",
    "        self.effluent = [\"Sumo__Plant__Effluent1__Q\",\n",
    "                    \"Sumo__Plant__Effluent1__SNOx\",\n",
    "                    \"Sumo__Plant__Effluent1__SNHx\",\n",
    "                    \"Sumo__Plant__Effluent1__SN_B\",\n",
    "                    \"Sumo__Plant__Effluent1__XN_B\",\n",
    "                    \"Sumo__Plant__Effluent1__SALK\",\n",
    "                    \"Sumo__Plant__Effluent1__XTSS\",\n",
    "                    \"Sumo__Plant__Effluent1__TKN\",\n",
    "                    \"Sumo__Plant__Effluent1__SKN\",\n",
    "                    \"Sumo__Plant__Effluent1__TN\",\n",
    "                    \"Sumo__Plant__Effluent1__TCOD\",\n",
    "                    \"Sumo__Plant__Effluent1__TBOD_5\",\n",
    "                    \"Sumo__Plant__Sludge1__Q\",\n",
    "                    \"Sumo__Plant__Sludge1__XTSS\"]\n",
    "\n",
    "        self.variables = self.sumo_time + self.influent + self.operation + self.effluent\n",
    "\n",
    "        self.resultdf = pd.DataFrame()\n",
    "        self.rewarddf = pd.DataFrame(columns=[\"Reward\"])\n",
    "        \n",
    "        ### Define RL Env Parameters ###\n",
    "        \n",
    "        # Define observation space\n",
    "        # define low observation bounds\n",
    "        lowobs = np.array([\n",
    "            # Inf Flow\n",
    "            #0,\n",
    "            # Inf NHx Load\n",
    "            #0,\n",
    "            # CSTR1 SNHx\n",
    "            0,\n",
    "            # CSTR2 SNHx\n",
    "            0,\n",
    "            # CSTR3 SNHx\n",
    "            0,\n",
    "            # CSTR4 SNHx\n",
    "            0,\n",
    "            # CSTR5 SNHx\n",
    "            0,\n",
    "            # CSTR1 SNOx\n",
    "            0,\n",
    "            # CSTR2 SNOx\n",
    "            0,\n",
    "            # CSTR3 SNOx\n",
    "            0,\n",
    "            # CSTR4 SNOx\n",
    "            0,\n",
    "            # CSTR5 SNOx\n",
    "            0,\n",
    "            # CSTR1 XTSS\n",
    "            0,\n",
    "        ]).astype(np.float32)\n",
    "        \n",
    "        # define high observation bounds\n",
    "        highobs = np.array([\n",
    "            # Inf Flow\n",
    "            #1,\n",
    "            # Inf NHx Load\n",
    "            #0,\n",
    "            # CSTR1 SNHx\n",
    "            1,\n",
    "            # CSTR2 SNHx\n",
    "            1,\n",
    "            # CSTR3 SNHx\n",
    "            1,\n",
    "            # CSTR4 SNHx\n",
    "            1,\n",
    "            # CSTR5 SNHx\n",
    "            1,\n",
    "            # CSTR1 SNOx\n",
    "            1,\n",
    "            # CSTR2 SNOx\n",
    "            1,\n",
    "            # CSTR3 SNOx\n",
    "            1,\n",
    "            # CSTR4 SNOx\n",
    "            1,\n",
    "            # CSTR5 SNOx\n",
    "            1,\n",
    "            # CSTR1 XTSS\n",
    "            1,\n",
    "        ]).astype(np.float32)\n",
    "            \n",
    "        self.observation_space = gym.spaces.Box(lowobs,highobs)\n",
    "        \n",
    "        # Define action space\n",
    "       \n",
    "        # define low action bounds\n",
    "        lowact = np.array([\n",
    "            # CSTR1 DO\n",
    "            0,\n",
    "            # CSTR2 DO\n",
    "            0,\n",
    "            # CSTR3 DO\n",
    "            0,\n",
    "            # CSTR4 DO\n",
    "            0,\n",
    "            # CSTR5 DO\n",
    "            0,\n",
    "            # recycle IMLR\n",
    "            0,\n",
    "            # RAS\n",
    "            0,\n",
    "        ]).astype(np.float32)\n",
    "        \n",
    "        # define high observation bounds\n",
    "        highact = np.array([\n",
    "            # CSTR1 DO\n",
    "            1,\n",
    "            # CSTR2 DO\n",
    "            1,\n",
    "            # CSTR3 DO\n",
    "            1,\n",
    "            # CSTR4 DO\n",
    "            1,\n",
    "            # CSTR5 DO\n",
    "            1,\n",
    "            # recycle IMLR\n",
    "            1,\n",
    "            # RAS\n",
    "            1,\n",
    "        ]).astype(np.float32)\n",
    "            \n",
    "        self.action_space = gym.spaces.Box(lowact,highact)\n",
    "        \n",
    "        # Define starting environment\n",
    "        self.state = np.zeros(11)\n",
    "        \n",
    "        # Start sim counter\n",
    "        self.i = 1344\n",
    "        self.total = 0\n",
    "        self.influent = 1344\n",
    "    \n",
    "    def step(self, action):\n",
    "        # Apply action\n",
    "        # This will be done by the RL agent which passes the new action into the SUMO_step function\n",
    "        ### This is where SUMO is called ###\n",
    "                                      \n",
    "        set_air1 = action[0]*3\n",
    "        set_air2 = action[1]*3\n",
    "        set_air3 = action[2]*3\n",
    "        set_air4 = action[3]*3\n",
    "        set_air5 = action[4]*3\n",
    "        set_recycle = action[5]*92230\n",
    "        set_ras = action[6]*36892\n",
    "        \n",
    "        data = SUMO_step(self, set_air1, set_air2, set_air3, set_air4, set_air5, set_recycle, set_ras)\n",
    "\n",
    "        # Increase simulation tracker by 1 timestep\n",
    "        self.i -= 1\n",
    "        self.total += 1\n",
    "        self.influent += 1\n",
    "\n",
    "        self.state = [\n",
    "            data[self.total]['results']['Sumo__Plant__CSTR1__SNHx']/50,\n",
    "            data[self.total]['results']['Sumo__Plant__CSTR2__SNHx']/50,\n",
    "            data[self.total]['results']['Sumo__Plant__CSTR3__SNHx']/50,\n",
    "            data[self.total]['results']['Sumo__Plant__CSTR4__SNHx']/50,\n",
    "            data[self.total]['results']['Sumo__Plant__CSTR5__SNHx']/50,\n",
    "            data[self.total]['results']['Sumo__Plant__CSTR1__SNOx']/50,\n",
    "            data[self.total]['results']['Sumo__Plant__CSTR2__SNOx']/50,\n",
    "            data[self.total]['results']['Sumo__Plant__CSTR3__SNOx']/50,\n",
    "            data[self.total]['results']['Sumo__Plant__CSTR4__SNOx']/50,\n",
    "            data[self.total]['results']['Sumo__Plant__CSTR5__SNOx']/50,\n",
    "            data[self.total]['results']['Sumo__Plant__CSTR1__XTSS']/10000\n",
    "        ]\n",
    "\n",
    "        # concatonate result to dataframe\n",
    "        self.resultdf = pd.concat([self.resultdf,(pd.DataFrame(data[self.total]['results'], index = [0]))])\n",
    "        # update time in first column\n",
    "        self.resultdf.iloc[-1][0] = self.resultdf.iloc[-1][0]*(self.total)\n",
    "\n",
    "        ### Calculate reward ###\n",
    "        reward = getreward(self, data)\n",
    "\n",
    "        # add reward\n",
    "        self.rewarddf.loc[self.total] = reward\n",
    "        \n",
    "        # Check if simulation is complete\n",
    "        # Will need to update run-time\n",
    "        if self.i <= 0: \n",
    "            done = True\n",
    "            # Export results if desired\n",
    "            self.resultdf.to_csv(\"RESULTS_FILE_NAME.csv\", index = True)\n",
    "            print(self.total)\n",
    "        else:\n",
    "            done = False\n",
    "\n",
    "        # Set placeholder for info\n",
    "        info = {}\n",
    "        \n",
    "        # Return step information\n",
    "        return self.state, reward, done, info\n",
    "    \n",
    "    def render(self):\n",
    "        # No renders as part of this environment\n",
    "        pass\n",
    "    \n",
    "    def reset(self):\n",
    "        # Reset starting environment\n",
    "        # Create active xml \n",
    "        # UNCOMMENT THIS SECTION DURING TRAINING/COMMENT DURING TESTING\n",
    "#         SUMOstate_active = \"ACTIVE_STATE_NAME\"\n",
    "#         path = 'YOUR_PATH\\\\Dynamita\\\\Sumo22\\\\PythonAPI\\\\YOUR_PATH'\n",
    "#         src = path + SUMOstate + \".xml\"\n",
    "#         dst = path + SUMOstate_active + \".xml\"\n",
    "\n",
    "#         shutil.copyfile(src, dst)\n",
    "        \n",
    "#         Path = os.path.join('Test_Training', 'Saved_Models', \n",
    "#                         'MODEL_NAME_HERE')\n",
    "#         model.save(Path)\n",
    "        \n",
    "#         # Reset state\n",
    "#         self.state = np.zeros(11)\n",
    "        \n",
    "        # Reset simulation length for tracking episodes (LEAVE UNCOMMENTED)\n",
    "        self.i = 1344\n",
    "        return self.state\n",
    "        \n",
    "    \n",
    "env = SUMOEnv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "12f380d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make Monitor Env\n",
    "env_monitor = Monitor(env)\n",
    "\n",
    "# Wrap environment for Stable Baselines\n",
    "env_wrap = DummyVecEnv([lambda: env_monitor])\n",
    "\n",
    "hyperparams = {\n",
    "        \"gamma\": 0.99,\n",
    "        \"learning_rate\": 0.001,\n",
    "        \"batch_size\": 200,\n",
    "        \"tau\": 0.005,\n",
    "        \"train_freq\": (4, \"episode\"),\n",
    "        \"target_policy_noise\": 0.05,\n",
    "        \"learning_starts\": 100000,\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "aa2235bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 0\n",
    "\n",
    "# Define RL model (UNCOMMENT DURING TRAINING)\n",
    "#model = TD3('MlpPolicy', env_wrap, **hyperparams, verbose=1, tensorboard_log=log_path, seed = seed)\n",
    "\n",
    "np.random.seed(seed)\n",
    "set_random_seed(seed)\n",
    "\n",
    "# Train Model (UNCOMMENT FOR TRAINING)\n",
    "#model.learn(total_timesteps = 150000, log_interval = 1) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0c9322d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# I do not think this section works, it was for the purpose of continuing training of an existing model\n",
    "\n",
    "# Path = os.path.join('Test_Training', 'Saved_Models', \n",
    "#                        'MODEL_NAME_HERE')\n",
    "\n",
    "# log_path = os.path.join('Test_Training', 'Test_Logs')\n",
    "\n",
    "# model = TD3.load(Path, env_wrap, **hyperparams, verbose=1, tensorboard_log=log_path, seed = seed)\n",
    "\n",
    "# model.load_replay_buffer(\"REPLAY_BUFFER_FILE\")\n",
    "\n",
    "# # Train Model\n",
    "# model.learn(total_timesteps = 100000, log_interval = 1) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1ec92dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# UNCOMMENT DURING TRAINING\n",
    "# Save final model\n",
    "# Path = os.path.join('Test_Training', 'Saved_Models', \n",
    "#                        'MODEL_NAME_HERE')\n",
    "# model.save(Path)\n",
    "# model.save_replay_buffer(\"REPLAY_BUFFER_FILE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "054b2433",
   "metadata": {},
   "outputs": [],
   "source": [
    "# UNCOMMENT FOR TESTING\n",
    "# Test the Trained Model\n",
    "\n",
    "Path = os.path.join('Test_Training', 'Saved_Models', \n",
    "                       'MODEL_NAME_HERE')\n",
    "model = TD3.load(Path)\n",
    "\n",
    "episodes = 26\n",
    "total = 0\n",
    "for episode in range (1, episodes+1):\n",
    "    obs = env_wrap.reset()\n",
    "    done = False\n",
    "    score = 0\n",
    "    \n",
    "    while not done:\n",
    "        action, _ = model.predict(obs)\n",
    "        print(action)\n",
    "        obs, reward, done, info = env_wrap.step(action)\n",
    "        print(obs)\n",
    "        score += reward\n",
    "    #print('Episode:{} Score:{}'.format(episode, score))\n",
    "    total += score\n",
    "\n",
    "print('Average score:{}'.format(total/episodes))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}