mirror of
https://github.com/wassname/ray.git
synced 2026-07-26 13:37:24 +08:00
[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:
@@ -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."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user