Debugging for policy gradients (#681)

* configuration option for tensorflow debugger

* add model checkpointing

* fix linting

* make it possible to run without checkpointing

* fix

* loading from checkpoint and expose debugger through cli

* todo for filters

* Fix typo.

examples/policy_gradient/examples/example.py

@@ -38,7 +38,8 @@ config = {"kl_coeff": 0.2,
           "num_agents": 5,
           "tensorboard_log_dir": "/tmp/ray",
           "full_trace_nth_sgd_batch": -1,
-          "full_trace_data_load": False}
+          "full_trace_data_load": False,
+          "model_checkpoint_file": "/tmp/iteration-%s.ckpt"}
 
 
 if __name__ == "__main__":
@@ -48,8 +49,13 @@ if __name__ == "__main__":
                       help="The gym environment to use.")
   parser.add_argument("--redis-address", default=None, type=str,
                       help="The Redis address of the cluster.")
+  parser.add_argument("--use-tf-debugger", default=False, type=bool,
+                      help="Run the script inside of tf-dbg.")
+  parser.add_argument("--load-checkpoint", default=None, type=str,
+                      help="Continue training from a checkpoint.")
 
   args = parser.parse_args()
+  config["use_tf_debugger"] = args.use_tf_debugger
 
   ray.init(redis_address=args.redis_address)
 
@@ -79,10 +85,21 @@ if __name__ == "__main__":
           config["tensorboard_log_dir"], mdp_name,
           str(datetime.today()).replace(" ", "_")),
       agent.sess.graph)
+
   global_step = 0
+
+  saver = tf.train.Saver(max_to_keep=None)
+  if args.load_checkpoint:
+    saver.restore(agent.sess, args.load_checkpoint)
+
   for j in range(config["max_iterations"]):
     iter_start = time.time()
-    print("== iteration", j)
+    print("\n== iteration", j)
+    if config["model_checkpoint_file"]:
+      checkpoint_path = saver.save(
+          agent.sess, config["model_checkpoint_file"] % j)
+      print("Checkpoint saved in file: %s" % checkpoint_path)
+    checkpointing_end = time.time()
     weights = ray.put(agent.get_weights())
     [a.load_weights.remote(weights) for a in agents]
     trajectory, total_reward, traj_len_mean = collect_samples(
@@ -113,7 +130,8 @@ if __name__ == "__main__":
     tuples_per_device = agent.load_data(
         trajectory, j == 0 and config["full_trace_data_load"])
     load_end = time.time()
-    rollouts_time = rollouts_end - iter_start
+    checkpointing_time = checkpointing_end - iter_start
+    rollouts_time = rollouts_end - checkpointing_end
     shuffle_time = shuffle_end - rollouts_end
     load_time = load_end - shuffle_end
     sgd_time = 0
@@ -168,6 +186,7 @@ if __name__ == "__main__":
       kl_coeff *= 0.5
     print("kl div:", kl)
     print("kl coeff:", kl_coeff)
+    print("checkpointing time:", checkpointing_time)
     print("rollouts time:", rollouts_time)
     print("shuffle time:", shuffle_time)
     print("load time:", load_time)

examples/policy_gradient/reinforce/agent.py

@@ -9,6 +9,7 @@ import tensorflow as tf
 import os
 
 from tensorflow.python.client import timeline
+from tensorflow.python import debug as tf_debug
 
 import ray
 
@@ -19,8 +20,13 @@ from reinforce.filter import MeanStdFilter
 from reinforce.rollout import rollouts, add_advantage_values
 from reinforce.utils import make_divisible_by, average_gradients
 
+# TODO(pcm): Make sure that both observation_filter and reward_filter
+# are correctly handled, i.e. (a) the values are accumulated accross
+# workers (if necessary), (b) they are passed between all the methods
+# correctly and no default arguments are used, and (c) they are saved
+# as part of the checkpoint so training can resume properly.
 
-# Each tower is a copy of the policy graph pinned to a specific device
+# Each tower is a copy of the policy graph pinned to a specific device.
 Tower = namedtuple("Tower", ["init_op", "grads", "policy"])
 
 
@@ -58,6 +64,9 @@ class Agent(object):
       config_proto = tf.ConfigProto(**config["tf_session_args"])
     self.preprocessor = preprocessor
     self.sess = tf.Session(config=config_proto)
+    if config["use_tf_debugger"] and not is_remote:
+      self.sess = tf_debug.LocalCLIDebugWrapperSession(self.sess)
+      self.sess.add_tensor_filter("has_inf_or_nan", tf_debug.has_inf_or_nan)
 
     # Defines the training inputs.
     self.kl_coeff = tf.placeholder(name="newkl", shape=(), dtype=tf.float32)