Scope summaries by worker

Author: dennybritz

PolicyGradient/a3c/estimator_test.py

@@ -47,14 +47,15 @@ def testPredict(self):
       pred = sess.run(estimator.predictions, feed_dict)
 
       # Assertions
-      self.assertTrue(loss > 0.0)
+      self.assertTrue(loss != 0.0)
       self.assertEqual(pred["probs"].shape, (1, len(VALID_ACTIONS)))
       self.assertEqual(pred["logits"].shape, (1, len(VALID_ACTIONS)))
 
   def testGradient(self):
     env = make_env()
     sp = StateProcessor()
     estimator = PolicyEstimator(len(VALID_ACTIONS))
+    grads = [g for g, _ in estimator.grads_and_vars]
 
     with self.test_session() as sess:
       sess.run(tf.initialize_all_variables())
@@ -64,16 +65,17 @@ def testGradient(self):
       processed_state = atari_helpers.atari_make_initial_state(state)
       processed_states = np.array([processed_state])
 
-      # Run feeds
+      # Run feeds to get gradients
       feed_dict = {
         estimator.states: processed_states,
         estimator.targets: [1.0],
         estimator.actions: [1]
       }
-      loss = sess.run(estimator.train_op, feed_dict)
+      grads_ = sess.run(grads, feed_dict)
 
-      # Assertions
-      self.assertTrue(loss > 0.0)
+      # Apply calculated gradients
+      grad_feed_dict = { k: v for k, v in zip(grads, grads_) }
+      _ = sess.run(estimator.train_op, grad_feed_dict)
 
 
 class ValueEstimatorTest(tf.test.TestCase):
@@ -99,13 +101,14 @@ def testPredict(self):
       pred = sess.run(estimator.predictions, feed_dict)
 
       # Assertions
-      self.assertTrue(loss > 0.0)
+      self.assertTrue(loss != 0.0)
       self.assertEqual(pred["logits"].shape, (1,))
 
   def testGradient(self):
     env = make_env()
     sp = StateProcessor()
     estimator = ValueEstimator()
+    grads = [g for g, _ in estimator.grads_and_vars]
 
     with self.test_session() as sess:
       sess.run(tf.initialize_all_variables())
@@ -120,10 +123,11 @@ def testGradient(self):
         estimator.states: processed_states,
         estimator.targets: [1.0],
       }
-      loss = sess.run(estimator.train_op, feed_dict)
+      grads_ = sess.run(grads, feed_dict)
 
-      # Assertions
-      self.assertTrue(loss > 0.0)
+      # Apply calculated gradients
+      grad_feed_dict = { k: v for k, v in zip(grads, grads_) }
+      _ = sess.run(estimator.train_op, grad_feed_dict)
 
 if __name__ == '__main__':
   unittest.main()

PolicyGradient/a3c/estimators.py

@@ -69,44 +69,41 @@ def __init__(self, num_outputs, reuse=False, trainable=True):
     with tf.variable_scope("policy_net"):
       self.logits = tf.contrib.layers.fully_connected(fc1, num_outputs, activation_fn=None)
       self.probs = tf.nn.softmax(self.logits)
+      self.probs = tf.clip_by_value(self.probs, 1e-6, 1.0)
 
       self.predictions = {
         "logits": self.logits,
         "probs": self.probs
       }
 
-      if not trainable:
-        return
-
       # We add cross-entropy to the loss to encourage exploration
-      self.cross_entropy = -tf.reduce_sum(self.probs * tf.log(self.probs), 1)
+      self.cross_entropy = -tf.reduce_sum(self.probs * tf.log(self.probs), 1, name="cross_entropy")
+      self.cross_entropy_mean = tf.reduce_mean(self.cross_entropy, name="cross_entropy_mean")
 
       # Get the predictions for the chosen actions only
       gather_indices = tf.range(batch_size) * tf.shape(self.probs)[1] + self.actions
       self.picked_action_probs = tf.gather(tf.reshape(self.probs, [-1]), gather_indices)
 
       self.losses = - (tf.log(self.picked_action_probs) * self.targets + 0.01 * self.cross_entropy)
-      self.loss = tf.reduce_sum(self.losses)
+      self.loss = tf.reduce_sum(self.losses, name="loss")
 
