Modified Reward Function, added local run functionality to TrainDispatcher.py, added functionality to save best model during training

Author: dkapur17

SBAgent/EvaluateExperiment.py

@@ -5,6 +5,7 @@
 import argparse
 import json
 import numpy as np
+import random
 from envs.utils.EnvBuilder import EnvBuilder
 from stable_baselines3 import PPO
 from tqdm import tqdm
@@ -18,6 +19,9 @@
 args = parser.parse_args()
 
 
+np.random.seed(42)
+random.seed(42)
+
 with open(args.experimentConfigFile, 'r') as f:
     experimentConfig = json.load(f)
 
@@ -28,7 +32,7 @@
 print(f"Running Evaluation on {experimentName}")
 
 env = EnvBuilder.buildEnvFromConfig(os.path.join('..', 'configs', configFileName), gui=args.gui)
-agent = PPO.load(os.path.join('models', modelName))
+agent = PPO.load(os.path.join('models', modelName, 'best_model'))
 
 totalTrials = args.trials
 successfulTrials = 0

SBAgent/EvaluateModel.py

@@ -7,22 +7,23 @@
 from envs.utils.EnvBuilder import EnvBuilder
 from stable_baselines3 import PPO
 from tqdm import tqdm
+from tabulate import tabulate
 
 parser = argparse.ArgumentParser()
 parser.add_argument("configFileName", help="Name of the environment config file.", type=str)
-parser.add_argument("inputModelName", help="(base|finetuned + )Name of the model to load.", type=str)
-parser.add_argument("--trials", default=100, help="Number of episodes to evaluate for.", type=int)
-parser.add_argument("--gui", action=argparse.BooleanOptionalAction, help="Whether or not to show GUI")
-
+parser.add_argument("inputModelPath", help="Path to model to evaluate", type=str)
+parser.add_argument("-t", "--trials", default=100, help="Number of episodes to evaluate for.", type=int)
+parser.add_argument('--gui', action='store_true', help='Enable GUI')
+parser.add_argument('--no-gui', action='store_false', dest='gui', help='Disable GUI')
 args = parser.parse_args()
 
 configFileName = args.configFileName
-modelName = args.outputModelName
+modelName = args.inputModelPath
 
 env = EnvBuilder.buildEnvFromConfig(os.path.join('..', 'configs', configFileName), gui=args.gui)
-agent = PPO.load(os.path.join('models', modelName))
+agent = PPO.load(modelName)
 
-totalTrials = args.trails
+totalTrials = args.trials
 successfulTrials = 0
 rewards = []
 durations = []
@@ -42,7 +43,7 @@
 
     if info['success']:
         successfulTrials += 1
