Add outflow to tensorboard

Author: eugenevinitsky

examples/train.py

@@ -159,6 +159,7 @@ def on_episode_start(info):
         episode = info["episode"]
         episode.user_data["avg_speed"] = []
         episode.user_data["energy"] = []
+        episode.user_data["outflow"] = []
 
     def on_episode_step(info):
         episode = info["episode"]
@@ -172,11 +173,14 @@ def on_episode_step(info):
 
     def on_episode_end(info):
         episode = info["episode"]
+        env = info["env"].get_unwrapped()[0]
         avg_speed = np.mean(episode.user_data["avg_speed"])
         avg_energy = np.mean(episode.user_data["avg_energy"])
 
         episode.custom_metrics["avg_speed"] = avg_speed
         episode.custom_metrics["avg_energy"] = avg_energy
+        episode.custom_metrics["outflow"] = env.k.vehicle.get_outflow_rate()
+
 
     config["callbacks"] = {"on_episode_start": tune.function(on_episode_start),
                            "on_episode_step": tune.function(on_episode_step),