-      tf.scalar_summary("policy_net/loss", self.loss)
-      tf.scalar_summary("policy_net/advantage_mean", tf.reduce_mean(self.targets))
-      tf.scalar_summary("policy_net/entropy_mean", tf.reduce_mean(self.cross_entropy))
-      tf.histogram_summary("policy_net/cross_entropy", self.cross_entropy)
-      tf.histogram_summary("policy_net/actions", self.actions)
+      tf.scalar_summary(self.loss.op.name, self.loss)
+      tf.scalar_summary(self.cross_entropy_mean.op.name, self.cross_entropy_mean)
+      tf.histogram_summary(self.cross_entropy.op.name, self.cross_entropy)
 
-      # Optimizer Parameters from original paper
-      self.optimizer = tf.train.AdamOptimizer(1e-4)
-      self.train_op = tf.contrib.layers.optimize_loss(
-        loss=self.loss,
-        global_step=tf.contrib.framework.get_global_step(),
-        learning_rate=1e-4,
-        optimizer=self.optimizer,
-        # clip_gradients=5.0,
-        summaries=tf.contrib.layers.optimizers.OPTIMIZER_SUMMARIES)
+      if trainable:
+        self.optimizer = tf.train.AdamOptimizer(1e-4)
+        self.grads_and_vars = self.optimizer.compute_gradients(self.loss)
+        self.grads_and_vars = [[grad, var] for grad, var in self.grads_and_vars if grad is not None]
+        self.train_op = self.optimizer.apply_gradients(self.grads_and_vars,
+          global_step=tf.contrib.framework.get_global_step())
 
-      # Merge summaries from this network and the shared network (but not the value net)
-      summary_ops = tf.get_collection(tf.GraphKeys.SUMMARIES)
-      self.summaries = tf.merge_summary([s for s in summary_ops if "policy_net" in s.name or "shared" in s.name])
+    # Merge summaries from this network and the shared network (but not the value net)
+    var_scope_name = tf.get_variable_scope().name
+    summary_ops = tf.get_collection(tf.GraphKeys.SUMMARIES)
+    sumaries = [s for s in summary_ops if "policy_net" in s.name or "shared" in s.name]
+    sumaries = [s for s in summary_ops if var_scope_name in s.name]
+    self.summaries = tf.merge_summary(sumaries)
 
 
 class ValueEstimator():
@@ -139,39 +136,36 @@ def __init__(self, reuse=False, trainable=True):
         inputs=fc1,
         num_outputs=1,
         activation_fn=None)
-      self.logits = tf.squeeze(self.logits, squeeze_dims=[1])
+      self.logits = tf.squeeze(self.logits, squeeze_dims=[1], name="logits")
 
       self.losses = tf.squared_difference(self.logits, self.targets)
-      self.loss = tf.reduce_sum(self.losses)
+      self.loss = tf.reduce_sum(self.losses, name="loss")
 
       self.predictions = {
         "logits": self.logits
       }
 
-      if not trainable:
-        return
-
-      # Optimizer Parameters from original paper
-      self.optimizer = tf.train.AdamOptimizer(1e-4)
-      self.train_op = tf.contrib.layers.optimize_loss(
-        loss=self.loss,
-        global_step=tf.contrib.framework.get_global_step(),
-        learning_rate=1e-4,
-        optimizer=self.optimizer,
-        # clip_gradients=5.0,
-        summaries=tf.contrib.layers.optimizers.OPTIMIZER_SUMMARIES)
-
       # Summaries
