[rllib] Add high-performance external application connector (#7641)

.travis.yml

@@ -253,8 +253,8 @@ matrix:
         - if [ $RAY_CI_RLLIB_FULL_AFFECTED != "1" ]; then exit; fi
         - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_A,examples_B --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/...
         - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_C --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/...
-        - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_E,examples_M,examples_P --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/...
-        - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_R,examples_T --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/...
+        - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_E,examples_L,examples_M,examples_P --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/...
+        - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_R,examples_S,examples_T --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/...
 
     # RLlib: tests_dir: Everything in rllib/tests/ directory (A-I).
     - os: linux

doc/source/rllib-env.rst

@@ -298,30 +298,96 @@ In this setup, the appropriate rewards for training lower-level agents must be p
 
 See this file for a runnable example: `hierarchical_training.py <https://github.com/ray-project/ray/blob/master/rllib/examples/hierarchical_training.py>`__.
 
-Interfacing with External Agents
+External Agents and Applications
 --------------------------------
 
 In many situations, it does not make sense for an environment to be "stepped" by RLlib. For example, if a policy is to be used in a web serving system, then it is more natural for an agent to query a service that serves policy decisions, and for that service to learn from experience over time. This case also naturally arises with **external simulators** that run independently outside the control of RLlib, but may still want to leverage RLlib for training.
 
 RLlib provides the `ExternalEnv <https://github.com/ray-project/ray/blob/master/rllib/env/external_env.py>`__ class for this purpose. Unlike other envs, ExternalEnv has its own thread of control. At any point, agents on that thread can query the current policy for decisions via ``self.get_action()`` and reports rewards via ``self.log_returns()``. This can be done for multiple concurrent episodes as well.
 
-ExternalEnv can be used to implement a simple REST policy `server <https://github.com/ray-project/ray/tree/master/rllib/examples/serving>`__ that learns over time using RLlib. In this example RLlib runs with ``num_workers=0`` to avoid port allocation issues, but in principle this could be scaled by increasing ``num_workers``.
-
 Logging off-policy actions
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-ExternalEnv also provides a ``self.log_action()`` call to support off-policy actions. This allows the client to make independent decisions, e.g., to compare two different policies, and for RLlib to still learn from those off-policy actions. Note that this requires the algorithm used to support learning from off-policy decisions (e.g., DQN).
-
-Data ingest
-~~~~~~~~~~~
-
-The ``log_action`` API of ExternalEnv can be used to ingest data from offline logs. The pattern would be as follows: First, some policy is followed to produce experience data which is stored in some offline storage system. Then, RLlib creates a number of workers that use a ExternalEnv to read the logs in parallel and ingest the experiences. After a round of training completes, the new policy can be deployed to collect more experiences.
-
-Note that envs can read from different partitions of the logs based on the ``worker_index`` attribute of the `env context <https://github.com/ray-project/ray/blob/master/rllib/env/env_context.py>`__ passed into the environment constructor.
+ExternalEnv provides a ``self.log_action()`` call to support off-policy actions. This allows the client to make independent decisions, e.g., to compare two different policies, and for RLlib to still learn from those off-policy actions. Note that this requires the algorithm used to support learning from off-policy decisions (e.g., DQN).
 
 .. seealso::
 
-    `Offline Datasets <rllib-offline.html>`__ provide higher-level interfaces for working with offline experience datasets.
+    `Offline Datasets <rllib-offline.html>`__ provide higher-level interfaces for working with off-policy experience datasets.
+
+External Application Clients
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+For applications that are running entirely outside the Ray cluster (i.e., cannot be packaged into a Python environment of any form), RLlib provides the ``PolicyServerInput`` application connector, which can be connected to over the network using ``PolicyClient`` instances.
+
+You can configure any Trainer to launch a policy server with the following config:
+
+.. code-block:: python
+
+    trainer_config = {
+        # An environment class is still required, but it doesn't need to be runnable.
+        # You only need to define its action and observation space attributes.
+        "env": YOUR_ENV_STUB,
+
+        # Use the policy server to generate experiences.
+        "input": (
+            lambda ioctx: PolicyServerInput(ioctx, SERVER_ADDRESS, SERVER_PORT)
+        ),
+        # Use the existing trainer process to run the server.
+        "num_workers": 0,
+        # Disable OPE, since the rollouts are coming from online clients.
+        "input_evaluation": [],
+    }
+
+Clients can then connect in either *local* or *remote* inference mode. In local inference mode, copies of the policy are downloaded from the server and cached on the client for a configurable period of time. This allows actions to be computed by the client without requiring a network round trip each time. In remote inference mode, each computed action requires a network call to the server.
+
+Example:
+
+.. code-block:: python
+
+    client = PolicyClient("http://localhost:9900", inference_mode="local")
+    episode_id = client.start_episode()
+    ...
+    action = client.get_action(episode_id, cur_obs)
+    ...
+    client.end_episode(episode_id, last_obs)
+
+To understand the difference between standard envs, external envs, and connecting with a ``PolicyClient``, refer to the following figure:
+
+.. https://docs.google.com/drawings/d/1hJvT9bVGHVrGTbnCZK29BYQIcYNRbZ4Dr6FOPMJDjUs/edit
+.. image:: rllib-external.svg
+
+Try it yourself by launching a `cartpole_server.py <https://github.com/ray-project/ray/blob/master/rllib/examples/serving/cartpole_server.py>`__, and connecting to it with any number of clients (`cartpole_client.py <https://github.com/ray-project/ray/blob/master/rllib/examples/serving/cartpole_client.py>`__):
+
+.. code-block:: bash
+
+    # Start the server by running:
+    >>> python rllib/examples/serving/cartpole_server.py --run=PPO
+    --
+    -- Starting policy server at localhost:9900
+    --
+
+    # To connect from a client with inference_mode="remote".
+    >>> python rllib/examples/serving/cartpole_client.py --inference-mode=remote
+    Total reward: 10.0
+    Total reward: 58.0
+    ...
+    Total reward: 200.0
+    ...
+
+    # To connect from a client with inference_mode="local" (faster).
+    >>> python rllib/examples/serving/cartpole_client.py --inference-mode=local
+    Querying server for new policy weights.
+    Generating new batch of experiences.
+    Total reward: 13.0
+    Total reward: 11.0
+    ...
+    Sending batch of 1000 steps back to server.
+    Querying server for new policy weights.
+    ...
+    Total reward: 200.0
+    ...
+
+For the best performance, when possible we recommend using ``inference_mode="local"`` when possible.
 
 Advanced Integrations
 ---------------------

doc/source/rllib-toc.rst

@@ -53,7 +53,7 @@ Training APIs
 
    -  `Stack Traces <rllib-training.html#stack-traces>`__
 
-*  `REST API <rllib-training.html#rest-api>`__
+*  `External Application API <rllib-training.html#external-application-api>`__
 
 Environments
 ------------
@@ -62,7 +62,10 @@ Environments
 * `OpenAI Gym <rllib-env.html#openai-gym>`__
 * `Vectorized <rllib-env.html#vectorized>`__
 * `Multi-Agent and Hierarchical <rllib-env.html#multi-agent-and-hierarchical>`__
-* `Interfacing with External Agents <rllib-env.html#interfacing-with-external-agents>`__
+* `External Agents and Applications <rllib-env.html#external-agents-and-applications>`__
+
+   -  `External Application Clients <rllib-env.html#external-application-clients>`__
+
 * `Advanced Integrations <rllib-env.html#advanced-integrations>`__
 
 Models, Preprocessors, and Action Distributions
@@ -185,8 +188,9 @@ try setting ``OMP_NUM_THREADS=1``. Similarly, check configured system limits wit
 If you encounter out-of-memory errors, consider setting ``redis_max_memory`` and ``object_store_memory`` in ``ray.init()`` to reduce memory usage.
 
 For debugging unexpected hangs or performance problems, you can run ``ray stack`` to dump
-the stack traces of all Ray workers on the current node, and ``ray timeline`` to dump
-a timeline visualization of tasks to a file.
+the stack traces of all Ray workers on the current node, ``ray timeline`` to dump
+a timeline visualization of tasks to a file, and ``ray memory`` to list all object
+references in the cluster.
 
 TensorFlow 2.0
 ~~~~~~~~~~~~~~

doc/source/rllib-training.rst

@@ -913,15 +913,13 @@ Stack Traces
 
 You can use the ``ray stack`` command to dump the stack traces of all the Python workers on a single node. This can be useful for debugging unexpected hangs or performance issues.
 
-REST API
---------
+External Application API
+------------------------
 
-In some cases (i.e., when interacting with an externally hosted simulator or production environment) it makes more sense to interact with RLlib as if were an independently running service, rather than RLlib hosting the simulations itself. This is possible via RLlib's external agents `interface <rllib-env.html#interfacing-with-external-agents>`__.
+In some cases (i.e., when interacting with an externally hosted simulator or production environment) it makes more sense to interact with RLlib as if it were an independently running service, rather than RLlib hosting the simulations itself. This is possible via RLlib's external applications interface `(full documentation) <rllib-env.html#external-agents-and-applications>`__.
 
-.. autoclass:: ray.rllib.utils.policy_client.PolicyClient
+.. autoclass:: ray.rllib.env.policy_client.PolicyClient
     :members:
 
-.. autoclass:: ray.rllib.utils.policy_server.PolicyServer
+.. autoclass:: ray.rllib.env.policy_server_input.PolicyServerInput
     :members:
-
-For a full client / server example that you can run, see the example `client script <https://github.com/ray-project/ray/blob/master/rllib/examples/serving/cartpole_client.py>`__ and also the corresponding `server script <https://github.com/ray-project/ray/blob/master/rllib/examples/serving/cartpole_server.py>`__, here configured to serve a policy for the toy CartPole-v0 environment.

doc/source/rllib.rst

@@ -105,6 +105,13 @@ RLlib provides ways to customize almost all aspects of training, including the `
 
 .. image:: rllib-components.svg
 
+Application Support
+~~~~~~~~~~~~~~~~~~~
+
+Beyond environments defined in Python, RLlib supports batch training on `offline datasets <rllib-offline.html>`__, and also provides a variety of integration strategies for `external applications <rllib-env.html#external-agents-and-applications>`__:
+
+.. image:: rllib-external.svg
+
 To learn more, proceed to the `table of contents <rllib-toc.html>`__.
 
 .. |tensorflow| image:: tensorflow.png

rllib/BUILD

@@ -1336,6 +1336,22 @@ py_test(
     args = ["--num-cpus=4"]
 )
 
+sh_test(
+    name = "examples/serving/test_local_inference",
+    tags = ["examples", "examples_L", "exclusive"],
+    size = "medium",
+    srcs = ["examples/serving/test_local_inference.sh"],
+    data = glob(["examples/serving/*.py"]),
+)
+
+sh_test(
+    name = "examples/serving/test_remote_inference",
+    tags = ["examples", "examples_R", "exclusive"],
+    size = "medium",
+    srcs = ["examples/serving/test_remote_inference.sh"],
+    data = glob(["examples/serving/*.py"]),
+)
+
 py_test(
     name = "examples/rock_paper_scissors_multiagent", main = "examples/rock_paper_scissors_multiagent.py",
     tags = ["examples", "examples_R"],

rllib/agents/dqn/dqn.py

@@ -86,7 +86,7 @@ DEFAULT_CONFIG = with_common_config({
     # Epsilon to add to the TD errors when updating priorities.
     "prioritized_replay_eps": 1e-6,
     # Whether to LZ4 compress observations
-    "compress_observations": True,
+    "compress_observations": False,
 
     # === Optimization ===
     # Learning rate for adam optimizer

rllib/env/__init__.py

@@ -2,11 +2,20 @@ from ray.rllib.env.base_env import BaseEnv
 from ray.rllib.env.dm_env_wrapper import DMEnv
 from ray.rllib.env.multi_agent_env import MultiAgentEnv
 from ray.rllib.env.external_env import ExternalEnv
-from ray.rllib.env.serving_env import ServingEnv
+from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv
 from ray.rllib.env.vector_env import VectorEnv
 from ray.rllib.env.env_context import EnvContext
+from ray.rllib.env.policy_client import PolicyClient
+from ray.rllib.env.policy_server_input import PolicyServerInput
 
 __all__ = [
-    "BaseEnv", "MultiAgentEnv", "ExternalEnv", "VectorEnv", "ServingEnv",
-    "EnvContext", "DMEnv"
+    "BaseEnv",
+    "MultiAgentEnv",
+    "ExternalEnv",
+    "ExternalMultiAgentEnv",
+    "VectorEnv",
+    "EnvContext",
+    "DMEnv",
+    "PolicyClient",
+    "PolicyServerInput",
 ]

rllib/env/policy_client.py

@@ -0,0 +1,294 @@
+"""REST client to interact with a policy server.
+
+This client supports both local and remote policy inference modes. Local
+inference is faster but causes more compute to be done on the client.
+"""
+
+import logging
+import threading
+import time
+
+import ray.cloudpickle as pickle
+from ray.rllib.evaluation.rollout_worker import RolloutWorker
+from ray.rllib.env import ExternalEnv, MultiAgentEnv, ExternalMultiAgentEnv
+from ray.rllib.utils.annotations import PublicAPI
+
+logger = logging.getLogger(__name__)
+logger.setLevel("INFO")  # TODO(ekl) seems to be needed for cartpole_client.py
+
+try:
+    import requests  # `requests` is not part of stdlib.
+except ImportError:
+    requests = None
+    logger.warning(
+        "Couldn't import `requests` library. Be sure to install it on"
+        " the client side.")
+
+
+@PublicAPI
+class PolicyClient:
+    """REST client to interact with a RLlib policy server."""
+
+    # Commands for local inference mode.
+    GET_WORKER_ARGS = "GET_WORKER_ARGS"
+    GET_WEIGHTS = "GET_WEIGHTS"
+    REPORT_SAMPLES = "REPORT_SAMPLES"
+
+    # Commands for remote inference mode.
+    START_EPISODE = "START_EPISODE"
+    GET_ACTION = "GET_ACTION"
+    LOG_ACTION = "LOG_ACTION"
+    LOG_RETURNS = "LOG_RETURNS"
+    END_EPISODE = "END_EPISODE"
+
+    @PublicAPI
+    def __init__(self, address, inference_mode="local", update_interval=10.0):
+        """Create a PolicyClient instance.
+
+        Args:
+            address (str): Server to connect to (e.g., "localhost:9090").
+            inference_mode (str): Whether to use 'local' or 'remote' policy
+                inference for computing actions.
+            update_interval (float): If using 'local' inference mode, the
+                policy is refreshed after this many seconds have passed.
+        """
+        self.address = address
+        if inference_mode == "local":
+            self.local = True
+            self._setup_local_rollout_worker(update_interval)
+        elif inference_mode == "remote":
+            self.local = False
+        else:
+            raise ValueError(
+                "inference_mode must be either 'local' or 'remote'")
+
+    @PublicAPI
+    def start_episode(self, episode_id=None, training_enabled=True):
+        """Record the start of an episode.
+
+        Arguments:
+            episode_id (str): Unique string id for the episode or None for
+                it to be auto-assigned.
+            training_enabled (bool): Whether to use experiences for this
+                episode to improve the policy.
+
+        Returns:
+            episode_id (str): Unique string id for the episode.
+        """
+
+        if self.local:
+            self._update_local_policy()
+            return self.env.start_episode(episode_id, training_enabled)
+
+        return self._send({
+            "episode_id": episode_id,
+            "command": PolicyClient.START_EPISODE,
+            "training_enabled": training_enabled,
+        })["episode_id"]
+
+    @PublicAPI
+    def get_action(self, episode_id, observation):
+        """Record an observation and get the on-policy action.
+
+        Arguments:
+            episode_id (str): Episode id returned from start_episode().
+            observation (obj): Current environment observation.
+
+        Returns:
+            action (obj): Action from the env action space.
+        """
+
+        if self.local:
+            self._update_local_policy()
+            return self.env.get_action(episode_id, observation)
+
+        return self._send({
+            "command": PolicyClient.GET_ACTION,
+            "observation": observation,
+            "episode_id": episode_id,
+        })["action"]
+
+    @PublicAPI
+    def log_action(self, episode_id, observation, action):
+        """Record an observation and (off-policy) action taken.
+
+        Arguments:
+            episode_id (str): Episode id returned from start_episode().
+            observation (obj): Current environment observation.
+            action (obj): Action for the observation.
+        """
+
+        if self.local:
+            self._update_local_policy()
+            return self.env.log_action(episode_id, observation, action)
+
+        self._send({
+            "command": PolicyClient.LOG_ACTION,
+            "observation": observation,
+            "action": action,
+            "episode_id": episode_id,
+        })
+
+    @PublicAPI
+    def log_returns(self, episode_id, reward, info=None):
+        """Record returns from the environment.
+
+        The reward will be attributed to the previous action taken by the
+        episode. Rewards accumulate until the next action. If no reward is
+        logged before the next action, a reward of 0.0 is assumed.
+
+        Arguments:
+            episode_id (str): Episode id returned from start_episode().
+            reward (float): Reward from the environment.
+        """
+
+        if self.local:
+            self._update_local_policy()
+            return self.env.log_returns(episode_id, reward, info)
+
+        self._send({
+            "command": PolicyClient.LOG_RETURNS,
+            "reward": reward,
+            "info": info,
+            "episode_id": episode_id,
+        })
+
+    @PublicAPI
+    def end_episode(self, episode_id, observation):
+        """Record the end of an episode.
+
+        Arguments:
+            episode_id (str): Episode id returned from start_episode().
+            observation (obj): Current environment observation.
+        """
+
+        if self.local:
+            self._update_local_policy()
+            return self.env.end_episode(episode_id, observation)
+
+        self._send({
+            "command": PolicyClient.END_EPISODE,
+            "observation": observation,
+            "episode_id": episode_id,
+        })
+
+    def _send(self, data):
+        payload = pickle.dumps(data)
+        response = requests.post(self.address, data=payload)
+        if response.status_code != 200:
+            logger.error("Request failed {}: {}".format(response.text, data))
+        response.raise_for_status()
+        parsed = pickle.loads(response.content)
+        return parsed
+
+    def _setup_local_rollout_worker(self, update_interval):
+        self.update_interval = update_interval
+        self.last_updated = 0
+
+        logger.info("Querying server for rollout worker settings.")
+        kwargs = self._send({
+            "command": PolicyClient.GET_WORKER_ARGS,
+        })["worker_args"]
+
+        (self.rollout_worker,
+         self.inference_thread) = create_embedded_rollout_worker(
+             kwargs, self._send)
+        self.env = self.rollout_worker.env
+
+    def _update_local_policy(self):
+        assert self.inference_thread.is_alive()
+        if time.time() - self.last_updated > self.update_interval:
+            logger.info("Querying server for new policy weights.")
+            resp = self._send({
+                "command": PolicyClient.GET_WEIGHTS,
+            })
+            weights = resp["weights"]
+            global_vars = resp["global_vars"]
+            logger.info(
+                "Updating rollout worker weights and global vars {}.".format(
+                    global_vars))
+            self.rollout_worker.set_weights(weights, global_vars)
+            self.last_updated = time.time()
+
+
+class _LocalInferenceThread(threading.Thread):
+    """Thread that handles experience generation (worker.sample() loop)."""
+
+    def __init__(self, rollout_worker, send_fn):
+        super().__init__()
+        self.daemon = True
+        self.rollout_worker = rollout_worker
+        self.send_fn = send_fn
+
+    def run(self):
+        try:
+            while True:
+                logger.info("Generating new batch of experiences.")
+                samples = self.rollout_worker.sample()
+                metrics = self.rollout_worker.get_metrics()
+                logger.info("Sending batch of {} steps back to server.".format(
+                    samples.count))
+                self.send_fn({
+                    "command": PolicyClient.REPORT_SAMPLES,
+                    "samples": samples,
+                    "metrics": metrics,
+                })
+        except Exception as e:
+            logger.info("Error: inference worker thread died!", e)
+
+
+def auto_wrap_external(real_env_creator):
+    """Wrap an environment in the ExternalEnv interface if needed.
+
+    Args:
+        real_env_creator (fn): Create an env given the env_config.
+    """
+
+    def wrapped_creator(env_config):
+        real_env = real_env_creator(env_config)
+        if not (isinstance(real_env, ExternalEnv)
+                or isinstance(real_env, ExternalMultiAgentEnv)):
+            logger.info(
+                "The env you specified is not a type of ExternalEnv. "
+                "Attempting to convert it automatically to ExternalEnv.")
+
+            if isinstance(real_env, MultiAgentEnv):
+                external_cls = MultiAgentEnv
+            else:
+                external_cls = ExternalEnv
+
+            class ExternalEnvWrapper(external_cls):
+                def __init__(self, real_env):
+                    super().__init__(real_env.action_space,
+                                     real_env.observation_space)
+
+                def run(self):
+                    # Since we are calling methods on this class in the
+                    # client, run doesn't need to do anything.
+                    time.sleep(999999)
+
+            return ExternalEnvWrapper(real_env)
+
+    return wrapped_creator
+
+
+def create_embedded_rollout_worker(kwargs, send_fn):
+    """Create a local rollout worker and a thread that samples from it.
+
+    Arguments:
+        kwargs (dict): args for the RolloutWorker constructor.
+        send_fn (fn): function to send a JSON request to the server.
+    """
+
+    # Since the server acts as an input datasource, we have to reset the
+    # input config to the default, which runs env rollouts.
+    kwargs = kwargs.copy()
+    del kwargs["input_creator"]
+    logger.info("Creating rollout worker with kwargs={}".format(kwargs))
+    real_env_creator = kwargs["env_creator"]
+    kwargs["env_creator"] = auto_wrap_external(real_env_creator)
+
+    rollout_worker = RolloutWorker(**kwargs)
+    inference_thread = _LocalInferenceThread(rollout_worker, send_fn)
+    inference_thread.start()
+    return rollout_worker, inference_thread

rllib/env/policy_server_input.py

@@ -0,0 +1,191 @@
+import logging
+import six.moves.queue as queue
+import threading
+import traceback
+
+from http.server import SimpleHTTPRequestHandler, HTTPServer
+from socketserver import ThreadingMixIn
+
+import ray.cloudpickle as pickle
+from ray.rllib.offline.input_reader import InputReader
+from ray.rllib.env.policy_client import PolicyClient, \
+    create_embedded_rollout_worker
+from ray.rllib.utils.annotations import override, PublicAPI
+
+logger = logging.getLogger(__name__)
+logger.setLevel("INFO")  # TODO(ekl) this is needed for cartpole_server.py
+
+
+class PolicyServerInput(ThreadingMixIn, HTTPServer, InputReader):
+    """REST policy server that acts as an offline data source.
+
+    This launches a multi-threaded server that listens on the specified host
+    and port to serve policy requests and forward experiences to RLlib. For
+    high performance experience collection, it implements InputReader.
+
+    For an example, run `examples/cartpole_server.py` along
+    with `examples/cartpole_client.py --inference-mode=local|remote`.
+
+    Examples:
+        >>> pg = PGTrainer(
+        ...     env="CartPole-v0", config={
+        ...         "input": lambda ioctx:
+        ...             PolicyServerInput(ioctx, addr, port),
+        ...         "num_workers": 0,  # Run just 1 server, in the trainer.
+        ...     }
+        >>> while True:
+                pg.train()
+
+        >>> client = PolicyClient("localhost:9900", inference_mode="local")
+        >>> eps_id = client.start_episode()
+        >>> action = client.get_action(eps_id, obs)
+        >>> ...
+        >>> client.log_returns(eps_id, reward)
+        >>> ...
+        >>> client.log_returns(eps_id, reward)
+    """
+
+    @PublicAPI
+    def __init__(self, ioctx, address, port):
+        """Create a PolicyServerInput.
+
+        This class implements rllib.offline.InputReader, and can be used with
+        any Trainer by configuring
+
+            {"num_workers": 0,
+             "input": lambda ioctx: PolicyServerInput(ioctx, addr, port)}
+
+        Note that by setting num_workers: 0, the trainer will only create one
+        rollout worker / PolicyServerInput. Clients can connect to the launched
+        server using rllib.env.PolicyClient.
+
+        Args:
+            ioctx (IOContext): IOContext provided by RLlib.
+            address (str): Server addr (e.g., "localhost").
+            port (int): Server port (e.g., 9900).
+        """
+
+        self.rollout_worker = ioctx.worker
+        self.samples_queue = queue.Queue()
+        self.metrics_queue = queue.Queue()
+
+        def get_metrics():
+            completed = []
+            while True:
+                try:
+                    completed.append(self.metrics_queue.get_nowait())
+                except queue.Empty:
+                    break
+            return completed
+
+        # Forwards client-reported rewards directly into the local rollout
+        # worker. This is a bit of a hack since it is patching the get_metrics
+        # function of the sampler.
+        self.rollout_worker.sampler.get_metrics = get_metrics
+
+        handler = _make_handler(self.rollout_worker, self.samples_queue,
+                                self.metrics_queue)
+        HTTPServer.__init__(self, (address, port), handler)
+        logger.info("")
+        logger.info("Starting connector server at {}:{}".format(address, port))
+        logger.info("")
+        thread = threading.Thread(name="server", target=self.serve_forever)
+        thread.start()
+
+    @override(InputReader)
+    def next(self):
+        return self.samples_queue.get()
+
+
+def _make_handler(rollout_worker, samples_queue, metrics_queue):
+    # Only used in remote inference mode. We must create a new rollout worker
+    # then since the original worker doesn't have the env properly wrapped in
+    # an ExternalEnv interface.
+    child_rollout_worker = None
+    inference_thread = None
+    lock = threading.Lock()
+
+    def setup_child_rollout_worker():
+        nonlocal lock
+        nonlocal child_rollout_worker
+        nonlocal inference_thread
+
+        with lock:
+            if child_rollout_worker is None:
+                (child_rollout_worker,
+                 inference_thread) = create_embedded_rollout_worker(
+                     rollout_worker.creation_args(), report_data)
+                child_rollout_worker.set_weights(rollout_worker.get_weights())
+
+    def report_data(data):
+        nonlocal child_rollout_worker
+
+        batch = data["samples"]
+        batch.decompress_if_needed()
+        samples_queue.put(batch)
+        for rollout_metric in data["metrics"]:
+            metrics_queue.put(rollout_metric)
+
+        if child_rollout_worker is not None:
+            child_rollout_worker.set_weights(rollout_worker.get_weights(),
+                                             rollout_worker.get_global_vars())
+
+    class Handler(SimpleHTTPRequestHandler):
+        def __init__(self, *a, **kw):
+            super().__init__(*a, **kw)
+
+        def do_POST(self):
+            content_len = int(self.headers.get("Content-Length"), 0)
+            raw_body = self.rfile.read(content_len)
+            parsed_input = pickle.loads(raw_body)
+            try:
+                response = self.execute_command(parsed_input)
+                self.send_response(200)
+                self.end_headers()
+                self.wfile.write(pickle.dumps(response))
+            except Exception:
+                self.send_error(500, traceback.format_exc())
+
+        def execute_command(self, args):
+            command = args["command"]
+            response = {}
+            # Local inference commands:
+            if command == PolicyClient.GET_WORKER_ARGS:
+                logger.info("Sending worker creation args to client.")
+                response["worker_args"] = rollout_worker.creation_args()
+            elif command == PolicyClient.GET_WEIGHTS:
+                logger.info("Sending worker weights to client.")
+                response["weights"] = rollout_worker.get_weights()
+                response["global_vars"] = rollout_worker.get_global_vars()
+            elif command == PolicyClient.REPORT_SAMPLES:
+                logger.info("Got sample batch of size {} from client.".format(
+                    args["samples"].count))
+                report_data(args)
+            # Remote inference commands:
+            elif command == PolicyClient.START_EPISODE:
+                setup_child_rollout_worker()
+                assert inference_thread.is_alive()
+                response["episode_id"] = (
+                    child_rollout_worker.env.start_episode(
+                        args["episode_id"], args["training_enabled"]))
+            elif command == PolicyClient.GET_ACTION:
+                assert inference_thread.is_alive()
+                response["action"] = child_rollout_worker.env.get_action(
+                    args["episode_id"], args["observation"])
+            elif command == PolicyClient.LOG_ACTION:
+                assert inference_thread.is_alive()
+                child_rollout_worker.env.log_action(
+                    args["episode_id"], args["observation"], args["action"])
+            elif command == PolicyClient.LOG_RETURNS:
+                assert inference_thread.is_alive()
+                child_rollout_worker.env.log_returns(
+                    args["episode_id"], args["reward"], args["info"])
+            elif command == PolicyClient.END_EPISODE:
+                assert inference_thread.is_alive()
+                child_rollout_worker.env.end_episode(args["episode_id"],
+                                                     args["observation"])
+            else:
+                raise ValueError("Unknown command: {}".format(command))
+            return response
+
+    return Handler

rllib/env/serving_env.py

@@ -1,4 +0,0 @@
-from ray.rllib.env.external_env import ExternalEnv
-
-# renamed to ExternalEnv in 0.6
-ServingEnv = ExternalEnv

rllib/evaluation/rollout_worker.py

@@ -240,6 +240,8 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker):
                 to ensure each remote worker has unique exploration behavior.
             _fake_sampler (bool): Use a fake (inf speed) sampler for testing.
         """
+        self._original_kwargs = locals().copy()
+        del self._original_kwargs["self"]
 
         global _global_worker
         _global_worker = self
@@ -279,6 +281,7 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker):
         self.compress_observations = compress_observations
         self.preprocessing_enabled = True
         self.last_batch = None
+        self.global_vars = None
         self._fake_sampler = _fake_sampler
 
         self.env = _validate_env(env_creator(env_context))
@@ -769,6 +772,11 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker):
     @DeveloperAPI
     def set_global_vars(self, global_vars):
         self.foreach_policy(lambda p, _: p.on_global_var_update(global_vars))
+        self.global_vars = global_vars
+
+    @DeveloperAPI
+    def get_global_vars(self):
+        return self.global_vars
 
     @DeveloperAPI
     def export_policy_model(self, export_dir, policy_id=DEFAULT_POLICY_ID):
@@ -786,6 +794,11 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker):
     def stop(self):
         self.async_env.stop()
 
+    @DeveloperAPI
+    def creation_args(self):
+        """Returns the args used to create this worker."""
+        return self._original_kwargs
+
     def _build_policy_map(self, policy_dict, policy_config):
         policy_map = {}
         preprocessors = {}

rllib/examples/serving/cartpole_client.py

@@ -1,18 +1,23 @@
-"""Example of querying a policy server. Copy this file for your use case.
+#!/usr/bin/env python
+"""Example of training with a policy server. Copy this file for your use case.
 
 To try this out, in two separate shells run:
-    $ python cartpole_server.py
-    $ python cartpole_client.py
+    $ python cartpole_server.py --run=[PPO|DQN]
+    $ python cartpole_client.py --inference-mode=local|remote
+
+Local inference mode offloads inference to the client for better performance.
 """
 
 import argparse
 import gym
 
-from ray.rllib.utils.policy_client import PolicyClient
+from ray.rllib.env.policy_client import PolicyClient
 
 parser = argparse.ArgumentParser()
 parser.add_argument(
     "--no-train", action="store_true", help="Whether to disable training.")
+parser.add_argument(
+    "--inference-mode", type=str, required=True, choices=["local", "remote"])
 parser.add_argument(
     "--off-policy",
     action="store_true",
@@ -26,7 +31,8 @@ parser.add_argument(
 if __name__ == "__main__":
     args = parser.parse_args()
     env = gym.make("CartPole-v0")
-    client = PolicyClient("http://localhost:9900")
+    client = PolicyClient(
+        "http://localhost:9900", inference_mode=args.inference_mode)
 
     eid = client.start_episode(training_enabled=not args.no_train)
     obs = env.reset()

rllib/examples/serving/cartpole_server.py

@@ -1,66 +1,84 @@
+#!/usr/bin/env python
 """Example of running a policy server. Copy this file for your use case.
 
 To try this out, in two separate shells run:
     $ python cartpole_server.py
-    $ python cartpole_client.py
+    $ python cartpole_client.py --inference-mode=local|remote
 """
 
+import argparse
 import os
-from gym import spaces
-import numpy as np
 
 import ray
 from ray.rllib.agents.dqn import DQNTrainer
-from ray.rllib.env.external_env import ExternalEnv
-from ray.rllib.utils.policy_server import PolicyServer
+from ray.rllib.agents.ppo import PPOTrainer
+from ray.rllib.env.policy_server_input import PolicyServerInput
 from ray.tune.logger import pretty_print
-from ray.tune.registry import register_env
 
 SERVER_ADDRESS = "localhost"
 SERVER_PORT = 9900
-CHECKPOINT_FILE = "last_checkpoint.out"
-
-
-class CartpoleServing(ExternalEnv):
-    def __init__(self):
-        ExternalEnv.__init__(
-            self, spaces.Discrete(2),
-            spaces.Box(low=-10, high=10, shape=(4, ), dtype=np.float32))
-
-    def run(self):
-        print("Starting policy server at {}:{}".format(SERVER_ADDRESS,
-                                                       SERVER_PORT))
-        server = PolicyServer(self, SERVER_ADDRESS, SERVER_PORT)
-        server.serve_forever()
+CHECKPOINT_FILE = "last_checkpoint_{}.out"
 
+parser = argparse.ArgumentParser()
+parser.add_argument("--run", type=str, default="DQN")
 
 if __name__ == "__main__":
+    args = parser.parse_args()
     ray.init()
-    register_env("srv", lambda _: CartpoleServing())
 
-    # We use DQN since it supports off-policy actions, but you can choose and
-    # configure any agent.
-    dqn = DQNTrainer(
-        env="srv",
-        config={
-            # Use a single process to avoid needing to set up a load balancer
-            "num_workers": 0,
-            # Configure the agent to run short iterations for debugging
-            "exploration_fraction": 0.01,
-            "learning_starts": 100,
-            "timesteps_per_iteration": 200,
-        })
+    env = "CartPole-v0"
+    connector_config = {
+        # Use the connector server to generate experiences.
+        "input": (
+            lambda ioctx: PolicyServerInput( \
+                ioctx, SERVER_ADDRESS, SERVER_PORT)
+        ),
+        # Use a single worker process to run the server.
+        "num_workers": 0,
+        # Disable OPE, since the rollouts are coming from online clients.
+        "input_evaluation": [],
+    }
+
+    if args.run == "DQN":
+        # Example of using DQN (supports off-policy actions).
+        trainer = DQNTrainer(
+            env=env,
+            config=dict(
+                connector_config, **{
+                    "exploration_config": {
+                        "type": "EpsilonGreedy",
+                        "initial_epsilon": 1.0,
+                        "final_epsilon": 0.02,
+                        "epsilon_timesteps": 1000,
+                    },
+                    "learning_starts": 100,
+                    "timesteps_per_iteration": 200,
+                    "log_level": "INFO",
+                }))
+    elif args.run == "PPO":
+        # Example of using PPO (does NOT support off-policy actions).
+        trainer = PPOTrainer(
+            env=env,
+            config=dict(
+                connector_config, **{
+                    "sample_batch_size": 1000,
+                    "train_batch_size": 4000,
+                }))
+    else:
+        raise ValueError("--run must be DQN or PPO")
+
+    checkpoint_path = CHECKPOINT_FILE.format(args.run)
 
     # Attempt to restore from checkpoint if possible.
-    if os.path.exists(CHECKPOINT_FILE):
-        checkpoint_path = open(CHECKPOINT_FILE).read()
+    if os.path.exists(checkpoint_path):
+        checkpoint_path = open(checkpoint_path).read()
         print("Restoring from checkpoint path", checkpoint_path)
-        dqn.restore(checkpoint_path)
+        trainer.restore(checkpoint_path)
 
     # Serving and training loop
     while True:
-        print(pretty_print(dqn.train()))
-        checkpoint_path = dqn.save()
-        print("Last checkpoint", checkpoint_path)
-        with open(CHECKPOINT_FILE, "w") as f:
-            f.write(checkpoint_path)
+        print(pretty_print(trainer.train()))
+        checkpoint = trainer.save()
+        print("Last checkpoint", checkpoint)
+        with open(checkpoint_path, "w") as f:
+            f.write(checkpoint)

rllib/examples/serving/test.sh

@@ -1,12 +0,0 @@
-#!/bin/bash
-
-pkill -f cartpole_server.py
-(python cartpole_server.py 2>&1 | grep -v 200) &
-pid=$!
-
-while ! curl localhost:9900; do   
-  sleep 1
-done
-
-python cartpole_client.py --stop-at-reward=100
-kill $pid

rllib/examples/serving/test_local_inference.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+rm -f last_checkpoint.out
+pkill -f cartpole_server.py
+sleep 1
+
+if [ -f cartpole_server.py ]; then
+    basedir="."
+else
+    basedir="rllib/examples/serving"  # In bazel.
+fi
+
+(python $basedir/cartpole_server.py --run=PPO 2>&1 | grep -v 200) &
+pid=$!
+
+echo "Waiting for server to start"
+while ! curl localhost:9900; do   
+  sleep 1
+done
+
+sleep 2
+python $basedir/cartpole_client.py --stop-at-reward=100 --inference-mode=local
+kill $pid

rllib/examples/serving/test_remote_inference.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+rm -f last_checkpoint.out
+pkill -f cartpole_server.py
+sleep 1
+
+if [ -f cartpole_server.py ]; then
+    basedir="."
+else
+    basedir="rllib/examples/serving"  # In bazel.
+fi
+
+(python $basedir/cartpole_server.py --run=DQN 2>&1 | grep -v 200) &
+pid=$!
+
+echo "Waiting for server to start"
+while ! curl localhost:9900; do   
+  sleep 1
+done
+
+sleep 2
+python $basedir/cartpole_client.py --stop-at-reward=100 --inference-mode=remote
+kill $pid
+

rllib/utils/policy_client.py

@@ -1,7 +1,10 @@
+"""DEPRECATED: Please use rllib.env.PolicyClient instead."""
+
 import logging
 import pickle
 
 from ray.rllib.utils.annotations import PublicAPI
+from ray.rllib.utils.deprecation import deprecation_warning
 
 logger = logging.getLogger(__name__)
 
@@ -16,7 +19,7 @@ except ImportError:
 
 @PublicAPI
 class PolicyClient:
-    """REST client to interact with a RLlib policy server."""
+    """DEPRECATED: Please use rllib.env.PolicyClient instead."""
 
     START_EPISODE = "START_EPISODE"
     GET_ACTION = "GET_ACTION"
@@ -26,6 +29,8 @@ class PolicyClient:
 
     @PublicAPI
     def __init__(self, address):
+        deprecation_warning(
+            "rllib.utils.PolicyServer", new="rllib.env.PolicyServerInput")
         self._address = address
 
     @PublicAPI

rllib/utils/policy_server.py

@@ -1,3 +1,5 @@
+"""DEPRECATED: Please use rllib.env.PolicyServerInput instead."""
+
 import pickle
 import traceback
 
@@ -6,44 +8,17 @@ from socketserver import ThreadingMixIn
 
 from ray.rllib.utils.annotations import PublicAPI
 from ray.rllib.utils.policy_client import PolicyClient
+from ray.rllib.utils.deprecation import deprecation_warning
 
 
 @PublicAPI
 class PolicyServer(ThreadingMixIn, HTTPServer):
-    """REST server than can be launched from a ExternalEnv.
-
-    This launches a multi-threaded server that listens on the specified host
-    and port to serve policy requests and forward experiences to RLlib.
-
-    Examples:
-        >>> class CartpoleServing(ExternalEnv):
-               def __init__(self):
-                   ExternalEnv.__init__(
-                       self, spaces.Discrete(2),
-                       spaces.Box(
-                           low=-10,
-                           high=10,
-                           shape=(4,),
-                           dtype=np.float32))
-               def run(self):
-                   server = PolicyServer(self, "localhost", 8900)
-                   server.serve_forever()
-        >>> register_env("srv", lambda _: CartpoleServing())
-        >>> pg = PGTrainer(env="srv", config={"num_workers": 0})
-        >>> while True:
-                pg.train()
-
-        >>> client = PolicyClient("localhost:8900")
-        >>> eps_id = client.start_episode()
-        >>> action = client.get_action(eps_id, obs)
-        >>> ...
-        >>> client.log_returns(eps_id, reward)
-        >>> ...
-        >>> client.log_returns(eps_id, reward)
-    """
+    """DEPRECATED: Please use rllib.env.PolicyServerInput instead."""
 
     @PublicAPI
     def __init__(self, external_env, address, port):
+        deprecation_warning(
+            "rllib.utils.PolicyClient", new="rllib.env.PolicyClient")
         handler = _make_handler(external_env)
         HTTPServer.__init__(self, (address, port), handler)