[rllib] Document "v2" APIs (#2316)

* re

* wip

* wip

* a3c working

* torch support

* pg works

* lint

* rm v2

* consumer id

* clean up pg

* clean up more

* fix python 2.7

* tf session management

* docs

* dqn wip

* fix compile

* dqn

* apex runs

* up

* impotrs

* ddpg

* quotes

* fix tests

* fix last r

* fix tests

* lint

* pass checkpoint restore

* kwar

* nits

* policy graph

* fix yapf

* com

* class

* pyt

* vectorization

* update

* test cpe

* unit test

* fix ddpg2

* changes

* wip

* args

* faster test

* common

* fix

* add alg option

* batch mode and policy serving

* multi serving test

* todo

* wip

* serving test

* doc async env

* num envs

* comments

* thread

* remove init hook

* update

* fix ppo

* comments1

* fix

* updates

* add jenkins tests

* fix

* fix pytorch

* fix

* fixes

* fix a3c policy

* fix squeeze

* fix trunc on apex

* fix squeezing for real

* update

* remove horizon test for now

* multiagent wip

* update

* fix race condition

* fix ma

* t

* doc

* st

* wip

* example

* wip

* working

* cartpole

* wip

* batch wip

* fix bug

* make other_batches None default

* working

* debug

* nit

* warn

* comments

* fix ppo

* fix obs filter

* update

* wip

* tf

* update

* fix

* cleanup

* cleanup

* spacing

* model

* fix

* dqn

* fix ddpg

* doc

* keep names

* update

* fix

* com

* docs

* clarify model outputs

* Update torch_policy_graph.py

* fix obs filter

* pass thru worker index

* fix

* rename

* vlad torch comments

* fix log action

* debug name

* fix lstm

* remove unused ddpg net

* remove conv net

* revert lstm

* wip

* wip

* cast

* wip

* works

* fix a3c

* works

* lstm util test

* doc

* clean up

* update

* fix lstm check

* move to end

* fix sphinx

* fix cmd

* remove bad doc

* envs

* vec

* doc prep

* models

* rl

* alg

* up

* clarify

* copy

* async sa

* fix

* comments

* fix a3c conf

* tune lstm

* fix reshape

* fix

* back to 16

* tuned a3c update

* update

* tuned

* optional

* merge

* wip

* fix up

* move pg class

* rename env

* wip

* update

* tip

* alg

* readme

* fix catalog

* readme

* doc

* context

* remove prep

* comma

* add env

* link to paper

* paper

* update

* rnn

* update

* wip

* clean up ev creation

* fix

* fix

* fix

* fix lint

* up

* no comma

* ma

* Update run_multi_node_tests.sh

* fix

* sphinx is stupid

* sphinx is stupid

* clarify torch graph

* no horizon

* fix config

* sb

* Update test_optimizers.py

python/ray/rllib/agents/__init__.py

@@ -0,0 +1,3 @@
+from ray.rllib.agents.agent import Agent, with_common_config
+
+__all__ = ["Agent", "with_common_config"]

python/ray/rllib/agents/a3c/__init__.py

@@ -0,0 +1,3 @@
+from ray.rllib.agents.a3c.a3c import A3CAgent, DEFAULT_CONFIG
+
+__all__ = ["A3CAgent", "DEFAULT_CONFIG"]

python/ray/rllib/agents/a3c/a3c.py

@@ -0,0 +1,132 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import pickle
+import os
+
+import ray
+from ray.rllib.agents.agent import Agent, with_common_config
+from ray.rllib.optimizers import AsyncGradientsOptimizer
+from ray.rllib.utils import FilterManager
+from ray.rllib.evaluation.metrics import collect_metrics
+from ray.tune.trial import Resources
+
+DEFAULT_CONFIG = with_common_config({
+    # Size of rollout batch
+    "sample_batch_size": 10,
+    # Use PyTorch as backend - no LSTM support
+    "use_pytorch": False,
+    # GAE(gamma) parameter
+    "lambda": 1.0,
+    # Max global norm for each gradient calculated by worker
+    "grad_clip": 40.0,
+    # Learning rate
+    "lr": 0.0001,
+    # Value Function Loss coefficient
+    "vf_loss_coeff": 0.5,
+    # Entropy coefficient
+    "entropy_coeff": -0.01,
+    # Whether to place workers on GPUs
+    "use_gpu_for_workers": False,
+    # Whether to emit extra summary stats
+    "summarize": False,
+    # Workers sample async
+    "sample_async": True,
+    # Model and preprocessor options
+    "model": {
+        # Use LSTM model. Requires TF.
+        "use_lstm": False,
+        # Max seq length for LSTM training.
+        "max_seq_len": 20,
+        # (Image statespace) - Converts image to Channels = 1
+        "grayscale": True,
+        # (Image statespace) - Each pixel
+        "zero_mean": False,
+        # (Image statespace) - Converts image to (dim, dim, C)
+        "dim": 80,
+        # (Image statespace) - Converts image shape to (C, dim, dim)
+        "channel_major": False,
+    },
+    # Configure TF for single-process operation
+    "tf_session_args": {
+        "intra_op_parallelism_threads": 1,
+        "inter_op_parallelism_threads": 1,
+        "gpu_options": {
+            "allow_growth": True,
+        },
+    },
+    # Arguments to pass to the rllib optimizer
+    "optimizer": {
+        # Number of gradients applied for each `train` step
+        "grads_per_step": 100,
+    },
+})
+
+
+class A3CAgent(Agent):
+    """A3C implementations in TensorFlow and PyTorch."""
+
+    _agent_name = "A3C"
+    _default_config = DEFAULT_CONFIG
+
+    @classmethod
+    def default_resource_request(cls, config):
+        cf = dict(cls._default_config, **config)
+        return Resources(
+            cpu=1,
+            gpu=0,
+            extra_cpu=cf["num_workers"],
+            extra_gpu=cf["use_gpu_for_workers"] and cf["num_workers"] or 0)
+
+    def _init(self):
+        if self.config["use_pytorch"]:
+            from ray.rllib.agents.a3c.a3c_torch_policy import \
+                A3CTorchPolicyGraph
+            policy_cls = A3CTorchPolicyGraph
+        else:
+            from ray.rllib.agents.a3c.a3c_tf_policy import A3CPolicyGraph
+            policy_cls = A3CPolicyGraph
+
+        self.local_evaluator = self.make_local_evaluator(
+            self.env_creator, policy_cls)
+        self.remote_evaluators = self.make_remote_evaluators(
+            self.env_creator, policy_cls, self.config["num_workers"],
+            {"num_gpus": 1 if self.config["use_gpu_for_workers"] else 0})
+        self.optimizer = AsyncGradientsOptimizer(
+            self.config["optimizer"], self.local_evaluator,
+            self.remote_evaluators)
+
+    def _train(self):
+        self.optimizer.step()
+        FilterManager.synchronize(
+            self.local_evaluator.filters, self.remote_evaluators)
+        result = collect_metrics(self.local_evaluator, self.remote_evaluators)
+        result = result._replace(
+            info=self.optimizer.stats())
+        return result
+
+    def _stop(self):
+        # workaround for https://github.com/ray-project/ray/issues/1516
+        for ev in self.remote_evaluators:
+            ev.__ray_terminate__.remote()
+
+    def _save(self, checkpoint_dir):
+        checkpoint_path = os.path.join(checkpoint_dir,
+                                       "checkpoint-{}".format(self.iteration))
+        agent_state = ray.get(
+            [a.save.remote() for a in self.remote_evaluators])
+        extra_data = {
+            "remote_state": agent_state,
+            "local_state": self.local_evaluator.save()
+        }
+        pickle.dump(extra_data, open(checkpoint_path + ".extra_data", "wb"))
+        return checkpoint_path
+
+    def _restore(self, checkpoint_path):
+        extra_data = pickle.load(open(checkpoint_path + ".extra_data", "rb"))
+        ray.get([
+            a.restore.remote(o)
+            for a, o in zip(self.remote_evaluators, extra_data["remote_state"])
+        ])
+        self.local_evaluator.restore(extra_data["local_state"])

python/ray/rllib/agents/a3c/a3c_tf_policy.py

@@ -0,0 +1,140 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import gym
+
+import ray
+from ray.rllib.utils.error import UnsupportedSpaceException
+from ray.rllib.evaluation.postprocessing import compute_advantages
+from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph
+from ray.rllib.models.misc import linear, normc_initializer
+from ray.rllib.models.catalog import ModelCatalog
+
+
+class A3CLoss(object):
+    def __init__(
+            self, action_dist, actions, advantages, v_target, vf,
+            vf_loss_coeff=0.5, entropy_coeff=-0.01):
+        log_prob = action_dist.logp(actions)
+
+        # The "policy gradients" loss
+        self.pi_loss = - tf.reduce_sum(log_prob * advantages)
+
+        delta = vf - v_target
+        self.vf_loss = 0.5 * tf.reduce_sum(tf.square(delta))
+        self.entropy = tf.reduce_sum(action_dist.entropy())
+        self.total_loss = (self.pi_loss +
+                           self.vf_loss * vf_loss_coeff +
+                           self.entropy * entropy_coeff)
+
+
+class A3CPolicyGraph(TFPolicyGraph):
+    def __init__(self, observation_space, action_space, config):
+        config = dict(ray.rllib.agents.a3c.a3c.DEFAULT_CONFIG, **config)
+        self.config = config
+        self.sess = tf.get_default_session()
+
+        # Setup the policy
+        self.observations = tf.placeholder(
+            tf.float32, [None] + list(observation_space.shape))
+        dist_class, logit_dim = ModelCatalog.get_action_dist(
+            action_space, self.config["model"])
+        self.model = ModelCatalog.get_model(
+            self.observations, logit_dim, self.config["model"])
+        action_dist = dist_class(self.model.outputs)
+        self.vf = tf.reshape(
+            linear(self.model.last_layer, 1, "value", normc_initializer(1.0)),
+            [-1])
+        self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
+                                          tf.get_variable_scope().name)
+        is_training = tf.placeholder_with_default(True, ())
+
+        # Setup the policy loss
+        if isinstance(action_space, gym.spaces.Box):
+            ac_size = action_space.shape[0]
+            actions = tf.placeholder(tf.float32, [None, ac_size], name="ac")
+        elif isinstance(action_space, gym.spaces.Discrete):
+            actions = tf.placeholder(tf.int64, [None], name="ac")
+        else:
+            raise UnsupportedSpaceException(
+                "Action space {} is not supported for A3C.".format(
+                    action_space))
+        advantages = tf.placeholder(tf.float32, [None], name="advantages")
+        v_target = tf.placeholder(tf.float32, [None], name="v_target")
+        self.loss = A3CLoss(
+            action_dist, actions, advantages, v_target, self.vf,
+            self.config["vf_loss_coeff"], self.config["entropy_coeff"])
+
+        # Initialize TFPolicyGraph
+        loss_in = [
+            ("obs", self.observations),
+            ("actions", actions),
+            ("advantages", advantages),
+            ("value_targets", v_target),
+        ]
+        for i, ph in enumerate(self.model.state_in):
+            loss_in.append(("state_in_{}".format(i), ph))
+        self.state_in = self.model.state_in
+        self.state_out = self.model.state_out
+        TFPolicyGraph.__init__(
+            self, observation_space, action_space, self.sess,
+            obs_input=self.observations, action_sampler=action_dist.sample(),
+            loss=self.loss.total_loss, loss_inputs=loss_in,
+            is_training=is_training, state_inputs=self.state_in,
+            state_outputs=self.state_out,
+            seq_lens=self.model.seq_lens,
+            max_seq_len=self.config["model"]["max_seq_len"])
+
+        if self.config.get("summarize"):
+            bs = tf.to_float(tf.shape(self.observations)[0])
+            tf.summary.scalar("model/policy_graph", self.loss.pi_loss / bs)
+            tf.summary.scalar("model/value_loss", self.loss.vf_loss / bs)
+            tf.summary.scalar("model/entropy", self.loss.entropy / bs)
+            tf.summary.scalar("model/grad_gnorm", tf.global_norm(self._grads))
+            tf.summary.scalar("model/var_gnorm", tf.global_norm(self.var_list))
+            self.summary_op = tf.summary.merge_all()
+
+        self.sess.run(tf.global_variables_initializer())
+
+    def extra_compute_action_fetches(self):
+        return {"vf_preds": self.vf}
+
+    def value(self, ob, *args):
+        feed_dict = {self.observations: [ob]}
+        assert len(args) == len(self.state_in), (args, self.state_in)
+        for k, v in zip(self.state_in, args):
+            feed_dict[k] = v
+        vf = self.sess.run(self.vf, feed_dict)
+        return vf[0]
+
+    def optimizer(self):
+        return tf.train.AdamOptimizer(self.config["lr"])
+
+    def gradients(self, optimizer):
+        grads = tf.gradients(self.loss.total_loss, self.var_list)
+        self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"])
+        clipped_grads = list(zip(self.grads, self.var_list))
+        return clipped_grads
+
+    def extra_compute_grad_fetches(self):
+        if self.config.get("summarize"):
+            return {"summary": self.summary_op}
+        else:
+            return {}
+
+    def get_initial_state(self):
+        return self.model.state_init
+
+    def postprocess_trajectory(self, sample_batch, other_agent_batches=None):
+        completed = sample_batch["dones"][-1]
+        if completed:
+            last_r = 0.0
+        else:
+            next_state = []
+            for i in range(len(self.state_in)):
+                next_state.append([sample_batch["state_out_{}".format(i)][-1]])
+            last_r = self.value(sample_batch["new_obs"][-1], *next_state)
+        return compute_advantages(
+            sample_batch, last_r, self.config["gamma"], self.config["lambda"])

python/ray/rllib/agents/a3c/a3c_torch_policy.py

