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

* rllib v0

* fix imports

* lint

* comments

* update docs

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

@@ -0,0 +1,287 @@
+# Code in this file is copied and adapted from
+# https://github.com/openai/evolution-strategies-starter.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from collections import namedtuple
+import gym
+import numpy as np
+import os
+import time
+
+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", "num_workers", "stepsize"
+])
+
+Result = namedtuple("Result", [
+    "noise_inds_n", "returns_n2", "sign_returns_n2", "lengths_n2",
+    "eval_return", "eval_length", "ob_sum", "ob_sumsq", "ob_count"
+])
+
+
+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."""
+  seed = 123
+  count = 250000000
+  noise = np.random.RandomState(seed).randn(count).astype(np.float32)
+  return noise
+
+
+class SharedNoiseTable(object):
+  def __init__(self, noise):
+    self.noise = noise
+    assert self.noise.dtype == np.float32
+
+  def get(self, i, dim):
+    return self.noise[i:i + dim]
+
+  def sample_index(self, stream, dim):
+    return stream.randint(0, len(self.noise) - dim + 1)
+
+
+@ray.remote
+class Worker(object):
+  def __init__(self, config, policy_params, env_name, noise,
+               min_task_runtime=0.2):
+    self.min_task_runtime = min_task_runtime
+    self.config = config
+    self.policy_params = policy_params
+    self.noise = SharedNoiseTable(noise)
+
+    self.env = gym.make(env_name)
+    self.sess = utils.make_session(single_threaded=True)
+    self.policy = policies.MujocoPolicy(self.env.observation_space,
+                                        self.env.action_space,
+                                        **policy_params)
+    tf_util.initialize()
+
+    self.rs = np.random.RandomState()
+
+    assert self.policy.needs_ob_stat == (self.config.calc_obstat_prob != 0)
+
+  def rollout_and_update_ob_stat(self, timestep_limit, task_ob_stat):
+    if (self.policy.needs_ob_stat and self.config.calc_obstat_prob != 0 and
+            self.rs.rand() < self.config.calc_obstat_prob):
+      rollout_rews, rollout_len, obs = self.policy.rollout(
+          self.env, timestep_limit=timestep_limit, save_obs=True,
+          random_stream=self.rs)
+      task_ob_stat.increment(obs.sum(axis=0), np.square(obs).sum(axis=0),
+                             len(obs))
+    else:
+      rollout_rews, rollout_len = self.policy.rollout(
+          self.env, timestep_limit=timestep_limit, random_stream=self.rs)
+    return rollout_rews, rollout_len
+
+  def do_rollouts(self, params, ob_mean, ob_std, timestep_limit=None):
+    # Set the network weights.
+    self.policy.set_trainable_flat(params)
+
+    if self.policy.needs_ob_stat:
+      self.policy.set_ob_stat(ob_mean, ob_std)
+
+    if self.config.eval_prob != 0:
+      raise NotImplementedError("Eval rollouts are not implemented.")
+
+    noise_inds, returns, sign_returns, lengths = [], [], [], []
+    # We set eps=0 because we're incrementing only.
+    task_ob_stat = utils.RunningStat(self.env.observation_space.shape, eps=0)
+
+    # Perform some rollouts with noise.
+    task_tstart = time.time()
+    while (len(noise_inds) == 0 or
+           time.time() - task_tstart < self.min_task_runtime):
+      noise_idx = self.noise.sample_index(self.rs, self.policy.num_params)
+      perturbation = self.config.noise_stdev * self.noise.get(
+          noise_idx, self.policy.num_params)
+
+      # These two sampling steps could be done in parallel on different actors
+      # letting us update twice as frequently.
+      self.policy.set_trainable_flat(params + perturbation)
+      rews_pos, len_pos = self.rollout_and_update_ob_stat(timestep_limit,
+                                                          task_ob_stat)
+
+      self.policy.set_trainable_flat(params - perturbation)
+      rews_neg, len_neg = self.rollout_and_update_ob_stat(timestep_limit,
+                                                          task_ob_stat)
+
+      noise_inds.append(noise_idx)
+      returns.append([rews_pos.sum(), rews_neg.sum()])
+      sign_returns.append([np.sign(rews_pos).sum(), np.sign(rews_neg).sum()])
+      lengths.append([len_pos, len_neg])
+
+      return Result(
+          noise_inds_n=np.array(noise_inds),
+          returns_n2=np.array(returns, dtype=np.float32),
+          sign_returns_n2=np.array(sign_returns, dtype=np.float32),
+          lengths_n2=np.array(lengths, dtype=np.int32),
+          eval_return=None,
+          eval_length=None,
+          ob_sum=(None if task_ob_stat.count == 0 else task_ob_stat.sum),
+          ob_sumsq=(None if task_ob_stat.count == 0 else task_ob_stat.sumsq),
+          ob_count=task_ob_stat.count)
+
+
+class EvolutionStrategies(Algorithm):
+  def __init__(self, env_name, config):
+    Algorithm.__init__(self, env_name, config)
+
+    policy_params = {
+        "ac_bins": "continuous:",
+        "ac_noise_std": 0.01,
+        "nonlin_type": "tanh",
+        "hidden_dims": [256, 256],
+        "connection_type": "ff"
+    }
+
+    # Create the shared noise table.
+    print("Creating shared noise table.")
+    noise_id = create_shared_noise.remote()
+    self.noise = SharedNoiseTable(ray.get(noise_id))
+
+    # Create the actors.
+    print("Creating actors.")
+    self.workers = [Worker.remote(config, policy_params, env_name, noise_id)
+                    for _ in range(config.num_workers)]
+
+    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)
+
+    self.episodes_so_far = 0
+    self.timesteps_so_far = 0
+    self.tstart = time.time()
+    self.iteration = 0
+
+  def train(self):
+    config = self.config
+
+    step_tstart = time.time()
+    theta = self.policy.get_trainable_flat()
+    assert theta.dtype == np.float32
+
+    # Put the current policy weights in the object store.
+    theta_id = ray.put(theta)
+    # Use the actors to do rollouts, note that we pass in the ID of the policy
+    # weights.
+    rollout_ids = [worker.do_rollouts.remote(
+        theta_id,
+        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)
+
+    curr_task_results = []
+    ob_count_this_batch = 0
+    # Loop over the results
+    for result in results:
+      assert result.eval_length is None, "We aren't doing eval rollouts."
+      assert result.noise_inds_n.ndim == 1
+      assert result.returns_n2.shape == (len(result.noise_inds_n), 2)
+      assert result.lengths_n2.shape == (len(result.noise_inds_n), 2)
+      assert result.returns_n2.dtype == np.float32
+
+      result_num_eps = result.lengths_n2.size
+      result_num_timesteps = result.lengths_n2.sum()
+      self.episodes_so_far += result_num_eps
+      self.timesteps_so_far += result_num_timesteps
+
+      curr_task_results.append(result)
+      # Update ob stats.
+      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.
+    noise_inds_n = np.concatenate([r.noise_inds_n for
+                                   r in curr_task_results])
+    returns_n2 = np.concatenate([r.returns_n2 for r in curr_task_results])
+    lengths_n2 = np.concatenate([r.lengths_n2 for r in curr_task_results])
+    assert noise_inds_n.shape[0] == returns_n2.shape[0] == lengths_n2.shape[0]
+    # Process the returns.
+    if config.return_proc_mode == "centered_rank":
+      proc_returns_n2 = utils.compute_centered_ranks(returns_n2)
+    else:
+      raise NotImplementedError(config.return_proc_mode)
+
+    # Compute and take a step.
+    g, count = utils.batched_weighted_sum(
+        proc_returns_n2[:, 0] - proc_returns_n2[:, 1],
+        (self.noise.get(idx, self.policy.num_params) for idx in noise_inds_n),
+        batch_size=500)
+    g /= returns_n2.size
+    assert (g.shape == (self.policy.num_params,) and g.dtype == np.float32 and
+            count == len(noise_inds_n))
+    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 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(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", self.episodes_so_far)
+    tlogger.record_tabular("TimestepsThisIter", lengths_n2.sum())
+    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 - self.tstart)
+    tlogger.dump_tabular()
+
+    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))
+
+    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/optimizers.py

@@ -0,0 +1,57 @@
+# Code in this file is copied and adapted from
+# https://github.com/openai/evolution-strategies-starter.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+
+class Optimizer(object):
+  def __init__(self, pi):
+    self.pi = pi
+    self.dim = pi.num_params
+    self.t = 0
+
+  def update(self, globalg):
+    self.t += 1
+    step = self._compute_step(globalg)
+    theta = self.pi.get_trainable_flat()
+    ratio = np.linalg.norm(step) / np.linalg.norm(theta)
+    self.pi.set_trainable_flat(theta + step)
+    return ratio
+
+  def _compute_step(self, globalg):
+    raise NotImplementedError
+
+
+class SGD(Optimizer):
+  def __init__(self, pi, stepsize, momentum=0.9):
+    Optimizer.__init__(self, pi)
+    self.v = np.zeros(self.dim, dtype=np.float32)
+    self.stepsize, self.momentum = stepsize, momentum
+
+  def _compute_step(self, globalg):
+    self.v = self.momentum * self.v + (1. - self.momentum) * globalg
+    step = -self.stepsize * self.v
+    return step
+
+
+class Adam(Optimizer):
+  def __init__(self, pi, stepsize, beta1=0.9, beta2=0.999, epsilon=1e-08):
+    Optimizer.__init__(self, pi)
+    self.stepsize = stepsize
+    self.beta1 = beta1
+    self.beta2 = beta2
+    self.epsilon = epsilon
+    self.m = np.zeros(self.dim, dtype=np.float32)
+    self.v = np.zeros(self.dim, dtype=np.float32)
+
+  def _compute_step(self, globalg):
+    a = self.stepsize * (np.sqrt(1 - self.beta2 ** self.t) /
+                         (1 - self.beta1 ** self.t))
+    self.m = self.beta1 * self.m + (1 - self.beta1) * globalg
+    self.v = self.beta2 * self.v + (1 - self.beta2) * (globalg * globalg)
+    step = -a * self.m / (np.sqrt(self.v) + self.epsilon)
+    return step

python/ray/rllib/evolution_strategies/policies.py

@@ -0,0 +1,241 @@
+# Code in this file is copied and adapted from
+# https://github.com/openai/evolution-strategies-starter.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+import pickle
+
+import h5py
+import numpy as np
+import tensorflow as tf
+
+from ray.rllib.evolution_strategies import tf_util as U
+
+logger = logging.getLogger(__name__)
+
+
+class Policy:
+  def __init__(self, *args, **kwargs):
+    self.args, self.kwargs = args, kwargs
+    self.scope = self._initialize(*args, **kwargs)
+    self.all_variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
+                                           self.scope.name)
+
+    self.trainable_variables = tf.get_collection(
+        tf.GraphKeys.TRAINABLE_VARIABLES, self.scope.name)
+    self.num_params = sum(int(np.prod(v.get_shape().as_list()))
+                          for v in self.trainable_variables)
+    self._setfromflat = U.SetFromFlat(self.trainable_variables)
+    self._getflat = U.GetFlat(self.trainable_variables)
+
+    logger.info('Trainable variables ({} parameters)'.format(self.num_params))
+    for v in self.trainable_variables:
+      shp = v.get_shape().as_list()
+      logger.info('- {} shape:{} size:{}'.format(v.name, shp, np.prod(shp)))
+    logger.info('All variables')
+    for v in self.all_variables:
+      shp = v.get_shape().as_list()
+      logger.info('- {} shape:{} size:{}'.format(v.name, shp, np.prod(shp)))
+
+    placeholders = [tf.placeholder(v.value().dtype, v.get_shape().as_list())
+                    for v in self.all_variables]
+    self.set_all_vars = U.function(
+        inputs=placeholders,
+        outputs=[],
+        updates=[tf.group(*[v.assign(p) for v, p
+                 in zip(self.all_variables, placeholders)])]
+    )
+
+  def _initialize(self, *args, **kwargs):
+    raise NotImplementedError
+
+  def save(self, filename):
+    assert filename.endswith('.h5')
+    with h5py.File(filename, 'w') as f:
+      for v in self.all_variables:
+        f[v.name] = v.eval()
+      # TODO: It would be nice to avoid pickle, but it's convenient to pass
+      # Python objects to _initialize (like Gym spaces or numpy arrays).
+      f.attrs['name'] = type(self).__name__
+      f.attrs['args_and_kwargs'] = np.void(pickle.dumps((self.args,
+                                                         self.kwargs),
+                                                        protocol=-1))
+
+  @classmethod
+  def Load(cls, filename, extra_kwargs=None):
+    with h5py.File(filename, 'r') as f:
+      args, kwargs = pickle.loads(f.attrs['args_and_kwargs'].tostring())
+      if extra_kwargs:
+        kwargs.update(extra_kwargs)
+      policy = cls(*args, **kwargs)
+      policy.set_all_vars(*[f[v.name][...] for v in policy.all_variables])
+    return policy
+
+  # === Rollouts/training ===
+
+  def rollout(self, env, render=False, timestep_limit=None, save_obs=False,
+              random_stream=None):
+    """Do a rollout.
+
+    If random_stream is provided, the rollout will take noisy actions with
+    noise drawn from that stream. Otherwise, no action noise will be added.
+    """
+    env_timestep_limit = env.spec.tags.get("wrapper_config.TimeLimit"
+                                           ".max_episode_steps")
+    timestep_limit = (env_timestep_limit if timestep_limit is None
+                      else min(timestep_limit, env_timestep_limit))
+    rews = []
+    t = 0
+    if save_obs:
+      obs = []
+    ob = env.reset()
+    for _ in range(timestep_limit):
+      ac = self.act(ob[None], random_stream=random_stream)[0]
+      if save_obs:
+        obs.append(ob)
+      ob, rew, done, _ = env.step(ac)
+      rews.append(rew)
+      t += 1
+      if render:
+        env.render()
+      if done:
+        break
+    rews = np.array(rews, dtype=np.float32)
+    if save_obs:
+      return rews, t, np.array(obs)
+    return rews, t
+
+  def act(self, ob, random_stream=None):
+    raise NotImplementedError
+
+  def set_trainable_flat(self, x):
+    self._setfromflat(x)
+
+  def get_trainable_flat(self):
+    return self._getflat()
+
+  @property
+  def needs_ob_stat(self):
+    raise NotImplementedError
+
+  def set_ob_stat(self, ob_mean, ob_std):
+    raise NotImplementedError
+
+
+def bins(x, dim, num_bins, name):
+  scores = U.dense(x, dim * num_bins, name, U.normc_initializer(0.01))
+  scores_nab = tf.reshape(scores, [-1, dim, num_bins])
+  return tf.argmax(scores_nab, 2)
+
+
+class MujocoPolicy(Policy):
+  def _initialize(self, ob_space, ac_space, ac_bins, ac_noise_std, nonlin_type,
+                  hidden_dims, connection_type):
+    self.ac_space = ac_space
+    self.ac_bins = ac_bins
+    self.ac_noise_std = ac_noise_std
+    self.hidden_dims = hidden_dims
+    self.connection_type = connection_type
+
+    assert len(ob_space.shape) == len(self.ac_space.shape) == 1
+    assert (np.all(np.isfinite(self.ac_space.low)) and
+            np.all(np.isfinite(self.ac_space.high))), "Action bounds required"
+
+    self.nonlin = {'tanh': tf.tanh,
+                   'relu': tf.nn.relu,
+                   'lrelu': U.lrelu,
+                   'elu': tf.nn.elu}[nonlin_type]
+
+    with tf.variable_scope(type(self).__name__) as scope:
+      # Observation normalization.
+      ob_mean = tf.get_variable(
+          'ob_mean', ob_space.shape, tf.float32,
+          tf.constant_initializer(np.nan), trainable=False)
+      ob_std = tf.get_variable(
+          'ob_std', ob_space.shape, tf.float32,
+          tf.constant_initializer(np.nan), trainable=False)
+      in_mean = tf.placeholder(tf.float32, ob_space.shape)
+      in_std = tf.placeholder(tf.float32, ob_space.shape)
+      self._set_ob_mean_std = U.function([in_mean, in_std], [], updates=[
+          tf.assign(ob_mean, in_mean),
+          tf.assign(ob_std, in_std),
+      ])
+
+      # Policy network.
+      o = tf.placeholder(tf.float32, [None] + list(ob_space.shape))
+      a = self._make_net(tf.clip_by_value((o - ob_mean) / ob_std, -5.0, 5.0))
+      self._act = U.function([o], a)
+    return scope
+
+  def _make_net(self, o):
+    # Process observation.
+    if self.connection_type == 'ff':
+      x = o
+      for ilayer, hd in enumerate(self.hidden_dims):
+        x = self.nonlin(U.dense(x, hd, 'l{}'.format(ilayer),
+                                U.normc_initializer(1.0)))
+    else:
+      raise NotImplementedError(self.connection_type)
+
+    # Map to action.
+    adim = self.ac_space.shape[0]
+    ahigh = self.ac_space.high
+    alow = self.ac_space.low
+    assert isinstance(self.ac_bins, str)
+    ac_bin_mode, ac_bin_arg = self.ac_bins.split(':')
+
+    if ac_bin_mode == 'uniform':
+      # Uniformly spaced bins, from ac_space.low to ac_space.high.
+      num_ac_bins = int(ac_bin_arg)
+      aidx_na = bins(x, adim, num_ac_bins, 'out')
+      ac_range_1a = (ahigh - alow)[None, :]
+      a = (1. / (num_ac_bins - 1.) * tf.to_float(aidx_na) * ac_range_1a +
+           alow[None, :])
+
+    elif ac_bin_mode == 'custom':
+      # Custom bins specified as a list of values from -1 to 1.
+      # The bins are rescaled to ac_space.low to ac_space.high.
+      acvals_k = np.array(list(map(float, ac_bin_arg.split(','))),
+                          dtype=np.float32)
+      logger.info('Custom action values: ' + ' '.join('{:.3f}'.format(x)
+                                                      for x in acvals_k))
+      assert acvals_k.ndim == 1 and acvals_k[0] == -1 and acvals_k[-1] == 1
+      acvals_ak = ((ahigh - alow)[:, None] / (acvals_k[-1] - acvals_k[0]) *
+                   (acvals_k - acvals_k[0])[None, :] + alow[:, None])
+
+      aidx_na = bins(x, adim, len(acvals_k), 'out')  # Values in [0, k-1].
+      a = tf.gather_nd(
+          acvals_ak,
+          tf.concat([
+              tf.tile(np.arange(adim)[None, :, None],
+                      [tf.shape(aidx_na)[0], 1, 1]),
+              2,
+              tf.expand_dims(aidx_na, -1)
+          ])  # (n, a, 2)
+      )  # (n, a)
+    elif ac_bin_mode == 'continuous':
+      a = U.dense(x, adim, 'out', U.normc_initializer(0.01))
+    else:
+      raise NotImplementedError(ac_bin_mode)
+
+    return a
+
+  def act(self, ob, random_stream=None):
+    a = self._act(ob)
+    if random_stream is not None and self.ac_noise_std != 0:
+      a += random_stream.randn(*a.shape) * self.ac_noise_std
+    return a
+
+  @property
+  def needs_ob_stat(self):
+    return True
+
+  @property
+  def needs_ref_batch(self):
+    return False
+
+  def set_ob_stat(self, ob_mean, ob_std):
+    self._set_ob_mean_std(ob_mean, ob_std)

python/ray/rllib/evolution_strategies/tabular_logger.py

@@ -0,0 +1,223 @@
+# Code in this file is copied and adapted from
+# https://github.com/openai/evolution-strategies-starter.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from collections import OrderedDict
+import os
+import sys
+import time
+
+import tensorflow as tf
+from tensorflow.core.util import event_pb2
+from tensorflow.python import pywrap_tensorflow
+from tensorflow.python.util import compat
+
+DEBUG = 10
+INFO = 20
+WARN = 30
+ERROR = 40
+
+DISABLED = 50
+
+
+class TbWriter(object):
+  """Based on SummaryWriter, but changed to allow for a different prefix."""
+  def __init__(self, dir, prefix):
+    self.dir = dir
+    # Start at 1, because EvWriter automatically generates an object with
+    # step = 0.
+    self.step = 1
+    self.evwriter = pywrap_tensorflow.EventsWriter(
+        compat.as_bytes(os.path.join(dir, prefix)))
+
+  def write_values(self, key2val):
+    summary = tf.Summary(value=[tf.Summary.Value(tag=k, simple_value=float(v))
+                                for (k, v) in key2val.items()])
+    event = event_pb2.Event(wall_time=time.time(), summary=summary)
+    event.step = self.step
+    self.evwriter.WriteEvent(event)
+    self.evwriter.Flush()
+    self.step += 1
+
+  def close(self):
+    self.evwriter.Close()
+
+# API
+
+
+def start(dir):
+  if _Logger.CURRENT is not _Logger.DEFAULT:
+    sys.stderr.write("WARNING: You asked to start logging (dir=%s), but you "
+                     "never stopped the previous logger (dir=%s)."
+                     "\n" % (dir, _Logger.CURRENT.dir))
+  _Logger.CURRENT = _Logger(dir=dir)
+
+
+def stop():
+  if _Logger.CURRENT is _Logger.DEFAULT:
+    sys.stderr.write("WARNING: You asked to stop logging, but you never "
+                     "started any previous logger."
+                     "\n" % (dir, _Logger.CURRENT.dir))
+    return
+  _Logger.CURRENT.close()
+  _Logger.CURRENT = _Logger.DEFAULT
+
+
+def record_tabular(key, val):
+  """Log a value of some diagnostic.
+
+  Call this once for each diagnostic quantity, each iteration.
+  """
+  _Logger.CURRENT.record_tabular(key, val)
+
+
+def dump_tabular():
+  """Write all of the diagnostics from the current iteration."""
+  _Logger.CURRENT.dump_tabular()
+
+
+def log(*args, **kwargs):
+  """Write the sequence of args, with no separators.
+
+  This is written to the console and output files (if you've configured an
+  output file).
+  """
+  level = kwargs['level'] if 'level' in kwargs else INFO
+  _Logger.CURRENT.log(*args, level=level)
+
+
+def debug(*args):
+  log(*args, level=DEBUG)
+
+
+def info(*args):
+  log(*args, level=INFO)
+
+
+def warn(*args):
+  log(*args, level=WARN)
+
+
+def error(*args):
+  log(*args, level=ERROR)
+
+
+def set_level(level):
+  """
+  Set logging threshold on current logger.
+  """
+  _Logger.CURRENT.set_level(level)
+
+
+def get_dir():
+  """
+  Get directory that log files are being written to.
+  will be None if there is no output directory (i.e., if you didn't call start)
+  """
+  return _Logger.CURRENT.get_dir()
+
+
+def get_expt_dir():
+  sys.stderr.write("get_expt_dir() is Deprecated. Switch to get_dir()\n")
+  return get_dir()
+
+# Backend
+
+
+class _Logger(object):
+  # A logger with no output files. (See right below class definition) so that
+  # you can still log to the terminal without setting up any output files.
+  DEFAULT = None
+  # Current logger being used by the free functions above.
+  CURRENT = None
+
+  def __init__(self, dir=None):
+    self.name2val = OrderedDict()  # Values this iteration.
+    self.level = INFO
+    self.dir = dir
+    self.text_outputs = [sys.stdout]
+    if dir is not None:
+      os.makedirs(dir, exist_ok=True)
+      self.text_outputs.append(open(os.path.join(dir, "log.txt"), "w"))
+      self.tbwriter = TbWriter(dir=dir, prefix="events")
+    else:
+      self.tbwriter = None
+
+  # Logging API, forwarded
+
+  def record_tabular(self, key, val):
+    self.name2val[key] = val
+
+  def dump_tabular(self):
+    # Create strings for printing.
+    key2str = OrderedDict()
+    for (key, val) in self.name2val.items():
+      if hasattr(val, "__float__"):
+        valstr = "%-8.3g" % val
+      else:
+        valstr = val
+      key2str[self._truncate(key)] = self._truncate(valstr)
+    keywidth = max(map(len, key2str.keys()))
+    valwidth = max(map(len, key2str.values()))
+    # Write to all text outputs
+    self._write_text("-" * (keywidth + valwidth + 7), "\n")
+    for (key, val) in key2str.items():
+      self._write_text("| ", key, " " * (keywidth - len(key)), " | ", val,
+                       " " * (valwidth - len(val)), " |\n")
+    self._write_text("-" * (keywidth + valwidth + 7), "\n")
+    for f in self.text_outputs:
+      try:
+        f.flush()
+      except OSError:
+        sys.stderr.write('Warning! OSError when flushing.\n')
+    # Write to tensorboard
+    if self.tbwriter is not None:
+      self.tbwriter.write_values(self.name2val)
+      self.name2val.clear()
+
+  def log(self, *args, **kwargs):
+    level = kwargs['level'] if 'level' in kwargs else INFO
+    if self.level <= level:
+      self._do_log(*args)
+
+  # Configuration
+
+  def set_level(self, level):
+    self.level = level
+
+  def get_dir(self):
+    return self.dir
+
+  def close(self):
+    for f in self.text_outputs[1:]:
+      f.close()
+    if self.tbwriter:
+      self.tbwriter.close()
+
+  # Misc
+
+  def _do_log(self, *args):
+    self._write_text(*args + ('\n',))
+    for f in self.text_outputs:
+      try:
+        f.flush()
+      except OSError:
+        print('Warning! OSError when flushing.')
+
+  def _write_text(self, *strings):
+    for f in self.text_outputs:
+      for string in strings:
+        f.write(string)
+
+  def _truncate(self, s):
+    if len(s) > 33:
+      return s[:30] + "..."
+    else:
+      return s
+
+
+_Logger.DEFAULT = _Logger()
+_Logger.CURRENT = _Logger.DEFAULT

python/ray/rllib/evolution_strategies/tf_util.py

@@ -0,0 +1,288 @@
+# Code in this file is copied and adapted from
+# https://github.com/openai/evolution-strategies-starter.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import tensorflow as tf
+import functools
+import os
+
+# Tensorflow must be at least version 1.0.0 for the example to work.
+if int(tf.__version__.split(".")[0]) < 1:
+  raise Exception("Your Tensorflow version is less than 1.0.0. Please update "
+                  "Tensorflow to the latest version.")
+
+# ================================================================
+# Import all names into common namespace
+# ================================================================
+
+clip = tf.clip_by_value
+
+# Make consistent with numpy
+
+
+def sum(x, axis=None, keepdims=False):
+  return tf.reduce_sum(x, reduction_indices=None if axis is None else [axis],
+                       keep_dims=keepdims)
+
+
+def mean(x, axis=None, keepdims=False):
+  return tf.reduce_mean(x, reduction_indices=None if axis is None else [axis],
+                        keep_dims=keepdims)
+
+
+def var(x, axis=None, keepdims=False):
+  meanx = mean(x, axis=axis, keepdims=keepdims)
+  return mean(tf.square(x - meanx), axis=axis, keepdims=keepdims)
+
+
+def std(x, axis=None, keepdims=False):
+  return tf.sqrt(var(x, axis=axis, keepdims=keepdims))
+
+
+def max(x, axis=None, keepdims=False):
+  return tf.reduce_max(x, reduction_indices=None if axis is None else [axis],
+                       keep_dims=keepdims)
+
+
+def min(x, axis=None, keepdims=False):
+  return tf.reduce_min(x, reduction_indices=None if axis is None else [axis],
+                       keep_dims=keepdims)
+
+
+def concatenate(arrs, axis=0):
+  return tf.concat(arrs, axis)
+
+
+def argmax(x, axis=None):
+  return tf.argmax(x, dimension=axis)
+
+# Extras
+
+
+def l2loss(params):
+  if len(params) == 0:
+    return tf.constant(0.0)
+  else:
+    return tf.add_n([sum(tf.square(p)) for p in params])
+
+
+def lrelu(x, leak=0.2):
+  f1 = 0.5 * (1 + leak)
+  f2 = 0.5 * (1 - leak)
+  return f1 * x + f2 * abs(x)
+
+
+def categorical_sample_logits(X):
+  # https://github.com/tensorflow/tensorflow/issues/456
+  U = tf.random_uniform(tf.shape(X))
+  return argmax(X - tf.log(-tf.log(U)), axis=1)
+
+# Global session
+
+
+def get_session():
+  return tf.get_default_session()
+
+
+def single_threaded_session():
+  tf_config = tf.ConfigProto(inter_op_parallelism_threads=1,
+                             intra_op_parallelism_threads=1)
+  return tf.Session(config=tf_config)
+
+
+ALREADY_INITIALIZED = set()
+
+
+def initialize():
+  new_variables = set(tf.global_variables()) - ALREADY_INITIALIZED
+  get_session().run(tf.variables_initializer(new_variables))
+  ALREADY_INITIALIZED.update(new_variables)
+
+
+def eval(expr, feed_dict=None):
+  if feed_dict is None:
+    feed_dict = {}
+  return get_session().run(expr, feed_dict=feed_dict)
+
+
+def set_value(v, val):
+  get_session().run(v.assign(val))
+
+
+def load_state(fname):
+  saver = tf.train.Saver()
+  saver.restore(get_session(), fname)
+
+
+def save_state(fname):
+  os.makedirs(os.path.dirname(fname), exist_ok=True)
+  saver = tf.train.Saver()
+  saver.save(get_session(), fname)
+
+# Model components
+
+
+def normc_initializer(std=1.0):
+  def _initializer(shape, dtype=None, partition_info=None):
+    out = np.random.randn(*shape).astype(np.float32)
+    out *= std / np.sqrt(np.square(out).sum(axis=0, keepdims=True))
+    return tf.constant(out)
+  return _initializer
+
+
+def dense(x, size, name, weight_init=None, bias=True):
+  w = tf.get_variable(name + "/w", [x.get_shape()[1], size],
+                      initializer=weight_init)
+  ret = tf.matmul(x, w)
+  if bias:
+    b = tf.get_variable(name + "/b", [size],
+                        initializer=tf.zeros_initializer())
+    return ret + b
+  else:
+    return ret
+
+# Basic Stuff
+
+
+def function(inputs, outputs, updates=None, givens=None):
+  if isinstance(outputs, list):
+    return _Function(inputs, outputs, updates, givens=givens)
+  elif isinstance(outputs, dict):
+    f = _Function(inputs, outputs.values(), updates, givens=givens)
+    return lambda *inputs: dict(zip(outputs.keys(), f(*inputs)))
+  else:
+    f = _Function(inputs, [outputs], updates, givens=givens)
+    return lambda *inputs: f(*inputs)[0]
+
+
+class _Function(object):
+  def __init__(self, inputs, outputs, updates, givens, check_nan=False):
+    assert all(len(i.op.inputs) == 0 for i in inputs), ("inputs should all be "
+                                                        "placeholders")
+    self.inputs = inputs
+    updates = updates or []
+    self.update_group = tf.group(*updates)
+    self.outputs_update = list(outputs) + [self.update_group]
+    self.givens = {} if givens is None else givens
+    self.check_nan = check_nan
+
+  def __call__(self, *inputvals):
+    assert len(inputvals) == len(self.inputs)
+    feed_dict = dict(zip(self.inputs, inputvals))
+    feed_dict.update(self.givens)
+    results = get_session().run(self.outputs_update, feed_dict=feed_dict)[:-1]
+    if self.check_nan:
+      if any(np.isnan(r).any() for r in results):
+        raise RuntimeError("Nan detected")
+    return results
+
+
+# Graph traversal
+
+VARIABLES = {}
+
+# Flat vectors
+
+
+def var_shape(x):
+  out = [k.value for k in x.get_shape()]
+  assert all(isinstance(a, int) for a in out), ("shape function assumes that "
+                                                "shape is fully known")
+  return out
+
+
+def numel(x):
+  return intprod(var_shape(x))
+
+
+def intprod(x):
+  return int(np.prod(x))
+
+
+def flatgrad(loss, var_list):
+  grads = tf.gradients(loss, var_list)
+  return tf.concat([tf.reshape(grad, [numel(v)], 0)
+                    for (v, grad) in zip(var_list, grads)])
+
+
+class SetFromFlat(object):
+  def __init__(self, var_list, dtype=tf.float32):
+    assigns = []
+    shapes = list(map(var_shape, var_list))
+    total_size = np.sum([intprod(shape) for shape in shapes])
+
+    self.theta = theta = tf.placeholder(dtype, [total_size])
+    start = 0
+    assigns = []
+    for (shape, v) in zip(shapes, var_list):
+      size = intprod(shape)
+      assigns.append(tf.assign(v, tf.reshape(theta[start:start + size],
+                                             shape)))
+      start += size
+    assert start == total_size
+    self.op = tf.group(*assigns)
+
+  def __call__(self, theta):
+    get_session().run(self.op, feed_dict={self.theta: theta})
+
+
+class GetFlat(object):
+  def __init__(self, var_list):
+    self.op = tf.concat([tf.reshape(v, [numel(v)]) for v in var_list], 0)
+
+  def __call__(self):
+    return get_session().run(self.op)
+
+# Misc
+
+
+def scope_vars(scope, trainable_only):
+  """Get variables inside a scope. The scope can be specified as a string."""
+  return tf.get_collection((tf.GraphKeys.TRAINABLE_VARIABLES if trainable_only
+                            else tf.GraphKeys.GLOBAL_VARIABLES),
+                           scope=(scope if isinstance(scope, str)
+                                  else scope.name))
+
+
+def in_session(f):
+  @functools.wraps(f)
+  def newfunc(*args, **kwargs):
+    with tf.Session():
+      f(*args, **kwargs)
+  return newfunc
+
+
+# A mapping from name -> (placeholder, dtype, shape).
+_PLACEHOLDER_CACHE = {}
+
+
+def get_placeholder(name, dtype, shape):
+  print("calling get_placeholder", name)
+  if name in _PLACEHOLDER_CACHE:
+    out, dtype1, shape1 = _PLACEHOLDER_CACHE[name]
+    assert dtype1 == dtype and shape1 == shape
+    return out
+  else:
+    out = tf.placeholder(dtype=dtype, shape=shape, name=name)
+    _PLACEHOLDER_CACHE[name] = (out, dtype, shape)
+    return out
+
+
+def get_placeholder_cached(name):
+  return _PLACEHOLDER_CACHE[name][0]
+
+
+def flattenallbut0(x):
+  return tf.reshape(x, [-1, intprod(x.get_shape().as_list()[1:])])
+
+
+def reset():
+  global _PLACEHOLDER_CACHE
+  global VARIABLES
+  _PLACEHOLDER_CACHE = {}
+  VARIABLES = {}
+  tf.reset_default_graph()

python/ray/rllib/evolution_strategies/utils.py

@@ -0,0 +1,86 @@
+# Code in this file is copied and adapted from
+# https://github.com/openai/evolution-strategies-starter.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import tensorflow as tf
+
+
+def compute_ranks(x):
+  """Returns ranks in [0, len(x))
+
+  Note: This is different from scipy.stats.rankdata, which returns ranks in
+  [1, len(x)].
+  """
+  assert x.ndim == 1
+  ranks = np.empty(len(x), dtype=int)
+  ranks[x.argsort()] = np.arange(len(x))
+  return ranks
+
+
+def compute_centered_ranks(x):
+  y = compute_ranks(x.ravel()).reshape(x.shape).astype(np.float32)
+  y /= (x.size - 1)
+  y -= 0.5
+  return y
+
+
+def make_session(single_threaded):
+  if not single_threaded:
+    return tf.InteractiveSession()
+  return tf.InteractiveSession(
+      config=tf.ConfigProto(inter_op_parallelism_threads=1,
+                            intra_op_parallelism_threads=1))
+
+
+def itergroups(items, group_size):
+  assert group_size >= 1
+  group = []
+  for x in items:
+    group.append(x)
+    if len(group) == group_size:
+      yield tuple(group)
+      del group[:]
+  if group:
+    yield tuple(group)
+
+
+def batched_weighted_sum(weights, vecs, batch_size):
+  total = 0
+  num_items_summed = 0
+  for batch_weights, batch_vecs in zip(itergroups(weights, batch_size),
+                                       itergroups(vecs, batch_size)):
+    assert len(batch_weights) == len(batch_vecs) <= batch_size
+    total += np.dot(np.asarray(batch_weights, dtype=np.float32),
+                    np.asarray(batch_vecs, dtype=np.float32))
+    num_items_summed += len(batch_weights)
+  return total, num_items_summed
+
+
+class RunningStat(object):
+  def __init__(self, shape, eps):
+    self.sum = np.zeros(shape, dtype=np.float32)
+    self.sumsq = np.full(shape, eps, dtype=np.float32)
+    self.count = eps
+
+  def increment(self, s, ssq, c):
+    self.sum += s
+    self.sumsq += ssq
+    self.count += c
+
+  @property
+  def mean(self):
+    return self.sum / self.count
+
+  @property
+  def std(self):
+    return np.sqrt(np.maximum(self.sumsq / self.count - np.square(self.mean),
+                              1e-2))
+
+  def set_from_init(self, init_mean, init_std, init_count):
+    self.sum[:] = init_mean * init_count
+    self.sumsq[:] = (np.square(init_mean) + np.square(init_std)) * init_count
+    self.count = init_count

python/ray/rllib/evolution_strategies/viz.py

@@ -0,0 +1,42 @@
+# Code in this file is copied and adapted from
+# https://github.com/openai/evolution-strategies-starter.
+
+import click
+
+
+@click.command()
+@click.argument("env_id")
+@click.argument("policy_file")
+@click.option("--record", is_flag=True)
+@click.option("--stochastic", is_flag=True)
+@click.option("--extra_kwargs")
+def main(env_id, policy_file, record, stochastic, extra_kwargs):
+  import gym
+  from gym import wrappers
+  import tensorflow as tf
+  from policies import MujocoPolicy
+  import numpy as np
+
+  env = gym.make(env_id)
+  if record:
+    import uuid
+    env = wrappers.Monitor(env, "/tmp/" + str(uuid.uuid4()), force=True)
+
+  if extra_kwargs:
+    import json
+    extra_kwargs = json.loads(extra_kwargs)
+
+  with tf.Session():
+    pi = MujocoPolicy.Load(policy_file, extra_kwargs=extra_kwargs)
+    while True:
+      rews, t = pi.rollout(env, render=True,
+                           random_stream=np.random if stochastic else None)
+      print("return={:.4f} len={}".format(rews.sum(), t))
+
+    if record:
+      env.close()
+      return
+
+
+if __name__ == "__main__":
+  main()

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

@@ -0,0 +1,278 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from collections import namedtuple
+
+import gym.spaces
+import tensorflow as tf
+import os
+
+from tensorflow.python.client import timeline
+from tensorflow.python import debug as tf_debug
+
+import ray
+
+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
+# 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.
+Tower = namedtuple("Tower", ["init_op", "grads", "policy"])
+
+
+class Agent(object):
+  """
+  Agent class that holds the simulator environment and the policy.
+
+  Initializes the tensorflow graphs for both training and evaluation.
+  One common policy graph is initialized on '/cpu:0' and holds all the shared
+  network weights. When run as a remote agent, only this graph is used.
+
+  When the agent is initialized locally with multiple GPU devices, copies of
+  the policy graph are also placed on each GPU. These per-GPU graphs share the
+  common policy network weights but take device-local input tensors.
+
+  The idea here is that training data can be bulk-loaded onto these
+  device-local variables. Synchronous SGD can then be run in parallel over
+  this GPU-local data.
+  """
+
+  def __init__(self, name, batchsize, preprocessor, config, is_remote):
+    if is_remote:
+      os.environ["CUDA_VISIBLE_DEVICES"] = ""
+      devices = ["/cpu:0"]
+    else:
+      devices = config["devices"]
+    self.devices = devices
+    self.config = config
+    self.env = BatchedEnv(name, batchsize, preprocessor=preprocessor)
+    if preprocessor.shape is None:
+      preprocessor.shape = self.env.observation_space.shape
+    if is_remote:
+      config_proto = tf.ConfigProto()
+    else:
+      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)
+    self.observations = tf.placeholder(tf.float32,
+                                       shape=(None,) + preprocessor.shape)
+    self.advantages = tf.placeholder(tf.float32, shape=(None,))
+
+    action_space = self.env.action_space
+    if isinstance(action_space, gym.spaces.Box):
+      # The first half of the dimensions are the means, the second half are the
+      # standard deviations.
+      self.action_dim = action_space.shape[0]
+      self.action_shape = (self.action_dim,)
+      self.logit_dim = 2 * self.action_dim
+      self.actions = tf.placeholder(tf.float32, shape=(None, self.action_dim))
+      self.distribution_class = DiagGaussian
+    elif isinstance(action_space, gym.spaces.Discrete):
+      self.action_dim = action_space.n
+      self.action_shape = ()
+      self.logit_dim = self.action_dim
+      self.actions = tf.placeholder(tf.int64, shape=(None,))
+      self.distribution_class = Categorical
+    else:
+      raise NotImplemented("action space" + str(type(action_space)) +
+                           "currently not supported")
+    self.prev_logits = tf.placeholder(tf.float32, shape=(None, self.logit_dim))
+
+    data_splits = zip(
+        tf.split(self.observations, len(devices)),
+        tf.split(self.advantages, len(devices)),
+        tf.split(self.actions, len(devices)),
+        tf.split(self.prev_logits, len(devices)))
+
+    # Parallel SGD ops
+    self.towers = []
+    self.batch_index = tf.placeholder(tf.int32)
+    assert config["sgd_batchsize"] % len(devices) == 0, \
+        "Batch size must be evenly divisible by devices"
+    if is_remote:
+      self.batch_size = 1
+      self.per_device_batch_size = 1
+    else:
+      self.batch_size = config["sgd_batchsize"]
+      self.per_device_batch_size = int(self.batch_size / len(devices))
+    self.optimizer = tf.train.AdamOptimizer(self.config["sgd_stepsize"])
+    self.setup_common_policy(
+        self.observations, self.advantages, self.actions, self.prev_logits)
+    for device, (obs, adv, acts, plog) in zip(devices, data_splits):
+      self.towers.append(
+          self.setup_per_device_policy(device, obs, adv, acts, plog))
+
+    avg = average_gradients([t.grads for t in self.towers])
+    self.train_op = self.optimizer.apply_gradients(avg)
+
+    # Metric ops
+    with tf.name_scope("test_outputs"):
+      self.mean_loss = tf.reduce_mean(
+          tf.stack(values=[t.policy.loss for t in self.towers]), 0)
+      self.mean_kl = tf.reduce_mean(
+          tf.stack(values=[t.policy.mean_kl for t in self.towers]), 0)
+      self.mean_entropy = tf.reduce_mean(
+          tf.stack(values=[t.policy.mean_entropy for t in self.towers]), 0)
+
+    # References to the model weights
+    self.variables = ray.experimental.TensorFlowVariables(
+        self.common_policy.loss,
+        self.sess)
+    self.observation_filter = MeanStdFilter(preprocessor.shape, clip=None)
+    self.reward_filter = MeanStdFilter((), clip=5.0)
+    self.sess.run(tf.global_variables_initializer())
+
+  def setup_common_policy(self, observations, advantages, actions, prev_log):
+    with tf.variable_scope("tower"):
+      self.common_policy = ProximalPolicyLoss(
+          self.env.observation_space, self.env.action_space,
+          observations, advantages, actions, prev_log, self.logit_dim,
+          self.kl_coeff, self.distribution_class, self.config, self.sess)
+
+  def setup_per_device_policy(
+          self, device, observations, advantages, actions, prev_log):
+    with tf.device(device):
+      with tf.variable_scope("tower", reuse=True):
+        all_obs = tf.Variable(
+            observations, trainable=False, validate_shape=False,
+            collections=[])
+        all_adv = tf.Variable(
+            advantages, trainable=False, validate_shape=False, collections=[])
+        all_acts = tf.Variable(
+            actions, trainable=False, validate_shape=False, collections=[])
+        all_plog = tf.Variable(
+            prev_log, trainable=False, validate_shape=False, collections=[])
+        obs_slice = tf.slice(
+            all_obs,
+            [self.batch_index] + [0] * len(self.preprocessor.shape),
+            [self.per_device_batch_size] + [-1] * len(self.preprocessor.shape))
+        obs_slice.set_shape(observations.shape)
+        adv_slice = tf.slice(
+            all_adv, [self.batch_index], [self.per_device_batch_size])
+        acts_slice = tf.slice(
+            all_acts,
+            [self.batch_index] + [0] * len(self.action_shape),
+            [self.per_device_batch_size] + [-1] * len(self.action_shape))
+        plog_slice = tf.slice(
+            all_plog, [self.batch_index, 0], [self.per_device_batch_size, -1])
+        policy = ProximalPolicyLoss(
+            self.env.observation_space, self.env.action_space,
+            obs_slice, adv_slice, acts_slice, plog_slice, self.logit_dim,
+            self.kl_coeff, self.distribution_class, self.config, self.sess)
+        grads = self.optimizer.compute_gradients(
+            policy.loss, colocate_gradients_with_ops=True)
+
+      return Tower(
+          tf.group(
+              *[all_obs.initializer,
+                all_adv.initializer,
+                all_acts.initializer,
+                all_plog.initializer]),
+          grads,
+          policy)
+
+  def load_data(self, trajectories, full_trace):
+    """
+    Bulk loads the specified trajectories into device memory.
+
+    The data is split equally across all the devices.
+
+    Returns:
+      The number of tuples loaded per device.
+    """
+
+    truncated_obs = make_divisible_by(
+        trajectories["observations"], self.batch_size)
+    if full_trace:
+      run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
+    else:
+      run_options = tf.RunOptions(trace_level=tf.RunOptions.NO_TRACE)
+    run_metadata = tf.RunMetadata()
+    self.sess.run(
+        [t.init_op for t in self.towers],
+        feed_dict={
+            self.observations: truncated_obs,
+            self.advantages: make_divisible_by(
+                trajectories["advantages"], self.batch_size),
+            self.actions: make_divisible_by(
+                trajectories["actions"].squeeze(), self.batch_size),
+            self.prev_logits: make_divisible_by(
+                trajectories["logprobs"], self.batch_size),
+        },
+        options=run_options,
+        run_metadata=run_metadata)
+    if full_trace:
+      trace = timeline.Timeline(step_stats=run_metadata.step_stats)
+      trace_file = open("/tmp/ray/timeline-load.json", "w")
+      trace_file.write(trace.generate_chrome_trace_format())
+
+    tuples_per_device = len(truncated_obs) / len(self.devices)
+    assert tuples_per_device % self.per_device_batch_size == 0
+    return tuples_per_device
+
+  def run_sgd_minibatch(self, batch_index, kl_coeff, full_trace, file_writer):
+    """
+    Run a single step of SGD.
+
+    Runs a SGD step over the batch with index batch_index as created by
+    load_rollouts_data(), updating local weights.
+
+    Returns:
+      (mean_loss, mean_kl, mean_entropy) evaluated over the batch.
+    """
+
+    if full_trace:
+      run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
+    else:
+      run_options = tf.RunOptions(trace_level=tf.RunOptions.NO_TRACE)
+    run_metadata = tf.RunMetadata()
+
+    _, loss, kl, entropy = self.sess.run(
+        [self.train_op, self.mean_loss, self.mean_kl, self.mean_entropy],
+        feed_dict={
+            self.batch_index: batch_index,
+            self.kl_coeff: kl_coeff},
+        options=run_options,
+        run_metadata=run_metadata)
+
+    if full_trace:
+      trace = timeline.Timeline(step_stats=run_metadata.step_stats)
+      trace_file = open("/tmp/ray/timeline-sgd.json", "w")
+      trace_file.write(trace.generate_chrome_trace_format())
+      file_writer.add_run_metadata(
+          run_metadata, "sgd_train_{}".format(batch_index))
+
+    return loss, kl, entropy
+
+  def get_weights(self):
+    return self.variables.get_weights()
+
+  def load_weights(self, weights):
+    self.variables.set_weights(weights)
+
+  def compute_trajectory(self, gamma, lam, horizon):
+    trajectory = rollouts(
+        self.common_policy,
+        self.env, horizon, self.observation_filter, self.reward_filter)
+    add_advantage_values(trajectory, gamma, lam, self.reward_filter)
+    return trajectory
+
+
+RemoteAgent = ray.remote(Agent)

python/ray/rllib/policy_gradient/distributions.py

@@ -0,0 +1,69 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import numpy as np
+
+
+class Categorical(object):
+  def __init__(self, logits):
+    self.logits = logits
+
+  def logp(self, x):
+    return -tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.logits,
+                                                           labels=x)
+
+  def entropy(self):
+    a0 = self.logits - tf.reduce_max(self.logits, reduction_indices=[1],
+                                     keep_dims=True)
+    ea0 = tf.exp(a0)
+    z0 = tf.reduce_sum(ea0, reduction_indices=[1], keep_dims=True)
+    p0 = ea0 / z0
+    return tf.reduce_sum(p0 * (tf.log(z0) - a0), reduction_indices=[1])
+
+  def kl(self, other):
+    a0 = self.logits - tf.reduce_max(self.logits, reduction_indices=[1],
+                                     keep_dims=True)
+    a1 = other.logits - tf.reduce_max(other.logits, reduction_indices=[1],
+                                      keep_dims=True)
+    ea0 = tf.exp(a0)
+    ea1 = tf.exp(a1)
+    z0 = tf.reduce_sum(ea0, reduction_indices=[1], keep_dims=True)
+    z1 = tf.reduce_sum(ea1, reduction_indices=[1], keep_dims=True)
+    p0 = ea0 / z0
+    return tf.reduce_sum(p0 * (a0 - tf.log(z0) - a1 + tf.log(z1)),
+                         reduction_indices=[1])
+
+  def sample(self):
+    return tf.multinomial(self.logits, 1)
+
+
+class DiagGaussian(object):
+  def __init__(self, flat):
+    self.flat = flat
+    mean, logstd = tf.split(flat, 2, axis=1)
+    self.mean = mean
+    self.logstd = logstd
+    self.std = tf.exp(logstd)
+
+  def logp(self, x):
+    return (-0.5 * tf.reduce_sum(tf.square((x - self.mean) / self.std),
+                                 reduction_indices=[1]) -
+            0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) -
+            tf.reduce_sum(self.logstd, reduction_indices=[1]))
+
+  def kl(self, other):
+    assert isinstance(other, DiagGaussian)
+    return tf.reduce_sum(other.logstd - self.logstd +
+                         (tf.square(self.std) +
+                          tf.square(self.mean - other.mean)) /
+                         (2.0 * tf.square(other.std)) - 0.5,
+                         reduction_indices=[1])
+
+  def entropy(self):
+    return tf.reduce_sum(self.logstd + .5 * np.log(2.0 * np.pi * np.e),
+                         reduction_indices=[1])
+
+  def sample(self):
+    return self.mean + self.std * tf.random_normal(tf.shape(self.mean))

python/ray/rllib/policy_gradient/env.py

@@ -0,0 +1,65 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+import numpy as np
+
+
+class AtariPixelPreprocessor(object):
+  def __init__(self):
+    self.shape = (80, 80, 3)
+
+  def __call__(self, observation):
+    "Convert images from (210, 160, 3) to (3, 80, 80) by downsampling."
+    return (observation[25:-25:2, ::2, :][None] - 128) / 128
+
+
+class AtariRamPreprocessor(object):
+  def __init__(self):
+    self.shape = (128,)
+
+  def __call__(self, observation):
+    return (observation - 128) / 128
+
+
+class NoPreprocessor(object):
+  def __init__(self):
+    self.shape = None
+
+  def __call__(self, observation):
+    return observation
+
+
+class BatchedEnv(object):
+  """This holds multiple gym enviroments and performs steps on all of them."""
+  def __init__(self, name, batchsize, preprocessor=None):
+    self.envs = [gym.make(name) for _ in range(batchsize)]
+    self.observation_space = self.envs[0].observation_space
+    self.action_space = self.envs[0].action_space
+    self.batchsize = batchsize
+    self.preprocessor = preprocessor if preprocessor else lambda obs: obs[None]
+
+  def reset(self):
+    observations = [self.preprocessor(env.reset()) for env in self.envs]
+    self.shape = observations[0].shape
+    self.dones = [False for _ in range(self.batchsize)]
+    return np.vstack(observations)
+
+  def step(self, actions, render=False):
+    observations = []
+    rewards = []
+    for i, action in enumerate(actions):
+      if self.dones[i]:
+        observations.append(np.zeros(self.shape))
+        rewards.append(0.0)
+        continue
+      observation, reward, done, info = self.envs[i].step(
+          action if len(action) > 1 else action[0])
+      if render:
+        self.envs[0].render()
+      observations.append(self.preprocessor(observation))
+      rewards.append(reward)
+      self.dones[i] = done
+    return (np.vstack(observations), np.array(rewards, dtype="float32"),
+            np.array(self.dones))

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/filter.py

@@ -0,0 +1,133 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+
+class NoFilter(object):
+  def __init__(self):
+    pass
+
+  def __call__(self, x, update=True):
+    return np.asarray(x)
+
+
+# http://www.johndcook.com/blog/standard_deviation/
+class RunningStat(object):
+
+  def __init__(self, shape=None):
+    self._n = 0
+    self._M = np.zeros(shape)
+    self._S = np.zeros(shape)
+
+  def push(self, x):
+    x = np.asarray(x)
+    # Unvectorized update of the running statistics.
+    assert x.shape == self._M.shape, ("x.shape = {}, self.shape = {}"
+                                      .format(x.shape, self._M.shape))
+    n1 = self._n
+    self._n += 1
+    if self._n == 1:
+      self._M[...] = x
+    else:
+      delta = x - self._M
+      self._M[...] += delta / self._n
+      self._S[...] += delta * delta * n1 / self._n
+
+  def update(self, other):
+    n1 = self._n
+    n2 = other._n
+    n = n1 + n2
+    delta = self._M - other._M
+    delta2 = delta * delta
+    M = (n1 * self._M + n2 * other._M) / n
+    S = self._S + other._S + delta2 * n1 * n2 / n
+    self._n = n
+    self._M = M
+    self._S = S
+
+  @property
+  def n(self):
+    return self._n
+
+  @property
+  def mean(self):
+    return self._M
+
+  @property
+  def var(self):
+    return self._S / (self._n - 1) if self._n > 1 else np.square(self._M)
+
+  @property
+  def std(self):
+    return np.sqrt(self.var)
+
+  @property
+  def shape(self):
+    return self._M.shape
+
+
+class MeanStdFilter(object):
+  def __init__(self, shape, demean=True, destd=True, clip=10.0):
+    self.demean = demean
+    self.destd = destd
+    self.clip = clip
+
+    self.rs = RunningStat(shape)
+
+  def __call__(self, x, update=True):
+    x = np.asarray(x)
+    if update:
+      if len(x.shape) == len(self.rs.shape) + 1:
+        # The vectorized case.
+        for i in range(x.shape[0]):
+          self.rs.push(x[i])
+      else:
+        # The unvectorized case.
+        self.rs.push(x)
+    if self.demean:
+      x = x - self.rs.mean
+    if self.destd:
+      x = x / (self.rs.std + 1e-8)
+    if self.clip:
+      x = np.clip(x, -self.clip, self.clip)
+    return x
+
+
+def test_running_stat():
+  for shp in ((), (3,), (3, 4)):
+    li = []
+    rs = RunningStat(shp)
+    for _ in range(5):
+      val = np.random.randn(*shp)
+      rs.push(val)
+      li.append(val)
+      m = np.mean(li, axis=0)
+      assert np.allclose(rs.mean, m)
+      v = np.square(m) if (len(li) == 1) else np.var(li, ddof=1, axis=0)
+      assert np.allclose(rs.var, v)
+
+
+def test_combining_stat():
+  for shape in [(), (3,), (3, 4)]:
+    li = []
+    rs1 = RunningStat(shape)
+    rs2 = RunningStat(shape)
+    rs = RunningStat(shape)
+    for _ in range(5):
+      val = np.random.randn(*shape)
+      rs1.push(val)
+      rs.push(val)
+      li.append(val)
+    for _ in range(9):
+      rs2.push(val)
+      rs.push(val)
+      li.append(val)
+    rs1.update(rs2)
+    assert np.allclose(rs.mean, rs1.mean)
+    assert np.allclose(rs.std, rs1.std)
+
+
+test_running_stat()
+test_combining_stat()

python/ray/rllib/policy_gradient/loss.py

@@ -0,0 +1,52 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym.spaces
+import tensorflow as tf
+from ray.rllib.policy_gradient.models.visionnet import vision_net
+from ray.rllib.policy_gradient.models.fcnet import fc_net
+
+
+class ProximalPolicyLoss(object):
+
+  def __init__(
+          self, observation_space, action_space,
+          observations, advantages, actions, prev_logits, logit_dim,
+          kl_coeff, distribution_class, config, sess):
+    assert (isinstance(action_space, gym.spaces.Discrete) or
+            isinstance(action_space, gym.spaces.Box))
+    self.prev_dist = distribution_class(prev_logits)
+
+    # Saved so that we can compute actions given different observations
+    self.observations = observations
+
+    if len(observation_space.shape) > 1:
+      self.curr_logits = vision_net(observations, num_classes=logit_dim)
+    else:
+      assert len(observation_space.shape) == 1
+      self.curr_logits = fc_net(observations, num_classes=logit_dim)
+    self.curr_dist = distribution_class(self.curr_logits)
+    self.sampler = self.curr_dist.sample()
+
+    # Make loss functions.
+    self.ratio = tf.exp(self.curr_dist.logp(actions) -
+                        self.prev_dist.logp(actions))
+    self.kl = self.prev_dist.kl(self.curr_dist)
+    self.mean_kl = tf.reduce_mean(self.kl)
+    self.entropy = self.curr_dist.entropy()
+    self.mean_entropy = tf.reduce_mean(self.entropy)
+    self.surr1 = self.ratio * advantages
+    self.surr2 = tf.clip_by_value(self.ratio, 1 - config["clip_param"],
+                                  1 + config["clip_param"]) * advantages
+    self.surr = tf.minimum(self.surr1, self.surr2)
+    self.loss = tf.reduce_mean(-self.surr + kl_coeff * self.kl -
+                               config["entropy_coeff"] * self.entropy)
+    self.sess = sess
+
+  def compute_actions(self, observations):
+    return self.sess.run([self.sampler, self.curr_logits],
+                         feed_dict={self.observations: observations})
+
+  def loss(self):
+    return self.loss

python/ray/rllib/policy_gradient/models/fcnet.py

@@ -0,0 +1,38 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import tensorflow.contrib.slim as slim
+
+import numpy as np
+
+
+def normc_initializer(std=1.0):
+  def _initializer(shape, dtype=None, partition_info=None):
+    out = np.random.randn(*shape).astype(np.float32)
+    out *= std / np.sqrt(np.square(out).sum(axis=0, keepdims=True))
+    return tf.constant(out)
+  return _initializer
+
+
+def fc_net(inputs, num_classes=10, logstd=False):
+  with tf.name_scope("fc_net"):
+    fc1 = slim.fully_connected(inputs, 128,
+                               weights_initializer=normc_initializer(1.0),
+                               scope="fc1")
+    fc2 = slim.fully_connected(fc1, 128,
+                               weights_initializer=normc_initializer(1.0),
+                               scope="fc2")
+    fc3 = slim.fully_connected(fc2, 128,
+                               weights_initializer=normc_initializer(1.0),
+                               scope="fc3")
+    fc4 = slim.fully_connected(fc3, num_classes,
+                               weights_initializer=normc_initializer(0.01),
+                               activation_fn=None, scope="fc4")
+    if logstd:
+      logstd = tf.get_variable(name="logstd", shape=[num_classes],
+                               initializer=tf.zeros_initializer)
+      return tf.concat(1, [fc4, logstd])
+    else:
+      return fc4

python/ray/rllib/policy_gradient/models/visionnet.py

@@ -0,0 +1,16 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import tensorflow.contrib.slim as slim
+
+
+def vision_net(inputs, num_classes=10):
+  with tf.name_scope("vision_net"):
+    conv1 = slim.conv2d(inputs, 16, [8, 8], 4, scope="conv1")
+    conv2 = slim.conv2d(conv1, 32, [4, 4], 2, scope="conv2")
+    fc1 = slim.conv2d(conv2, 512, [10, 10], padding="VALID", scope="fc1")
+    fc2 = slim.conv2d(fc1, num_classes, [1, 1], activation_fn=None,
+                      normalizer_fn=None, scope="fc2")
+    return tf.squeeze(fc2, [1, 2])

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

@@ -0,0 +1,103 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import ray
+
+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(),
+             reward_filter=NoFilter()):
+  """Perform a batch of rollouts of a policy in an environment.
+
+  Args:
+    policy: The policy that will be rollout out. Can be an arbitrary object
+      that supports a compute_actions(observation) function.
+    env: The environment the rollout is computed in. Needs to support the
+      OpenAI gym API and needs to support batches of data.
+    horizon: Upper bound for the number of timesteps for each rollout in the
+      batch.
+    observation_filter: Function that is applied to each of the observations.
+    reward_filter: Function that is applied to each of the rewards.
+
+  Returns:
+    A trajectory, which is a dictionary with keys "observations", "rewards",
+      "orig_rewards", "actions", "logprobs", "dones". Each value is an array of
+      shape (num_timesteps, env.batchsize, shape).
+  """
+
+  observation = observation_filter(env.reset())
+  done = np.array(env.batchsize * [False])
+  t = 0
+  observations = []
+  raw_rewards = []  # Empirical rewards
+  actions = []
+  logprobs = []
+  dones = []
+
+  while not done.all() and t < horizon:
+    action, logprob = policy.compute_actions(observation)
+    observations.append(observation[None])
+    actions.append(action[None])
+    logprobs.append(logprob[None])
+    observation, raw_reward, done = env.step(action)
+    observation = observation_filter(observation)
+    raw_rewards.append(raw_reward[None])
+    dones.append(done[None])
+    t += 1
+
+  return {"observations": np.vstack(observations),
+          "raw_rewards": np.vstack(raw_rewards),
+          "actions": np.vstack(actions),
+          "logprobs": np.vstack(logprobs),
+          "dones": np.vstack(dones)}
+
+
+def add_advantage_values(trajectory, gamma, lam, reward_filter):
+  rewards = trajectory["raw_rewards"]
+  dones = trajectory["dones"]
+  advantages = np.zeros_like(rewards)
+  last_advantage = np.zeros(rewards.shape[1], dtype="float32")
+
+  for t in reversed(range(len(rewards))):
+    delta = rewards[t, :] * (1 - dones[t, :])
+    last_advantage = delta + gamma * lam * last_advantage
+    advantages[t, :] = last_advantage
+    reward_filter(advantages[t, :])
+
+  trajectory["advantages"] = advantages
+
+
+@ray.remote
+def compute_trajectory(policy, env, gamma, lam, horizon, observation_filter,
+                       reward_filter):
+  trajectory = rollouts(policy, env, horizon, observation_filter,
+                        reward_filter)
+  add_advantage_values(trajectory, gamma, lam, reward_filter)
+  return trajectory
+
+
+def collect_samples(agents, num_timesteps, gamma, lam, horizon,
+                    observation_filter=NoFilter(), reward_filter=NoFilter()):
+  num_timesteps_so_far = 0
+  trajectories = []
+  total_rewards = []
+  traj_len_means = []
+  while num_timesteps_so_far < num_timesteps:
+    trajectory_batch = ray.get(
+        [agent.compute_trajectory.remote(gamma, lam, horizon)
+         for agent in agents])
+    trajectory = concatenate(trajectory_batch)
+    trajectory = flatten(trajectory)
+    not_done = np.logical_not(trajectory["dones"])
+    total_rewards.append(
+        trajectory["raw_rewards"][not_done].sum(axis=0).mean() / len(agents))
+    traj_len_means.append(not_done.sum(axis=0).mean() / len(agents))
+    trajectory = {key: val[not_done] for key, val in trajectory.items()}
+    num_timesteps_so_far += len(trajectory["dones"])
+    trajectories.append(trajectory)
+  return (concatenate(trajectories), np.mean(total_rewards),
+          np.mean(traj_len_means))

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

@@ -0,0 +1,61 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import unittest
+import numpy as np
+import tensorflow as tf
+from numpy.testing import assert_allclose
+
+from ray.rllib.policy_gradient.distributions import Categorical
+from ray.rllib.policy_gradient.utils import flatten, concatenate
+
+
+class DistibutionsTest(unittest.TestCase):
+
+  def testCategorical(self):
+    num_samples = 100000
+    logits = tf.placeholder(tf.float32, shape=(None, 10))
+    z = 8 * (np.random.rand(10) - 0.5)
+    data = np.tile(z, (num_samples, 1))
+    c = Categorical(logits)
+    sample_op = c.sample()
+    sess = tf.Session()
+    sess.run(tf.global_variables_initializer())
+    samples = sess.run(sample_op, feed_dict={logits: data})
+    counts = np.zeros(10)
+    for sample in samples:
+      counts[sample] += 1.0
+    probs = np.exp(z) / np.sum(np.exp(z))
+    self.assertTrue(np.sum(np.abs(probs - counts / num_samples)) <= 0.01)
+
+
+class UtilsTest(unittest.TestCase):
+
+  def testFlatten(self):
+    d = {"s": np.array([[[1, -1], [2, -2]], [[3, -3], [4, -4]]]),
+         "a": np.array([[[5], [-5]], [[6], [-6]]])}
+    flat = flatten(d.copy(), start=0, stop=2)
+    assert_allclose(d["s"][0][0][:], flat["s"][0][:])
+    assert_allclose(d["s"][0][1][:], flat["s"][1][:])
+    assert_allclose(d["s"][1][0][:], flat["s"][2][:])
+    assert_allclose(d["s"][1][1][:], flat["s"][3][:])
+    assert_allclose(d["a"][0][0], flat["a"][0])
+    assert_allclose(d["a"][0][1], flat["a"][1])
+    assert_allclose(d["a"][1][0], flat["a"][2])
+    assert_allclose(d["a"][1][1], flat["a"][3])
+
+  def testConcatenate(self):
+    d1 = {"s": np.array([0, 1]), "a": np.array([2, 3])}
+    d2 = {"s": np.array([4, 5]), "a": np.array([6, 7])}
+    d = concatenate([d1, d2])
+    assert_allclose(d["s"], np.array([0, 1, 4, 5]))
+    assert_allclose(d["a"], np.array([2, 3, 6, 7]))
+
+    D = concatenate([d])
+    assert_allclose(D["s"], np.array([0, 1, 4, 5]))
+    assert_allclose(D["a"], np.array([2, 3, 6, 7]))
+
+
+if __name__ == "__main__":
+  unittest.main(verbosity=2)

python/ray/rllib/policy_gradient/utils.py

@@ -0,0 +1,88 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import tensorflow as tf
+
+
+def flatten(weights, start=0, stop=2):
+  """This methods reshapes all values in a dictionary.
+
+  The indices from start to stop will be flattened into a single index.
+
+  Args:
+    weights: A dictionary mapping keys to numpy arrays.
+    start: The starting index.
+    stop: The ending index.
+  """
+  for key, val in weights.items():
+    new_shape = val.shape[0:start] + (-1,) + val.shape[stop:]
+    weights[key] = val.reshape(new_shape)
+  return weights
+
+
+def concatenate(weights_list):
+  keys = weights_list[0].keys()
+  result = {}
+  for key in keys:
+    result[key] = np.concatenate([l[key] for l in weights_list])
+  return result
+
+
+def shuffle(trajectory):
+  permutation = np.random.permutation(trajectory["dones"].shape[0])
+  for key, val in trajectory.items():
+    trajectory[key] = val[permutation]
+  return trajectory
+
+
+def make_divisible_by(array, n):
+  return array[0:array.shape[0] - array.shape[0] % n]
+
+
+def average_gradients(tower_grads):
+  """
+  Average gradients across towers.
+
+  Calculate the average gradient for each shared variable across all towers.
+  Note that this function provides a synchronization point across all towers.
+
+  Args:
+    tower_grads: List of lists of (gradient, variable) tuples. The outer list
+      is over individual gradients. The inner list is over the gradient
+      calculation for each tower.
+
+  Returns:
+     List of pairs of (gradient, variable) where the gradient has been averaged
+     across all towers.
+
+  TODO(ekl): We could use NCCL if this becomes a bottleneck.
+  """
+
+  average_grads = []
+  for grad_and_vars in zip(*tower_grads):
+
+    # Note that each grad_and_vars looks like the following:
+    #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
+    grads = []
+    for g, _ in grad_and_vars:
+      if g is not None:
+        # Add 0 dimension to the gradients to represent the tower.
+        expanded_g = tf.expand_dims(g, 0)
+
+        # Append on a 'tower' dimension which we will average over below.
+        grads.append(expanded_g)
+
+    # Average over the 'tower' dimension.
+    grad = tf.concat(axis=0, values=grads)
+    grad = tf.reduce_mean(grad, 0)
+
+    # Keep in mind that the Variables are redundant because they are shared
+    # across towers. So .. we will just return the first tower's pointer to
+    # the Variable.
+    v = grad_and_vars[0][1]
+    grad_and_var = (grad, v)
+    average_grads.append(grad_and_var)
+
+  return average_grads