[rllib] Move policy gradient and evolution strategies algorithms from examples/ to ray/rllib/ (#694)

* rllib v0

* fix imports

* lint

* comments

* update docs

doc/source/example-evolution-strategies.rst

@@ -11,14 +11,14 @@ To run the application, first install some dependencies.
 
 You can view the `code for this example`_.
 
-.. _`code for this example`: https://github.com/ray-project/ray/tree/master/examples/evolution_strategies
+.. _`code for this example`: https://github.com/ray-project/ray/tree/master/python/ray/rllib/evolution_strategies
 
 The script can be run as follows. Note that the configuration is tuned to work
 on the ``Humanoid-v1`` gym environment.
 
 .. code-block:: bash
 
-  python examples/evolution_strategies/evolution_strategies.py
+  python/ray/rllib/evolution_strategies/example.py
 
 At the heart of this example, we define a ``Worker`` class. These workers have
 a method ``do_rollouts``, which will be used to perform simulate randomly

doc/source/example-policy-gradient.rst

@@ -23,7 +23,7 @@ Then you can run the example as follows.
 
 .. code-block:: bash
 
-  python ray/examples/policy_gradient/examples/example.py --environment=Pong-ram-v3
+  python/ray/rllib/policy_gradient/example.py --environment=Pong-ram-v3
 
 This will train an agent on the ``Pong-ram-v3`` Atari environment. You can also
 try passing in the ``Pong-v0`` environment or the ``CartPole-v0`` environment.
@@ -41,4 +41,4 @@ Many of the TensorBoard metrics are also printed to the console, but you might
 find it easier to visualize and compare between runs using the TensorBoard UI.
 
 .. _`TensorFlow with GPU support`: https://www.tensorflow.org/install/
-.. _`code for this example`: https://github.com/ray-project/ray/tree/master/examples/policy_gradient
+.. _`code for this example`: https://github.com/ray-project/ray/tree/master/python/ray/rllib/policy_gradient

examples/policy_gradient/examples/example.py

@@ -1,194 +0,0 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from datetime import datetime
-
-import argparse
-import time
-
-import ray
-import numpy as np
-import tensorflow as tf
-
-from reinforce.env import (NoPreprocessor, AtariRamPreprocessor,
-                           AtariPixelPreprocessor)
-from reinforce.agent import Agent, RemoteAgent
-from reinforce.rollout import collect_samples
-from reinforce.utils import shuffle
-
-
-config = {"kl_coeff": 0.2,
-          "num_sgd_iter": 30,
-          "max_iterations": 1000,
-          "sgd_stepsize": 5e-5,
-          # TODO(pcm): Expose the choice between gpus and cpus
-          # as a command line argument.
-          "devices": ["/cpu:%d" % i for i in range(4)],
-          "tf_session_args": {
-              "device_count": {"CPU": 4},
-              "log_device_placement": False,
-              "allow_soft_placement": True,
-          },
-          "sgd_batchsize": 128,  # total size across all devices
-          "entropy_coeff": 0.0,
-          "clip_param": 0.3,
-          "kl_target": 0.01,
-          "timesteps_per_batch": 40000,
-          "num_agents": 5,
-          "tensorboard_log_dir": "/tmp/ray",
-          "full_trace_nth_sgd_batch": -1,
-          "full_trace_data_load": False,
-          "model_checkpoint_file": "/tmp/iteration-%s.ckpt"}
-
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Run the policy gradient "
-                                               "algorithm.")
-  parser.add_argument("--environment", default="Pong-v0", type=str,
-                      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)
-
-  mdp_name = args.environment
-  if args.environment == "Pong-v0":
-    preprocessor = AtariPixelPreprocessor()
-  elif mdp_name == "Pong-ram-v3":
-    preprocessor = AtariRamPreprocessor()
-  elif mdp_name == "CartPole-v0":
-    preprocessor = NoPreprocessor()
-  elif mdp_name == "Walker2d-v1":
-    preprocessor = NoPreprocessor()
-  else:
-    print("No environment was chosen, so defaulting to Pong-v0.")
-    mdp_name = "Pong-v0"
-    preprocessor = AtariPixelPreprocessor()
-
-  print("Using the environment {}.".format(mdp_name))
-  agents = [RemoteAgent.remote(mdp_name, 1, preprocessor, config, True)
-            for _ in range(config["num_agents"])]
-  agent = Agent(mdp_name, 1, preprocessor, config, False)
-
-  kl_coeff = config["kl_coeff"]
-
-  file_writer = tf.summary.FileWriter(
-      "{}/trpo_{}_{}".format(
-          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("\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(
-        agents, config["timesteps_per_batch"], 0.995, 1.0, 2000)
-    print("total reward is ", total_reward)
-    print("trajectory length mean is ", traj_len_mean)
-    print("timesteps:", trajectory["dones"].shape[0])
-    traj_stats = tf.Summary(value=[
-        tf.Summary.Value(
-            tag="policy_gradient/rollouts/mean_reward",
-            simple_value=total_reward),
-        tf.Summary.Value(
-            tag="policy_gradient/rollouts/traj_len_mean",
-            simple_value=traj_len_mean)])
-    file_writer.add_summary(traj_stats, global_step)
-    global_step += 1
-    trajectory["advantages"] = ((trajectory["advantages"] -
-                                 trajectory["advantages"].mean()) /
-                                trajectory["advantages"].std())
-    rollouts_end = time.time()
-    print("Computing policy (iterations=" + str(config["num_sgd_iter"]) +
-          ", stepsize=" + str(config["sgd_stepsize"]) + "):")
-    names = ["iter", "loss", "kl", "entropy"]
-    print(("{:>15}" * len(names)).format(*names))
-    num_devices = len(config["devices"])
-    trajectory = shuffle(trajectory)
-    shuffle_end = time.time()
-    tuples_per_device = agent.load_data(
-        trajectory, j == 0 and config["full_trace_data_load"])
-    load_end = time.time()
-    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
-    for i in range(config["num_sgd_iter"]):
-      sgd_start = time.time()
-      batch_index = 0
-      num_batches = int(tuples_per_device) // int(agent.per_device_batch_size)
-      loss, kl, entropy = [], [], []
-      permutation = np.random.permutation(num_batches)
-      while batch_index < num_batches:
-        full_trace = (
-            i == 0 and j == 0 and
-            batch_index == config["full_trace_nth_sgd_batch"])
-        batch_loss, batch_kl, batch_entropy = agent.run_sgd_minibatch(
-            permutation[batch_index] * agent.per_device_batch_size,
-            kl_coeff, full_trace, file_writer)
-        loss.append(batch_loss)
-        kl.append(batch_kl)
-        entropy.append(batch_entropy)
-        batch_index += 1
-      loss = np.mean(loss)
-      kl = np.mean(kl)
-      entropy = np.mean(entropy)
-      sgd_end = time.time()
-      print("{:>15}{:15.5e}{:15.5e}{:15.5e}".format(i, loss, kl, entropy))
-
-      values = []
-      if i == config["num_sgd_iter"] - 1:
-        metric_prefix = "policy_gradient/sgd/final_iter/"
-        values.append(tf.Summary.Value(
-            tag=metric_prefix + "kl_coeff",
-            simple_value=kl_coeff))
-      else:
-        metric_prefix = "policy_gradient/sgd/intermediate_iters/"
-      values.extend([
-          tf.Summary.Value(
-              tag=metric_prefix + "mean_entropy",
-              simple_value=entropy),
-          tf.Summary.Value(
-              tag=metric_prefix + "mean_loss",
-              simple_value=loss),
-          tf.Summary.Value(
-              tag=metric_prefix + "mean_kl",
-              simple_value=kl)])
-      sgd_stats = tf.Summary(value=values)
-      file_writer.add_summary(sgd_stats, global_step)
-      global_step += 1
-      sgd_time += sgd_end - sgd_start
-    if kl > 2.0 * config["kl_target"]:
-      kl_coeff *= 1.5
-    elif kl < 0.5 * config["kl_target"]:
-      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)
-    print("sgd time:", sgd_time)
-    print("sgd examples/s:", len(trajectory["observations"]) / sgd_time)

examples/policy_gradient/setup.py

@@ -1,9 +0,0 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from setuptools import setup, find_packages
-
-setup(name="reinforce",
-      version="0.0.1",
-      packages=find_packages())

python/ray/rllib/common.py

@@ -0,0 +1,32 @@
+from collections import namedtuple
+
+
+TrainingResult = namedtuple("TrainingResult", [
+    "training_iteration",
+    "episode_reward_mean",
+    "episode_len_mean",
+])
+
+
+class Algorithm(object):
+  """All RLlib algorithms extend this base class.
+
+  Algorithm objects retain internal model state between calls to train(), so
+  you should create a new algorithm instance for each training session.
+
+  TODO(ekl): support checkpoint / restore of training state.
+  """
+
+  def __init__(self, env_name, config):
+    self.env_name = env_name
+    self.config = config
+
+  def train(self):
+    """
+    Runs one logical iteration of training.
+
+    Returns:
+      A TrainingResult that describes training progress.
+    """
+
+    raise NotImplementedError

python/ray/rllib/evolution_strategies/__init__.py

@@ -0,0 +1,4 @@
+from ray.rllib.evolution_strategies.evolution_strategies import (
+    EvolutionStrategies, DEFAULT_CONFIG)
+
+__all__ = ["EvolutionStrategies", "DEFAULT_CONFIG"]

python/ray/rllib/evolution_strategies/evolution_strategies.py

@@ -5,25 +5,26 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import argparse
 from collections import namedtuple
 import gym
 import numpy as np
 import os
-import ray
 import time
 
-import optimizers
-import policies
-import tabular_logger as tlogger
-import tf_util
-import utils
+import ray
+from ray.rllib.common import Algorithm, TrainingResult
+
+from ray.rllib.evolution_strategies import optimizers
+from ray.rllib.evolution_strategies import policies
+from ray.rllib.evolution_strategies import tabular_logger as tlogger
+from ray.rllib.evolution_strategies import tf_util
+from ray.rllib.evolution_strategies import utils
 
 
 Config = namedtuple("Config", [
     "l2coeff", "noise_stdev", "episodes_per_batch", "timesteps_per_batch",
     "calc_obstat_prob", "eval_prob", "snapshot_freq", "return_proc_mode",
-    "episode_cutoff_mode"
+    "episode_cutoff_mode", "num_workers", "stepsize"
 ])
 
 Result = namedtuple("Result", [
@@ -32,6 +33,20 @@ Result = namedtuple("Result", [
 ])
 
 
+DEFAULT_CONFIG = Config(
+    l2coeff=0.005,
+    noise_stdev=0.02,
+    episodes_per_batch=10000,
+    timesteps_per_batch=100000,
+    calc_obstat_prob=0.01,
+    eval_prob=0,
+    snapshot_freq=0,
+    return_proc_mode="centered_rank",
+    episode_cutoff_mode="env_default",
+    num_workers=10,
+    stepsize=.01)
+
+
 @ray.remote
 def create_shared_noise():
   """Create a large array of noise to be shared by all workers."""
@@ -135,72 +150,46 @@ class Worker(object):
           ob_count=task_ob_stat.count)
 
 
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Train an RL agent on Pong.")
-  parser.add_argument("--num-workers", default=10, type=int,
-                      help=("The number of actors to create in aggregate "
-                            "across the cluster."))
-  parser.add_argument("--env-name", default="Pendulum-v0", type=str,
-                      help="The name of the gym environment to use.")
-  parser.add_argument("--stepsize", default=0.01, type=float,
-                      help="The stepsize to use.")
-  parser.add_argument("--redis-address", default=None, type=str,
-                      help="The Redis address of the cluster.")
+class EvolutionStrategies(Algorithm):
+  def __init__(self, env_name, config):
+    Algorithm.__init__(self, env_name, config)
 
-  args = parser.parse_args()
-  num_workers = args.num_workers
-  env_name = args.env_name
-  stepsize = args.stepsize
+    policy_params = {
+        "ac_bins": "continuous:",
+        "ac_noise_std": 0.01,
+        "nonlin_type": "tanh",
+        "hidden_dims": [256, 256],
+        "connection_type": "ff"
+    }
 
-  ray.init(redis_address=args.redis_address,
-           num_workers=(0 if args.redis_address is None else None))
+    # Create the shared noise table.
+    print("Creating shared noise table.")
+    noise_id = create_shared_noise.remote()
+    self.noise = SharedNoiseTable(ray.get(noise_id))
 
-  config = Config(l2coeff=0.005,
-                  noise_stdev=0.02,
-                  episodes_per_batch=10000,
-                  timesteps_per_batch=100000,
-                  calc_obstat_prob=0.01,
-                  eval_prob=0,
-                  snapshot_freq=20,
-                  return_proc_mode="centered_rank",
-                  episode_cutoff_mode="env_default")
+    # Create the actors.
+    print("Creating actors.")
+    self.workers = [Worker.remote(config, policy_params, env_name, noise_id)
+                    for _ in range(config.num_workers)]
 
-  policy_params = {
-      "ac_bins": "continuous:",
-      "ac_noise_std": 0.01,
-      "nonlin_type": "tanh",
-      "hidden_dims": [256, 256],
-      "connection_type": "ff"
-  }
+    env = gym.make(env_name)
+    utils.make_session(single_threaded=False)
+    self.policy = policies.MujocoPolicy(
+        env.observation_space, env.action_space, **policy_params)
+    tf_util.initialize()
+    self.optimizer = optimizers.Adam(self.policy, config.stepsize)
+    self.ob_stat = utils.RunningStat(env.observation_space.shape, eps=1e-2)
 
-  # Create the shared noise table.
-  print("Creating shared noise table.")
-  noise_id = create_shared_noise.remote()
-  noise = SharedNoiseTable(ray.get(noise_id))
+    self.episodes_so_far = 0
+    self.timesteps_so_far = 0
+    self.tstart = time.time()
+    self.iteration = 0
 
-  # Create the actors.
-  print("Creating actors.")
-  workers = [Worker.remote(config, policy_params, env_name, noise_id)
-             for _ in range(num_workers)]
+  def train(self):
+    config = self.config
 
-  env = gym.make(env_name)
-  sess = utils.make_session(single_threaded=False)
-  policy = policies.MujocoPolicy(env.observation_space, env.action_space,
-                                 **policy_params)
-  tf_util.initialize()
-  optimizer = optimizers.Adam(policy, stepsize)
-
-  ob_stat = utils.RunningStat(env.observation_space.shape, eps=1e-2)
-
-  episodes_so_far = 0
-  timesteps_so_far = 0
-  tstart = time.time()
-
-  iteration = 0
-
-  while True:
     step_tstart = time.time()
-    theta = policy.get_trainable_flat()
+    theta = self.policy.get_trainable_flat()
     assert theta.dtype == np.float32
 
     # Put the current policy weights in the object store.
@@ -209,8 +198,9 @@ if __name__ == "__main__":
     # weights.
     rollout_ids = [worker.do_rollouts.remote(
         theta_id,
-        ob_stat.mean if policy.needs_ob_stat else None,
-        ob_stat.std if policy.needs_ob_stat else None) for worker in workers]
+        self.ob_stat.mean if self.policy.needs_ob_stat else None,
+        self.ob_stat.std if self.policy.needs_ob_stat else None)
+        for worker in self.workers]
 
     # Get the results of the rollouts.
     results = ray.get(rollout_ids)