@@ -0,0 +1,79 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+import ray
+from ray.rllib.models.pytorch.misc import var_to_np
+from ray.rllib.models.catalog import ModelCatalog
+from ray.rllib.evaluation.postprocessing import compute_advantages
+from ray.rllib.evaluation.torch_policy_graph import TorchPolicyGraph
+
+
+class A3CLoss(nn.Module):
+    def __init__(self, policy_model, vf_loss_coeff=0.5, entropy_coeff=-0.01):
+        nn.Module.__init__(self)
+        self.policy_model = policy_model
+        self.vf_loss_coeff = vf_loss_coeff
+        self.entropy_coeff = entropy_coeff
+
+    def forward(self, observations, actions, advantages, value_targets):
+        logits, values = self.policy_model(observations)
+        log_probs = F.log_softmax(logits, dim=1)
+        probs = F.softmax(logits, dim=1)
+        action_log_probs = log_probs.gather(1, actions.view(-1, 1))
+        entropy = -(log_probs * probs).sum(-1).sum()
+        pi_err = -advantages.dot(action_log_probs.reshape(-1))
+        value_err = F.mse_loss(values.reshape(-1), value_targets)
+        overall_err = sum([
+            pi_err,
+            self.vf_loss_coeff * value_err,
+            self.entropy_coeff * entropy,
+        ])
+        return overall_err
+
+
+class A3CTorchPolicyGraph(TorchPolicyGraph):
+    """A simple, non-recurrent PyTorch policy example."""
+
+    def __init__(self, obs_space, action_space, config):
+        config = dict(ray.rllib.agents.a3c.a3c.DEFAULT_CONFIG, **config)
+        self.config = config
+        _, self.logit_dim = ModelCatalog.get_action_dist(
+            action_space, self.config["model"])
+        self.model = ModelCatalog.get_torch_model(
+            obs_space.shape, self.logit_dim, self.config["model"])
+        loss = A3CLoss(
+            self.model, self.config["vf_loss_coeff"],
+            self.config["entropy_coeff"])
+        TorchPolicyGraph.__init__(
+            self, obs_space, action_space, self.model, loss,
+            loss_inputs=[
+                "obs", "actions", "advantages", "value_targets"])
+
+    def extra_action_out(self, model_out):
+        return {"vf_preds": var_to_np(model_out[1])}
+
+    def optimizer(self):
+        return torch.optim.Adam(
+            self.model.parameters(), lr=self.config["lr"])
+
+    def postprocess_trajectory(self, sample_batch, other_agent_batches=None):
+        completed = sample_batch["dones"][-1]
+        if completed:
+            last_r = 0.0
+        else:
+            last_r = self._value(sample_batch["new_obs"][-1])
+        return compute_advantages(
+            sample_batch, last_r, self.config["gamma"], self.config["lambda"])
+
+    def _value(self, obs):
+        with self.lock:
+            obs = torch.from_numpy(obs).float().unsqueeze(0)
+            res = self.model.hidden_layers(obs)
+            res = self.model.value_branch(res)
+            res = res.squeeze()
+            return var_to_np(res)

python/ray/rllib/agents/agent.py

@@ -0,0 +1,358 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import copy
+import json
+import numpy as np
+import os
+import pickle
+
+import tensorflow as tf
+from ray.rllib.evaluation.common_policy_evaluator import CommonPolicyEvaluator
+from ray.tune.registry import ENV_CREATOR, _global_registry
+from ray.tune.result import TrainingResult
+from ray.tune.trainable import Trainable
+
+COMMON_CONFIG = {
+    # Discount factor of the MDP
+    "gamma": 0.99,
+    # Number of steps after which the rollout gets cut
+    "horizon": None,
+    # Number of environments to evaluate vectorwise per worker.
+    "num_envs": 1,
+    # Number of actors used for parallelism
+    "num_workers": 2,
+    # Default sample batch size
+    "sample_batch_size": 200,
+    # Whether to rollout "complete_episodes" or "truncate_episodes"
+    "batch_mode": "truncate_episodes",
+    # Whether to use a background thread for sampling (slightly off-policy)
+    "sample_async": False,
+    # Which observation filter to apply to the observation
+    "observation_filter": "NoFilter",
+    # Whether to use rllib or deepmind preprocessors
+    "preprocessor_pref": "rllib",
+    # Arguments to pass to the env creator
+    "env_config": {},
+    # Arguments to pass to model
+    "model": {},
+    # Arguments to pass to the rllib optimizer
+    "optimizer": {},
+    # Override default TF session args if non-empty
+    "tf_session_args": {},
+    # Whether to LZ4 compress observations
+    "compress_observations": False,
+
+    # === Multiagent ===
+    "multiagent": {
+        "policy_graphs": {},
+        "policy_mapping_fn": None,
+    },
+}
+
+
+def with_common_config(extra_config):
+    """Returns the given config dict merged with common agent confs."""
+
+    config = copy.deepcopy(COMMON_CONFIG)
+    config.update(extra_config)
+    return config
+
+
+def _deep_update(original, new_dict, new_keys_allowed, whitelist):
+    """Updates original dict with values from new_dict recursively.
+    If new key is introduced in new_dict, then if new_keys_allowed is not
+    True, an error will be thrown. Further, for sub-dicts, if the key is
+    in the whitelist, then new subkeys can be introduced.
+
+    Args:
+        original (dict): Dictionary with default values.
+        new_dict (dict): Dictionary with values to be updated
+        new_keys_allowed (bool): Whether new keys are allowed.
+        whitelist (list): List of keys that correspond to dict values
+            where new subkeys can be introduced. This is only at
+            the top level.
+    """
+    for k, value in new_dict.items():
+        if k not in original and k != "env":
+            if not new_keys_allowed:
+                raise Exception(
+                    "Unknown config parameter `{}` ".format(k))
+        if type(original.get(k)) is dict:
+            if k in whitelist:
+                _deep_update(original[k], value, True, [])
+            else:
+                _deep_update(original[k], value, new_keys_allowed, [])
+        else:
+            original[k] = value
+    return original
+
+
+class Agent(Trainable):
+    """All RLlib agents extend this base class.
+
+    Agent objects retain internal model state between calls to train(), so
+    you should create a new agent instance for each training session.
+
+    Attributes:
+        env_creator (func): Function that creates a new training env.
+        config (obj): Algorithm-specific configuration data.
+        logdir (str): Directory in which training outputs should be placed.
+    """
+
+    _allow_unknown_configs = False
+    _allow_unknown_subkeys = [
+        "tf_session_args", "env_config", "model", "optimizer", "multiagent"]
+
+    def make_local_evaluator(self, env_creator, policy_graph):
+        """Convenience method to return configured local evaluator."""
+
+        return self._make_evaluator(
+            CommonPolicyEvaluator, env_creator, policy_graph, 0)
+
+    def make_remote_evaluators(
+            self, env_creator, policy_graph, count, remote_args):
+        """Convenience method to return a number of remote evaluators."""
+
+        cls = CommonPolicyEvaluator.as_remote(**remote_args).remote
+        return [
+            self._make_evaluator(cls, env_creator, policy_graph, i+1)
+            for i in range(count)]
+
+    def _make_evaluator(self, cls, env_creator, policy_graph, worker_index):
+        config = self.config
+
+        def session_creator():
+            return tf.Session(
+                config=tf.ConfigProto(**config["tf_session_args"]))
+
+        return cls(
+            env_creator,
+            self.config["multiagent"]["policy_graphs"] or policy_graph,
+            policy_mapping_fn=self.config["multiagent"]["policy_mapping_fn"],
+            tf_session_creator=(
+                session_creator if config["tf_session_args"] else None),
+            batch_steps=config["sample_batch_size"],
+            batch_mode=config["batch_mode"],
+            episode_horizon=config["horizon"],
+            preprocessor_pref=config["preprocessor_pref"],
+            sample_async=config["sample_async"],
+            compress_observations=config["compress_observations"],
+            num_envs=config["num_envs"],
+            observation_filter=config["observation_filter"],
+            env_config=config["env_config"],
+            model_config=config["model"],
+            policy_config=config,
+            worker_index=worker_index)
+
+    @classmethod
+    def resource_help(cls, config):
+        return (
+            "\n\nYou can adjust the resource requests of RLlib agents by "
+            "setting `num_workers` and other configs. See the "
+            "DEFAULT_CONFIG defined by each agent for more info.\n\n"
+            "The config of this agent is: " + json.dumps(config))
+
+    def __init__(
+            self, config=None, env=None, logger_creator=None):
+        """Initialize an RLLib agent.
+
+        Args:
+            config (dict): Algorithm-specific configuration data.
+            env (str): Name of the environment to use. Note that this can also
+                be specified as the `env` key in config.
+            logger_creator (func): Function that creates a ray.tune.Logger
+                object. If unspecified, a default logger is created.
+        """
+
+        config = config or {}
+
+        # Agents allow env ids to be passed directly to the constructor.
+        self._env_id = env or config.get("env")
+        Trainable.__init__(self, config, logger_creator)
+
+    def _setup(self):
+        env = self._env_id
+        if env:
+            self.config["env"] = env
+            if _global_registry.contains(ENV_CREATOR, env):
+                self.env_creator = _global_registry.get(ENV_CREATOR, env)
+            else:
+                import gym  # soft dependency
+                self.env_creator = lambda env_config: gym.make(env)
+        else:
+            self.env_creator = lambda env_config: None
+
+        # Merge the supplied config with the class default
+        merged_config = self._default_config.copy()
+        merged_config = _deep_update(merged_config, self.config,
+                                     self._allow_unknown_configs,
+                                     self._allow_unknown_subkeys)
+        self.config = merged_config
+
+        # TODO(ekl) setting the graph is unnecessary for PyTorch agents
+        with tf.Graph().as_default():
+            self._init()
+
+    def _init(self):
+        """Subclasses should override this for custom initialization."""
+
+        raise NotImplementedError
+
+    @property
+    def iteration(self):
+        """Current training iter, auto-incremented with each train() call."""
+
+        return self._iteration
+
+    @property
+    def _agent_name(self):
+        """Subclasses should override this to declare their name."""
+
+        raise NotImplementedError
+
+    @property
+    def _default_config(self):
+        """Subclasses should override this to declare their default config."""
+
+        raise NotImplementedError
+
+    def compute_action(self, observation, state=None):
+        """Computes an action using the current trained policy."""
+
+        if state is None:
+            state = []
+        obs = self.local_evaluator.filters["default"](
+            observation, update=False)
+        return self.local_evaluator.for_policy(
+            lambda p: p.compute_single_action(
+                obs, state, is_training=False)[0])
+
+
+class _MockAgent(Agent):
+    """Mock agent for use in tests"""
+
+    _agent_name = "MockAgent"
+    _default_config = {
+        "mock_error": False,
+        "persistent_error": False,
+    }
+
+    def _init(self):
+        self.info = None
+        self.restored = False
+
+    def _train(self):
+        if self.config["mock_error"] and self.iteration == 1 \
+                and (self.config["persistent_error"] or not self.restored):
+            raise Exception("mock error")
+        return TrainingResult(
+            episode_reward_mean=10, episode_len_mean=10,
+            timesteps_this_iter=10, info={})
+
+    def _save(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "mock_agent.pkl")
+        with open(path, 'wb') as f:
+            pickle.dump(self.info, f)
+        return path
+
+    def _restore(self, checkpoint_path):
+        with open(checkpoint_path, 'rb') as f:
+            info = pickle.load(f)
+        self.info = info
+        self.restored = True
+
+    def set_info(self, info):
+        self.info = info
+        return info
+
+    def get_info(self):
+        return self.info
+
+
+class _SigmoidFakeData(_MockAgent):
+    """Agent that returns sigmoid learning curves.
+
+    This can be helpful for evaluating early stopping algorithms."""
+
+    _agent_name = "SigmoidFakeData"
+    _default_config = {
+        "width": 100,
+        "height": 100,
+        "offset": 0,
+        "iter_time": 10,
+        "iter_timesteps": 1,
+    }
+
+    def _train(self):
+        i = max(0, self.iteration - self.config["offset"])
+        v = np.tanh(float(i) / self.config["width"])
+        v *= self.config["height"]
+        return TrainingResult(
+            episode_reward_mean=v, episode_len_mean=v,
+            timesteps_this_iter=self.config["iter_timesteps"],
+            time_this_iter_s=self.config["iter_time"], info={})
+
+
+class _ParameterTuningAgent(_MockAgent):
+
+    _agent_name = "ParameterTuningAgent"
+    _default_config = {
+        "reward_amt": 10,
+        "dummy_param": 10,
+        "dummy_param2": 15,
+        "iter_time": 10,
+        "iter_timesteps": 1
+    }
+
+    def _train(self):
+        return TrainingResult(
+            episode_reward_mean=self.config["reward_amt"] * self.iteration,
+            episode_len_mean=self.config["reward_amt"],
+            timesteps_this_iter=self.config["iter_timesteps"],
+            time_this_iter_s=self.config["iter_time"], info={})
+
+
+def get_agent_class(alg):
+    """Returns the class of a known agent given its name."""
+
+    if alg == "DDPG":
+        from ray.rllib.agents import ddpg
+        return ddpg.DDPGAgent
+    elif alg == "APEX_DDPG":
+        from ray.rllib.agents import ddpg
+        return ddpg.ApexDDPGAgent
+    elif alg == "PPO":
+        from ray.rllib.agents import ppo
+        return ppo.PPOAgent
+    elif alg == "ES":
+        from ray.rllib.agents import es
+        return es.ESAgent
+    elif alg == "DQN":
+        from ray.rllib.agents import dqn
+        return dqn.DQNAgent
+    elif alg == "APEX":
+        from ray.rllib.agents import dqn
+        return dqn.ApexAgent
+    elif alg == "A3C":
+        from ray.rllib.agents import a3c
+        return a3c.A3CAgent
+    elif alg == "BC":
+        from ray.rllib.agents import bc
+        return bc.BCAgent
+    elif alg == "PG":
+        from ray.rllib.agents import pg
+        return pg.PGAgent
+    elif alg == "script":
+        from ray.tune import script_runner
+        return script_runner.ScriptRunner
+    elif alg == "__fake":
+        return _MockAgent
+    elif alg == "__sigmoid_fake_data":
+        return _SigmoidFakeData
+    elif alg == "__parameter_tuning":
+        return _ParameterTuningAgent
+    else:
+        raise Exception(
+            ("Unknown algorithm {}.").format(alg))

python/ray/rllib/agents/bc/__init__.py

@@ -0,0 +1,3 @@
+from ray.rllib.agents.bc.bc import BCAgent, DEFAULT_CONFIG
+
+__all__ = ["BCAgent", "DEFAULT_CONFIG"]

python/ray/rllib/agents/bc/bc.py

@@ -0,0 +1,96 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import ray
+from ray.rllib.agents.agent import Agent
+from ray.rllib.agents.bc.bc_evaluator import BCEvaluator, \
+    GPURemoteBCEvaluator, RemoteBCEvaluator
+from ray.rllib.optimizers import AsyncGradientsOptimizer
+from ray.tune.result import TrainingResult
+from ray.tune.trial import Resources
+
+DEFAULT_CONFIG = {
+    # Number of workers (excluding master)
+    "num_workers": 1,
+    # Size of rollout batch
+    "batch_size": 100,
+    # Max global norm for each gradient calculated by worker
+    "grad_clip": 40.0,
+    # Learning rate
+    "lr": 0.0001,
+    # Whether to use a GPU for local optimization.
+    "gpu": False,
+    # Whether to place workers on GPUs
+    "use_gpu_for_workers": False,
+    # Model and preprocessor options
+    "model": {
+        # (Image statespace) - Converts image to Channels = 1
+        "grayscale": True,
+        # (Image statespace) - Each pixel
+        "zero_mean": False,
+        # (Image statespace) - Converts image to (dim, dim, C)
+        "dim": 80,
+        # (Image statespace) - Converts image shape to (C, dim, dim)
+        "channel_major": False
+    },
+    # Arguments to pass to the rllib optimizer
+    "optimizer": {
+        # Number of gradients applied for each `train` step
+        "grads_per_step": 100,
+    },
+    # Arguments to pass to the env creator
+    "env_config": {},
+}
+
+
+class BCAgent(Agent):
+    _agent_name = "BC"
+    _default_config = DEFAULT_CONFIG
+    _allow_unknown_configs = True
+
+    @classmethod
+    def default_resource_request(cls, config):
+        cf = dict(cls._default_config, **config)
+        if cf["use_gpu_for_workers"]:
+            num_gpus_per_worker = 1
+        else:
+            num_gpus_per_worker = 0
+        return Resources(
+            cpu=1, gpu=cf["gpu"] and 1 or 0,
+            extra_cpu=cf["num_workers"],
+            extra_gpu=num_gpus_per_worker * cf["num_workers"])
+
+    def _init(self):
+        self.local_evaluator = BCEvaluator(
+            self.env_creator, self.config, self.logdir)
+        if self.config["use_gpu_for_workers"]:
+            remote_cls = GPURemoteBCEvaluator
+        else:
+            remote_cls = RemoteBCEvaluator
+        self.remote_evaluators = [
+            remote_cls.remote(self.env_creator, self.config, self.logdir)
+            for _ in range(self.config["num_workers"])]
+        self.optimizer = AsyncGradientsOptimizer(
+            self.config["optimizer"], self.local_evaluator,
+            self.remote_evaluators)
+
+    def _train(self):
+        self.optimizer.step()
+        metric_lists = [re.get_metrics.remote() for re in
+                        self.remote_evaluators]
+        total_samples = 0
+        total_loss = 0
+        for metrics in metric_lists:
+            for m in ray.get(metrics):
+                total_samples += m["num_samples"]
+                total_loss += m["loss"]
+        result = TrainingResult(
+            mean_loss=total_loss / total_samples,
+            timesteps_this_iter=total_samples,
+        )
+        return result
+
+    def compute_action(self, observation):
+        action, info = self.local_evaluator.policy.compute(observation)
+        return action

