[rllib] Envs for vectorized execution, async execution, and policy serving (#2170)

## What do these changes do?

**Vectorized envs**: Users can either implement `VectorEnv`, or alternatively set `num_envs=N` to auto-vectorize gym envs (this vectorizes just the action computation part).

```
# CartPole-v0 on single core with 64x64 MLP:

# vector_width=1:
Actions per second 2720.1284458322966

# vector_width=8:
Actions per second 13773.035334888269

# vector_width=64:
Actions per second 37903.20472563333
```

**Async envs**: The more general form of `VectorEnv` is `AsyncVectorEnv`, which allows agents to execute out of lockstep. We use this as an adapter to support `ServingEnv`. Since we can convert any other form of env to `AsyncVectorEnv`, utils.sampler has been rewritten to run against this interface.

**Policy serving**: This provides an env which is not stepped. Rather, the env executes in its own thread, querying the policy for actions via `self.get_action(obs)`, and reporting results via `self.log_returns(rewards)`. We also support logging of off-policy actions via `self.log_action(obs, action)`. This is a more convenient API for some use cases, and also provides parallelizable support for policy serving (for example, if you start a HTTP server in the env) and ingest of offline logs (if the env reads from serving logs).

Any of these types of envs can be passed to RLlib agents. RLlib handles conversions internally in CommonPolicyEvaluator, for example:
 ```
        gym.Env => rllib.VectorEnv => rllib.AsyncVectorEnv
        rllib.ServingEnv => rllib.AsyncVectorEnv
```
This commit is contained in:
Eric Liang
2018-06-18 11:55:32 -07:00
committed by GitHub
parent 8560993b46
commit 7dee2c6735
28 changed files with 1218 additions and 342 deletions
@@ -8,12 +8,16 @@ import tensorflow as tf
import ray
from ray.rllib.models import ModelCatalog
from ray.rllib.optimizers import SampleBatch
from ray.rllib.optimizers.policy_evaluator import PolicyEvaluator
from ray.rllib.utils.atari_wrappers import wrap_deepmind
from ray.rllib.utils.async_vector_env import AsyncVectorEnv, _VectorEnvToAsync
from ray.rllib.utils.atari_wrappers import wrap_deepmind, is_atari
from ray.rllib.utils.compression import pack
from ray.rllib.utils.filter import get_filter
from ray.rllib.utils.sampler import AsyncSampler, SyncSampler
from ray.rllib.utils.serving_env import ServingEnv, _ServingEnvToAsync
from ray.rllib.utils.tf_policy_graph import TFPolicyGraph
from ray.rllib.utils.vector_env import VectorEnv
from ray.tune.registry import get_registry
from ray.tune.result import TrainingResult
@@ -53,10 +57,8 @@ def collect_metrics(local_evaluator, remote_evaluators):
class CommonPolicyEvaluator(PolicyEvaluator):
"""Policy evaluator implementation that operates on a rllib.PolicyGraph.
TODO: vector env
TODO: multi-agent
TODO: consumer buffering for multi-agent
TODO: complete episode batch mode
TODO: multi-gpu
Examples:
# Create a policy evaluator and using it to collect experiences.
@@ -89,9 +91,11 @@ class CommonPolicyEvaluator(PolicyEvaluator):
tf_session_creator=None,
batch_steps=100,
batch_mode="truncate_episodes",
episode_horizon=None,
preprocessor_pref="rllib",
sample_async=False,
compress_observations=False,
num_envs=1,
observation_filter="NoFilter",
registry=None,
env_config=None,
@@ -108,13 +112,20 @@ class CommonPolicyEvaluator(PolicyEvaluator):
This is optional and only useful with TFPolicyGraph.
batch_steps (int): The target number of env transitions to include
in each sample batch returned from this evaluator.
batch_mode (str): One of the following choices:
complete_episodes: each batch will be at least batch_steps
in size, and will include one or more complete episodes.
truncate_episodes: each batch will be around batch_steps
in size, and include transitions from one episode only.
pack_episodes: each batch will be exactly batch_steps in
size, and may include transitions from multiple episodes.
batch_mode (str): One of the following batch modes:
"truncate_episodes": Each call to sample() will return a batch
of exactly `batch_steps` in size. Episodes may be truncated
in order to meet this size requirement. When
`num_envs > 1`, episodes will be truncated to sequences of
`batch_size / num_envs` in length.
"complete_episodes": Each call to sample() will return a batch
of at least `batch_steps in size. Episodes will not be
truncated, but multiple episodes may be packed within one
batch to meet the batch size. Note that when
`num_envs > 1`, episode steps will be buffered until the
episode completes, and hence batches may contain
significant amounts of off-policy data.
episode_horizon (int): Whether to stop episodes at this horizon.
preprocessor_pref (str): Whether to prefer RLlib preprocessors
("rllib") or deepmind ("deepmind") when applicable.
sample_async (bool): Whether to compute samples asynchronously in
@@ -122,6 +133,9 @@ class CommonPolicyEvaluator(PolicyEvaluator):
to be slightly off-policy.
compress_observations (bool): If true, compress the observations
returned.
num_envs (int): If more than one, will create multiple envs
and vectorize the computation of actions. This has no effect if
if the env already implements VectorEnv.
observation_filter (str): Name of observation filter to use.
registry (tune.Registry): User-registered objects. Pass in the
value from tune.registry.get_registry() if you're having
@@ -135,9 +149,6 @@ class CommonPolicyEvaluator(PolicyEvaluator):
env_config = env_config or {}
policy_config = policy_config or {}
model_config = model_config or {}
assert batch_mode in [
"complete_episodes", "truncate_episodes", "pack_episodes"]
self.env_creator = env_creator
self.policy_graph = policy_graph
self.batch_steps = batch_steps
@@ -145,15 +156,25 @@ class CommonPolicyEvaluator(PolicyEvaluator):
self.compress_observations = compress_observations
self.env = env_creator(env_config)
is_atari = hasattr(self.env.unwrapped, "ale")
if is_atari and "custom_preprocessor" not in model_config and \
if isinstance(self.env, VectorEnv) or \
isinstance(self.env, ServingEnv) or \
isinstance(self.env, AsyncVectorEnv):
def wrap(env):
return env # we can't auto-wrap these env types
elif is_atari(self.env) and \
"custom_preprocessor" not in model_config and \
preprocessor_pref == "deepmind":
self.env = wrap_deepmind(self.env, dim=model_config.get("dim", 80))
def wrap(env):
return wrap_deepmind(env, dim=model_config.get("dim", 80))
else:
self.env = ModelCatalog.get_preprocessor_as_wrapper(
registry, self.env, model_config)
def wrap(env):
return ModelCatalog.get_preprocessor_as_wrapper(
registry, env, model_config)
self.env = wrap(self.env)
def make_env():
return wrap(env_creator(env_config))
self.vectorized = hasattr(self.env, "vector_reset")
self.policy_map = {}
if issubclass(policy_graph, TFPolicyGraph):
@@ -179,24 +200,41 @@ class CommonPolicyEvaluator(PolicyEvaluator):
observation_filter, self.env.observation_space.shape)
self.filters = {"obs_filter": self.obs_filter}
if self.vectorized:
raise NotImplementedError("Vector envs not yet supported")
else:
if batch_mode not in [
"pack_episodes", "truncate_episodes", "complete_episodes"]:
raise NotImplementedError("Batch mode not yet supported")
pack = batch_mode == "pack_episodes"
if batch_mode == "complete_episodes":
batch_steps = 999999
if sample_async:
self.sampler = AsyncSampler(
self.env, self.policy_map["default"], self.obs_filter,
batch_steps, pack=pack)
self.sampler.start()
# Always use vector env for consistency even if num_envs = 1
if not isinstance(self.env, AsyncVectorEnv):
if isinstance(self.env, ServingEnv):
self.vector_env = _ServingEnvToAsync(self.env)
else:
self.sampler = SyncSampler(
self.env, self.policy_map["default"], self.obs_filter,
batch_steps, pack=pack)
if not isinstance(self.env, VectorEnv):
self.env = VectorEnv.wrap(
make_env, [self.env], num_envs=num_envs)
self.vector_env = _VectorEnvToAsync(self.env)
else:
self.vector_env = self.env
if self.batch_mode == "truncate_episodes":
if batch_steps % num_envs != 0:
raise ValueError(
"In 'truncate_episodes' batch mode, `batch_steps` must be "
"evenly divisible by `num_envs`. Got {} and {}.".format(
batch_steps, num_envs))
batch_steps = batch_steps // num_envs
pack_episodes = True
elif self.batch_mode == "complete_episodes":
batch_steps = float("inf") # never cut episodes
pack_episodes = False # sampler will return 1 episode per poll
else:
raise ValueError(
"Unsupported batch mode: {}".format(self.batch_mode))
if sample_async:
self.sampler = AsyncSampler(
self.vector_env, self.policy_map["default"], self.obs_filter,
batch_steps, horizon=episode_horizon, pack=pack_episodes)
self.sampler.start()
else:
self.sampler = SyncSampler(
self.vector_env, self.policy_map["default"], self.obs_filter,
batch_steps, horizon=episode_horizon, pack=pack_episodes)
def sample(self):
"""Evaluate the current policies and return a batch of experiences.
@@ -205,8 +243,13 @@ class CommonPolicyEvaluator(PolicyEvaluator):
SampleBatch from evaluating the current policies.
"""
batch = self.policy_map["default"].postprocess_trajectory(
self.sampler.get_data())
batches = [self.sampler.get_data()]
steps_so_far = batches[0].count
while steps_so_far < self.batch_steps:
batch = self.sampler.get_data()
steps_so_far += batch.count
batches.append(batch)
batch = SampleBatch.concat_samples(batches)
if self.compress_observations:
batch["obs"] = [pack(o) for o in batch["obs"]]
@@ -214,11 +257,6 @@ class CommonPolicyEvaluator(PolicyEvaluator):
return batch
def apply(self, func):
"""Apply the given function to this evaluator instance."""
return func(self)
def for_policy(self, func):
"""Apply the given function to this evaluator's default policy."""