From 9bcaaaeaf5bd493cc2b9a0f4d9874d2a471faa6d Mon Sep 17 00:00:00 2001 From: Philipp Moritz Date: Mon, 19 Jun 2017 00:58:41 +0000 Subject: [PATCH] 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 | 25 +++++++++++++++++--- examples/policy_gradient/reinforce/agent.py | 11 ++++++++- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/examples/policy_gradient/examples/example.py b/examples/policy_gradient/examples/example.py index b262cbc0d..0e65ec887 100644 --- a/examples/policy_gradient/examples/example.py +++ b/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) diff --git a/examples/policy_gradient/reinforce/agent.py b/examples/policy_gradient/reinforce/agent.py index b2037edc9..380519742 100644 --- a/examples/policy_gradient/reinforce/agent.py +++ b/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)