python/ray/rllib/agents/bc/bc_evaluator.py

@@ -0,0 +1,63 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import pickle
+from six.moves import queue
+
+import ray
+from ray.rllib.agents.bc.experience_dataset import ExperienceDataset
+from ray.rllib.agents.bc.policy import BCPolicy
+from ray.rllib.evaluation.interface import PolicyEvaluator
+from ray.rllib.models import ModelCatalog
+
+
+class BCEvaluator(PolicyEvaluator):
+    def __init__(self, env_creator, config, logdir):
+        env = ModelCatalog.get_preprocessor_as_wrapper(env_creator(
+            config["env_config"]), config["model"])
+        self.dataset = ExperienceDataset(config["dataset_path"])
+        self.policy = BCPolicy(env.observation_space, env.action_space, config)
+        self.config = config
+        self.logdir = logdir
+        self.metrics_queue = queue.Queue()
+
+    def sample(self):
+        return self.dataset.sample(self.config["batch_size"])
+
+    def compute_gradients(self, samples):
+        gradient, info = self.policy.compute_gradients(samples)
+        self.metrics_queue.put(
+            {"num_samples": info["num_samples"], "loss": info["loss"]})
+        return gradient, {}
+
+    def apply_gradients(self, grads):
+        self.policy.apply_gradients(grads)
+
+    def get_weights(self):
+        return self.policy.get_weights()
+
+    def set_weights(self, params):
+        self.policy.set_weights(params)
+
+    def save(self):
+        weights = self.get_weights()
+        return pickle.dumps({
+            "weights": weights})
+
+    def restore(self, objs):
+        objs = pickle.loads(objs)
+        self.set_weights(objs["weights"])
+
+    def get_metrics(self):
+        completed = []
+        while True:
+            try:
+                completed.append(self.metrics_queue.get_nowait())
+            except queue.Empty:
+                break
+        return completed
+
+
+RemoteBCEvaluator = ray.remote(BCEvaluator)
+GPURemoteBCEvaluator = ray.remote(num_gpus=1)(BCEvaluator)

python/ray/rllib/agents/bc/experience_dataset.py

@@ -0,0 +1,33 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import itertools
+import pickle
+
+import numpy as np
+
+
+class ExperienceDataset(object):
+    def __init__(self, dataset_path):
+        """Create dataset of experience to imitate.
+
+        Parameters
+        ----------
+        dataset_path:
+          Path of file containing the database as pickled list of trajectories,
+          each trajectory being a list of steps,
+          each step containing the observation and action as its first two
+            elements.
+          The file must be available on each machine used by a BCEvaluator.
+        """
+        self._dataset = list(itertools.chain.from_iterable(
+            pickle.load(open(dataset_path, "rb"))))
+
+    def sample(self, batch_size):
+        indexes = np.random.choice(len(self._dataset), batch_size)
+        samples = {
+            'observations': [self._dataset[i][0] for i in indexes],
+            'actions': [self._dataset[i][1] for i in indexes]
+        }
+        return samples

python/ray/rllib/agents/bc/policy.py

@@ -0,0 +1,105 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import gym
+
+import ray
+from ray.rllib.models.catalog import ModelCatalog
+
+
+class BCPolicy(object):
+    def __init__(self, obs_space, action_space, config):
+        self.local_steps = 0
+        self.config = config
+        self.summarize = config.get("summarize")
+        self._setup_graph(obs_space, action_space)
+        self.setup_loss(action_space)
+        self.setup_gradients()
+        self.initialize()
+
+    def _setup_graph(self, obs_space, ac_space):
+        self.x = tf.placeholder(tf.float32, [None] + list(obs_space.shape))
+        dist_class, self.logit_dim = ModelCatalog.get_action_dist(
+            ac_space, self.config["model"])
+        self._model = ModelCatalog.get_model(
+            self.x, self.logit_dim, self.config["model"])
+        self.logits = self._model.outputs
+        self.curr_dist = dist_class(self.logits)
+        self.sample = self.curr_dist.sample()
+        self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
+                                          tf.get_variable_scope().name)
+
+    def setup_loss(self, action_space):
+        if isinstance(action_space, gym.spaces.Box):
+            self.ac = tf.placeholder(tf.float32,
+                                     [None] + list(action_space.shape),
+                                     name="ac")
+        elif isinstance(action_space, gym.spaces.Discrete):
+            self.ac = tf.placeholder(tf.int64, [None], name="ac")
+        else:
+            raise NotImplementedError(
+                "action space" + str(type(action_space)) +
+                "currently not supported")
+        log_prob = self.curr_dist.logp(self.ac)
+        self.pi_loss = - tf.reduce_sum(log_prob)
+        self.loss = self.pi_loss
+
+    def setup_gradients(self):
+        grads = tf.gradients(self.loss, self.var_list)
+        self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"])
+        grads_and_vars = list(zip(self.grads, self.var_list))
+        opt = tf.train.AdamOptimizer(self.config["lr"])
+        self._apply_gradients = opt.apply_gradients(grads_and_vars)
+
+    def initialize(self):
+        if self.summarize:
+            bs = tf.to_float(tf.shape(self.x)[0])
+            tf.summary.scalar("model/policy_loss", self.pi_loss / bs)
+            tf.summary.scalar("model/grad_gnorm", tf.global_norm(self.grads))
+            tf.summary.scalar("model/var_gnorm", tf.global_norm(self.var_list))
+            self.summary_op = tf.summary.merge_all()
+
+        # TODO(rliaw): Can consider exposing these parameters
+        self.sess = tf.Session(graph=self.g, config=tf.ConfigProto(
+            intra_op_parallelism_threads=1, inter_op_parallelism_threads=2,
+            gpu_options=tf.GPUOptions(allow_growth=True)))
+        self.variables = ray.experimental.TensorFlowVariables(self.loss,
+                                                              self.sess)
+        self.sess.run(tf.global_variables_initializer())
+
+    def compute_gradients(self, samples):
+        info = {}
+        feed_dict = {
+            self.x: samples["observations"],
+            self.ac: samples["actions"]
+        }
+        self.grads = [g for g in self.grads if g is not None]
+        self.local_steps += 1
+        if self.summarize:
+            loss, grad, summ = self.sess.run(
+                [self.loss, self.grads, self.summary_op], feed_dict=feed_dict)
+            info["summary"] = summ
+        else:
+            loss, grad = self.sess.run([self.loss, self.grads],
+                                       feed_dict=feed_dict)
+        info["num_samples"] = len(samples)
+        info["loss"] = loss
+        return grad, info
+
+    def apply_gradients(self, grads):
+        feed_dict = {self.grads[i]: grads[i]
+                     for i in range(len(grads))}
+        self.sess.run(self._apply_gradients, feed_dict=feed_dict)
+
+    def get_weights(self):
+        weights = self.variables.get_weights()
+        return weights
+
+    def set_weights(self, weights):
+        self.variables.set_weights(weights)
+
+    def compute(self, ob, *args):
+        action = self.sess.run(self.sample, {self.x: [ob]})
+        return action, None

python/ray/rllib/agents/ddpg/README.md

