fix(a3c): use episode outcome for log coloring

 - previously the win was inferred from the total episode reward, but it's not always the case that a winning episode has a > 0 reward, especially with long episodes.
 - pass the win info along and use it to control color instead of total reward

Author: justindujardin

libraries/mathy_python/mathy/agents/a3c/util.py

@@ -1,9 +1,9 @@
-from typing import Dict
-from colr import color
 import datetime
-
-
+import multiprocessing
 from dataclasses import dataclass, field
+from typing import Dict
+
+from colr import color
 
 
 @dataclass
@@ -35,14 +35,15 @@ def truncate(value):
 
 
 def record(
-    episode,
-    episode_reward,
-    worker_idx,
-    global_ep_reward,
-    result_queue,
+    episode: int,
+    is_win: bool,
+    episode_reward: float,
+    worker_idx: int,
+    global_ep_reward: float,
+    result_queue: multiprocessing.Queue,
     losses: EpisodeLosses,
-    num_steps,
-    env_name,
+    num_steps: int,
+    env_name: str,
 ):
     """Helper function to store score and print statistics.
   Arguments:
@@ -56,10 +57,12 @@ def record(
   """
 
     now = datetime.datetime.now().strftime("%H:%M:%S")
+    # Clamp to range -2, 2
+    episode_reward = min(2.0, max(-2.0, episode_reward))
 
     global_ep_reward = global_ep_reward * 0.99 + episode_reward * 0.01
 
-    fore = "green" if episode_reward > 0.0 else "red"
+    fore = "green" if is_win else "red"
     heading = "{:<8} {:<3} {:<8} {:<10}".format(
         now, f"w{worker_idx}", f"ep: {episode}", f"steps: {num_steps}"
     )

libraries/mathy_python/mathy/agents/a3c/worker.py

@@ -246,7 +246,7 @@ def run_episode(self, episode_memory: EpisodeMemory) -> float:
                 continue
 
             # Take an env step
-            observation, reward, done, _ = env.step(action)
+            observation, reward, done, last_obs_info = env.step(action)
 
             observation = MathyObservation(
                 nodes=observation.nodes,
@@ -262,10 +262,7 @@ def run_episode(self, episode_memory: EpisodeMemory) -> float:
             )
             ep_reward += reward
             episode_memory.store(
-                observation=last_observation,
-                action=action,
-                reward=reward,
-                value=value,
+                observation=last_observation, action=action, reward=reward, value=value,
             )
             if time_count == self.args.update_gradients_every or done:
                 if done and self.args.print_training and self.worker_idx == 0:
@@ -275,7 +272,9 @@ def run_episode(self, episode_memory: EpisodeMemory) -> float:
                 self.maybe_write_histograms()
                 time_count = 0
                 if done:
-                    self.finish_episode(ep_reward, ep_steps, env.state)
+                    self.finish_episode(
+                        last_obs_info.get("win", False), ep_reward, ep_steps, env.state
+                    )
 
             ep_steps += 1
             time_count += 1
@@ -383,15 +382,24 @@ def update_global_network(
 
         if done:
             episode_memory.clear()
+        else:
+            episode_memory.clear_except_window(self.args.prediction_window_size)
 
-    def finish_episode(self, episode_reward, episode_steps, last_state: MathyEnvState):
+    def finish_episode(
+        self,
+        is_win: bool,
+        episode_reward: float,
+        episode_steps: int,
+        last_state: MathyEnvState,
+    ):
         env_name = self.teacher.get_env(self.worker_idx, self.iteration)
 
         # Only observe/track the most-greedy worker (high epsilon exploration
         # stats are unlikely to be consistent or informative)
         if self.worker_idx == 0:
             A3CWorker.global_moving_average_reward = record(
                 A3CWorker.global_episode,
+                is_win,
                 episode_reward,
                 self.worker_idx,
                 A3CWorker.global_moving_average_reward,
@@ -490,7 +498,9 @@ def compute_policy_value_loss(
 
         policy_loss *= advantage
         policy_loss = tf.reduce_mean(policy_loss)
-
+        if self.args.normalize_pi_loss:
+            policy_loss /= sequence_length
+        # Scale the policy loss down by the seq_len to make invariant to length
         total_loss = value_loss + policy_loss + entropy_loss + rp_loss
         prefix = self.tb_prefix
         tf.summary.scalar(f"{prefix}/policy_loss", data=policy_loss, step=step)

libraries/mathy_python/mathy/envs/gym/mathy_gym_env.py

@@ -66,8 +66,10 @@ def action_size(self) -> int:
     def step(self, action):
         self.state, transition, change = self.mathy.get_next_state(self.state, action)
         observation = self._observe(self.state)
-        info = {"transition": transition}
         done = is_terminal_transition(transition)
+        info = {"transition": transition, "done": done}
+        if done:
+            info["win"] = transition.reward > 0.0
         self.last_action = action
         self.last_change = change
         self.last_reward = round(float(transition.reward), 4)