Update evaluate and train to evaluate per episode

ultra/tests/test_evaluate.py

@@ -404,6 +404,7 @@ def run_experiment(scenario_info, num_agents, log_dir, headless=True):
                 eval_rate=10,
                 eval_episodes=1,
                 max_episode_steps=2,
+                episode_count=1,
                 policy_classes=agent_classes,
                 scenario_info=scenario_info,
                 timestep_sec=0.1,

ultra/ultra/evaluate.py

@@ -49,6 +49,7 @@ def evaluation_check(
     eval_rate,
     eval_episodes,
     max_episode_steps,
+    episode_count,
     scenario_info,
     timestep_sec,
     headless,
@@ -58,8 +59,8 @@ def evaluation_check(
     agent_ids_to_evaluate = [
         agent_id
         for agent_id in agent_ids
-        if (episode.get_itr(agent_id) + 1) % eval_rate == 0
-        and episode.last_eval_iterations[agent_id] != episode.get_itr(agent_id)
+        if (episode_count + 1) % eval_rate == 0
+        and episode.last_eval_iterations[agent_id] != episode_count
     ]
 
     # Skip evaluation if there are no agents needing an evaluation.
@@ -71,9 +72,7 @@ def evaluation_check(
 
     for agent_id in agent_ids_to_evaluate:
         # Get the checkpoint directory for the current agent and save its model.
-        checkpoint_directory = episode.checkpoint_dir(
-            agent_id, episode.get_itr(agent_id)
-        )
+        checkpoint_directory = episode.checkpoint_dir(agent_id, episode_count)
         agents[agent_id].save(checkpoint_directory)
 
         # Perform the evaluation on this agent and save the data.
@@ -97,11 +96,14 @@ def evaluation_check(
             )[0]
         )
         episode.eval_count += 1
-        episode.last_eval_iterations[agent_id] = episode.get_itr(agent_id)
+        episode.last_eval_iterations[agent_id] = episode_count
 
     # Put the evaluation data for all agents into the episode and record the TensorBoard.
     episode.info[episode.active_tag] = evaluation_data
     episode.record_tensorboard()
+    episode.gap_mode()
+    episode.calculate_gap()
+    episode.record_tensorboard()
     episode.train_mode()
 
 
@@ -119,6 +121,7 @@ def evaluate(
     headless,
     timestep_sec,
     log_dir,
+    explore=False,
 ):
     torch.set_num_threads(1)
 

ultra/ultra/train.py

@@ -108,6 +108,8 @@ def train(
     # policy_classes list, transform it to an etag of "dqn-v0:ppo-v0".
     etag = ":".join([policy_class.split(":")[-1] for policy_class in policy_classes])
 
+    episode_count = 0
+    old_episode = None
     for episode in episodes(num_episodes, etag=etag, log_dir=log_dir):
         # Reset the environment and retrieve the initial observations.
         observations = env.reset()
@@ -136,19 +138,6 @@ def train(
             if any([episode.get_itr(agent_id) >= 1000000 for agent_id in agents]):
                 finished = True
                 break
-
-            # Perform the evaluation check.
-            evaluation_check(
-                agents=agents,
-                agent_ids=agent_ids,
-                policy_classes=agent_classes,
-                episode=episode,
-                log_dir=log_dir,
-                max_episode_steps=max_episode_steps,
-                **eval_info,
-                **env.info,
-            )
-
             # Request and perform actions on each agent that received an observation.
             actions = {
                 agent_id: agents[agent_id].act(observation, explore=True)
@@ -184,10 +173,26 @@ def train(
             total_step += 1
             observations = next_observations
 
-        # Normalize the data and record this episode on tensorboard.
-        episode.record_episode()
-        episode.record_tensorboard()
+        episode.record_episode(old_episode, eval_info["eval_rate"], count=episode_count)
+        old_episode = episode
 
+        if (episode_count + 1) % eval_info["eval_rate"] == 0:
+            episode.record_tensorboard()
+            old_episode = None
+
+        if eval_info["eval_episodes"] != 0:
+            evaluation_check(
+                agents=agents,
+                agent_ids=agent_ids,
+                policy_classes=agent_classes,
+                episode=episode,
+                log_dir=log_dir,
+                max_episode_steps=max_episode_steps,
+                episode_count=episode_count,
+                **eval_info,
+                **env.info,
+            )
+        episode_count += 1
         if finished:
             break
 
@@ -231,9 +236,9 @@ def train(
     )
     parser.add_argument(
         "--eval-rate",
-        help="Evaluation rate based on number of observations",
+        help="Evaluation rate based on number of episodes",
         type=int,
-        default=10000,
+        default=100,
     )
     parser.add_argument(
         "--seed",
@@ -283,7 +288,7 @@ def train(
                 num_episodes=int(args.episodes),
                 max_episode_steps=int(args.max_episode_steps),
                 eval_info={
-                    "eval_rate": float(args.eval_rate),
+                    "eval_rate": int(args.eval_rate),
                     "eval_episodes": int(args.eval_episodes),
                 },
                 timestep_sec=float(args.timestep),

ultra/ultra/utils/episode.py

@@ -44,8 +44,8 @@
 class LogInfo:
     def __init__(self):
         self.data = {
-            "env_score": 0,
-            "episode_reward": 0,
+            "env_score": 0.0,
+            "episode_reward": 0.0,
             "dist_center": 0,
             "goal_dist": 0,
             "speed": 0,
@@ -66,7 +66,7 @@ def __init__(self):
         }
 
     def add(self, infos, rewards):
-        self.data["env_score"] += int(infos["logs"]["env_score"])
+        self.data["env_score"] += infos["logs"]["env_score"]
         self.data["speed"] += infos["logs"]["speed"]
         self.data["max_speed_violation"] += (
             1 if infos["logs"]["speed"] > infos["logs"]["closest_wp"].speed_limit else 0
@@ -80,7 +80,7 @@ def add(self, infos, rewards):
         self.data["ego_linear_jerk"] += infos["logs"]["linear_jerk"]
         self.data["ego_angular_jerk"] += infos["logs"]["angular_jerk"]
         self.data["episode_reward"] += rewards
-        self.data["final_pos"] = infos["logs"]["position"]
+        self.data["final_pos"] = infos["logs"]["position"][:2]
         self.data["start_pos"] = infos["logs"]["start"].position
         self.data["dist_travelled"] = math.sqrt(
             (self.data["final_pos"][1] - self.data["start_pos"][1]) ** 2
@@ -104,14 +104,13 @@ def step(self):
 
     def normalize(self):
         steps = self.data["episode_length"]
-        self.data["env_score"] /= steps
         self.data["dist_center"] /= steps
         self.data["speed"] /= steps
         self.data["ego_linear_jerk"] /= steps
         self.data["ego_angular_jerk"] /= steps
-        self.data["ego_num_violations"] /= steps
-        self.data["social_num_violations"] /= steps
-        self.data["max_speed_violation"] /= steps
+        # self.data["ego_num_violations"] /= steps
+        # self.data["social_num_violations"] /= steps
+        # self.data["max_speed_violation"] /= steps
 
 
 class Episode:
@@ -186,6 +185,19 @@ def train_mode(self):
     def eval_mode(self):
         self.active_tag = "Evaluation"
 
+    def gap_mode(self):
+        self.active_tag = "Gap"
+
+    def calculate_gap(self):
+        gap_info = self.info["Gap"]
+        for agent_id, agent_info in self.info["Train"].items():
+            for key in agent_info.data:
+                if np.isscalar(gap_info[agent_id].data[key]):
+                    gap_info[agent_id].data[key] = (
+                        self.info["Train"][agent_id].data[key]
+                        - self.info["Evaluation"][agent_id].data[key]
+                    )
+
     def reset(self, mode="Train"):
         self.start_time = time.time()
         self.timestep_sec = 0.1
@@ -238,10 +250,21 @@ def record_step(
         # Increment this episode's step count.
         self.steps += 1
 
-    def record_episode(self):
+    def record_episode(self, old_episode=None, eval_rate=None, count=None):
         for _, agent_info in self.info[self.active_tag].items():
             agent_info.normalize()
 
+        if (old_episode is not None) and (eval_rate is not None):
+            count = count % eval_rate
+            for agent_id, agent_info in self.info[self.active_tag].items():
+                for key in agent_info.data:
+                    if np.isscalar(agent_info.data[key]):
+                        agent_info.data[key] = (
+                            agent_info.data[key]
+                            + old_episode.info[self.active_tag][agent_id].data[key]
+                            * count
+                        ) / (count + 1)
+
     def initialize_tb_writer(self):
         if self.tb_writer is None:
             self.tb_writer = SummaryWriter(
@@ -333,7 +356,7 @@ def episodes(n, etag=None, log_dir=None):
                     for agent_id, agent_info in e.info[e.active_tag].items()
                 ]
                 row = (
-                    f"{e.index}/{n}",
+                    f"{e.index + 1}/{n}",
                     f"{e.sim2wall_ratio:.2f}",
                     f"{e.steps}",
                     f"{e.steps_per_second:.2f}",