[rllib] Full checkpoint/restore for all algorithms (#875)

* wip

* working for all but dqn

* update

* add train

* rename

* update

* Update test

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()

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

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]

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]

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()
 

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]

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)

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()))