mirror of
https://github.com/wassname/ray.git
synced 2026-07-13 12:01:30 +08:00
[rllib] Part 1 of multiagent support: make sampler path support multiagent envs (#2268)
This refactors the RLlib sampler to support multi-agent environments. The main changes were: AsyncVectorEnv now produces dicts of env_id -> agent_id -> value rather than env_id -> value. This lets it model both vectorized and multi-agent envs (or both). The sampler class operates over the above nested dict structure for all envs. Single agent envs just return a dict with one agent_id=single_agent. When sample() is called on a policy evaluator, in the single agent case we return a SampleBatch, otherwise we return a MultiAgentBatch (which is a list of sample batches per policy). Left for another PR: Exposing multi-agent in the public interfaces. Optimizations such as evaluating multiple policies in one TF run.
This commit is contained in:
Eric Liang
+198
-141
@@ -7,13 +7,16 @@ 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
|
||||
from ray.rllib.optimizers.sample_batch import MultiAgentSampleBatchBuilder, \
|
||||
MultiAgentBatch
|
||||
from ray.rllib.utils.async_vector_env import AsyncVectorEnv
|
||||
|
||||
|
||||
CompletedRollout = namedtuple("CompletedRollout",
|
||||
["episode_length", "episode_reward"])
|
||||
RolloutMetrics = namedtuple(
|
||||
"RolloutMetrics", ["episode_length", "episode_reward", "agent_rewards"])
|
||||
|
||||
PolicyEvalData = namedtuple(
|
||||
"PolicyEvalData", ["env_id", "agent_id", "obs", "rnn_state"])
|
||||
|
||||
|
||||
class SyncSampler(object):
|
||||
@@ -26,26 +29,23 @@ class SyncSampler(object):
|
||||
thread."""
|
||||
|
||||
def __init__(
|
||||
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, env, policies, policy_mapping_fn, obs_filters,
|
||||
num_local_steps, horizon=None, pack=False):
|
||||
self.async_vector_env = AsyncVectorEnv.wrap_async(env)
|
||||
self.num_local_steps = num_local_steps
|
||||
self.horizon = horizon
|
||||
self.policy = policy
|
||||
self._obs_filter = obs_filter
|
||||
self.rollout_provider = _env_runner(self.async_vector_env, self.policy,
|
||||
self.num_local_steps, self.horizon,
|
||||
self._obs_filter, pack)
|
||||
self.policies = policies
|
||||
self.policy_mapping_fn = policy_mapping_fn
|
||||
self._obs_filters = obs_filters
|
||||
self.rollout_provider = _env_runner(
|
||||
self.async_vector_env, self.policies, self.policy_mapping_fn,
|
||||
self.num_local_steps, self.horizon, self._obs_filters, pack)
|
||||
self.metrics_queue = queue.Queue()
|
||||
|
||||
def get_data(self):
|
||||
while True:
|
||||
item = next(self.rollout_provider)
|
||||
if isinstance(item, CompletedRollout):
|
||||
if isinstance(item, RolloutMetrics):
|
||||
self.metrics_queue.put(item)
|
||||
else:
|
||||
return item
|
||||
@@ -67,23 +67,20 @@ class AsyncSampler(threading.Thread):
|
||||
accumulate and the gradient can be calculated on up to 5 batches."""
|
||||
|
||||
def __init__(
|
||||
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
|
||||
self, env, policies, policy_mapping_fn, obs_filters,
|
||||
num_local_steps, horizon=None, pack=False):
|
||||
for _, f in obs_filters.items():
|
||||
assert getattr(f, "is_concurrent", False), \
|
||||
"Observation Filter must support concurrent updates."
|
||||
self.async_vector_env = AsyncVectorEnv.wrap_async(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.policy = policy
|
||||
self._obs_filter = obs_filter
|
||||
self.policies = policies
|
||||
self.policy_mapping_fn = policy_mapping_fn
|
||||
self._obs_filters = obs_filters
|
||||
self.daemon = True
|
||||
self.pack = pack
|
||||
|
||||
@@ -95,15 +92,15 @@ class AsyncSampler(threading.Thread):
|
||||
raise e
|
||||
|
||||
def _run(self):
|
||||
rollout_provider = _env_runner(self.async_vector_env, self.policy,
|
||||
self.num_local_steps, self.horizon,
|
||||
self._obs_filter, self.pack)
|
||||
rollout_provider = _env_runner(
|
||||
self.async_vector_env, self.policies, self.policy_mapping_fn,
|
||||
self.num_local_steps, self.horizon, self._obs_filters, self.pack)
|
||||
while True:
|
||||
# The timeout variable exists because apparently, if one worker
|
||||
# dies, the other workers won't die with it, unless the timeout is
|
||||
# set to some large number. This is an empirical observation.
|
||||
item = next(rollout_provider)
|
||||
if isinstance(item, CompletedRollout):
|
||||
if isinstance(item, RolloutMetrics):
|
||||
self.metrics_queue.put(item)
|
||||
else:
|
||||
self.queue.put(item, timeout=600.0)
|
||||
@@ -115,8 +112,9 @@ class AsyncSampler(threading.Thread):
|
||||
if isinstance(rollout, BaseException):
|
||||
raise rollout
|
||||
|
||||
# We can't auto-concat rollouts in vector mode
|
||||
if self.async_vector_env.num_envs > 1:
|
||||
# We can't auto-concat rollouts in these modes
|
||||
if self.async_vector_env.num_envs > 1 or \
|
||||
isinstance(rollout, MultiAgentBatch):
|
||||
return rollout
|
||||
|
||||
# Auto-concat rollouts; TODO(ekl) is this important for A3C perf?
|
||||
@@ -141,23 +139,22 @@ class AsyncSampler(threading.Thread):
|
||||
|
||||
|
||||
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
|
||||
certain length, the thread runner appends the policy to the queue. Yields
|
||||
when `timestep_limit` is surpassed, environment terminates, or
|
||||
`num_local_steps` is reached.
|
||||
async_vector_env, policies, policy_mapping_fn, num_local_steps,
|
||||
horizon, obs_filters, pack):
|
||||
"""This implements the common experience collection logic.
|
||||
|
||||
Args:
|
||||
async_vector_env: env implementing AsyncVectorEnv.
|
||||
policy: Policy used to interact with environment. Also sets fields
|
||||
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.
|
||||
obs_filter: Filter used to process observations.
|
||||
pack: Whether to pack multiple episodes into each batch. This
|
||||
async_vector_env (AsyncVectorEnv): env implementing AsyncVectorEnv.
|
||||
policies (dict): Map of policy ids to PolicyGraph instances.
|
||||
policy_mapping_fn (func): Function that maps agent ids to policy ids.
|
||||
This is called when an agent first enters the environment. The
|
||||
agent is then "bound" to the returned policy for the episode.
|
||||
num_local_steps (int): Number of episode steps before `SampleBatch` is
|
||||
yielded. Set to infinity to yield complete episodes.
|
||||
horizon (int): Horizon of the episode.
|
||||
obs_filters (dict): Map of policy id to filter used to process
|
||||
observations for the policy.
|
||||
pack (bool): Whether to pack multiple episodes into each batch. This
|
||||
guarantees batches will be exactly `num_local_steps` in size.
|
||||
|
||||
Yields:
|
||||
@@ -181,110 +178,131 @@ def _env_runner(
|
||||
if batch_builder_pool:
|
||||
return batch_builder_pool.pop()
|
||||
else:
|
||||
return SampleBatchBuilder()
|
||||
return MultiAgentSampleBatchBuilder(policies)
|
||||
|
||||
episodes = defaultdict(
|
||||
lambda: _Episode(policy.get_initial_state(), get_batch_builder))
|
||||
active_episodes = defaultdict(
|
||||
lambda: _MultiAgentEpisode(
|
||||
policies, policy_mapping_fn, get_batch_builder))
|
||||
|
||||
while True:
|
||||
# Get observations from ready envs
|
||||
# Get observations from all ready agents
|
||||
unfiltered_obs, rewards, dones, infos, off_policy_actions = \
|
||||
async_vector_env.poll()
|
||||
ready_eids = []
|
||||
ready_obs = []
|
||||
ready_rnn_states = []
|
||||
|
||||
# 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)
|
||||
# Map of policy_id to list of PolicyEvalData
|
||||
to_eval = defaultdict(list)
|
||||
|
||||
if episode.last_observation is None:
|
||||
episode.last_observation = filtered_obs
|
||||
continue # This is the initial observation after a reset
|
||||
# For each environment
|
||||
for env_id, agent_obs in unfiltered_obs.items():
|
||||
new_episode = env_id not in active_episodes
|
||||
episode = active_episodes[env_id]
|
||||
if not new_episode:
|
||||
episode.length += 1
|
||||
episode.batch_builder.count += 1
|
||||
episode.add_agent_rewards(rewards[env_id])
|
||||
|
||||
episode.length += 1
|
||||
episode.total_reward += rewards[eid]
|
||||
|
||||
# Handle episode terminations
|
||||
if dones[eid] or episode.length >= horizon:
|
||||
done = True
|
||||
yield CompletedRollout(episode.length, episode.total_reward)
|
||||
# Check episode termination conditions
|
||||
if dones[env_id]["__all__"] or episode.length >= horizon:
|
||||
all_done = True
|
||||
yield RolloutMetrics(
|
||||
episode.length, episode.total_reward,
|
||||
dict(episode.agent_rewards))
|
||||
else:
|
||||
done = False
|
||||
all_done = False
|
||||
|
||||
if infos[eid].get("training_enabled", True):
|
||||
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)
|
||||
# For each agent in the environment
|
||||
for agent_id, raw_obs in agent_obs.items():
|
||||
policy_id = episode.policy_for(agent_id)
|
||||
filtered_obs = obs_filters[policy_id](raw_obs)
|
||||
agent_done = bool(all_done or dones[env_id].get(agent_id))
|
||||
if not agent_done:
|
||||
to_eval[policy_id].append(
|
||||
PolicyEvalData(
|
||||
env_id, agent_id, filtered_obs,
|
||||
episode.rnn_state_for(agent_id)))
|
||||
|
||||
# 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 \
|
||||
last_observation = episode.last_observation_for(agent_id)
|
||||
episode.set_last_observation(agent_id, filtered_obs)
|
||||
|
||||
# Record transition info if applicable
|
||||
if last_observation is not None and \
|
||||
infos[env_id][agent_id].get("training_enabled", True):
|
||||
episode.batch_builder.add_values(
|
||||
agent_id,
|
||||
policy_id,
|
||||
t=episode.length - 1,
|
||||
obs=last_observation,
|
||||
actions=episode.last_action_for(agent_id),
|
||||
rewards=rewards[env_id][agent_id],
|
||||
dones=agent_done,
|
||||
infos=infos[env_id][agent_id],
|
||||
new_obs=filtered_obs,
|
||||
**episode.last_pi_info_for(agent_id))
|
||||
|
||||
# Cut the batch if we're not packing multiple episodes into one,
|
||||
# or if we've exceeded the requested batch size.
|
||||
if episode.batch_builder.has_pending_data():
|
||||
if (all_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 crosses episode boundaries
|
||||
episode.batch_builder.postprocess_batch_so_far(
|
||||
policy.postprocess_trajectory)
|
||||
yield episode.batch_builder.build_and_reset()
|
||||
elif all_done:
|
||||
# Make sure postprocessor stays within one episode
|
||||
episode.batch_builder.postprocess_batch_so_far()
|
||||
|
||||
if done:
|
||||
if all_done:
|
||||
# Handle episode termination
|
||||
batch_builder_pool.append(episode.batch_builder)
|
||||
del episodes[eid]
|
||||
resetted_obs = async_vector_env.try_reset(eid)
|
||||
del active_episodes[env_id]
|
||||
resetted_obs = async_vector_env.try_reset(env_id)
|
||||
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
|
||||
# Creates a new episode
|
||||
episode = active_episodes[env_id]
|
||||
for agent_id, raw_obs in resetted_obs.items():
|
||||
policy_id = episode.policy_for(agent_id)
|
||||
filtered_obs = obs_filters[policy_id](raw_obs)
|
||||
episode.set_last_observation(agent_id, filtered_obs)
|
||||
to_eval[policy_id].append(
|
||||
PolicyEvalData(
|
||||
env_id, agent_id, filtered_obs,
|
||||
episode.rnn_state_for(agent_id)))
|
||||
|
||||
if not ready_eids:
|
||||
continue # No actions to take
|
||||
# Map of env_id -> agent_id -> action
|
||||
action_dict = defaultdict(dict)
|
||||
|
||||
# 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
|
||||
# TODO(ekl) fuse all policy evaluation into one TF run
|
||||
for policy_id, eval_data in to_eval.items():
|
||||
rnn_in_cols = _to_column_format([t.rnn_state for t in eval_data])
|
||||
actions, rnn_out_cols, pi_info_cols = \
|
||||
policies[policy_id].compute_actions(
|
||||
[t.obs for t in eval_data], rnn_in_cols, is_training=True)
|
||||
# Add RNN state info
|
||||
for f_i, column in enumerate(rnn_in_cols):
|
||||
pi_info_cols["state_in_{}".format(f_i)] = column
|
||||
for f_i, column in enumerate(rnn_out_cols):
|
||||
pi_info_cols["state_out_{}".format(f_i)] = column
|
||||
# Save output rows
|
||||
for i, action in enumerate(actions):
|
||||
env_id = eval_data[i].env_id
|
||||
agent_id = eval_data[i].agent_id
|
||||
action_dict[env_id][agent_id] = action
|
||||
episode = active_episodes[env_id]
|
||||
episode.set_rnn_state(agent_id, [c[i] for c in rnn_out_cols])
|
||||
episode.set_last_pi_info(
|
||||
agent_id, {k: v[i] for k, v in pi_info_cols.items()})
|
||||
if env_id in off_policy_actions and \
|
||||
agent_id in off_policy_actions[env_id]:
|
||||
episode.set_last_action(
|
||||
agent_id, off_policy_actions[env_id][agent_id])
|
||||
else:
|
||||
episode.set_last_action(agent_id, action)
|
||||
|
||||
# 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()}
|
||||
async_vector_env.send_actions(dict(action_dict))
|
||||
|
||||
|
||||
def _to_column_format(rnn_state_rows):
|
||||
@@ -293,18 +311,57 @@ def _to_column_format(rnn_state_rows):
|
||||
[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
|
||||
class _MultiAgentEpisode(object):
|
||||
def __init__(self, policies, policy_mapping_fn, batch_builder_factory):
|
||||
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
|
||||
self.agent_rewards = defaultdict(float)
|
||||
self._policies = policies
|
||||
self._policy_mapping_fn = policy_mapping_fn
|
||||
self._agent_to_policy = {}
|
||||
self._agent_to_rnn_state = {}
|
||||
self._agent_to_last_obs = {}
|
||||
self._agent_to_last_action = {}
|
||||
self._agent_to_last_pi_info = {}
|
||||
|
||||
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
|
||||
def add_agent_rewards(self, reward_dict):
|
||||
for agent_id, reward in reward_dict.items():
|
||||
self.agent_rewards[agent_id] += reward
|
||||
self.total_reward += reward
|
||||
|
||||
def policy_for(self, agent_id):
|
||||
if agent_id not in self._agent_to_policy:
|
||||
self._agent_to_policy[agent_id] = self._policy_mapping_fn(agent_id)
|
||||
return self._agent_to_policy[agent_id]
|
||||
|
||||
def rnn_state_for(self, agent_id):
|
||||
if agent_id not in self._agent_to_rnn_state:
|
||||
policy = self._policies[self.policy_for(agent_id)]
|
||||
self._agent_to_rnn_state[agent_id] = policy.get_initial_state()
|
||||
return self._agent_to_rnn_state[agent_id]
|
||||
|
||||
def last_observation_for(self, agent_id):
|
||||
return self._agent_to_last_obs.get(agent_id)
|
||||
|
||||
def last_action_for(self, agent_id):
|
||||
action = self._agent_to_last_action[agent_id]
|
||||
# Concatenate tuple actions
|
||||
if isinstance(action, list):
|
||||
action = np.concatenate(action, axis=0).flatten()
|
||||
return action
|
||||
|
||||
def last_pi_info_for(self, agent_id):
|
||||
return self._agent_to_last_pi_info[agent_id]
|
||||
|
||||
def set_rnn_state(self, agent_id, rnn_state):
|
||||
self._agent_to_rnn_state[agent_id] = rnn_state
|
||||
|
||||
def set_last_observation(self, agent_id, obs):
|
||||
self._agent_to_last_obs[agent_id] = obs
|
||||
|
||||
def set_last_action(self, agent_id, action):
|
||||
self._agent_to_last_action[agent_id] = action
|
||||
|
||||
def set_last_pi_info(self, agent_id, pi_info):
|
||||
self._agent_to_last_pi_info[agent_id] = pi_info
|
||||
|
||||
Reference in New Issue
Block a user