diff --git a/python/ray/rllib/a3c/a3c.py b/python/ray/rllib/a3c/a3c.py index 28acc320a..c83245ca2 100644 --- a/python/ray/rllib/a3c/a3c.py +++ b/python/ray/rllib/a3c/a3c.py @@ -3,6 +3,7 @@ from __future__ import division from __future__ import print_function import numpy as np +import pickle import tensorflow as tf import six.moves.queue as queue import os @@ -134,11 +135,24 @@ class A3C(Algorithm): avg_length = np.mean(episode_lengths) if episode_lengths else None res = TrainingResult( self.experiment_id.hex, self.iteration, - avg_reward, avg_length, None, dict()) + avg_reward, avg_length, dict()) return res + def save(self): + checkpoint_path = os.path.join( + self.logdir, "checkpoint-{}".format(self.iteration)) + objects = [ + self.parameters, + self.iteration] + pickle.dump(objects, open(checkpoint_path, "wb")) + return checkpoint_path + def restore(self, checkpoint_path): - raise NotImplementedError # TODO(ekl) + objects = pickle.load(open(checkpoint_path, "rb")) + self.parameters = objects[0] + self.policy.set_weights(self.parameters) + self.iteration = objects[1] def compute_action(self, observation): - raise NotImplementedError # TODO(ekl) + actions = self.policy.compute_actions(observation)[0] + return actions.argmax() diff --git a/python/ray/rllib/common.py b/python/ray/rllib/common.py index 5410edcbf..88edb8c0a 100644 --- a/python/ray/rllib/common.py +++ b/python/ray/rllib/common.py @@ -50,7 +50,6 @@ TrainingResult = namedtuple("TrainingResult", [ "training_iteration", "episode_reward_mean", "episode_len_mean", - "latest_checkpoint", "info" ]) @@ -90,8 +89,7 @@ class Algorithm(object): self.__class__.__name__, datetime.today().strftime("%Y-%m-%d_%H-%M-%S")) if upload_dir.startswith("file"): - self.logdir = "file://" + tempfile.mkdtemp( - prefix=prefix, dir="/tmp/ray") + self.logdir = tempfile.mkdtemp(prefix=prefix, dir="/tmp/ray") else: self.logdir = os.path.join(upload_dir, prefix) log_path = os.path.join(self.logdir, "config.json") @@ -110,11 +108,19 @@ class Algorithm(object): raise NotImplementedError - def restore(self, checkpoint_path): - """Restores training state from a given checkpoint. + def save(self): + """Saves the current model state to a checkpoint. - These checkpoints are returned from calls to train() in the - checkpoint_path field of TrainingResult. + Returns: + Checkpoint path that may be passed to restore(). + """ + + raise NotImplementedError + + def restore(self, checkpoint_path): + """Restores training state from a given model checkpoint. + + These checkpoints are returned from calls to save(). """ raise NotImplementedError diff --git a/python/ray/rllib/dqn/dqn.py b/python/ray/rllib/dqn/dqn.py index b0dd7f641..c6e43598a 100644 --- a/python/ray/rllib/dqn/dqn.py +++ b/python/ray/rllib/dqn/dqn.py @@ -6,6 +6,8 @@ import time import gym import numpy as np +import pickle +import os import tensorflow as tf from ray.rllib.common import Algorithm, TrainingResult @@ -103,10 +105,17 @@ DEFAULT_CONFIG = dict( class DQN(Algorithm): def __init__(self, env_name, config, upload_dir=None): config.update({"alg": "DQN"}) + Algorithm.__init__(self, env_name, config, upload_dir=upload_dir) - env = gym.make(env_name) + + with tf.Graph().as_default(): + self._init() + + def _init(self): + config = self.config + env = gym.make(self.env_name) # TODO(ekl): replace this with RLlib preprocessors - if "NoFrameskip" in env_name: + if "NoFrameskip" in self.env_name: env = ScaledFloatFrame(wrap_dqn(env)) self.env = env @@ -153,6 +162,7 @@ class DQN(Algorithm): self.num_timesteps = 0 self.num_iterations = 0 self.file_writer = tf.summary.FileWriter(self.logdir, self.sess.graph) + self.saver = tf.train.Saver(max_to_keep=None) def train(self): config = self.config @@ -235,12 +245,41 @@ class DQN(Algorithm): res = TrainingResult( self.experiment_id.hex, self.num_iterations, mean_100ep_reward, - mean_100ep_length, None, info) + mean_100ep_length, info) self.num_iterations += 1 return res + def save(self): + checkpoint_path = self.saver.save( + self.sess, + os.path.join(self.logdir, "checkpoint"), + global_step=self.num_iterations) + extra_data = [ + self.replay_buffer, + self.beta_schedule, + self.exploration, + self.episode_rewards, + self.episode_lengths, + self.saved_mean_reward, + self.obs, + self.num_timesteps, + self.num_iterations] + pickle.dump(extra_data, open(checkpoint_path + ".extra_data", "wb")) + return checkpoint_path + def restore(self, checkpoint_path): - raise NotImplementedError # TODO(ekl) + self.saver.restore(self.sess, checkpoint_path) + extra_data = pickle.load(open(checkpoint_path + ".extra_data", "rb")) + self.replay_buffer = extra_data[0] + self.beta_schedule = extra_data[1] + self.exploration = extra_data[2] + self.episode_rewards = extra_data[3] + self.episode_lengths = extra_data[4] + self.saved_mean_reward = extra_data[5] + self.obs = extra_data[6] + self.num_timesteps = extra_data[7] + self.num_iterations = extra_data[8] def compute_action(self, observation): - raise NotImplementedError # TODO(ekl) + return self.dqn_graph.act( + self.sess, np.array(observation)[None], 0.0)[0] diff --git a/python/ray/rllib/evolution_strategies/evolution_strategies.py b/python/ray/rllib/evolution_strategies/evolution_strategies.py index 0667df8d6..ea14c29ba 100644 --- a/python/ray/rllib/evolution_strategies/evolution_strategies.py +++ b/python/ray/rllib/evolution_strategies/evolution_strategies.py @@ -9,8 +9,11 @@ from collections import namedtuple import gym import numpy as np import os +import pickle import time +import tensorflow as tf + import ray from ray.rllib.common import Algorithm, TrainingResult from ray.rllib.models import ModelCatalog @@ -163,22 +166,27 @@ class EvolutionStrategies(Algorithm): Algorithm.__init__(self, env_name, config, upload_dir=upload_dir) + with tf.Graph().as_default(): + self._init() + + def _init(self): + policy_params = { "ac_noise_std": 0.01 } - env = gym.make(env_name) + env = gym.make(self.env_name) preprocessor = ModelCatalog.get_preprocessor( - env_name, env.observation_space.shape) + self.env_name, env.observation_space.shape) preprocessor_shape = preprocessor.transform_shape( env.observation_space.shape) - utils.make_session(single_threaded=False) + self.sess = utils.make_session(single_threaded=False) self.policy = policies.GenericPolicy( env.observation_space, env.action_space, preprocessor, **policy_params) tf_util.initialize() - self.optimizer = optimizers.Adam(self.policy, config["stepsize"]) + self.optimizer = optimizers.Adam(self.policy, self.config["stepsize"]) self.ob_stat = utils.RunningStat(preprocessor_shape, eps=1e-2) # Create the shared noise table. @@ -189,8 +197,8 @@ class EvolutionStrategies(Algorithm): # Create the actors. print("Creating actors.") self.workers = [ - Worker.remote(config, policy_params, env_name, noise_id) - for _ in range(config["num_workers"])] + Worker.remote(self.config, policy_params, self.env_name, noise_id) + for _ in range(self.config["num_workers"])] self.episodes_so_far = 0 self.timesteps_so_far = 0 @@ -327,14 +335,32 @@ class EvolutionStrategies(Algorithm): "time_elapsed": step_tend - self.tstart } res = TrainingResult(self.experiment_id.hex, self.iteration, - returns_n2.mean(), lengths_n2.mean(), None, info) + returns_n2.mean(), lengths_n2.mean(), info) self.iteration += 1 return res + def save(self): + checkpoint_path = os.path.join( + self.logdir, "checkpoint-{}".format(self.iteration)) + weights = self.policy.get_trainable_flat() + objects = [ + weights, + self.ob_stat, + self.episodes_so_far, + self.timesteps_so_far, + self.iteration] + pickle.dump(objects, open(checkpoint_path, "wb")) + return checkpoint_path + def restore(self, checkpoint_path): - raise NotImplementedError # TODO(ekl) + objects = pickle.load(open(checkpoint_path, "rb")) + self.policy.set_trainable_flat(objects[0]) + self.ob_stat = objects[1] + self.episodes_so_far = objects[2] + self.timesteps_so_far = objects[3] + self.iteration = objects[4] def compute_action(self, observation): - raise NotImplementedError # TODO(ekl) + return self.policy.act([observation])[0] diff --git a/python/ray/rllib/policy_gradient/agent.py b/python/ray/rllib/policy_gradient/agent.py index 18b0f38db..456568f64 100644 --- a/python/ray/rllib/policy_gradient/agent.py +++ b/python/ray/rllib/policy_gradient/agent.py @@ -3,6 +3,7 @@ from __future__ import division from __future__ import print_function import gym.spaces +import pickle import tensorflow as tf import os @@ -178,6 +179,14 @@ class Agent(object): extra_feed_dict={self.kl_coeff: kl_coeff}, file_writer=file_writer if full_trace else None) + def save(self): + return pickle.dumps([self.observation_filter, self.reward_filter]) + + def restore(self, objs): + objs = pickle.loads(objs) + self.observation_filter = objs[0] + self.reward_filter = objs[1] + def get_weights(self): return self.variables.get_weights() diff --git a/python/ray/rllib/policy_gradient/policy_gradient.py b/python/ray/rllib/policy_gradient/policy_gradient.py index 46d721744..0adb4f224 100644 --- a/python/ray/rllib/policy_gradient/policy_gradient.py +++ b/python/ray/rllib/policy_gradient/policy_gradient.py @@ -6,6 +6,7 @@ import os import time import numpy as np +import pickle import tensorflow as tf from tensorflow.python import debug as tf_debug @@ -69,10 +70,9 @@ DEFAULT_CONFIG = { # If this is True, the TensorFlow debugger is invoked if an Inf or NaN # is detected "tf_debug_inf_or_nan": False, - # If True, we write checkpoints and tensorflow logging - "write_logs": True, - # Name of the model checkpoint file - "model_checkpoint_file": "checkpoint"} + # If True, we write tensorflow logs and checkpoints + "write_logs": True +} class PolicyGradient(Algorithm): @@ -81,18 +81,20 @@ class PolicyGradient(Algorithm): Algorithm.__init__(self, env_name, config, upload_dir=upload_dir) + with tf.Graph().as_default(): + self._init() + + def _init(self): self.global_step = 0 self.j = 0 - self.kl_coeff = config["kl_coeff"] + self.kl_coeff = self.config["kl_coeff"] self.model = Agent(self.env_name, 1, self.config, self.logdir, False) self.agents = [ RemoteAgent.remote( self.env_name, 1, self.config, self.logdir, True) - for _ in range(config["num_agents"])] + for _ in range(self.config["num_agents"])] self.start_time = time.time() - # TF does not support to write logs to S3 at the moment - write_tf_logs = config["write_logs"] and self.logdir.startswith("file") - if write_tf_logs: + if self.config["write_logs"]: self.file_writer = tf.summary.FileWriter( self.logdir, self.model.sess.graph) else: @@ -213,21 +215,9 @@ class PolicyGradient(Algorithm): elif kl < 0.5 * config["kl_target"]: self.kl_coeff *= 0.5 - checkpointing_start = time.time() - if config["model_checkpoint_file"]: - checkpoint_path = self.saver.save( - model.sess, - os.path.join( - self.logdir, config["model_checkpoint_file"]), - global_step=j) - else: - checkpoint_path = None - checkpointing_time = time.time() - checkpointing_start - info = { "kl_divergence": kl, "kl_coefficient": self.kl_coeff, - "checkpointing_time": checkpointing_time, "rollouts_time": rollouts_time, "shuffle_time": shuffle_time, "load_time": load_time, @@ -237,7 +227,6 @@ class PolicyGradient(Algorithm): print("kl div:", kl) print("kl coeff:", self.kl_coeff) - print("checkpointing time:", checkpointing_time) print("rollouts time:", rollouts_time) print("shuffle time:", shuffle_time) print("load time:", load_time) @@ -246,14 +235,35 @@ class PolicyGradient(Algorithm): print("total time so far:", time.time() - self.start_time) result = TrainingResult( - self.experiment_id.hex, j, total_reward, traj_len_mean, - checkpoint_path, info) + self.experiment_id.hex, j, total_reward, traj_len_mean, info) return result + def save(self): + checkpoint_path = self.saver.save( + self.model.sess, + os.path.join(self.logdir, "checkpoint"), + global_step=self.j) + agent_state = ray.get([a.save.remote() for a in self.agents]) + extra_data = [ + self.model.save(), + self.global_step, + self.j, + self.kl_coeff, + agent_state] + pickle.dump(extra_data, open(checkpoint_path + ".extra_data", "wb")) + return checkpoint_path + def restore(self, checkpoint_path): self.saver.restore(self.model.sess, checkpoint_path) + extra_data = pickle.load(open(checkpoint_path + ".extra_data", "rb")) + self.model.restore(extra_data[0]) + self.global_step = extra_data[1] + self.j = extra_data[2] + self.kl_coeff = extra_data[3] + ray.get([ + a.restore.remote(o) + for (a, o) in zip(self.agents, extra_data[4])]) def compute_action(self, observation): - return self.model.common_policy.compute_actions( - observation[np.newaxis, :])[0][0] + return self.model.common_policy.compute([observation])[0][0] diff --git a/python/ray/rllib/test.sh b/python/ray/rllib/test/test.sh similarity index 100% rename from python/ray/rllib/test.sh rename to python/ray/rllib/test/test.sh diff --git a/python/ray/rllib/test/test_checkpoint_restore.py b/python/ray/rllib/test/test_checkpoint_restore.py new file mode 100755 index 000000000..149b8aab0 --- /dev/null +++ b/python/ray/rllib/test/test_checkpoint_restore.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import ray +import random + +from ray.rllib.dqn import DQN, DEFAULT_CONFIG as DQN_CONFIG +from ray.rllib.evolution_strategies import ( + EvolutionStrategies, DEFAULT_CONFIG as ES_CONFIG) +from ray.rllib.policy_gradient import ( + PolicyGradient, DEFAULT_CONFIG as PG_CONFIG) +from ray.rllib.a3c import A3C, DEFAULT_CONFIG as A3C_CONFIG + +ray.init() +for (cls, default_config) in [ + (DQN, DQN_CONFIG), + # TODO(ekl) this fails with multiple ES instances in a process + (EvolutionStrategies, ES_CONFIG), + (PolicyGradient, PG_CONFIG), + (A3C, A3C_CONFIG)]: + config = default_config.copy() + config["num_sgd_iter"] = 5 + config["episodes_per_batch"] = 100 + config["timesteps_per_batch"] = 1000 + alg1 = cls('CartPole-v0', config) + alg2 = cls('CartPole-v0', config) + + for _ in range(3): + res = alg1.train() + print("current status: " + str(res)) + + # Sync the models + alg2.restore(alg1.save()) + + for _ in range(10): + obs = [ + random.random(), random.random(), random.random(), random.random()] + a1 = alg1.compute_action(obs) + a2 = alg2.compute_action(obs) + print("Checking computed actions", obs, a1, a2) + + # TODO(ekl) this fails for stochastic policies + assert(a1 == a2) diff --git a/python/ray/rllib/train.py b/python/ray/rllib/train.py index ff092f6cc..156cf4d7a 100755 --- a/python/ray/rllib/train.py +++ b/python/ray/rllib/train.py @@ -29,6 +29,10 @@ parser.add_argument("--config", default="{}", type=str, help="The configuration options of the algorithm.") parser.add_argument("--upload-dir", default="file:///tmp/ray", type=str, help="Where the traces are stored.") +parser.add_argument("--checkpoint-freq", default=sys.maxsize, type=int, + help="How many iterations between checkpoints.") +parser.add_argument("--restore", default="", type=str, + help="If specified, restores state from this checkpoint.") if __name__ == "__main__": @@ -66,6 +70,9 @@ if __name__ == "__main__": result_logger = ray.rllib.common.RLLibLogger( os.path.join(alg.logdir, "result.json")) + if args.restore: + alg.restore(args.restore) + for i in range(args.num_iterations): result = alg.train() @@ -76,3 +83,6 @@ if __name__ == "__main__": result_logger.write("\n") print("current status: {}".format(result)) + + if (i + 1) % args.checkpoint_freq == 0: + print("checkpoint path: {}".format(alg.save()))