@@ -227,13 +217,13 @@ if __name__ == "__main__":
 
       result_num_eps = result.lengths_n2.size
       result_num_timesteps = result.lengths_n2.sum()
-      episodes_so_far += result_num_eps
-      timesteps_so_far += result_num_timesteps
+      self.episodes_so_far += result_num_eps
+      self.timesteps_so_far += result_num_timesteps
 
       curr_task_results.append(result)
       # Update ob stats.
-      if policy.needs_ob_stat and result.ob_count > 0:
-        ob_stat.increment(result.ob_sum, result.ob_sumsq, result.ob_count)
+      if self.policy.needs_ob_stat and result.ob_count > 0:
+        self.ob_stat.increment(result.ob_sum, result.ob_sumsq, result.ob_count)
         ob_count_this_batch += result.ob_count
 
     # Assemble the results.
@@ -251,44 +241,47 @@ if __name__ == "__main__":
     # Compute and take a step.
     g, count = utils.batched_weighted_sum(
         proc_returns_n2[:, 0] - proc_returns_n2[:, 1],
-        (noise.get(idx, policy.num_params) for idx in noise_inds_n),
+        (self.noise.get(idx, self.policy.num_params) for idx in noise_inds_n),
         batch_size=500)
     g /= returns_n2.size
-    assert (g.shape == (policy.num_params,) and g.dtype == np.float32 and
+    assert (g.shape == (self.policy.num_params,) and g.dtype == np.float32 and
             count == len(noise_inds_n))
-    update_ratio = optimizer.update(-g + config.l2coeff * theta)
+    update_ratio = self.optimizer.update(-g + config.l2coeff * theta)
 
     # Update ob stat (we're never running the policy in the master, but we
     # might be snapshotting the policy).
-    if policy.needs_ob_stat:
-      policy.set_ob_stat(ob_stat.mean, ob_stat.std)
+    if self.policy.needs_ob_stat:
+      self.policy.set_ob_stat(self.ob_stat.mean, self.ob_stat.std)
 
     step_tend = time.time()
     tlogger.record_tabular("EpRewMean", returns_n2.mean())
     tlogger.record_tabular("EpRewStd", returns_n2.std())
     tlogger.record_tabular("EpLenMean", lengths_n2.mean())
 
-    tlogger.record_tabular("Norm",
-                           float(np.square(policy.get_trainable_flat()).sum()))
+    tlogger.record_tabular(
+        "Norm", float(np.square(self.policy.get_trainable_flat()).sum()))
     tlogger.record_tabular("GradNorm", float(np.square(g).sum()))
     tlogger.record_tabular("UpdateRatio", float(update_ratio))
 
     tlogger.record_tabular("EpisodesThisIter", lengths_n2.size)
-    tlogger.record_tabular("EpisodesSoFar", episodes_so_far)
+    tlogger.record_tabular("EpisodesSoFar", self.episodes_so_far)
     tlogger.record_tabular("TimestepsThisIter", lengths_n2.sum())
-    tlogger.record_tabular("TimestepsSoFar", timesteps_so_far)
+    tlogger.record_tabular("TimestepsSoFar", self.timesteps_so_far)
 
     tlogger.record_tabular("ObCount", ob_count_this_batch)
 
     tlogger.record_tabular("TimeElapsedThisIter", step_tend - step_tstart)
-    tlogger.record_tabular("TimeElapsed", step_tend - tstart)
+    tlogger.record_tabular("TimeElapsed", step_tend - self.tstart)
     tlogger.dump_tabular()
 
-    if config.snapshot_freq != 0 and iteration % config.snapshot_freq == 0:
-        filename = os.path.join("/tmp",
-                                "snapshot_iter{:05d}.h5".format(iteration))
-        assert not os.path.exists(filename)
-        policy.save(filename)
-        tlogger.log("Saved snapshot {}".format(filename))
+    if (config.snapshot_freq != 0 and
+            self.iteration % config.snapshot_freq == 0):
+      filename = os.path.join(
+          "/tmp", "snapshot_iter{:05d}.h5".format(self.iteration))
+      assert not os.path.exists(filename)
+      self.policy.save(filename)
+      tlogger.log("Saved snapshot {}".format(filename))
 
-    iteration += 1
+    res = TrainingResult(self.iteration, returns_n2.mean(), lengths_n2.mean())
+    self.iteration += 1
+    return res

python/ray/rllib/evolution_strategies/example.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+
+import ray
+from ray.rllib.evolution_strategies import EvolutionStrategies, DEFAULT_CONFIG
+
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(description="Train an RL agent on Pong.")
+  parser.add_argument("--num-workers", default=10, type=int,
+                      help=("The number of actors to create in aggregate "
+                            "across the cluster."))
+  parser.add_argument("--env-name", default="Pendulum-v0", type=str,
+                      help="The name of the gym environment to use.")
+  parser.add_argument("--stepsize", default=0.01, type=float,
+                      help="The stepsize to use.")
+  parser.add_argument("--redis-address", default=None, type=str,
+                      help="The Redis address of the cluster.")
+
+  args = parser.parse_args()
+  num_workers = args.num_workers
+  env_name = args.env_name
+  stepsize = args.stepsize
+
+  ray.init(redis_address=args.redis_address,
+           num_workers=(0 if args.redis_address is None else None))
+
+  config = DEFAULT_CONFIG._replace(
+      num_workers=num_workers,
+      stepsize=stepsize)
+
+  alg = EvolutionStrategies(env_name, config)
+  while True:
+    result = alg.train()
+    print("current status: {}".format(result))

python/ray/rllib/evolution_strategies/policies.py

@@ -12,7 +12,7 @@ import h5py
 import numpy as np
 import tensorflow as tf
 
-import tf_util as U
+from ray.rllib.evolution_strategies import tf_util as U
 
 logger = logging.getLogger(__name__)
 

python/ray/rllib/example.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python
+"""Demonstrates the RLlib algorithm API through a simple bakeoff."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import ray
+import ray.rllib.evolution_strategies as es
+import ray.rllib.policy_gradient as pg
+
+
+if __name__ == "__main__":
+  ray.init()
+
+  # TODO(ekl): get the algorithms working on a common set of envs
+  env_name = "CartPole-v0"
+  alg1 = es.EvolutionStrategies(env_name, es.DEFAULT_CONFIG)
+  alg2 = pg.PolicyGradient(env_name, pg.DEFAULT_CONFIG)
+
+  while True:
+    r1 = alg1.train()
+    r2 = alg2.train()
+    print("evolution strategies: {}".format(r1))
+    print("policy gradient: {}".format(r2))

python/ray/rllib/policy_gradient/__init__.py

@@ -0,0 +1,4 @@
+from ray.rllib.policy_gradient.policy_gradient import (
+    PolicyGradient, DEFAULT_CONFIG)
+
+__all__ = ["PolicyGradient", "DEFAULT_CONFIG"]

python/ray/rllib/policy_gradient/agent.py

@@ -13,12 +13,13 @@ from tensorflow.python import debug as tf_debug
 
 import ray
 
-from reinforce.distributions import Categorical, DiagGaussian
-from reinforce.env import BatchedEnv
-from reinforce.policy import ProximalPolicyLoss
-from reinforce.filter import MeanStdFilter
-from reinforce.rollout import rollouts, add_advantage_values
-from reinforce.utils import make_divisible_by, average_gradients
+from ray.rllib.policy_gradient.distributions import Categorical, DiagGaussian
+from ray.rllib.policy_gradient.env import BatchedEnv
+from ray.rllib.policy_gradient.loss import ProximalPolicyLoss
+from ray.rllib.policy_gradient.filter import MeanStdFilter
+from ray.rllib.policy_gradient.rollout import rollouts, add_advantage_values
+from ray.rllib.policy_gradient.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

python/ray/rllib/policy_gradient/example.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+
+import ray
+from ray.rllib.policy_gradient import PolicyGradient, DEFAULT_CONFIG
+
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(description="Run the policy gradient "
+                                               "algorithm.")
+  parser.add_argument("--environment", default="Pong-v0", type=str,
+                      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 = DEFAULT_CONFIG.copy()
+  config["use_tf_debugger"] = args.use_tf_debugger
+  if args.load_checkpoint:
+    config["load_checkpoint"] = args.load_checkpoint
+
+  ray.init(redis_address=args.redis_address)
+
+  alg = PolicyGradient(args.environment, config)
+  result = alg.train()
+  while result.training_iteration < config["max_iterations"]:
+    print("\n== iteration", result.training_iteration)
+    result = alg.train()
+    print("current status: {}".format(result))

python/ray/rllib/policy_gradient/loss.py

@@ -4,8 +4,8 @@ from __future__ import print_function
 
 import gym.spaces
 import tensorflow as tf
-from reinforce.models.visionnet import vision_net
-from reinforce.models.fcnet import fc_net
+from ray.rllib.policy_gradient.models.visionnet import vision_net
+from ray.rllib.policy_gradient.models.fcnet import fc_net
 
 
 class ProximalPolicyLoss(object):

python/ray/rllib/policy_gradient/policy_gradient.py

@@ -0,0 +1,188 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from datetime import datetime
+import time
+
+import numpy as np
+import tensorflow as tf
+
+import ray
+from ray.rllib.common import Algorithm, TrainingResult
+from ray.rllib.policy_gradient.agent import Agent, RemoteAgent
+from ray.rllib.policy_gradient.env import (
+    NoPreprocessor, AtariRamPreprocessor, AtariPixelPreprocessor)
+from ray.rllib.policy_gradient.rollout import collect_samples
+from ray.rllib.policy_gradient.utils import shuffle
+
+
+DEFAULT_CONFIG = {
+    "kl_coeff": 0.2,
+    "num_sgd_iter": 30,
+    "max_iterations": 1000,
+    "sgd_stepsize": 5e-5,
+    # TODO(pcm): Expose the choice between gpus and cpus
+    # as a command line argument.
+    "devices": ["/cpu:%d" % i for i in range(4)],
+    "tf_session_args": {
+        "device_count": {"CPU": 4},
+        "log_device_placement": False,
+        "allow_soft_placement": True,
+    },
+    "sgd_batchsize": 128,  # total size across all devices
+    "entropy_coeff": 0.0,
+    "clip_param": 0.3,
+    "kl_target": 0.01,
+    "timesteps_per_batch": 40000,
+    "num_agents": 5,
+    "tensorboard_log_dir": "/tmp/ray",
+    "full_trace_nth_sgd_batch": -1,
+    "full_trace_data_load": False,
+    "use_tf_debugger": False,
+    "model_checkpoint_file": "/tmp/iteration-%s.ckpt"}
+
+
+class PolicyGradient(Algorithm):
+  def __init__(self, env_name, config):
+    Algorithm.__init__(self, env_name, config)
+
+    # TODO(ekl) the preprocessor should be associated with the env elsewhere
+    if self.env_name == "Pong-v0":
+      preprocessor = AtariPixelPreprocessor()
+    elif self.env_name == "Pong-ram-v3":
+      preprocessor = AtariRamPreprocessor()
+    elif self.env_name == "CartPole-v0":
+      preprocessor = NoPreprocessor()
+    elif self.env_name == "Walker2d-v1":
+      preprocessor = NoPreprocessor()
+    else:
+      preprocessor = AtariPixelPreprocessor()
+
+    self.preprocessor = preprocessor
+    self.global_step = 0
+    self.j = 0
+    self.kl_coeff = config["kl_coeff"]
+    self.model = Agent(
+        self.env_name, 1, self.preprocessor, self.config, False)
+    self.agents = [
+        RemoteAgent.remote(
+            self.env_name, 1, self.preprocessor, self.config, True)
+        for _ in range(config["num_agents"])]
+
+  def train(self):
+    agents = self.agents
+    config = self.config
+    model = self.model
+    j = self.j
+    self.j += 1
+
+    saver = tf.train.Saver(max_to_keep=None)
+    if "load_checkpoint" in config:
+      saver.restore(model.sess, config["load_checkpoint"])
+
+    file_writer = tf.summary.FileWriter(
+        "{}/trpo_{}_{}".format(
+            config["tensorboard_log_dir"], self.env_name,
+            str(datetime.today()).replace(" ", "_")),
+        model.sess.graph)
+    iter_start = time.time()
+    if config["model_checkpoint_file"]:
+      checkpoint_path = saver.save(
+          model.sess, config["model_checkpoint_file"] % j)
+      print("Checkpoint saved in file: %s" % checkpoint_path)
+    checkpointing_end = time.time()
+    weights = ray.put(model.get_weights())
+    [a.load_weights.remote(weights) for a in agents]
+    trajectory, total_reward, traj_len_mean = collect_samples(
+        agents, config["timesteps_per_batch"], 0.995, 1.0, 2000)
+    print("total reward is ", total_reward)
+    print("trajectory length mean is ", traj_len_mean)
+    print("timesteps:", trajectory["dones"].shape[0])
+    traj_stats = tf.Summary(value=[
+        tf.Summary.Value(
+            tag="policy_gradient/rollouts/mean_reward",
+            simple_value=total_reward),
+        tf.Summary.Value(
+            tag="policy_gradient/rollouts/traj_len_mean",
+            simple_value=traj_len_mean)])
+    file_writer.add_summary(traj_stats, self.global_step)
+    self.global_step += 1
+    trajectory["advantages"] = ((trajectory["advantages"] -
+                                 trajectory["advantages"].mean()) /
+                                trajectory["advantages"].std())
+    rollouts_end = time.time()
+    print("Computing policy (iterations=" + str(config["num_sgd_iter"]) +
+          ", stepsize=" + str(config["sgd_stepsize"]) + "):")
+    names = ["iter", "loss", "kl", "entropy"]
+    print(("{:>15}" * len(names)).format(*names))
+    trajectory = shuffle(trajectory)
+    shuffle_end = time.time()
+    tuples_per_device = model.load_data(
+        trajectory, j == 0 and config["full_trace_data_load"])
+    load_end = time.time()
+    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
+    for i in range(config["num_sgd_iter"]):
+      sgd_start = time.time()
+      batch_index = 0
+      num_batches = int(tuples_per_device) // int(model.per_device_batch_size)
+      loss, kl, entropy = [], [], []
+      permutation = np.random.permutation(num_batches)
+      while batch_index < num_batches:
+        full_trace = (
+            i == 0 and j == 0 and
+            batch_index == config["full_trace_nth_sgd_batch"])
+        batch_loss, batch_kl, batch_entropy = model.run_sgd_minibatch(
+            permutation[batch_index] * model.per_device_batch_size,
+            self.kl_coeff, full_trace, file_writer)
+        loss.append(batch_loss)
+        kl.append(batch_kl)
+        entropy.append(batch_entropy)
+        batch_index += 1
+      loss = np.mean(loss)
+      kl = np.mean(kl)
+      entropy = np.mean(entropy)
+      sgd_end = time.time()
+      print("{:>15}{:15.5e}{:15.5e}{:15.5e}".format(i, loss, kl, entropy))
+
+      values = []
+      if i == config["num_sgd_iter"] - 1:
+        metric_prefix = "policy_gradient/sgd/final_iter/"
+        values.append(tf.Summary.Value(
+            tag=metric_prefix + "kl_coeff",
+            simple_value=self.kl_coeff))
+      else:
+        metric_prefix = "policy_gradient/sgd/intermediate_iters/"
+      values.extend([
+          tf.Summary.Value(
+              tag=metric_prefix + "mean_entropy",
+              simple_value=entropy),
+          tf.Summary.Value(
+              tag=metric_prefix + "mean_loss",
+              simple_value=loss),
+          tf.Summary.Value(
+              tag=metric_prefix + "mean_kl",
+              simple_value=kl)])
+      sgd_stats = tf.Summary(value=values)
+      file_writer.add_summary(sgd_stats, self.global_step)
+      self.global_step += 1
+      sgd_time += sgd_end - sgd_start
+    if kl > 2.0 * config["kl_target"]:
+      self.kl_coeff *= 1.5
+    elif kl < 0.5 * config["kl_target"]:
+      self.kl_coeff *= 0.5
+
+    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)
+    print("sgd time:", sgd_time)
+    print("sgd examples/s:", len(trajectory["observations"]) / sgd_time)
+
+    return TrainingResult(j, total_reward, traj_len_mean)

python/ray/rllib/policy_gradient/rollout.py

@@ -5,8 +5,8 @@ from __future__ import print_function
 import numpy as np
 import ray
 
-from reinforce.filter import NoFilter
-from reinforce.utils import flatten, concatenate
+from ray.rllib.policy_gradient.filter import NoFilter
+from ray.rllib.policy_gradient.utils import flatten, concatenate
 
 
 def rollouts(policy, env, horizon, observation_filter=NoFilter(),

python/ray/rllib/policy_gradient/test/test.py

@@ -7,8 +7,8 @@ import numpy as np
 import tensorflow as tf
 from numpy.testing import assert_allclose
 
-from reinforce.distributions import Categorical
-from reinforce.utils import flatten, concatenate
+from ray.rllib.policy_gradient.distributions import Categorical
+from ray.rllib.policy_gradient.utils import flatten, concatenate
 
 
 class DistibutionsTest(unittest.TestCase):