Update evaluate and train to evaluate per episode

ultra/tests/test_evaluate.py

@@ -81,6 +81,7 @@ def run_experiment():
             policy_class = "ultra.baselines.sac:sac-v0"
             log_dir = "tests/output_eval_check_logs"
 
+            episode_count = 0
             for episode in episodes(1, etag=policy_class, log_dir=log_dir):
                 observations = env.reset()
                 state = observations[AGENT_ID]
@@ -95,19 +96,6 @@ def run_experiment():
                         dill.dump(spec, spec_output, pickle.HIGHEST_PROTOCOL)
 
                 while not dones["__all__"]:
-                    evaluation_check(
-                        agent=agent,
-                        agent_id=AGENT_ID,
-                        episode=episode,
-                        eval_rate=10,
-                        eval_episodes=1,
-                        max_episode_steps=2,
-                        policy_class=policy_class,
-                        scenario_info=("00", "eval_test"),
-                        timestep_sec=0.1,
-                        headless=True,
-                        log_dir=log_dir,
-                    )
                     action = agent.act(state, explore=True)
                     observations, rewards, dones, infos = env.step({AGENT_ID: action})
                     next_state = observations[AGENT_ID]
@@ -130,6 +118,21 @@ def run_experiment():
                     )
                     total_step += 1
                     state = next_state
+                evaluation_check(
+                    agent=agent,
+                    agent_id=AGENT_ID,
+                    episode=episode,
+                    eval_rate=10,
+                    eval_episodes=1,
+                    max_episode_steps=2,
+                    episode_count=episode_count,
+                    policy_class=policy_class,
+                    scenario_info=("00", "eval_test"),
+                    timestep_sec=0.1,
+                    headless=True,
+                    log_dir=log_dir,
+                )
+                episode_count += 1
 
             env.close()
 
@@ -182,7 +185,6 @@ def test_evaluate_agent(self):
                 agent_id="AGENT_001",
                 policy_class=policy_class,
                 seed=seed,
-                itr_count=0,
                 checkpoint_dir=model,
                 scenario_info=("00", "eval_test"),
                 num_episodes=1,

ultra/ultra/evaluate.py

@@ -53,6 +53,7 @@ def evaluation_check(
     eval_rate,
     eval_episodes,
     max_episode_steps,
+    episode_count,
     scenario_info,
     timestep_sec,
     headless,
@@ -61,10 +62,12 @@ def evaluation_check(
     agent_itr = episode.get_itr(agent_id)
 
     print(
-        f"Agent iteration : {agent_itr}, Eval rate : {eval_rate}, last_eval_iter : {episode.last_eval_iteration}"
+        f"Agent iteration : {agent_itr}, Episode count : {episode_count}, Eval rate : {eval_rate}, last_eval_iter : {episode.last_eval_iteration}"
     )
-    if (agent_itr + 1) % eval_rate == 0 and episode.last_eval_iteration != agent_itr:
-        checkpoint_dir = episode.checkpoint_dir(agent_itr)
+    if (
+        episode_count + 1
+    ) % eval_rate == 0 and episode.last_eval_iteration != episode_count:
+        checkpoint_dir = episode.checkpoint_dir(episode_count)
         agent.save(checkpoint_dir)
         episode.eval_mode()
         episode.info[episode.active_tag][agent_id] = ray.get(
@@ -74,7 +77,6 @@ def evaluation_check(
                     agent_id=agent_id,
                     policy_class=policy_class,
                     seed=episode.eval_count,
-                    itr_count=agent_itr,
                     checkpoint_dir=checkpoint_dir,
                     scenario_info=scenario_info,
                     num_episodes=eval_episodes,
@@ -86,7 +88,10 @@ def evaluation_check(
             ]
         )[0]
         episode.eval_count += 1
-        episode.last_eval_iteration = agent_itr
+        episode.last_eval_iteration = episode_count
+        episode.record_tensorboard()
+        episode.gap_mode()
+        episode.calculate_gap()
         episode.record_tensorboard()
         episode.train_mode()
 
@@ -98,14 +103,14 @@ def evaluate(
     seed,
     agent_id,
     policy_class,
-    itr_count,
     checkpoint_dir,
     scenario_info,
     num_episodes,
     max_episode_steps,
     headless,
     timestep_sec,
     log_dir,
+    explore=False,
 ):
 
     torch.set_num_threads(1)
@@ -137,7 +142,7 @@ def evaluate(
 
         episode.reset(mode="Evaluation")
         while not dones["__all__"]:
-            action = agent.act(state, explore=False)
+            action = agent.act(state, explore=True)
             observations, rewards, dones, infos = env.step({agent_id: action})
 
             next_state = observations[agent_id]
@@ -238,7 +243,6 @@ def evaluate(
         ):
             model = sorted_models[episode.index]
             print("model: ", model)
-            episode_count = model.split("/")[-1]
             episode.eval_mode()
             episode.info[episode.active_tag][AGENT_ID] = ray.get(
                 [
@@ -247,7 +251,6 @@ def evaluate(
                         agent_id=AGENT_ID,
                         policy_class=policy_class,
                         seed=episode.eval_count,
-                        itr_count=0,
                         checkpoint_dir=model,
                         scenario_info=(args.task, args.level),
                         num_episodes=int(args.episodes),

ultra/ultra/train.py

@@ -75,6 +75,8 @@ def train(
 
     agent = spec.build_agent()
 
+    episode_count = 0
+    old_episode = None
     for episode in episodes(num_episodes, etag=policy_class, log_dir=log_dir):
         observations = env.reset()
         state = observations[AGENT_ID]
@@ -93,16 +95,6 @@ def train(
             if episode.get_itr(AGENT_ID) >= 1000000:
                 finished = True
                 break
-            evaluation_check(
-                agent=agent,
-                agent_id=AGENT_ID,
-                policy_class=policy_class,
-                episode=episode,
-                log_dir=log_dir,
-                max_episode_steps=max_episode_steps,
-                **eval_info,
-                **env.info,
-            )
             action = agent.act(state, explore=True)
             observations, rewards, dones, infos = env.step({AGENT_ID: action})
             next_state = observations[AGENT_ID]
@@ -124,8 +116,25 @@ def train(
             total_step += 1
             state = next_state
 
-        episode.record_episode()
-        episode.record_tensorboard()
+        episode.record_episode(old_episode, eval_info["eval_rate"])
+        old_episode = episode
+
+        if (episode_count + 1) % eval_info["eval_rate"] == 0:
+            episode.record_tensorboard()
+            old_episode = None
+
+        evaluation_check(
+            agent=agent,
+            agent_id=AGENT_ID,
+            policy_class=policy_class,
+            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
 
@@ -169,9 +178,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",
@@ -209,7 +218,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

@@ -35,8 +35,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,
@@ -57,7 +57,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
@@ -71,7 +71,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
@@ -95,14 +95,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:
@@ -177,6 +176,18 @@ 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:
+                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
@@ -219,10 +230,20 @@ def record_step(self, agent_id, infos, rewards, total_step=0, loss_output=None):
         self.steps += 1
         self.agents_itr[agent_id] += 1
 
-    def record_episode(self):
+    def record_episode(self, old_episode=None, eval_rate=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):
+            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]
+                            * eval_rate
+                        ) / eval_rate
+
     def initialize_tb_writer(self):
         if self.tb_writer is None:
             self.tb_writer = SummaryWriter(
@@ -314,7 +335,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}",