-      tf.scalar_summary("value_net/loss", self.loss)
-      tf.scalar_summary("value_net/max_value", tf.reduce_max(self.logits))
-      tf.scalar_summary("value_net/min_value", tf.reduce_min(self.logits))
-      tf.scalar_summary("value_net/mean_value", tf.reduce_mean(self.logits))
-      tf.scalar_summary("value_net/reward_max", tf.reduce_max(self.targets))
-      tf.scalar_summary("value_net/reward_min", tf.reduce_min(self.targets))
-      tf.scalar_summary("value_net/reward_mean", tf.reduce_mean(self.targets))
-      tf.histogram_summary("value_net/reward_targets", self.targets)
-      tf.histogram_summary("value_net/values", self.logits)
-
-      # Merge summaries from this network and the shared network (but not the policy net)
-      summary_ops = tf.get_collection(tf.GraphKeys.SUMMARIES)
-      self.summaries = tf.merge_summary([s for s in summary_ops if "value_net" in s.name or "shared" in s.name])
+      prefix = tf.get_variable_scope().name
+      tf.scalar_summary(self.loss.name, self.loss)
+      tf.scalar_summary("{}/max_value".format(prefix), tf.reduce_max(self.logits))
+      tf.scalar_summary("{}/min_value".format(prefix), tf.reduce_min(self.logits))
+      tf.scalar_summary("{}/mean_value".format(prefix), tf.reduce_mean(self.logits))
+      tf.scalar_summary("{}/reward_max".format(prefix), tf.reduce_max(self.targets))
+      tf.scalar_summary("{}/reward_min".format(prefix), tf.reduce_min(self.targets))
+      tf.scalar_summary("{}/reward_mean".format(prefix), tf.reduce_mean(self.targets))
+      tf.histogram_summary("{}/reward_targets".format(prefix), self.targets)
+      tf.histogram_summary("{}/values".format(prefix), self.logits)
+
+      if trainable:
+        self.optimizer = tf.train.AdamOptimizer(1e-4)
+        self.grads_and_vars = self.optimizer.compute_gradients(self.loss)
+        self.grads_and_vars = [[grad, var] for grad, var in self.grads_and_vars if grad is not None]
+        self.train_op = self.optimizer.apply_gradients(self.grads_and_vars,
+          global_step=tf.contrib.framework.get_global_step())
+
+    var_scope_name = tf.get_variable_scope().name
+    summary_ops = tf.get_collection(tf.GraphKeys.SUMMARIES)
+    sumaries = [s for s in summary_ops if "policy_net" in s.name or "shared" in s.name]
+    sumaries = [s for s in summary_ops if var_scope_name in s.name]
+    self.summaries = tf.merge_summary(sumaries)

PolicyGradient/a3c/worker.py

@@ -165,20 +165,30 @@ def update(self, transitions, sess):
       policy_targets.append(policy_target)
       value_targets.append(reward)
 
+    # Calculate the gradients
     feed_dict = {
-      self.global_policy_net.states: np.array(states),
-      self.global_policy_net.targets: policy_targets,
-      self.global_policy_net.actions: actions,
-      self.global_value_net.states: np.array(states),
-      self.global_value_net.targets: value_targets,
+      self.policy_net.states: np.array(states),
+      self.policy_net.targets: policy_targets,
+      self.policy_net.actions: actions,
+      self.value_net.states: np.array(states),
+      self.value_net.targets: value_targets,
     }
 
-    # Apply policy net update
-    global_step, pnet_loss, pnet_summaries, vnet_loss, vnet_summaries = sess.run(
-      [self.global_step, self.global_policy_net.train_op,
-      self.global_policy_net.summaries, self.global_value_net.train_op,
-      self.global_value_net.summaries],
-      feed_dict)
+    # Calculate the local gradients
+    pnet_loss, vnet_loss, pnet_grads, vnet_grads, pnet_summaries, vnet_summaries = sess.run([
+      self.policy_net.loss,
+      self.value_net.loss,
+      [g for g, _ in self.policy_net.grads_and_vars],
+      [g for g, _ in self.value_net.grads_and_vars],
+      self.policy_net.summaries,
+      self.value_net.summaries
+    ], feed_dict)
+
+    # Apply the gradients to the global nets
+    pnet_grad_ops = [g for g, _ in self.global_policy_net.grads_and_vars]
+    vnet_grad_ops = [g for g, _ in self.global_value_net.grads_and_vars]
+    grad_feed_dict = { k: v for k, v in zip(pnet_grad_ops + vnet_grad_ops, pnet_grads + vnet_grads)}
+    global_step, _, _, = sess.run([self.global_step, self.global_policy_net.train_op, self.global_value_net.train_op], grad_feed_dict)
 
     # Write summaries
     if self.summary_writer is not None:

PolicyGradient/a3c/worker_test.py

@@ -86,7 +86,7 @@ def testRunNStepsAndUpdate(self):
       state = self.sp.process(self.env.reset())
       processed_state = atari_helpers.atari_make_initial_state(state)
       w.state = processed_state
-      transitions = w.run_n_steps(10, sess)
+      transitions, local_t, global_t = w.run_n_steps(10, sess)
       policy_net_loss, value_net_loss, policy_net_summaries, value_net_summaries = w.update(transitions, sess)
 
     self.assertEqual(len(transitions), 10)