added tensorboard to A2C

Author: hill-a

data/cartpole.gif

@@ -67,6 +67,7 @@ def __init__(self, policy, env, gamma=0.99, n_steps=5, vf_coef=0.25, ent_coef=0.
         self.initial_state = None
         self.learning_rate_schedule = None
         self.writer = None
+        self.summary = None
 
         # if we are loading, it is possible the environment is not known, however the obs and action space are known
         if _init_setup_model:
@@ -107,12 +108,29 @@ def setup_model(self):
                     self.vf_loss = mse(tf.squeeze(train_model.value_fn), self.rewards_ph)
                     loss = self.pg_loss - self.entropy * self.ent_coef + self.vf_loss * self.vf_coef
 
+                    tf.summary.scalar('entropy_loss', self.entropy)
+                    tf.summary.scalar('policy_gradient_loss', self.pg_loss)
+                    tf.summary.scalar('value_function_loss', self.vf_loss)
+                    tf.summary.scalar('loss', loss)
+
                     self.params = find_trainable_variables("model")
                     grads = tf.gradients(loss, self.params)
                     if self.max_grad_norm is not None:
                         grads, _ = tf.clip_by_global_norm(grads, self.max_grad_norm)
                     grads = list(zip(grads, self.params))
 
+                with tf.variable_scope("info", reuse=False):
+                    tf.summary.scalar('rewards', tf.reduce_mean(self.rewards_ph))
+                    tf.summary.histogram('rewards', self.rewards_ph)
+                    tf.summary.scalar('learning_rate', tf.reduce_mean(self.learning_rate))
+                    tf.summary.histogram('learning_rate', self.learning_rate)
+                    tf.summary.scalar('advs', tf.reduce_mean(self.advs_ph))
+                    tf.summary.histogram('advs', self.advs_ph)
+                    if len(self.env.observation_space.shape) == 3:
+                        tf.summary.image('observation', train_model.obs_ph)
+                    else:
+                        tf.summary.histogram('observation', train_model.obs_ph)
+
                 trainer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate_ph, decay=self.alpha,
                                                     epsilon=self.epsilon)
                 self.apply_backprop = trainer.apply_gradients(grads)
@@ -125,10 +143,12 @@ def setup_model(self):
                 self.initial_state = step_model.initial_state
                 tf.global_variables_initializer().run(session=self.sess)
 
+                self.summary = tf.summary.merge_all()
+
             if self.tensorboard_log is not None:
                 self.writer = tf.summary.FileWriter(self.tensorboard_log, graph=self.graph)
 
-    def _train_step(self, obs, states, rewards, masks, actions, values):
+    def _train_step(self, obs, states, rewards, masks, actions, values, update):
         """
         applies a training step to the model
 
@@ -138,6 +158,7 @@ def _train_step(self, obs, states, rewards, masks, actions, values):
         :param masks: ([bool]) Whether or not the episode is over (used for reccurent policies)
         :param actions: ([float]) The actions taken
         :param values: ([float]) The logits values
+        :param update: (int) the current step iteration
         :return: (float, float, float) policy loss, value loss, policy entropy
         """
         advs = rewards - values
@@ -152,8 +173,22 @@ def _train_step(self, obs, states, rewards, masks, actions, values):
             td_map[self.train_model.states_ph] = states
             td_map[self.train_model.masks_ph] = masks
 
-        policy_loss, value_loss, policy_entropy, _ = self.sess.run(
-            [self.pg_loss, self.vf_loss, self.entropy, self.apply_backprop], td_map)
+        if update % 10 == 9:
+            run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
+            run_metadata = tf.RunMetadata()
+            summary, policy_loss, value_loss, policy_entropy, _ = self.sess.run(
+                [self.summary, self.pg_loss, self.vf_loss, self.entropy, self.apply_backprop],
+                td_map, options=run_options, run_metadata=run_metadata)
+
+            if self.writer is not None:
+                self.writer.add_run_metadata(run_metadata, 'step%d' % (update * (self.n_batch + 1)))
+        else:
+            summary, policy_loss, value_loss, policy_entropy, _ = self.sess.run(
+                [self.summary, self.pg_loss, self.vf_loss, self.entropy, self.apply_backprop], td_map)
+
+        if self.writer is not None:
+            self.writer.add_summary(summary, update * (self.n_batch + 1))
+
         return policy_loss, value_loss, policy_entropy
 
     def learn(self, total_timesteps, callback=None, seed=None, log_interval=100):
@@ -168,7 +203,7 @@ def learn(self, total_timesteps, callback=None, seed=None, log_interval=100):
             t_start = time.time()
             for update in range(1, total_timesteps // self.n_batch + 1):
                 obs, states, rewards, masks, actions, values = runner.run()
-                _, value_loss, policy_entropy = self._train_step(obs, states, rewards, masks, actions, values)
+                _, value_loss, policy_entropy = self._train_step(obs, states, rewards, masks, actions, values, update)
                 n_seconds = time.time() - t_start
                 fps = int((update * self.n_batch) / n_seconds)