wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-10 01:23:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/rllib/utils/sampler.py
+160 -72
View File
@@ -2,12 +2,14 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from collections import namedtuple
from collections import defaultdict, namedtuple
import numpy as np
import six.moves.queue as queue
import threading
from ray.rllib.optimizers.sample_batch import SampleBatchBuilder
from ray.rllib.utils.vector_env import VectorEnv
from ray.rllib.utils.async_vector_env import AsyncVectorEnv, _VectorEnvToAsync
CompletedRollout = namedtuple("CompletedRollout",
@@ -22,17 +24,20 @@ class SyncSampler(object):
This class provides data on invocation, rather than on a separate
thread."""
_async = False
def __init__(
self, env, policy, obs_filter, num_local_steps, horizon=None,
pack=False):
self, env, policy, obs_filter, num_local_steps,
horizon=None, pack=False):
if not isinstance(env, AsyncVectorEnv):
if not isinstance(env, VectorEnv):
env = VectorEnv.wrap(make_env=None, existing_envs=[env])
env = _VectorEnvToAsync(env)
self.async_vector_env = env
self.num_local_steps = num_local_steps
self.horizon = horizon
self.env = env
self.policy = policy
self._obs_filter = obs_filter
self.rollout_provider = _env_runner(self.env, self.policy,
self.rollout_provider = _env_runner(self.async_vector_env, self.policy,
self.num_local_steps, self.horizon,
self._obs_filter, pack)
self.metrics_queue = queue.Queue()
@@ -60,28 +65,29 @@ class AsyncSampler(threading.Thread):
Note that batch_size is only a unit of measure here. Batches can
accumulate and the gradient can be calculated on up to 5 batches."""
_async = True
def __init__(
self, env, policy, obs_filter, num_local_steps, horizon=None,
pack=False):
self, env, policy, obs_filter, num_local_steps,
horizon=None, pack=False):
assert getattr(
obs_filter, "is_concurrent",
False), ("Observation Filter must support concurrent updates.")
if not isinstance(env, AsyncVectorEnv):
if not isinstance(env, VectorEnv):
env = VectorEnv.wrap(make_env=None, existing_envs=[env])
env = _VectorEnvToAsync(env)
self.async_vector_env = env
threading.Thread.__init__(self)
self.queue = queue.Queue(5)
self.metrics_queue = queue.Queue()
self.num_local_steps = num_local_steps
self.horizon = horizon
self.env = env
self.policy = policy
self._obs_filter = obs_filter
self.started = False
self.daemon = True
self.pack = pack
def run(self):
self.started = True
try:
self._run()
except BaseException as e:
@@ -89,7 +95,7 @@ class AsyncSampler(threading.Thread):
raise e
def _run(self):
rollout_provider = _env_runner(self.env, self.policy,
rollout_provider = _env_runner(self.async_vector_env, self.policy,
self.num_local_steps, self.horizon,
self._obs_filter, self.pack)
while True:
@@ -103,15 +109,17 @@ class AsyncSampler(threading.Thread):
self.queue.put(item, timeout=600.0)
def get_data(self):
"""Gets currently accumulated data.
Returns:
rollout (SampleBatch): trajectory data (unprocessed)
"""
assert self.started, "Sampler never started running!"
rollout = self.queue.get(timeout=600.0)
# Propagate errors
if isinstance(rollout, BaseException):
raise rollout
# We can't auto-concat rollouts in vector mode
if self.async_vector_env.num_envs > 1:
return rollout
# Auto-concat rollouts; TODO(ekl) is this important for A3C perf?
while not rollout["dones"][-1]:
try:
part = self.queue.get_nowait()
@@ -132,7 +140,8 @@ class AsyncSampler(threading.Thread):
return completed
def _env_runner(env, policy, num_local_steps, horizon, obs_filter, pack):
def _env_runner(
async_vector_env, policy, num_local_steps, horizon, obs_filter, pack):
"""This implements the logic of the thread runner.
It continually runs the policy, and as long as the rollout exceeds a
@@ -141,9 +150,9 @@ def _env_runner(env, policy, num_local_steps, horizon, obs_filter, pack):
`num_local_steps` is reached.
Args:
env: Environment generated by env_creator
async_vector_env: env implementing AsyncVectorEnv.
policy: Policy used to interact with environment. Also sets fields
to be included in `SampleBatch`
to be included in `SampleBatch`.
num_local_steps: Number of steps before `SampleBatch` is yielded. Set
to infinity to yield complete episodes.
horizon: Horizon of the episode.
@@ -155,67 +164,146 @@ def _env_runner(env, policy, num_local_steps, horizon, obs_filter, pack):
rollout (SampleBatch): Object containing state, action, reward,
terminal condition, and other fields as dictated by `policy`.
"""
last_observation = obs_filter(env.reset())
try:
horizon = horizon if horizon else env.spec.max_episode_steps
if not horizon:
horizon = async_vector_env.get_unwrapped().spec.max_episode_steps
except Exception:
print("Warning, no horizon specified, assuming infinite")
if not horizon:
horizon = 999999
last_features = policy.get_initial_state()
features = last_features
length = 0
rewards = 0
rollout_number = 0
horizon = float("inf")
# Pool of batch builders, which can be shared across episodes to pack
# trajectory data.
batch_builder_pool = []
def get_batch_builder():
if batch_builder_pool:
return batch_builder_pool.pop()
else:
return SampleBatchBuilder()
episodes = defaultdict(
lambda: _Episode(policy.get_initial_state(), get_batch_builder))
while True:
batch_builder = SampleBatchBuilder()
# Get observations from ready envs
unfiltered_obs, rewards, dones, _, off_policy_actions = \
async_vector_env.poll()
ready_eids = []
ready_obs = []
ready_rnn_states = []
for _ in range(num_local_steps):
# Assume batch size one for now
action, features, pi_info = policy.compute_single_action(
last_observation, last_features, is_training=True)
for i, state_value in enumerate(last_features):
pi_info["state_in_{}".format(i)] = state_value
for i, state_value in enumerate(features):
pi_info["state_out_{}".format(i)] = state_value
observation, reward, terminal, info = env.step(action)
observation = obs_filter(observation)
# Process and record the new observations
for eid, raw_obs in unfiltered_obs.items():
episode = episodes[eid]
filtered_obs = obs_filter(raw_obs)
ready_eids.append(eid)
ready_obs.append(filtered_obs)
ready_rnn_states.append(episode.rnn_state)
length += 1
rewards += reward
if length >= horizon:
terminal = True
if episode.last_observation is None:
episode.last_observation = filtered_obs
continue # This is the initial observation after a reset
# Concatenate multiagent actions
if isinstance(action, list):
action = np.concatenate(action, axis=0).flatten()
episode.length += 1
episode.total_reward += rewards[eid]
# Collect the experience.
batch_builder.add_values(
obs=last_observation,
actions=action,
rewards=reward,
dones=terminal,
new_obs=observation,
**pi_info)
# Handle episode terminations
if dones[eid] or episode.length >= horizon:
done = True
yield CompletedRollout(episode.length, episode.total_reward)
else:
done = False
last_observation = observation
last_features = features
episode.batch_builder.add_values(
obs=episode.last_observation,
actions=episode.last_action_flat(),
rewards=rewards[eid],
dones=done,
new_obs=filtered_obs,
**episode.last_pi_info)
if terminal:
yield CompletedRollout(length, rewards)
# Cut the batch if we're not packing multiple episodes into one,
# or if we've exceeded the requested batch size.
if (done and not pack) or \
episode.batch_builder.count >= num_local_steps:
yield episode.batch_builder.build_and_reset(
policy.postprocess_trajectory)
elif done:
# Make sure postprocessor never goes across episode boundaries
episode.batch_builder.postprocess_batch_so_far(
policy.postprocess_trajectory)
if (length >= horizon or
not env.metadata.get("semantics.autoreset")):
last_observation = obs_filter(env.reset())
last_features = policy.get_initial_state()
rollout_number += 1
length = 0
rewards = 0
if not pack:
break
if done:
# Handle episode termination
batch_builder_pool.append(episode.batch_builder)
del episodes[eid]
resetted_obs = async_vector_env.try_reset(eid)
if resetted_obs is None:
# Reset not supported, drop this env from the ready list
assert horizon == float("inf"), \
"Setting episode horizon requires reset() support."
ready_eids.pop()
ready_obs.pop()
ready_rnn_states.pop()
else:
# Reset successful, put in the new obs as ready
episode = episodes[eid]
episode.last_observation = obs_filter(resetted_obs)
ready_obs[-1] = episode.last_observation
ready_rnn_states[-1] = episode.rnn_state
else:
episode.last_observation = filtered_obs
# Once we have enough experience, yield it, and have the ThreadRunner
# place it on a queue.
yield batch_builder.build()
if not ready_eids:
continue # No actions to take
# Compute action for ready envs
ready_rnn_state_cols = _to_column_format(ready_rnn_states)
actions, new_rnn_state_cols, pi_info_cols = policy.compute_actions(
ready_obs, ready_rnn_state_cols, is_training=True)
# Add RNN state info
for f_i, column in enumerate(ready_rnn_state_cols):
pi_info_cols["state_in_{}".format(f_i)] = column
for f_i, column in enumerate(new_rnn_state_cols):
pi_info_cols["state_out_{}".format(f_i)] = column
# Return computed actions to ready envs. We also send to envs that have
# taken off-policy actions; those envs are free to ignore the action.
async_vector_env.send_actions(dict(zip(ready_eids, actions)))
# Store the computed action info
for i, eid in enumerate(ready_eids):
episode = episodes[eid]
if eid in off_policy_actions:
episode.last_action = off_policy_actions[eid]
else:
episode.last_action = actions[i]
episode.rnn_state = [column[i] for column in new_rnn_state_cols]
episode.last_pi_info = {
k: column[i] for k, column in pi_info_cols.items()}
def _to_column_format(rnn_state_rows):
num_cols = len(rnn_state_rows[0])
return [
[row[i] for row in rnn_state_rows] for i in range(num_cols)]
class _Episode(object):
def __init__(self, init_rnn_state, batch_builder_factory):
self.rnn_state = init_rnn_state
self.batch_builder = batch_builder_factory()
self.last_action = None
self.last_observation = None
self.last_pi_info = None
self.total_reward = 0.0
self.length = 0
def last_action_flat(self):
# Concatenate multiagent actions
if isinstance(self.last_action, list):
return np.concatenate(self.last_action, axis=0).flatten()
return self.last_action