-    if info['reason'] == "collision":
+    elif info['reason'] == "collision":
         nCollisions +=1
     else:
         incompleteDistances.append(np.linalg.norm(obs[:(obs.shape[0]//2)]))
@@ -52,17 +53,13 @@
 
 env.close()
 
+evaluationResults = {
+    'Success Rate': f"{successfulTrials/totalTrials * 100:.2f}%",
+    'Collision Rate': f"{nCollisions/totalTrials * 100:.2f}%",
+    'Mean Incompletion Distance': f"{sum(incompleteDistances)/len(incompleteDistances):.2f}m" if len(incompleteDistances) > 0 else "N/A",
+    'Mean Reward': f"{sum(rewards)/len(rewards):.2f}",
+    'Mean Episode Length': f"{sum(durations)/len(durations)}"
+}
 
-
-
-print(f"---------------------------------------------------------")
-print(f"EVALUATION STATISTICS")
-print()
-print(f"Success Rate: {successfulTrials/totalTrials * 100:.2f}%")
-print(f"Mean Reward: {sum(rewards)/len(rewards):.2f}")
-print(f"Minimum Reward: {min(rewards):.2f}")
-print(f"Maximium Reward: {max(rewards):.2f}")
-print(f"Mean Episode Duration: {sum(durations)/len(durations):.2f} steps")
-print(f"Shortest Episode: {min(durations)} steps")
-print(f"Longest Episode: {max(durations)} steps")
-print(f"---------------------------------------------------------")
+evaluationTable = [[k, v] for k,v in evaluationResults.items()]
+print(tabulate(evaluationTable, headers=["Metric", "Value"], tablefmt='github'))

SBAgent/TrainModel.py

@@ -4,7 +4,7 @@
 import os
 import argparse
 from stable_baselines3 import PPO
-from stable_baselines3.common.callbacks import CheckpointCallback
+from stable_baselines3.common.callbacks import EvalCallback
 from envs.utils.EnvBuilder import EnvBuilder
 
 
@@ -22,12 +22,11 @@
 
 env = EnvBuilder.buildEnvFromConfig(os.path.join('..', 'configs', configFileName), gui=False)
 
-checkpoint_callback = CheckpointCallback(
-  save_freq=1000000,
-  save_path=os.path.join("checkpoints", modelName),
-  name_prefix=f"chkpt",
-)
+eval_callback = EvalCallback(env, best_model_save_path=os.path.join('models', modelName), 
+                                log_path=os.path.join('sbEvalLogs', modelName), 
+                                eval_freq=100_000, deterministic=True, render=False)
+
+agent = PPO("MlpPolicy", env, verbose=1, tensorboard_log=os.path.join('logs', modelName))
+agent.learn(n_steps, callback=eval_callback, tb_log_name="train_logs")
+
 
-agent = PPO('MlpPolicy', env, verbose=1, tensorboard_log=os.path.join('logs', modelName))
-agent.learn(n_steps, callback=checkpoint_callback, tb_log_name="train_logs")
-agent.save(os.path.join('models', modelName))

SBAgent/checkpoints/base/.hidden renamed to SBAgent/sbEvalLogs/base/.hidden

@@ -8,29 +8,42 @@
 parser = argparse.ArgumentParser()
 parser.add_argument("experimentConfigFile", help="Experiment Config File Path")
 parser.add_argument("-s", "--steps", default=2_000_000, help="Number of timesteps to train for", type=int)
+parser.add_argument('--local', action='store_true', help='Run on Local Machine')
 args = parser.parse_args()
 
-with open('trainScriptTemplate.sh', 'r') as f:
-    script = ''.join(f.readlines())
 
-with open(args.experimentConfigFile, 'r') as f:
-    experimentConfig = json.load(f)
+if args.local:
+    with open(args.experimentConfigFile, 'r') as f:
+        experimentConfig = json.load(f)
 
-experimentName = experimentConfig["name"]
-envConfig = experimentConfig["trainParameters"]["config"]
-modelName = experimentConfig["trainParameters"]["outputModelName"]
+    experimentName = experimentConfig["name"]
+    envConfig = experimentConfig["trainParameters"]["config"]
+    modelName = experimentConfig["trainParameters"]["outputModelName"]
 
-script = script.replace("{outputFile}", f"jobOutputs/{experimentName}_train_output.txt")
-script = script.replace("{jobName}", f"{experimentName}_train")
-script = script.replace("{configFile}", envConfig)
-script = script.replace("{outputModelName}", modelName)
-script = script.replace("{steps}", str(args.steps))
+    os.chdir('SBAgent')
+    os.system(f"python TrainModel.py {envConfig} {modelName} --steps {args.steps}")
+else:
+    with open('trainScriptTemplate.sh', 'r') as f:
+        script = ''.join(f.readlines())
 
-tmp = tempfile.NamedTemporaryFile()
+    with open(args.experimentConfigFile, 'r') as f:
+        experimentConfig = json.load(f)
 
-with open(tmp.name, 'w') as f:
-    f.write(script)
+    experimentName = experimentConfig["name"]
+    envConfig = experimentConfig["trainParameters"]["config"]
+    modelName = experimentConfig["trainParameters"]["outputModelName"]
 
-print(f"Dispatching Train Job for {experimentName}")
+    script = script.replace("{outputFile}", f"jobOutputs/{experimentName}_train_output.txt")
+    script = script.replace("{jobName}", f"{experimentName}_train")
+    script = script.replace("{configFile}", envConfig)
+    script = script.replace("{outputModelName}", modelName)
+    script = script.replace("{steps}", str(args.steps))
 
-os.system(f"sbatch {tmp.name}")
+    tmp = tempfile.NamedTemporaryFile()
+
+    with open(tmp.name, 'w') as f:
+        f.write(script)
+
+    print(f"Dispatching Train Job for {experimentName}")
+
+    os.system(f"sbatch {tmp.name}")

SBAgent/checkpoints/finetuned/.hidden renamed to SBAgent/sbEvalLogs/finetuned/.hidden

@@ -13,7 +13,7 @@ class ObstacleAviary(BaseSingleAgentAviary):
 
     CLOSE_TO_FINISH_REWARD = 5
     SUCCESS_REWARD = 1000
-    COLLISION_PENALTY = -100
+    COLLISION_PENALTY = -1000
 
     SUCCESS_EPSILON = 0.1
 
@@ -226,8 +226,8 @@ def _computeReward(self):
         if np.linalg.norm(self.targetPos - pos) < ObstacleAviary.SUCCESS_EPSILON:
             return ObstacleAviary.SUCCESS_REWARD
 
-        if np.linalg.norm(self.targetPos - pos) < ObstacleAviary.MINOR_SAFETY_BOUND_RADIUS:
-            return ObstacleAviary.CLOSE_TO_FINISH_REWARD
+        # if np.linalg.norm(self.targetPos - pos) < ObstacleAviary.MINOR_SAFETY_BOUND_RADIUS:
+        #     return ObstacleAviary.CLOSE_TO_FINISH_REWARD
 
         offsetToClosestObstacle = self._computeOffsetToClosestObstacle()
 
@@ -239,7 +239,8 @@ def _computeReward(self):
         majorBoundBreach = distToClosestObstacle < ObstacleAviary.MAJOR_SAFETY_BOUND_RADIUS
         minorBoundBreach = distToClosestObstacle < ObstacleAviary.MINOR_SAFETY_BOUND_RADIUS
 
-        return 0.5*np.linalg.norm(pos - self.initPos) - 2*np.linalg.norm(self.targetPos - pos) - 10*majorBoundBreach - 2*minorBoundBreach
+        # return 0.5*np.linalg.norm(pos - self.initPos) -*np.linalg.norm(self.targetPos - pos) - 10*majorBoundBreach - 2*minorBoundBreach
+        return -2*np.linalg.norm(self.targetPos - pos) - 1*majorBoundBreach - 0.1*minorBoundBreach
 
 
     def _computeOffsetToClosestObstacle(self):