@@ -0,0 +1 @@
+Implementation of deep deterministic policy gradients (https://arxiv.org/abs/1509.02971), including an Ape-X variant.

python/ray/rllib/agents/ddpg/__init__.py

@@ -0,0 +1,8 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from ray.rllib.agents.ddpg.apex import ApexDDPGAgent
+from ray.rllib.agents.ddpg.ddpg import DDPGAgent, DEFAULT_CONFIG
+
+__all__ = ["DDPGAgent", "ApexDDPGAgent", "DEFAULT_CONFIG"]

python/ray/rllib/agents/ddpg/apex.py

@@ -0,0 +1,50 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from ray.rllib.agents.ddpg.ddpg import DDPGAgent, DEFAULT_CONFIG as DDPG_CONFIG
+from ray.utils import merge_dicts
+
+APEX_DDPG_DEFAULT_CONFIG = merge_dicts(
+    DDPG_CONFIG,
+    {
+        "optimizer_class": "AsyncSamplesOptimizer",
+        "optimizer":
+            merge_dicts(
+                DDPG_CONFIG["optimizer"], {
+                    "max_weight_sync_delay": 400,
+                    "num_replay_buffer_shards": 4,
+                    "debug": False
+                }),
+        "n_step": 3,
+        "num_workers": 32,
+        "buffer_size": 2000000,
+        "learning_starts": 50000,
+        "train_batch_size": 512,
+        "sample_batch_size": 50,
+        "max_weight_sync_delay": 400,
+        "target_network_update_freq": 500000,
+        "timesteps_per_iteration": 25000,
+        "per_worker_exploration": True,
+        "worker_side_prioritization": True,
+    },
+)
+
+
+class ApexDDPGAgent(DDPGAgent):
+    """DDPG variant that uses the Ape-X distributed policy optimizer.
+
+    By default, this is configured for a large single node (32 cores). For
+    running in a large cluster, increase the `num_workers` config var.
+    """
+
+    _agent_name = "APEX_DDPG"
+    _default_config = APEX_DDPG_DEFAULT_CONFIG
+
+    def update_target_if_needed(self):
+        # Ape-X updates based on num steps trained, not sampled
+        if self.optimizer.num_steps_trained - self.last_target_update_ts > \
+                self.config["target_network_update_freq"]:
+            self.local_evaluator.for_policy(lambda p: p.update_target())
+            self.last_target_update_ts = self.optimizer.num_steps_trained
+            self.num_target_updates += 1

python/ray/rllib/agents/ddpg/ddpg.py

@@ -0,0 +1,130 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from ray.rllib.agents.agent import with_common_config
+from ray.rllib.agents.dqn.dqn import DQNAgent
+from ray.rllib.agents.ddpg.ddpg_policy_graph import DDPGPolicyGraph
+from ray.rllib.utils.schedules import ConstantSchedule, LinearSchedule
+
+OPTIMIZER_SHARED_CONFIGS = [
+    "buffer_size", "prioritized_replay", "prioritized_replay_alpha",
+    "prioritized_replay_beta", "prioritized_replay_eps", "sample_batch_size",
+    "train_batch_size", "learning_starts", "clip_rewards"
+]
+
+DEFAULT_CONFIG = with_common_config({
+    # === Model ===
+    # Hidden layer sizes of the policy network
+    "actor_hiddens": [64, 64],
+    # Hidden layers activation of the policy network
+    "actor_hidden_activation": "relu",
+    # Hidden layer sizes of the critic network
+    "critic_hiddens": [64, 64],
+    # Hidden layers activation of the critic network
+    "critic_hidden_activation": "relu",
+    # N-step Q learning
+    "n_step": 1,
+
+    # === Exploration ===
+    # Max num timesteps for annealing schedules. Exploration is annealed from
+    # 1.0 to exploration_fraction over this number of timesteps scaled by
+    # exploration_fraction
+    "schedule_max_timesteps": 100000,
+    # Number of env steps to optimize for before returning
+    "timesteps_per_iteration": 1000,
+    # Fraction of entire training period over which the exploration rate is
+    # annealed
+    "exploration_fraction": 0.1,
+    # Final value of random action probability
+    "exploration_final_eps": 0.02,
+    # OU-noise scale
+    "noise_scale": 0.1,
+    # theta
+    "exploration_theta": 0.15,
+    # sigma
+    "exploration_sigma": 0.2,
+    # Update the target network every `target_network_update_freq` steps.
+    "target_network_update_freq": 0,
+    # Update the target by \tau * policy + (1-\tau) * target_policy
+    "tau": 0.002,
+
+    # === Replay buffer ===
+    # Size of the replay buffer. Note that if async_updates is set, then
+    # each worker will have a replay buffer of this size.
+    "buffer_size": 50000,
+    # If True prioritized replay buffer will be used.
+    "prioritized_replay": True,
+    # Alpha parameter for prioritized replay buffer.
+    "prioritized_replay_alpha": 0.6,
+    # Beta parameter for sampling from prioritized replay buffer.
+    "prioritized_replay_beta": 0.4,
+    # Epsilon to add to the TD errors when updating priorities.
+    "prioritized_replay_eps": 1e-6,
+    # Whether to clip rewards to [-1, 1] prior to adding to the replay buffer.
+    "clip_rewards": True,
+
+    # === Optimization ===
+    # Learning rate for adam optimizer
+    "actor_lr": 1e-4,
+    "critic_lr": 1e-3,
+    # If True, use huber loss instead of squared loss for critic network
+    # Conventionally, no need to clip gradients if using a huber loss
+    "use_huber": False,
+    # Threshold of a huber loss
+    "huber_threshold": 1.0,
+    # Weights for L2 regularization
+    "l2_reg": 1e-6,
+    # If not None, clip gradients during optimization at this value
+    "grad_norm_clipping": None,
+    # How many steps of the model to sample before learning starts.
+    "learning_starts": 1500,
+    # Update the replay buffer with this many samples at once. Note that this
+    # setting applies per-worker if num_workers > 1.
+    "sample_batch_size": 1,
+    # Size of a batched sampled from replay buffer for training. Note that
+    # if async_updates is set, then each worker returns gradients for a
+    # batch of this size.
+    "train_batch_size": 256,
+
+    # === Parallelism ===
+    # Whether to use a GPU for local optimization.
+    "gpu": False,
+    # Number of workers for collecting samples with. This only makes sense
+    # to increase if your environment is particularly slow to sample, or if
+    # you"re using the Async or Ape-X optimizers.
+    "num_workers": 0,
+    # Whether to allocate GPUs for workers (if > 0).
+    "num_gpus_per_worker": 0,
+    # Whether to allocate CPUs for workers (if > 0).
+    "num_cpus_per_worker": 1,
+    # Optimizer class to use.
+    "optimizer_class": "SyncReplayOptimizer",
+    # Whether to use a distribution of epsilons across workers for exploration.
+    "per_worker_exploration": False,
+    # Whether to compute priorities on workers.
+    "worker_side_prioritization": False,
+})
+
+
+class DDPGAgent(DQNAgent):
+    """DDPG implementation in TensorFlow."""
+    _agent_name = "DDPG"
+    _default_config = DEFAULT_CONFIG
+    _policy_graph = DDPGPolicyGraph
+
+    def _make_exploration_schedule(self, worker_index):
+        # Override DQN's schedule to take into account `noise_scale`
+        if self.config["per_worker_exploration"]:
+            assert self.config["num_workers"] > 1, \
+                "This requires multiple workers"
+            return ConstantSchedule(
+                self.config["noise_scale"] * 0.4 **
+                (1 + worker_index / float(self.config["num_workers"] - 1) * 7))
+        else:
+            return LinearSchedule(
+                schedule_timesteps=int(self.config["exploration_fraction"] *
+                                       self.config["schedule_max_timesteps"]),
+                initial_p=self.config["noise_scale"] * 1.0,
+                final_p=self.config["noise_scale"] *
+                self.config["exploration_final_eps"])

python/ray/rllib/agents/ddpg/ddpg_policy_graph.py

@@ -0,0 +1,356 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from gym.spaces import Box
+import numpy as np
+import tensorflow as tf
+import tensorflow.contrib.layers as layers
+
+import ray
+from ray.rllib.agents.dqn.dqn_policy_graph import _huber_loss, \
+    _minimize_and_clip, _scope_vars, _postprocess_dqn
+from ray.rllib.models import ModelCatalog
+from ray.rllib.utils.error import UnsupportedSpaceException
+from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph
+
+
+A_SCOPE = "a_func"
+P_SCOPE = "p_func"
+P_TARGET_SCOPE = "target_p_func"
+Q_SCOPE = "q_func"
+Q_TARGET_SCOPE = "target_q_func"
+
+
+class PNetwork(object):
+    """Maps an observations (i.e., state) to an action where each entry takes
+    value from (0, 1) due to the sigmoid function."""
+
+    def __init__(
+            self, model, dim_actions, hiddens=[64, 64], activation="relu"):
+        action_out = model.last_layer
+        activation = tf.nn.__dict__[activation]
+        for hidden in hiddens:
+            action_out = layers.fully_connected(
+                action_out, num_outputs=hidden, activation_fn=activation)
+        # Use sigmoid layer to bound values within (0, 1)
+        # shape of action_scores is [batch_size, dim_actions]
+        self.action_scores = layers.fully_connected(
+            action_out, num_outputs=dim_actions, activation_fn=tf.nn.sigmoid)
+
+
+class ActionNetwork(object):
+    """Acts as a stochastic policy for inference, but a deterministic policy
+    for training, thus ignoring the batch_size issue when constructing a
+    stochastic action."""
+
+    def __init__(
+            self, p_values, low_action, high_action, stochastic, eps,
+            theta=0.15, sigma=0.2):
+
+        # shape is [None, dim_action]
+        deterministic_actions = (
+            (high_action - low_action) * p_values + low_action)
+
+        exploration_sample = tf.get_variable(
+            name="ornstein_uhlenbeck",
+            dtype=tf.float32,
+            initializer=low_action.size * [.0],
+            trainable=False)
+        normal_sample = tf.random_normal(
+            shape=[low_action.size], mean=0.0, stddev=1.0)
+        exploration_value = tf.assign_add(
+            exploration_sample,
+            theta * (.0 - exploration_sample) + sigma * normal_sample)
+        stochastic_actions = deterministic_actions + eps * (
+            high_action - low_action) * exploration_value
+
+        self.actions = tf.cond(
+            stochastic, lambda: stochastic_actions,
+            lambda: deterministic_actions)
+
+
+class QNetwork(object):
+    def __init__(
+            self, model, action_inputs,
+            hiddens=[64, 64], activation="relu"):
+        q_out = tf.concat([model.last_layer, action_inputs], axis=1)
+        activation = tf.nn.__dict__[activation]
+        for hidden in hiddens:
+            q_out = layers.fully_connected(
+                q_out, num_outputs=hidden, activation_fn=activation)
+        self.value = layers.fully_connected(
+            q_out, num_outputs=1, activation_fn=None)
+
+
+class ActorCriticLoss(object):
+    def __init__(
+            self, q_t, q_tp1, q_tp0, importance_weights, rewards, done_mask,
+            gamma=0.99, n_step=1, use_huber=False, huber_threshold=1.0):
+
+        q_t_selected = tf.squeeze(q_t, axis=len(q_t.shape) - 1)
+
+        q_tp1_best = tf.squeeze(
+            input=q_tp1, axis=len(q_tp1.shape) - 1)
+        q_tp1_best_masked = (1.0 - done_mask) * q_tp1_best
+
+        # compute RHS of bellman equation
+        q_t_selected_target = rewards + gamma**n_step * q_tp1_best_masked
+
+        # compute the error (potentially clipped)
+        self.td_error = q_t_selected - tf.stop_gradient(q_t_selected_target)
+        if use_huber:
+            errors = _huber_loss(self.td_error, huber_threshold)
+        else:
+            errors = 0.5 * tf.square(self.td_error)
+
+        self.critic_loss = tf.reduce_mean(importance_weights * errors)
+
+        # for policy gradient
+        self.actor_loss = -1.0 * tf.reduce_mean(q_tp0)
+        self.total_loss = self.actor_loss + self.critic_loss
+
+
+class DDPGPolicyGraph(TFPolicyGraph):
+    def __init__(self, observation_space, action_space, config):
+        config = dict(ray.rllib.agents.ddpg.ddpg.DEFAULT_CONFIG, **config)
+        if not isinstance(action_space, Box):
+            raise UnsupportedSpaceException(
+                "Action space {} is not supported for DDPG.".format(
+                    action_space))
+
+        self.config = config
+        self.cur_epsilon = 1.0
+        dim_actions = action_space.shape[0]
+        low_action = action_space.low
+        high_action = action_space.high
+        self.actor_optimizer = tf.train.AdamOptimizer(
+            learning_rate=config["actor_lr"])
+        self.critic_optimizer = tf.train.AdamOptimizer(
+            learning_rate=config["critic_lr"])
+
+        def _build_q_network(obs, actions):
+            return QNetwork(
+                ModelCatalog.get_model(obs, 1, config["model"]),
+                actions,
+                config["critic_hiddens"],
+                config["critic_hidden_activation"]).value
+
+        def _build_p_network(obs):
+            return PNetwork(
+                ModelCatalog.get_model(obs, 1, config["model"]),
+                dim_actions,
+                config["actor_hiddens"],
+                config["actor_hidden_activation"]).action_scores
+
+        def _build_action_network(p_values, stochastic, eps):
+            return ActionNetwork(
+                p_values,
+                low_action,
+                high_action,
+                stochastic,
+                eps,
+                config["exploration_theta"],
+                config["exploration_sigma"]).actions
+
+        # Action inputs
+        self.stochastic = tf.placeholder(tf.bool, (), name="stochastic")
+        self.eps = tf.placeholder(tf.float32, (), name="eps")
+        self.cur_observations = tf.placeholder(
+            tf.float32, shape=(None, ) + observation_space.shape)
+
+        # Actor: P (policy) network
+        with tf.variable_scope(P_SCOPE) as scope:
+            p_values = _build_p_network(self.cur_observations)
+            self.p_func_vars = _scope_vars(scope.name)
+
+        # Action outputs
+        with tf.variable_scope(A_SCOPE):
+            self.output_actions = _build_action_network(
+                p_values, self.stochastic, self.eps)
+
+        with tf.variable_scope(A_SCOPE, reuse=True):
+            exploration_sample = tf.get_variable(name="ornstein_uhlenbeck")
+            self.reset_noise_op = tf.assign(exploration_sample,
+                                            dim_actions * [.0])
+
+        # Replay inputs
+        self.obs_t = tf.placeholder(
+            tf.float32,
+            shape=(None, ) + observation_space.shape,
+            name="observation")
+        self.act_t = tf.placeholder(
+            tf.float32, shape=(None, ) + action_space.shape, name="action")
+        self.rew_t = tf.placeholder(tf.float32, [None], name="reward")
+        self.obs_tp1 = tf.placeholder(
+            tf.float32, shape=(None, ) + observation_space.shape)
+        self.done_mask = tf.placeholder(tf.float32, [None], name="done")
+        self.importance_weights = tf.placeholder(
+            tf.float32, [None], name="weight")
+
+        # p network evaluation
+        with tf.variable_scope(P_SCOPE, reuse=True) as scope:
+            self.p_t = _build_p_network(self.obs_t)
+
+        # target p network evaluation
+        with tf.variable_scope(P_TARGET_SCOPE) as scope:
+            p_tp1 = _build_p_network(self.obs_tp1)
+            target_p_func_vars = _scope_vars(scope.name)
+
+        # Action outputs
+        with tf.variable_scope(A_SCOPE, reuse=True):
+            deterministic_flag = tf.constant(value=False, dtype=tf.bool)
+            zero_eps = tf.constant(value=.0, dtype=tf.float32)
+            output_actions = _build_action_network(
+                self.p_t, deterministic_flag, zero_eps)
+
+            output_actions_estimated = _build_action_network(
+                p_tp1, deterministic_flag, zero_eps)
+
+        # q network evaluation
+        with tf.variable_scope(Q_SCOPE) as scope:
+            q_t = _build_q_network(self.obs_t, self.act_t)
+            self.q_func_vars = _scope_vars(scope.name)
+        with tf.variable_scope(Q_SCOPE, reuse=True):
+            q_tp0 = _build_q_network(self.obs_t, output_actions)
+
+        # target q network evalution
+        with tf.variable_scope(Q_TARGET_SCOPE) as scope:
+            q_tp1 = _build_q_network(self.obs_tp1, output_actions_estimated)
+            target_q_func_vars = _scope_vars(scope.name)
+
+        self.loss = ActorCriticLoss(
+            q_t, q_tp1, q_tp0, self.importance_weights, self.rew_t,
+            self.done_mask, config["gamma"], config["n_step"],
+            config["use_huber"], config["huber_threshold"])
+
+        if config["l2_reg"] is not None:
+            for var in self.p_func_vars:
+                if "bias" not in var.name:
+                    self.loss.actor_loss += (
+                        config["l2_reg"] * 0.5 * tf.nn.l2_loss(var))
+            for var in self.q_func_vars:
+                if "bias" not in var.name:
+                    self.loss.critic_loss += (
+                        config["l2_reg"] * 0.5 * tf.nn.l2_loss(var))
+
+        # update_target_fn will be called periodically to copy Q network to
+        # target Q network
+        self.tau_value = config.get("tau")
+        self.tau = tf.placeholder(tf.float32, (), name="tau")
+        update_target_expr = []
+        for var, var_target in zip(
+                sorted(self.q_func_vars, key=lambda v: v.name),
+                sorted(target_q_func_vars, key=lambda v: v.name)):
+            update_target_expr.append(
+                var_target.assign(self.tau * var +
+                                  (1.0 - self.tau) * var_target))
+        for var, var_target in zip(
+                sorted(self.p_func_vars, key=lambda v: v.name),
+                sorted(target_p_func_vars, key=lambda v: v.name)):
+            update_target_expr.append(
+                var_target.assign(self.tau * var +
+                                  (1.0 - self.tau) * var_target))
+        self.update_target_expr = tf.group(*update_target_expr)
+
+        self.sess = tf.get_default_session()
+        self.loss_inputs = [
+            ("obs", self.obs_t),
+            ("actions", self.act_t),
+            ("rewards", self.rew_t),
+            ("new_obs", self.obs_tp1),
+            ("dones", self.done_mask),
+            ("weights", self.importance_weights),
+        ]
+        self.is_training = tf.placeholder_with_default(True, ())
+        TFPolicyGraph.__init__(
+            self, observation_space, action_space, self.sess,
+            obs_input=self.cur_observations,
+            action_sampler=self.output_actions, loss=self.loss.total_loss,
+            loss_inputs=self.loss_inputs, is_training=self.is_training)
+        self.sess.run(tf.global_variables_initializer())
+
+        # Note that this encompasses both the policy and Q-value networks and
+        # their corresponding target networks
+        self.variables = ray.experimental.TensorFlowVariables(
+            tf.group(q_tp0, q_tp1), self.sess)
+
+        # Hard initial update
+        self.update_target(tau=1.0)
+
+    def gradients(self, optimizer):
+        if self.config["grad_norm_clipping"] is not None:
+            actor_grads_and_vars = _minimize_and_clip(
+                self.actor_optimizer,
+                self.loss.actor_loss,
+                var_list=self.p_func_vars,
+                clip_val=self.config["grad_norm_clipping"])
+            critic_grads_and_vars = _minimize_and_clip(
+                self.critic_optimizer,
+                self.loss.critic_loss,
+                var_list=self.q_func_vars,
+                clip_val=self.config["grad_norm_clipping"])
+        else:
+            actor_grads_and_vars = self.actor_optimizer.compute_gradients(
+                self.loss.actor_loss, var_list=self.p_func_vars)
+            critic_grads_and_vars = self.critic_optimizer.compute_gradients(
+                self.loss.critic_loss, var_list=self.q_func_vars)
+        actor_grads_and_vars = [
+            (g, v) for (g, v) in actor_grads_and_vars if g is not None]
+        critic_grads_and_vars = [
+            (g, v) for (g, v) in critic_grads_and_vars if g is not None]
+        grads_and_vars = actor_grads_and_vars + critic_grads_and_vars
+        return grads_and_vars
+
+    def extra_compute_action_feed_dict(self):
+        return {
+            self.stochastic: True,
+            self.eps: self.cur_epsilon,
+        }
+
+    def extra_compute_grad_fetches(self):
+        return {
+            "td_error": self.loss.td_error,
+        }
+
+    def postprocess_trajectory(self, sample_batch, other_agent_batches=None):
+        return _postprocess_dqn(self, sample_batch)
+
+    def compute_td_error(self, obs_t, act_t, rew_t, obs_tp1, done_mask,
+                         importance_weights):
+        td_err = self.sess.run(
+            self.loss.td_error,
+            feed_dict={
+                self.obs_t: [np.array(ob) for ob in obs_t],
+                self.act_t: act_t,
+                self.rew_t: rew_t,
+                self.obs_tp1: [np.array(ob) for ob in obs_tp1],
+                self.done_mask: done_mask,
+                self.importance_weights: importance_weights
+            })
+        return td_err
+
+    def reset_noise(self, sess):
+        sess.run(self.reset_noise_op)
+
+    # support both hard and soft sync
+    def update_target(self, tau=None):
+        return self.sess.run(
+            self.update_target_expr,
+            feed_dict={self.tau: tau or self.tau_value})
+
+    def set_epsilon(self, epsilon):
+        self.cur_epsilon = epsilon
+
+    def get_weights(self):
+        return self.variables.get_weights()
+
+    def set_weights(self, weights):
+        self.variables.set_weights(weights)
+
+    def get_state(self):
+        return [TFPolicyGraph.get_state(self), self.cur_epsilon]
+
+    def set_state(self, state):
+        TFPolicyGraph.set_state(self, state[0])
+        self.set_epsilon(state[1])

python/ray/rllib/agents/dqn/README.md

@@ -0,0 +1 @@
+Code in this package is adapted from https://github.com/openai/baselines/tree/master/baselines/deepq.

python/ray/rllib/agents/dqn/__init__.py

@@ -0,0 +1,8 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from ray.rllib.agents.dqn.apex import ApexAgent
+from ray.rllib.agents.dqn.dqn import DQNAgent, DEFAULT_CONFIG
+
+__all__ = ["ApexAgent", "DQNAgent", "DEFAULT_CONFIG"]

python/ray/rllib/agents/dqn/apex.py

@@ -0,0 +1,61 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from ray.rllib.agents.dqn.dqn import DQNAgent, DEFAULT_CONFIG as DQN_CONFIG
+from ray.tune.trial import Resources
+from ray.utils import merge_dicts
+
+APEX_DEFAULT_CONFIG = merge_dicts(
+    DQN_CONFIG,
+    {
+        "optimizer_class": "AsyncSamplesOptimizer",
+        "optimizer":
+            merge_dicts(
+                DQN_CONFIG["optimizer"], {
+                    "max_weight_sync_delay": 400,
+                    "num_replay_buffer_shards": 4,
+                    "debug": False
+                }),
+        "n_step": 3,
+        "gpu": True,
+        "num_workers": 32,
+        "buffer_size": 2000000,
+        "learning_starts": 50000,
+        "train_batch_size": 512,
+        "sample_batch_size": 50,
+        "max_weight_sync_delay": 400,
+        "target_network_update_freq": 500000,
+        "timesteps_per_iteration": 25000,
+        "per_worker_exploration": True,
+        "worker_side_prioritization": True,
+    },
+)
+
+
+class ApexAgent(DQNAgent):
+    """DQN variant that uses the Ape-X distributed policy optimizer.
+
+    By default, this is configured for a large single node (32 cores). For
+    running in a large cluster, increase the `num_workers` config var.
+    """
+
+    _agent_name = "APEX"
+    _default_config = APEX_DEFAULT_CONFIG
+
+    @classmethod
+    def default_resource_request(cls, config):
+        cf = dict(cls._default_config, **config)
+        return Resources(
+            cpu=1 + cf["optimizer"]["num_replay_buffer_shards"],
+            gpu=cf["gpu"] and 1 or 0,
+            extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"],
+            extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"])
+
+    def update_target_if_needed(self):
+        # Ape-X updates based on num steps trained, not sampled
+        if self.optimizer.num_steps_trained - self.last_target_update_ts > \
+                self.config["target_network_update_freq"]:
+            self.local_evaluator.for_policy(lambda p: p.update_target())
+            self.last_target_update_ts = self.optimizer.num_steps_trained
+            self.num_target_updates += 1

python/ray/rllib/agents/dqn/common/wrappers.py

@@ -0,0 +1,20 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from ray.rllib.models import ModelCatalog
+from ray.rllib.utils.atari_wrappers import wrap_deepmind
+
+
+def wrap_dqn(env, options, random_starts):
+    """Apply a common set of wrappers for DQN."""
+
+    is_atari = hasattr(env.unwrapped, "ale")
+
+    # Override atari default to use the deepmind wrappers.
+    # TODO(ekl) this logic should be pushed to the catalog.
+    if is_atari and "custom_preprocessor" not in options:
+        return wrap_deepmind(
+            env, random_starts=random_starts, dim=options.get("dim", 80))
+
+    return ModelCatalog.get_preprocessor_as_wrapper(env, options)

python/ray/rllib/agents/dqn/dqn.py

@@ -0,0 +1,222 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import pickle
+import os
+
+import ray
+from ray.rllib import optimizers
+from ray.rllib.agents.agent import Agent, with_common_config
+from ray.rllib.agents.dqn.dqn_policy_graph import DQNPolicyGraph
+from ray.rllib.evaluation.metrics import collect_metrics
+from ray.rllib.utils.schedules import ConstantSchedule, LinearSchedule
+from ray.tune.trial import Resources
+
+
+OPTIMIZER_SHARED_CONFIGS = [
+    "buffer_size", "prioritized_replay", "prioritized_replay_alpha",
+    "prioritized_replay_beta", "prioritized_replay_eps", "sample_batch_size",
+    "train_batch_size", "learning_starts", "clip_rewards"]
+
+DEFAULT_CONFIG = with_common_config({
+    # === Model ===
+    # Whether to use dueling dqn
+    "dueling": True,
+    # Whether to use double dqn
+    "double_q": True,
+    # Hidden layer sizes of the state and action value networks
+    "hiddens": [256],
+    # N-step Q learning
+    "n_step": 1,
+    # Whether to use rllib or deepmind preprocessors
+    "preprocessor_pref": "deepmind",
+
+    # === Exploration ===
+    # Max num timesteps for annealing schedules. Exploration is annealed from
+    # 1.0 to exploration_fraction over this number of timesteps scaled by
+    # exploration_fraction
+    "schedule_max_timesteps": 100000,
+    # Number of env steps to optimize for before returning
+    "timesteps_per_iteration": 1000,
+    # Fraction of entire training period over which the exploration rate is
+    # annealed
+    "exploration_fraction": 0.1,
+    # Final value of random action probability
+    "exploration_final_eps": 0.02,
+    # Update the target network every `target_network_update_freq` steps.
+    "target_network_update_freq": 500,
+
+    # === Replay buffer ===
+    # Size of the replay buffer. Note that if async_updates is set, then
+    # each worker will have a replay buffer of this size.
+    "buffer_size": 50000,
+    # If True prioritized replay buffer will be used.
+    "prioritized_replay": True,
+    # Alpha parameter for prioritized replay buffer.
+    "prioritized_replay_alpha": 0.6,
+    # Beta parameter for sampling from prioritized replay buffer.
+    "prioritized_replay_beta": 0.4,
+    # Epsilon to add to the TD errors when updating priorities.
+    "prioritized_replay_eps": 1e-6,
+    # Whether to clip rewards to [-1, 1] prior to adding to the replay buffer.
+    "clip_rewards": True,
+    # Whether to LZ4 compress observations
+    "compress_observations": True,
+
+    # === Optimization ===
+    # Learning rate for adam optimizer
+    "lr": 5e-4,
+    # If not None, clip gradients during optimization at this value
+    "grad_norm_clipping": 40,
+    # How many steps of the model to sample before learning starts.
+    "learning_starts": 1000,
+    # Update the replay buffer with this many samples at once. Note that
+    # this setting applies per-worker if num_workers > 1.
+    "sample_batch_size": 4,
+    # Size of a batched sampled from replay buffer for training. Note that
+    # if async_updates is set, then each worker returns gradients for a
+    # batch of this size.
+    "train_batch_size": 32,
+
+    # === Parallelism ===
+    # Whether to use a GPU for local optimization.
+    "gpu": False,
+    # Number of workers for collecting samples with. This only makes sense
+    # to increase if your environment is particularly slow to sample, or if
+    # you"re using the Async or Ape-X optimizers.
+    "num_workers": 0,
+    # Whether to allocate GPUs for workers (if > 0).
+    "num_gpus_per_worker": 0,
+    # Whether to allocate CPUs for workers (if > 0).
+    "num_cpus_per_worker": 1,
+    # Optimizer class to use.
+    "optimizer_class": "SyncReplayOptimizer",
+    # Whether to use a distribution of epsilons across workers for exploration.
+    "per_worker_exploration": False,
+    # Whether to compute priorities on workers.
+    "worker_side_prioritization": False,
+})
+
+
+class DQNAgent(Agent):
+    """DQN implementation in TensorFlow."""
+
+    _agent_name = "DQN"
+    _default_config = DEFAULT_CONFIG
+    _policy_graph = DQNPolicyGraph
+
+    @classmethod
+    def default_resource_request(cls, config):
+        cf = dict(cls._default_config, **config)
+        return Resources(
+            cpu=1, gpu=cf["gpu"] and 1 or 0,
+            extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"],
+            extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"])
+
+    def _init(self):
+        # Update effective batch size to include n-step
+        adjusted_batch_size = (
+            self.config["sample_batch_size"] + self.config["n_step"] - 1)
+        self.config["sample_batch_size"] = adjusted_batch_size
+
+        self.exploration0 = self._make_exploration_schedule(0)
+        self.explorations = [
+            self._make_exploration_schedule(i)
+            for i in range(self.config["num_workers"])]
+
+        for k in OPTIMIZER_SHARED_CONFIGS:
+            if k not in self.config["optimizer"]:
+                self.config["optimizer"][k] = self.config[k]
+
+        self.local_evaluator = self.make_local_evaluator(
+            self.env_creator, self._policy_graph)
+        self.remote_evaluators = self.make_remote_evaluators(
+            self.env_creator, self._policy_graph, self.config["num_workers"],
+            {"num_cpus": self.config["num_cpus_per_worker"],
+             "num_gpus": self.config["num_gpus_per_worker"]})
+        self.optimizer = getattr(optimizers, self.config["optimizer_class"])(
+            self.config["optimizer"], self.local_evaluator,
+            self.remote_evaluators)
+
+        self.last_target_update_ts = 0
+        self.num_target_updates = 0
+
+    def _make_exploration_schedule(self, worker_index):
+        # Use either a different `eps` per worker, or a linear schedule.
+        if self.config["per_worker_exploration"]:
+            assert self.config["num_workers"] > 1, \
+                "This requires multiple workers"
+            return ConstantSchedule(
+                0.4 ** (
+                    1 + worker_index / float(
+                        self.config["num_workers"] - 1) * 7))
+        return LinearSchedule(
+            schedule_timesteps=int(
+                self.config["exploration_fraction"] *
+                self.config["schedule_max_timesteps"]),
+            initial_p=1.0,
+            final_p=self.config["exploration_final_eps"])
+
+    @property
+    def global_timestep(self):
+        return self.optimizer.num_steps_sampled
+
+    def update_target_if_needed(self):
+        if self.global_timestep - self.last_target_update_ts > \
+                self.config["target_network_update_freq"]:
+            self.local_evaluator.foreach_policy(lambda p, _: p.update_target())
+            self.last_target_update_ts = self.global_timestep
+            self.num_target_updates += 1
+
+    def _train(self):
+        start_timestep = self.global_timestep
+
+        while (self.global_timestep - start_timestep <
+               self.config["timesteps_per_iteration"]):
+            self.optimizer.step()
+            self.update_target_if_needed()
+
+        exp_vals = [self.exploration0.value(self.global_timestep)]
+        self.local_evaluator.foreach_policy(
+            lambda p, _: p.set_epsilon(exp_vals[0]))
+        for i, e in enumerate(self.remote_evaluators):
+            exp_val = self.explorations[i].value(self.global_timestep)
+            e.foreach_policy.remote(lambda p, _: p.set_epsilon(exp_val))
+            exp_vals.append(exp_val)
+
+        result = collect_metrics(
+            self.local_evaluator, self.remote_evaluators)
+        return result._replace(
+            info=dict({
+                "min_exploration": min(exp_vals),
+                "max_exploration": max(exp_vals),
+                "num_target_updates": self.num_target_updates,
+            }, **self.optimizer.stats()))
+
+    def _stop(self):
+        # workaround for https://github.com/ray-project/ray/issues/1516
+        for ev in self.remote_evaluators:
+            ev.__ray_terminate__.remote()
+
+    def _save(self, checkpoint_dir):
+        checkpoint_path = os.path.join(
+            checkpoint_dir, "checkpoint-{}".format(self.iteration))
+        extra_data = [
+            self.local_evaluator.save(),
+            ray.get([e.save.remote() for e in self.remote_evaluators]),
+            self.optimizer.save(),
+            self.num_target_updates,
+            self.last_target_update_ts]
+        pickle.dump(extra_data, open(checkpoint_path + ".extra_data", "wb"))
+        return checkpoint_path
+
+    def _restore(self, checkpoint_path):
+        extra_data = pickle.load(open(checkpoint_path + ".extra_data", "rb"))
+        self.local_evaluator.restore(extra_data[0])
+        ray.get([
+            e.restore.remote(d) for (d, e)
+            in zip(extra_data[1], self.remote_evaluators)])
+        self.optimizer.restore(extra_data[2])
+        self.num_target_updates = extra_data[3]
+        self.last_target_update_ts = extra_data[4]

python/ray/rllib/agents/dqn/dqn_policy_graph.py

@@ -0,0 +1,336 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from gym.spaces import Discrete
+import numpy as np
+import tensorflow as tf
+import tensorflow.contrib.layers as layers
+
+import ray
+from ray.rllib.models import ModelCatalog
+from ray.rllib.evaluation.sample_batch import SampleBatch
+from ray.rllib.utils.error import UnsupportedSpaceException
+from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph
+
+
+Q_SCOPE = "q_func"
+Q_TARGET_SCOPE = "target_q_func"
+
+
+class QNetwork(object):
+    def __init__(self, model, num_actions, dueling=False, hiddens=[256]):
+        with tf.variable_scope("action_value"):
+            action_out = model.last_layer
+            for hidden in hiddens:
+                action_out = layers.fully_connected(
+                    action_out, num_outputs=hidden, activation_fn=tf.nn.relu)
+            action_scores = layers.fully_connected(
+                action_out, num_outputs=num_actions, activation_fn=None)
+
+        if dueling:
+            with tf.variable_scope("state_value"):
+                state_out = model.last_layer
+                for hidden in hiddens:
+                    state_out = layers.fully_connected(
+                        state_out, num_outputs=hidden,
+                        activation_fn=tf.nn.relu)
+                state_score = layers.fully_connected(
+                    state_out, num_outputs=1, activation_fn=None)
+            action_scores_mean = tf.reduce_mean(action_scores, 1)
+            action_scores_centered = action_scores - tf.expand_dims(
+                action_scores_mean, 1)
+            self.value = state_score + action_scores_centered
+        else:
+            self.value = action_scores
+
+
+class QValuePolicy(object):
+    def __init__(self, q_values, observations, num_actions, stochastic, eps):
+        deterministic_actions = tf.argmax(q_values, axis=1)
+        batch_size = tf.shape(observations)[0]
+        random_actions = tf.random_uniform(
+            tf.stack([batch_size]), minval=0, maxval=num_actions,
+            dtype=tf.int64)
+        chose_random = tf.random_uniform(
+            tf.stack([batch_size]), minval=0, maxval=1, dtype=tf.float32) < eps
+        stochastic_actions = tf.where(
+            chose_random, random_actions, deterministic_actions)
+        self.action = tf.cond(
+            stochastic, lambda: stochastic_actions,
+            lambda: deterministic_actions)
+
+
+class QLoss(object):
+    def __init__(
+            self, q_t_selected, q_tp1_best, importance_weights, rewards,
+            done_mask, gamma=0.99, n_step=1):
+
+        q_tp1_best_masked = (1.0 - done_mask) * q_tp1_best
+
+        # compute RHS of bellman equation
+        q_t_selected_target = rewards + gamma ** n_step * q_tp1_best_masked
+
+        # compute the error (potentially clipped)
+        self.td_error = q_t_selected - tf.stop_gradient(q_t_selected_target)
+        self.loss = tf.reduce_mean(
+            importance_weights * _huber_loss(self.td_error))
+
+
+class DQNPolicyGraph(TFPolicyGraph):
+    def __init__(self, observation_space, action_space, config):
+        config = dict(ray.rllib.agents.dqn.dqn.DEFAULT_CONFIG, **config)
+        if not isinstance(action_space, Discrete):
+            raise UnsupportedSpaceException(
+                "Action space {} is not supported for DQN.".format(
+                    action_space))
+
+        self.config = config
+        self.cur_epsilon = 1.0
+        num_actions = action_space.n
+
+        def _build_q_network(obs):
+            return QNetwork(
+                ModelCatalog.get_model(obs, 1, config["model"]),
+                num_actions, config["dueling"], config["hiddens"]).value
+
+        # Action inputs
+        self.stochastic = tf.placeholder(tf.bool, (), name="stochastic")
+        self.eps = tf.placeholder(tf.float32, (), name="eps")
+        self.cur_observations = tf.placeholder(
+            tf.float32, shape=(None,) + observation_space.shape)
+
+        # Action Q network
+        with tf.variable_scope(Q_SCOPE) as scope:
+            q_values = _build_q_network(self.cur_observations)
+            self.q_func_vars = _scope_vars(scope.name)
+
+        # Action outputs
+        self.output_actions = QValuePolicy(
+            q_values,
+            self.cur_observations,
+            num_actions,
+            self.stochastic,
+            self.eps).action
+
+        # Replay inputs
+        self.obs_t = tf.placeholder(
+            tf.float32, shape=(None,) + observation_space.shape)
+        self.act_t = tf.placeholder(tf.int32, [None], name="action")
+        self.rew_t = tf.placeholder(tf.float32, [None], name="reward")
+        self.obs_tp1 = tf.placeholder(
+            tf.float32, shape=(None,) + observation_space.shape)
+        self.done_mask = tf.placeholder(tf.float32, [None], name="done")
+        self.importance_weights = tf.placeholder(
+            tf.float32, [None], name="weight")
+
+        # q network evaluation
+        with tf.variable_scope(Q_SCOPE, reuse=True):
+            q_t = _build_q_network(self.obs_t)
+
+        # target q network evalution
+        with tf.variable_scope(Q_TARGET_SCOPE) as scope:
+            q_tp1 = _build_q_network(self.obs_tp1)
+            self.target_q_func_vars = _scope_vars(scope.name)
+
+        # q scores for actions which we know were selected in the given state.
+        q_t_selected = tf.reduce_sum(
+            q_t * tf.one_hot(self.act_t, num_actions), 1)
+
+        # compute estimate of best possible value starting from state at t + 1
+        if config["double_q"]:
+            with tf.variable_scope(Q_SCOPE, reuse=True):
+                q_tp1_using_online_net = _build_q_network(self.obs_tp1)
+            q_tp1_best_using_online_net = tf.argmax(q_tp1_using_online_net, 1)
+            q_tp1_best = tf.reduce_sum(
+                q_tp1 * tf.one_hot(
+                    q_tp1_best_using_online_net, num_actions), 1)
+        else:
+            q_tp1_best = tf.reduce_max(q_tp1, 1)
+
+        self.loss = QLoss(
+            q_t_selected, q_tp1_best, self.importance_weights,
+            self.rew_t, self.done_mask, config["gamma"], config["n_step"])
+
+        # update_target_fn will be called periodically to copy Q network to
+        # target Q network
+        update_target_expr = []
+        for var, var_target in zip(
+            sorted(self.q_func_vars, key=lambda v: v.name),
+                sorted(self.target_q_func_vars, key=lambda v: v.name)):
+            update_target_expr.append(var_target.assign(var))
+        self.update_target_expr = tf.group(*update_target_expr)
+
+        # initialize TFPolicyGraph
+        self.sess = tf.get_default_session()
+        self.loss_inputs = [
+            ("obs", self.obs_t),
+            ("actions", self.act_t),
+            ("rewards", self.rew_t),
+            ("new_obs", self.obs_tp1),
+            ("dones", self.done_mask),
+            ("weights", self.importance_weights),
+        ]
+        self.is_training = tf.placeholder_with_default(True, ())
+        TFPolicyGraph.__init__(
+            self, observation_space, action_space, self.sess,
+            obs_input=self.cur_observations,
+            action_sampler=self.output_actions, loss=self.loss.loss,
+            loss_inputs=self.loss_inputs, is_training=self.is_training)
+        self.sess.run(tf.global_variables_initializer())
+
+    def optimizer(self):
+        return tf.train.AdamOptimizer(learning_rate=self.config["lr"])
+
+    def gradients(self, optimizer):
+        if self.config["grad_norm_clipping"] is not None:
+            grads_and_vars = _minimize_and_clip(
+                optimizer, self.loss.loss, var_list=self.q_func_vars,
+                clip_val=self.config["grad_norm_clipping"])
+        else:
+            grads_and_vars = optimizer.compute_gradients(
+                self.loss.loss, var_list=self.q_func_vars)
+        grads_and_vars = [
+            (g, v) for (g, v) in grads_and_vars if g is not None]
+        return grads_and_vars
+
+    def extra_compute_action_feed_dict(self):
+        return {
+            self.stochastic: True,
+            self.eps: self.cur_epsilon,
+        }
+
+    def extra_compute_grad_fetches(self):
+        return {
+            "td_error": self.loss.td_error,
+        }
+
+    def postprocess_trajectory(self, sample_batch, other_agent_batches=None):
+        return _postprocess_dqn(self, sample_batch)
+
+    def compute_td_error(
+            self, obs_t, act_t, rew_t, obs_tp1, done_mask, importance_weights):
+        td_err = self.sess.run(
+            self.loss.td_error,
+            feed_dict={
+                self.obs_t: [np.array(ob) for ob in obs_t],
+                self.act_t: act_t,
+                self.rew_t: rew_t,
+                self.obs_tp1: [np.array(ob) for ob in obs_tp1],
+                self.done_mask: done_mask,
+                self.importance_weights: importance_weights
+            })
+        return td_err
+
+    def update_target(self):
+        return self.sess.run(self.update_target_expr)
+
+    def set_epsilon(self, epsilon):
+        self.cur_epsilon = epsilon
+
+    def get_state(self):
+        return [TFPolicyGraph.get_state(self), self.cur_epsilon]
+
+    def set_state(self, state):
+        TFPolicyGraph.set_state(self, state[0])
+        self.set_epsilon(state[1])
+
+
+def adjust_nstep(n_step, gamma, obs, actions, rewards, new_obs, dones):
+    """Rewrites the given trajectory fragments to encode n-step rewards.
+
+    reward[i] = (
+        reward[i] * gamma**0 +
+        reward[i+1] * gamma**1 +
+        ... +
+        reward[i+n_step-1] * gamma**(n_step-1))
+
+    The ith new_obs is also adjusted to point to the (i+n_step-1)'th new obs.
+
+    If the episode finishes, the reward will be truncated. After this rewrite,
+    all the arrays will be shortened by (n_step - 1).
+    """
+    for i in range(len(rewards) - n_step + 1):
+        if dones[i]:
+            continue  # episode end
+        for j in range(1, n_step):
+            new_obs[i] = new_obs[i + j]
+            rewards[i] += gamma ** j * rewards[i + j]
+            if dones[i + j]:
+                break  # episode end
+    # truncate ends of the trajectory
+    new_len = len(obs) - n_step + 1
+    for arr in [obs, actions, rewards, new_obs, dones]:
+        del arr[new_len:]
+
+
+def _postprocess_dqn(policy_graph, sample_batch):
+    obs, actions, rewards, new_obs, dones = [
+        list(x) for x in sample_batch.columns(
+            ["obs", "actions", "rewards", "new_obs", "dones"])]
+
+    # N-step Q adjustments
+    if policy_graph.config["n_step"] > 1:
+        adjust_nstep(
+            policy_graph.config["n_step"], policy_graph.config["gamma"],
+            obs, actions, rewards, new_obs, dones)
+
+    batch = SampleBatch({
+        "obs": obs, "actions": actions, "rewards": rewards,
+        "new_obs": new_obs, "dones": dones,
+        "weights": np.ones_like(rewards)})
+
+    # Prioritize on the worker side
+    if batch.count > 0 and policy_graph.config["worker_side_prioritization"]:
+        td_errors = policy_graph.compute_td_error(
+            batch["obs"], batch["actions"], batch["rewards"],
+            batch["new_obs"], batch["dones"], batch["weights"])
+        new_priorities = (
+            np.abs(td_errors) + policy_graph.config["prioritized_replay_eps"])
+        batch.data["weights"] = new_priorities
+
+    return batch
+
+
+def _huber_loss(x, delta=1.0):
+    """Reference: https://en.wikipedia.org/wiki/Huber_loss"""
+    return tf.where(
+        tf.abs(x) < delta,
+        tf.square(x) * 0.5,
+        delta * (tf.abs(x) - 0.5 * delta))
+
+
+def _minimize_and_clip(optimizer, objective, var_list, clip_val=10):
+    """Minimized `objective` using `optimizer` w.r.t. variables in
+    `var_list` while ensure the norm of the gradients for each
+    variable is clipped to `clip_val`
+    """
+    gradients = optimizer.compute_gradients(objective, var_list=var_list)
+    for i, (grad, var) in enumerate(gradients):
+        if grad is not None:
+            gradients[i] = (tf.clip_by_norm(grad, clip_val), var)
+    return gradients
+
+
+def _scope_vars(scope, trainable_only=False):
+    """
+    Get variables inside a scope
+    The scope can be specified as a string
+
+    Parameters
+    ----------
+    scope: str or VariableScope
+      scope in which the variables reside.
+    trainable_only: bool
+      whether or not to return only the variables that were marked as
+      trainable.
+
+    Returns
+    -------
+    vars: [tf.Variable]
+      list of variables in `scope`.
+    """
+    return tf.get_collection(
+        tf.GraphKeys.TRAINABLE_VARIABLES
+        if trainable_only else tf.GraphKeys.VARIABLES,
+        scope=scope if isinstance(scope, str) else scope.name)

python/ray/rllib/agents/es/__init__.py

@@ -0,0 +1,3 @@
+from ray.rllib.agents.es.es import (ESAgent, DEFAULT_CONFIG)
+
+__all__ = ["ESAgent", "DEFAULT_CONFIG"]

python/ray/rllib/agents/es/es.py

@@ -0,0 +1,333 @@
+# 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 numpy as np
+import os
+import pickle
+import time
+
+import ray
+from ray.rllib.agents import Agent
+from ray.tune.trial import Resources
+
+from ray.rllib.agents.es import optimizers
+from ray.rllib.agents.es import policies
+from ray.rllib.agents.es import tabular_logger as tlogger
+from ray.rllib.agents.es import utils
+
+
+Result = namedtuple("Result", [
+    "noise_indices", "noisy_returns", "sign_noisy_returns", "noisy_lengths",
+    "eval_returns", "eval_lengths"
+])
+
+
+DEFAULT_CONFIG = {
+    'l2_coeff': 0.005,
+    'noise_stdev': 0.02,
+    'episodes_per_batch': 1000,
+    'timesteps_per_batch': 10000,
+    'eval_prob': 0.003,
+    'return_proc_mode': "centered_rank",
+    'num_workers': 10,
+    'stepsize': 0.01,
+    'observation_filter': "MeanStdFilter",
+    'noise_size': 250000000,
+    'env_config': {},
+}
+
+
+@ray.remote
+def create_shared_noise(count):
+    """Create a large array of noise to be shared by all workers."""
+    seed = 123
+    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, dim):
+        return np.random.randint(0, len(self.noise) - dim + 1)
+
+
+@ray.remote
+class Worker(object):
+    def __init__(self, config, policy_params, env_creator, 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 = env_creator(config["env_config"])
+        from ray.rllib import models
+        self.preprocessor = models.ModelCatalog.get_preprocessor(self.env)
+
+        self.sess = utils.make_session(single_threaded=True)
+        self.policy = policies.GenericPolicy(
+            self.sess, self.env.action_space, self.preprocessor,
+            config["observation_filter"], **policy_params)
+
+    def rollout(self, timestep_limit, add_noise=True):
+        rollout_rewards, rollout_length = policies.rollout(
+            self.policy, self.env, timestep_limit=timestep_limit,
+            add_noise=add_noise)
+        return rollout_rewards, rollout_length
+
+    def do_rollouts(self, params, timestep_limit=None):
+        # Set the network weights.
+        self.policy.set_weights(params)
+
+        noise_indices, returns, sign_returns, lengths = [], [], [], []
+        eval_returns, eval_lengths = [], []
+
+        # Perform some rollouts with noise.
+        task_tstart = time.time()
+        while (len(noise_indices) == 0 or
+               time.time() - task_tstart < self.min_task_runtime):
+
+            if np.random.uniform() < self.config["eval_prob"]:
+                # Do an evaluation run with no perturbation.
+                self.policy.set_weights(params)
+                rewards, length = self.rollout(timestep_limit, add_noise=False)
+                eval_returns.append(rewards.sum())
+                eval_lengths.append(length)
+            else:
+                # Do a regular run with parameter perturbations.
+                noise_index = self.noise.sample_index(self.policy.num_params)
+
+                perturbation = self.config["noise_stdev"] * self.noise.get(
+                    noise_index, 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_weights(params + perturbation)
+                rewards_pos, lengths_pos = self.rollout(timestep_limit)
+
+                self.policy.set_weights(params - perturbation)
+                rewards_neg, lengths_neg = self.rollout(timestep_limit)
+
+                noise_indices.append(noise_index)
+                returns.append([rewards_pos.sum(), rewards_neg.sum()])
+                sign_returns.append(
+                    [np.sign(rewards_pos).sum(), np.sign(rewards_neg).sum()])
+                lengths.append([lengths_pos, lengths_neg])
+
+        return Result(
+            noise_indices=noise_indices,
+            noisy_returns=returns,
+            sign_noisy_returns=sign_returns,
+            noisy_lengths=lengths,
+            eval_returns=eval_returns,
+            eval_lengths=eval_lengths)
+
+
+class ESAgent(Agent):
+    """Large-scale implementation of Evolution Strategies in Ray."""
+
+    _agent_name = "ES"
+    _default_config = DEFAULT_CONFIG
+
+    @classmethod
+    def default_resource_request(cls, config):
+        cf = dict(cls._default_config, **config)
+        return Resources(cpu=1, gpu=0, extra_cpu=cf["num_workers"])
+
+    def _init(self):
+        policy_params = {
+            "action_noise_std": 0.01
+        }
+
+        env = self.env_creator(self.config["env_config"])
+        from ray.rllib import models
+        preprocessor = models.ModelCatalog.get_preprocessor(env)
+
+        self.sess = utils.make_session(single_threaded=False)
+        self.policy = policies.GenericPolicy(
+            self.sess, env.action_space, preprocessor,
+            self.config["observation_filter"], **policy_params)
+        self.optimizer = optimizers.Adam(self.policy, self.config["stepsize"])
+
+        # Create the shared noise table.
+        print("Creating shared noise table.")
+        noise_id = create_shared_noise.remote(self.config["noise_size"])
+        self.noise = SharedNoiseTable(ray.get(noise_id))
+
+        # Create the actors.
+        print("Creating actors.")
+        self.workers = [
+            Worker.remote(
+                self.config, policy_params, self.env_creator, noise_id)
+            for _ in range(self.config["num_workers"])]
+
+        self.episodes_so_far = 0
+        self.timesteps_so_far = 0
+        self.tstart = time.time()
+
+    def _collect_results(self, theta_id, min_episodes, min_timesteps):
+        num_episodes, num_timesteps = 0, 0
+        results = []
+        while num_episodes < min_episodes or num_timesteps < min_timesteps:
+            print(
+                "Collected {} episodes {} timesteps so far this iter".format(
+                    num_episodes, num_timesteps))
+            rollout_ids = [worker.do_rollouts.remote(theta_id)
+                           for worker in self.workers]
+            # Get the results of the rollouts.
+            for result in ray.get(rollout_ids):
+                results.append(result)
+                # Update the number of episodes and the number of timesteps
+                # keeping in mind that result.noisy_lengths is a list of lists,
+                # where the inner lists have length 2.
+                num_episodes += sum(len(pair) for pair
+                                    in result.noisy_lengths)
+                num_timesteps += sum(sum(pair) for pair
+                                     in result.noisy_lengths)
+        return results, num_episodes, num_timesteps
+
+    def _train(self):
+        config = self.config
+
+        step_tstart = time.time()
+        theta = self.policy.get_weights()
+        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.
+        results, num_episodes, num_timesteps = self._collect_results(
+            theta_id,
+            config["episodes_per_batch"],
+            config["timesteps_per_batch"])
+
+        all_noise_indices = []
+        all_training_returns = []
+        all_training_lengths = []
+        all_eval_returns = []
+        all_eval_lengths = []
+
+        # Loop over the results.
+        for result in results:
+            all_eval_returns += result.eval_returns
+            all_eval_lengths += result.eval_lengths
+
+            all_noise_indices += result.noise_indices
+            all_training_returns += result.noisy_returns
+            all_training_lengths += result.noisy_lengths
+
+        assert len(all_eval_returns) == len(all_eval_lengths)
+        assert (len(all_noise_indices) == len(all_training_returns) ==
+                len(all_training_lengths))
+
+        self.episodes_so_far += num_episodes
+        self.timesteps_so_far += num_timesteps
+
+        # Assemble the results.
+        eval_returns = np.array(all_eval_returns)
+        eval_lengths = np.array(all_eval_lengths)
+        noise_indices = np.array(all_noise_indices)
+        noisy_returns = np.array(all_training_returns)
+        noisy_lengths = np.array(all_training_lengths)
+
+        # Process the returns.
+        if config["return_proc_mode"] == "centered_rank":
+            proc_noisy_returns = utils.compute_centered_ranks(noisy_returns)
+        else:
+            raise NotImplementedError(config["return_proc_mode"])
+
+        # Compute and take a step.
+        g, count = utils.batched_weighted_sum(
+            proc_noisy_returns[:, 0] - proc_noisy_returns[:, 1],
+            (self.noise.get(index, self.policy.num_params)
+             for index in noise_indices),
+            batch_size=500)
+        g /= noisy_returns.size
+        assert (
+            g.shape == (self.policy.num_params,) and
+            g.dtype == np.float32 and
+            count == len(noise_indices))
+        # Compute the new weights theta.
+        theta, update_ratio = self.optimizer.update(
+            -g + config["l2_coeff"] * theta)
+        # Set the new weights in the local copy of the policy.
+        self.policy.set_weights(theta)
+
+        step_tend = time.time()
+        tlogger.record_tabular("EvalEpRewMean", eval_returns.mean())
+        tlogger.record_tabular("EvalEpRewStd", eval_returns.std())
+        tlogger.record_tabular("EvalEpLenMean", eval_lengths.mean())
+
+        tlogger.record_tabular("EpRewMean", noisy_returns.mean())
+        tlogger.record_tabular("EpRewStd", noisy_returns.std())
+        tlogger.record_tabular("EpLenMean", noisy_lengths.mean())
+
+        tlogger.record_tabular("Norm", float(np.square(theta).sum()))
+        tlogger.record_tabular("GradNorm", float(np.square(g).sum()))
+        tlogger.record_tabular("UpdateRatio", float(update_ratio))
+
+        tlogger.record_tabular("EpisodesThisIter", noisy_lengths.size)
+        tlogger.record_tabular("EpisodesSoFar", self.episodes_so_far)
+        tlogger.record_tabular("TimestepsThisIter", noisy_lengths.sum())
+        tlogger.record_tabular("TimestepsSoFar", self.timesteps_so_far)
+
+        tlogger.record_tabular("TimeElapsedThisIter", step_tend - step_tstart)
+        tlogger.record_tabular("TimeElapsed", step_tend - self.tstart)
+        tlogger.dump_tabular()
+
+        info = {
+            "weights_norm": np.square(theta).sum(),
+            "grad_norm": np.square(g).sum(),
+            "update_ratio": update_ratio,
+            "episodes_this_iter": noisy_lengths.size,
+            "episodes_so_far": self.episodes_so_far,
+            "timesteps_this_iter": noisy_lengths.sum(),
+            "timesteps_so_far": self.timesteps_so_far,
+            "time_elapsed_this_iter": step_tend - step_tstart,
+            "time_elapsed": step_tend - self.tstart
+        }
+
+        result = ray.tune.result.TrainingResult(
+            episode_reward_mean=eval_returns.mean(),
+            episode_len_mean=eval_lengths.mean(),
+            timesteps_this_iter=noisy_lengths.sum(),
+            info=info)
+
+        return result
+
+    def _stop(self):
+        # workaround for https://github.com/ray-project/ray/issues/1516
+        for w in self.workers:
+            w.__ray_terminate__.remote()
+
+    def _save(self, checkpoint_dir):
+        checkpoint_path = os.path.join(
+            checkpoint_dir, "checkpoint-{}".format(self.iteration))
+        weights = self.policy.get_weights()
+        objects = [
+            weights,
+            self.episodes_so_far,
+            self.timesteps_so_far]
+        pickle.dump(objects, open(checkpoint_path, "wb"))
+        return checkpoint_path
+
+    def _restore(self, checkpoint_path):
+        objects = pickle.load(open(checkpoint_path, "rb"))
+        self.policy.set_weights(objects[0])
+        self.episodes_so_far = objects[1]
+        self.timesteps_so_far = objects[2]
+
+    def compute_action(self, observation):
+        return self.policy.compute(observation, update=False)[0]

python/ray/rllib/agents/es/optimizers.py

@@ -0,0 +1,56 @@
+# 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_weights()
+        ratio = np.linalg.norm(step) / np.linalg.norm(theta)
+        return theta + step, 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/agents/es/policies.py

@@ -0,0 +1,80 @@
+# 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 gym
+import numpy as np
+import tensorflow as tf
+
+import ray
+from ray.rllib.models import ModelCatalog
+from ray.rllib.utils.filter import get_filter
+
+
+def rollout(policy, env, timestep_limit=None, add_noise=False):
+    """Do a rollout.
+
+    If add_noise is True, the rollout will take noisy actions with
+    noise drawn from that stream. Otherwise, no action noise will be added.
+    """
+    env_timestep_limit = env.spec.max_episode_steps
+    timestep_limit = (env_timestep_limit if timestep_limit is None
+                      else min(timestep_limit, env_timestep_limit))
+    rews = []
+    t = 0
+    observation = env.reset()
+    for _ in range(timestep_limit or 999999):
+        ac = policy.compute(observation, add_noise=add_noise)[0]
+        observation, rew, done, _ = env.step(ac)
+        rews.append(rew)
+        t += 1
+        if done:
+            break
+    rews = np.array(rews, dtype=np.float32)
+    return rews, t
+
+
+class GenericPolicy(object):
+    def __init__(self, sess, action_space, preprocessor,
+                 observation_filter, action_noise_std):
+        self.sess = sess
+        self.action_space = action_space
+        self.action_noise_std = action_noise_std
+        self.preprocessor = preprocessor
+        self.observation_filter = get_filter(
+            observation_filter, self.preprocessor.shape)
+        self.inputs = tf.placeholder(
+            tf.float32, [None] + list(self.preprocessor.shape))
+
+        # Policy network.
+        dist_class, dist_dim = ModelCatalog.get_action_dist(
+            self.action_space, dist_type="deterministic")
+        model = ModelCatalog.get_model(self.inputs, dist_dim)
+        dist = dist_class(model.outputs)
+        self.sampler = dist.sample()
+
+        self.variables = ray.experimental.TensorFlowVariables(
+            model.outputs, self.sess)
+
+        self.num_params = sum(np.prod(variable.shape.as_list())
+                              for _, variable
+                              in self.variables.variables.items())
+        self.sess.run(tf.global_variables_initializer())
+
+    def compute(self, observation, add_noise=False, update=True):
+        observation = self.preprocessor.transform(observation)
+        observation = self.observation_filter(observation[None], update=update)
+        action = self.sess.run(self.sampler,
+                               feed_dict={self.inputs: observation})
+        if add_noise and isinstance(self.action_space, gym.spaces.Box):
+            action += np.random.randn(*action.shape) * self.action_noise_std
+        return action
+
+    def set_weights(self, x):
+        self.variables.set_flat(x)
+
+    def get_weights(self):
+        return self.variables.get_flat()

python/ray/rllib/agents/es/tabular_logger.py

@@ -0,0 +1,225 @@
+# 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/agents/es/utils.py

@@ -0,0 +1,59 @@
+# 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.Session()
+    return tf.Session(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

python/ray/rllib/agents/pg/__init__.py

@@ -0,0 +1,3 @@
+from ray.rllib.agents.pg.pg import PGAgent, DEFAULT_CONFIG
+
+__all__ = ["PGAgent", "DEFAULT_CONFIG"]

python/ray/rllib/agents/pg/pg.py

@@ -0,0 +1,54 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from ray.rllib.agents.agent import Agent, with_common_config
+from ray.rllib.agents.pg.pg_policy_graph import PGPolicyGraph
+from ray.rllib.evaluation.metrics import collect_metrics
+from ray.rllib.optimizers import SyncSamplesOptimizer
+from ray.tune.trial import Resources
+
+
+DEFAULT_CONFIG = with_common_config({
+    # No remote workers by default
+    "num_workers": 0,
+    # Learning rate
+    "lr": 0.0004,
+    # Override model config
+    "model": {
+        # Use LSTM model.
+        "use_lstm": False,
+        # Max seq length for LSTM training.
+        "max_seq_len": 20,
+    },
+})
+
+
+class PGAgent(Agent):
+    """Simple policy gradient agent.
+
+    This is an example agent to show how to implement algorithms in RLlib.
+    In most cases, you will probably want to use the PPO agent instead.
+    """
+
+    _agent_name = "PG"
+    _default_config = DEFAULT_CONFIG
+
+    @classmethod
+    def default_resource_request(cls, config):
+        cf = dict(cls._default_config, **config)
+        return Resources(cpu=1, gpu=0, extra_cpu=cf["num_workers"])
+
+    def _init(self):
+        self.local_evaluator = self.make_local_evaluator(
+            self.env_creator, PGPolicyGraph)
+        self.remote_evaluators = self.make_remote_evaluators(
+            self.env_creator, PGPolicyGraph, self.config["num_workers"], {})
+        self.optimizer = SyncSamplesOptimizer(
+            self.config["optimizer"], self.local_evaluator,
+            self.remote_evaluators)
+
+    def _train(self):
+        self.optimizer.step()
+        return collect_metrics(
+            self.optimizer.local_evaluator, self.optimizer.remote_evaluators)

python/ray/rllib/agents/pg/pg_policy_graph.py

@@ -0,0 +1,64 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+import ray
+from ray.rllib.models.catalog import ModelCatalog
+from ray.rllib.evaluation.postprocessing import compute_advantages
+from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph
+
+
+class PGLoss(object):
+    def __init__(self, action_dist, actions, advantages):
+        self.loss = -tf.reduce_mean(action_dist.logp(actions) * advantages)
+
+
+class PGPolicyGraph(TFPolicyGraph):
+    def __init__(self, obs_space, action_space, config):
+        config = dict(ray.rllib.agents.pg.pg.DEFAULT_CONFIG, **config)
+        self.config = config
+
+        # Setup policy
+        obs = tf.placeholder(tf.float32, shape=[None] + list(obs_space.shape))
+        dist_class, self.logit_dim = ModelCatalog.get_action_dist(
+            action_space, self.config["model"])
+        self.model = ModelCatalog.get_model(
+            obs, self.logit_dim, options=self.config["model"])
+        action_dist = dist_class(self.model.outputs)  # logit for each action
+
+        # Setup policy loss
+        actions = ModelCatalog.get_action_placeholder(action_space)
+        advantages = tf.placeholder(tf.float32, [None], name="adv")
+        loss = PGLoss(action_dist, actions, advantages).loss
+
+        # Initialize TFPolicyGraph
+        sess = tf.get_default_session()
+        loss_in = [
+            ("obs", obs),
+            ("actions", actions),
+            ("advantages", advantages),
+        ]
+
+        # LSTM support
+        for i, ph in enumerate(self.model.state_in):
+            loss_in.append(("state_in_{}".format(i), ph))
+
+        is_training = tf.placeholder_with_default(True, ())
+        TFPolicyGraph.__init__(
+            self, obs_space, action_space, sess, obs_input=obs,
+            action_sampler=action_dist.sample(), loss=loss,
+            loss_inputs=loss_in, is_training=is_training,
+            state_inputs=self.model.state_in,
+            state_outputs=self.model.state_out,
+            seq_lens=self.model.seq_lens,
+            max_seq_len=config["model"]["max_seq_len"])
+        sess.run(tf.global_variables_initializer())
+
+    def postprocess_trajectory(self, sample_batch, other_agent_batches=None):
+        return compute_advantages(
+            sample_batch, 0.0, self.config["gamma"], use_gae=False)
+
+    def get_initial_state(self):
+        return self.model.state_init

python/ray/rllib/agents/ppo/__init__.py

@@ -0,0 +1,3 @@
+from ray.rllib.agents.ppo.ppo import (PPOAgent, DEFAULT_CONFIG)
+
+__all__ = ["PPOAgent", "DEFAULT_CONFIG"]

python/ray/rllib/agents/ppo/ppo.py

@@ -0,0 +1,140 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import numpy as np
+import pickle
+
+import ray
+from ray.rllib.agents import Agent, with_common_config
+from ray.rllib.agents.ppo.ppo_tf_policy import PPOTFPolicyGraph
+from ray.rllib.evaluation.metrics import collect_metrics
+from ray.rllib.utils import FilterManager
+from ray.rllib.optimizers.multi_gpu_optimizer import LocalMultiGPUOptimizer
+from ray.tune.trial import Resources
+
+DEFAULT_CONFIG = with_common_config({
+    # If true, use the Generalized Advantage Estimator (GAE)
+    # with a value function, see https://arxiv.org/pdf/1506.02438.pdf.
+    "use_gae": True,
+    # GAE(lambda) parameter
+    "lambda": 1.0,
+    # Initial coefficient for KL divergence
+    "kl_coeff": 0.2,
+    # Number of timesteps collected for each SGD round
+    "timesteps_per_batch": 4000,
+    # Number of SGD iterations in each outer loop
+    "num_sgd_iter": 30,
+    # Stepsize of SGD
+    "sgd_stepsize": 5e-5,
+    # Total SGD batch size across all devices for SGD
+    "sgd_batchsize": 128,
+    # Coefficient of the value function loss
+    "vf_loss_coeff": 1.0,
+    # Coefficient of the entropy regularizer
+    "entropy_coeff": 0.0,
+    # PPO clip parameter
+    "clip_param": 0.3,
+    # Target value for KL divergence
+    "kl_target": 0.01,
+    # Number of GPUs to use for SGD
+    "num_gpus": 0,
+    # Whether to allocate GPUs for workers (if > 0).
+    "num_gpus_per_worker": 0,
+    # Whether to allocate CPUs for workers (if > 0).
+    "num_cpus_per_worker": 1,
+    # Whether to rollout "complete_episodes" or "truncate_episodes"
+    "batch_mode": "complete_episodes",
+    # Which observation filter to apply to the observation
+    "observation_filter": "MeanStdFilter",
+})
+
+
+class PPOAgent(Agent):
+    """Multi-GPU optimized implementation of PPO in TensorFlow."""
+
+    _agent_name = "PPO"
+    _default_config = DEFAULT_CONFIG
+
+    @classmethod
+    def default_resource_request(cls, config):
+        cf = dict(cls._default_config, **config)
+        return Resources(
+            cpu=1,
+            gpu=cf["num_gpus"],
+            extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"],
+            extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"])
+
+    def _init(self):
+        self.local_evaluator = self.make_local_evaluator(
+            self.env_creator, PPOTFPolicyGraph)
+        self.remote_evaluators = self.make_remote_evaluators(
+            self.env_creator, PPOTFPolicyGraph, self.config["num_workers"],
+            {"num_cpus": self.config["num_cpus_per_worker"],
+             "num_gpus": self.config["num_gpus_per_worker"]})
+        self.optimizer = LocalMultiGPUOptimizer(
+            {"sgd_batch_size": self.config["sgd_batchsize"],
+             "sgd_stepsize": self.config["sgd_stepsize"],
+             "num_sgd_iter": self.config["num_sgd_iter"],
+             "timesteps_per_batch": self.config["timesteps_per_batch"]},
+            self.local_evaluator, self.remote_evaluators)
+
+    def _train(self):
+        def postprocess_samples(batch):
+            # Divide by the maximum of value.std() and 1e-4
+            # to guard against the case where all values are equal
+            value = batch["advantages"]
+            standardized = (value - value.mean()) / max(1e-4, value.std())
+            batch.data["advantages"] = standardized
+            batch.shuffle()
+            dummy = np.zeros_like(batch["advantages"])
+            if not self.config["use_gae"]:
+                batch.data["value_targets"] = dummy
+                batch.data["vf_preds"] = dummy
+        extra_fetches = self.optimizer.step(postprocess_fn=postprocess_samples)
+        kl = np.array(extra_fetches["kl"]).mean(axis=1)[-1]
+        total_loss = np.array(extra_fetches["total_loss"]).mean(axis=1)[-1]
+        policy_loss = np.array(extra_fetches["policy_loss"]).mean(axis=1)[-1]
+        vf_loss = np.array(extra_fetches["vf_loss"]).mean(axis=1)[-1]
+        entropy = np.array(extra_fetches["entropy"]).mean(axis=1)[-1]
+
+        newkl = self.local_evaluator.for_policy(lambda pi: pi.update_kl(kl))
+
+        info = {
+            "kl_divergence": kl,
+            "kl_coefficient": newkl,
+            "total_loss": total_loss,
+            "policy_loss": policy_loss,
+            "vf_loss": vf_loss,
+            "entropy": entropy,
+        }
+
+        FilterManager.synchronize(
+            self.local_evaluator.filters, self.remote_evaluators)
+        res = collect_metrics(self.local_evaluator, self.remote_evaluators)
+        res = res._replace(info=info)
+        return res
+
+    def _stop(self):
+        # workaround for https://github.com/ray-project/ray/issues/1516
+        for ev in self.remote_evaluators:
+            ev.__ray_terminate__.remote()
+
+    def _save(self, checkpoint_dir):
+        checkpoint_path = os.path.join(checkpoint_dir,
+                                       "checkpoint-{}".format(self.iteration))
+        agent_state = ray.get(
+            [a.save.remote() for a in self.remote_evaluators])
+        extra_data = [
+            self.local_evaluator.save(),
+            agent_state]
+        pickle.dump(extra_data, open(checkpoint_path + ".extra_data", "wb"))
+        return checkpoint_path
+
+    def _restore(self, checkpoint_path):
+        extra_data = pickle.load(open(checkpoint_path + ".extra_data", "rb"))
+        self.local_evaluator.restore(extra_data[0])
+        ray.get([
+            a.restore.remote(o)
+                for (a, o) in zip(self.remote_evaluators, extra_data[1])])

python/ray/rllib/agents/ppo/ppo_tf_policy.py

@@ -0,0 +1,199 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from ray.rllib.evaluation.postprocessing import compute_advantages
+from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph
+from ray.rllib.models.catalog import ModelCatalog
+
+
+class PPOLoss(object):
+    def __init__(
+            self, action_space, value_targets, advantages, actions, logprobs,
+            vf_preds, curr_action_dist, value_fn, cur_kl_coeff,
+            entropy_coeff=0, clip_param=0.1, vf_loss_coeff=1.0, use_gae=True):
+        """Constructs the loss for Proximal Policy Objective.
+
+        Arguments:
+            action_space: Environment observation space specification.
+            value_targets (Placeholder): Placeholder for target values; used
+                for GAE.
+            actions (Placeholder): Placeholder for actions taken
+                from previous model evaluation.
+            advantages (Placeholder): Placeholder for calculated advantages
+                from previous model evaluation.
+            logprobs (Placeholder): Placeholder for logits output from
+                previous model evaluation.
+            vf_preds (Placeholder): Placeholder for value function output
+                from previous model evaluation.
+            curr_action_dist (ActionDistribution): ActionDistribution
+                of the current model.
+            value_fn (Tensor): Current value function output Tensor.
+            cur_kl_coeff (Variable): Variable holding the current PPO KL
+                coefficient.
+            entropy_coeff (float): Coefficient of the entropy regularizer.
+            clip_param (float): Clip parameter
+            vf_loss_coeff (float): Coefficient of the value function loss
+            use_gae (bool): If true, use the Generalized Advantage Estimator.
+        """
+        dist_cls, _ = ModelCatalog.get_action_dist(action_space)
+        prev_dist = dist_cls(logprobs)
+        # Make loss functions.
+        logp_ratio = tf.exp(
+            curr_action_dist.logp(actions) - prev_dist.logp(actions))
+        action_kl = prev_dist.kl(curr_action_dist)
+        self.mean_kl = tf.reduce_mean(action_kl)
+
+        curr_entropy = curr_action_dist.entropy()
+        self.mean_entropy = tf.reduce_mean(curr_entropy)
+
+        surrogate_loss = tf.minimum(
+            advantages * logp_ratio,
+            advantages * tf.clip_by_value(
+                logp_ratio, 1 - clip_param, 1 + clip_param))
+        self.mean_policy_loss = tf.reduce_mean(-surrogate_loss)
+
+        if use_gae:
+            vf_loss1 = tf.square(value_fn - value_targets)
+            vf_clipped = vf_preds + tf.clip_by_value(
+                value_fn - vf_preds, -clip_param, clip_param)
+            vf_loss2 = tf.square(vf_clipped - value_targets)
+            vf_loss = tf.minimum(vf_loss1, vf_loss2)
+            self.mean_vf_loss = tf.reduce_mean(vf_loss)
+            loss = tf.reduce_mean(
+                -surrogate_loss + cur_kl_coeff*action_kl +
+                vf_loss_coeff*vf_loss - entropy_coeff*curr_entropy)
+        else:
+            self.mean_vf_loss = tf.constant(0.0)
+            loss = tf.reduce_mean(
+                -surrogate_loss + cur_kl_coeff*action_kl -
+                entropy_coeff*curr_entropy)
+        self.loss = loss
+
+
+class PPOTFPolicyGraph(TFPolicyGraph):
+    def __init__(self, observation_space, action_space,
+                 config, existing_inputs=None):
+        """
+        Arguments:
+            observation_space: Environment observation space specification.
+            action_space: Environment action space specification.
+            config (dict): Configuration values for PPO graph.
+            existing_inputs (list): Optional list of tuples that specify the
+                placeholders upon which the graph should be built upon.
+        """
+        self.sess = tf.get_default_session()
+        self.action_space = action_space
+        self.config = config
+        self.kl_coeff_val = self.config["kl_coeff"]
+        self.kl_target = self.config["kl_target"]
+        dist_cls, logit_dim = ModelCatalog.get_action_dist(
+            action_space)
+
+        if existing_inputs:
+            self.loss_in = existing_inputs
+            obs_ph, value_targets_ph, adv_ph, act_ph, \
+                logprobs_ph, vf_preds_ph = [ph for _, ph in existing_inputs]
+        else:
+            obs_ph = tf.placeholder(
+                tf.float32, name="obs", shape=(None,)+observation_space.shape)
+            # Targets of the value function.
+            value_targets_ph = tf.placeholder(
+                tf.float32, name="value_targets", shape=(None,))
+            # Advantage values in the policy gradient estimator.
+            adv_ph = tf.placeholder(
+                tf.float32, name="advantages", shape=(None,))
+            act_ph = ModelCatalog.get_action_placeholder(action_space)
+            # Log probabilities from the policy before the policy update.
+            logprobs_ph = tf.placeholder(
+                tf.float32, name="logprobs", shape=(None, logit_dim))
+            # Value function predictions before the policy update.
+            vf_preds_ph = tf.placeholder(
+                tf.float32, name="vf_preds", shape=(None,))
+            self.loss_in = [
+                ("obs", obs_ph),
+                ("value_targets", value_targets_ph),
+                ("advantages", adv_ph),
+                ("actions", act_ph),
+                ("logprobs", logprobs_ph),
+                ("vf_preds", vf_preds_ph)
+            ]
+        # TODO(ekl) feed RNN states in here
+
+        # KL Coefficient
+        self.kl_coeff = tf.get_variable(
+            initializer=tf.constant_initializer(self.kl_coeff_val),
+            name="kl_coeff", shape=(), trainable=False, dtype=tf.float32)
+
+        self.logits = ModelCatalog.get_model(
+            obs_ph, logit_dim, self.config["model"]).outputs
+        curr_action_dist = dist_cls(self.logits)
+        self.sampler = curr_action_dist.sample()
+        if self.config["use_gae"]:
+            vf_config = self.config["model"].copy()
+            # Do not split the last layer of the value function into
+            # mean parameters and standard deviation parameters and
+            # do not make the standard deviations free variables.
+            vf_config["free_log_std"] = False
+            with tf.variable_scope("value_function"):
+                self.value_function = ModelCatalog.get_model(
+                    obs_ph, 1, vf_config).outputs
+            self.value_function = tf.reshape(self.value_function, [-1])
+        else:
+            self.value_function = tf.constant("NA")
+
+        self.loss_obj = PPOLoss(
+            action_space, value_targets_ph, adv_ph, act_ph,
+            logprobs_ph, vf_preds_ph,
+            curr_action_dist, self.value_function, self.kl_coeff,
+            entropy_coeff=self.config["entropy_coeff"],
+            clip_param=self.config["clip_param"],
+            vf_loss_coeff=self.config["kl_target"],
+            use_gae=self.config["use_gae"])
+        self.is_training = tf.placeholder_with_default(True, ())
+
+        TFPolicyGraph.__init__(
+            self, observation_space, action_space,
+            self.sess, obs_input=obs_ph,
+            action_sampler=self.sampler, loss=self.loss_obj.loss,
+            loss_inputs=self.loss_in,
+            is_training=self.is_training)
+
+    def copy(self, existing_inputs):
+        """Creates a copy of self using existing input placeholders."""
+        return PPOTFPolicyGraph(
+            None, self.action_space, self.config,
+            existing_inputs=existing_inputs)
+
+    def extra_compute_action_fetches(self):
+        return {"vf_preds": self.value_function, "logprobs": self.logits}
+
+    def extra_apply_grad_fetches(self):
+        return {
+            "total_loss": self.loss_obj.loss,
+            "policy_loss": self.loss_obj.mean_policy_loss,
+            "vf_loss": self.loss_obj.mean_vf_loss,
+            "kl": self.loss_obj.mean_kl,
+            "entropy": self.loss_obj.mean_entropy
+        }
+
+    def update_kl(self, sampled_kl):
+        if sampled_kl > 2.0 * self.kl_target:
+            self.kl_coeff_val *= 1.5
+        elif sampled_kl < 0.5 * self.kl_target:
+            self.kl_coeff_val *= 0.5
+        self.kl_coeff.load(self.kl_coeff_val, session=self.sess)
+        return self.kl_coeff_val
+
+    def postprocess_trajectory(self, sample_batch, other_agent_batches=None):
+        last_r = 0.0
+        batch = compute_advantages(
+            sample_batch, last_r, self.config["gamma"],
+            self.config["lambda"], use_gae=self.config["use_gae"])
+        return batch
+
+    def gradients(self, optimizer):
+        return optimizer.compute_gradients(
+            self._loss, colocate_gradients_with_ops=True)

python/ray/rllib/agents/ppo/rollout.py

@@ -0,0 +1,34 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import ray
+from ray.rllib.evaluation.sample_batch import SampleBatch
+
+
+def collect_samples(agents, timesteps_per_batch):
+    num_timesteps_so_far = 0
+    trajectories = []
+    # This variable maps the object IDs of trajectories that are currently
+    # computed to the agent that they are computed on; we start some initial
+    # tasks here.
+
+    agent_dict = {}
+
+    for agent in agents:
+        fut_sample = agent.sample.remote()
+        agent_dict[fut_sample] = agent
+
+    while num_timesteps_so_far < timesteps_per_batch:
+        # TODO(pcm): Make wait support arbitrary iterators and remove the
+        # conversion to list here.
+        [fut_sample], _ = ray.wait(list(agent_dict))
+        agent = agent_dict.pop(fut_sample)
+        # Start task with next trajectory and record it in the dictionary.
+        fut_sample2 = agent.sample.remote()
+        agent_dict[fut_sample2] = agent
+
+        next_sample = ray.get(fut_sample)
+        num_timesteps_so_far += next_sample.count
+        trajectories.append(next_sample)
+    return SampleBatch.concat_samples(trajectories)

python/ray/rllib/agents/ppo/test/test.py

@@ -0,0 +1,62 @@
+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.models.action_dist import Categorical
+from ray.rllib.agents.ppo.utils import flatten, concatenate
+
+
+# TODO(ekl): move to rllib/models dir
+class DistributionsTest(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/agents/ppo/utils.py

@@ -0,0 +1,36 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+
+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["actions"].shape[0])
+    for key, val in trajectory.items():
+        trajectory[key] = val[permutation]
+    return trajectory