From aeb5be77332606e97af7f569a34b8d211cc91558 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Sat, 15 Aug 2020 15:09:00 +0200 Subject: [PATCH] [RLlib] Trajectory View API (part 2.5): Actual implementations (not used yet) of a SampleCollector. (#10112) --- rllib/evaluation/episode.py | 10 +- .../multi_agent_sample_collector.py | 247 +++++++++ .../evaluation/per_policy_sample_collector.py | 487 ++++++++++++++++++ rllib/evaluation/rollout_worker.py | 149 +++--- rllib/evaluation/sample_batch_builder.py | 4 +- rllib/evaluation/sample_collector.py | 38 +- rllib/evaluation/tests/test_trajectories.py | 70 --- rllib/policy/view_requirement.py | 12 +- rllib/rollout.py | 2 +- 9 files changed, 854 insertions(+), 165 deletions(-) create mode 100644 rllib/evaluation/multi_agent_sample_collector.py create mode 100644 rllib/evaluation/per_policy_sample_collector.py delete mode 100644 rllib/evaluation/tests/test_trajectories.py diff --git a/rllib/evaluation/episode.py b/rllib/evaluation/episode.py index b09171c4e..dbb33ac2c 100644 --- a/rllib/evaluation/episode.py +++ b/rllib/evaluation/episode.py @@ -92,11 +92,17 @@ class MultiAgentEpisode: self._agent_reward_history = defaultdict(list) @DeveloperAPI - def policy_for(self, agent_id: AgentID = _DUMMY_AGENT_ID) -> Policy: - """Returns the policy for the specified agent. + def policy_for(self, agent_id: AgentID = _DUMMY_AGENT_ID) -> PolicyID: + """Returns and stores the policy ID for the specified agent. If the agent is new, the policy mapping fn will be called to bind the agent to a policy for the duration of the episode. + + Args: + agent_id (AgentID): The agent ID to lookup the policy ID for. + + Returns: + PolicyID: The policy ID for the specified agent. """ if agent_id not in self._agent_to_policy: diff --git a/rllib/evaluation/multi_agent_sample_collector.py b/rllib/evaluation/multi_agent_sample_collector.py new file mode 100644 index 000000000..6b546e810 --- /dev/null +++ b/rllib/evaluation/multi_agent_sample_collector.py @@ -0,0 +1,247 @@ +import logging +from typing import Dict, Optional + +from ray.rllib.agents.callbacks import DefaultCallbacks +from ray.rllib.env.base_env import _DUMMY_AGENT_ID +from ray.rllib.evaluation.episode import MultiAgentEpisode +from ray.rllib.evaluation.per_policy_sample_collector import \ + _PerPolicySampleCollector +from ray.rllib.evaluation.sample_collector import _SampleCollector +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import override +from ray.rllib.utils.debug import summarize +from ray.rllib.utils.types import AgentID, EnvID, EpisodeID, PolicyID, \ + TensorType +from ray.util.debug import log_once + +logger = logging.getLogger(__name__) + + +class _MultiAgentSampleCollector(_SampleCollector): + """Builds SampleBatches for each policy (and agent) in a multi-agent env. + + Note: This is an experimental class only used when + `config._use_trajectory_view_api` = True. + Once `_use_trajectory_view_api` becomes the default in configs: + This class will deprecate the `SampleBatchBuilder` class. + + Input data is collected in central per-policy buffers, which + efficiently pre-allocate memory (over n timesteps) and re-use the same + memory even for succeeding agents and episodes. + Input_dicts for action computations, SampleBatches for postprocessing, and + train_batch dicts are - if possible - created from the central per-policy + buffers via views to avoid copying of data). + """ + + def __init__( + self, + policy_map: Dict[PolicyID, Policy], + callbacks: DefaultCallbacks, + # TODO: (sven) make `num_agents` flexibly grow in size. + num_agents: int = 100, + num_timesteps=None, + time_major: Optional[bool] = False): + """Initializes a _MultiAgentSampleCollector object. + + Args: + policy_map (Dict[PolicyID,Policy]): Maps policy ids to policy + instances. + callbacks (DefaultCallbacks): RLlib callbacks (configured in the + Trainer config dict). Used for trajectory postprocessing event. + num_agents (int): The max number of agent slots to pre-allocate + in the buffer. + num_timesteps (int): The max number of timesteps to pre-allocate + in the buffer. + time_major (Optional[bool]): Whether to preallocate buffers and + collect samples in time-major fashion (TxBx...). + """ + + self.policy_map = policy_map + self.callbacks = callbacks + if num_agents == float("inf") or num_agents is None: + num_agents = 1000 + self.num_agents = int(num_agents) + + # Collect SampleBatches per-policy in PolicyTrajectories objects. + self.rollout_sample_collectors = {} + for pid, policy in policy_map.items(): + # Figure out max-shifts (before and after). + view_reqs = policy.training_view_requirements + max_shift_before = 0 + max_shift_after = 0 + for vr in view_reqs.values(): + shift = force_list(vr.shift) + if max_shift_before > shift[0]: + max_shift_before = shift[0] + if max_shift_after < shift[-1]: + max_shift_after = shift[-1] + # Figure out num_timesteps and num_agents. + kwargs = {"time_major": time_major} + if policy.is_recurrent(): + kwargs["num_timesteps"] = \ + policy.config["model"]["max_seq_len"] + kwargs["time_major"] = True + elif num_timesteps is not None: + kwargs["num_timesteps"] = num_timesteps + + self.rollout_sample_collectors[pid] = _PerPolicySampleCollector( + num_agents=self.num_agents, + shift_before=-max_shift_before, + shift_after=max_shift_after, + **kwargs) + + # Internal agent-to-policy map. + self.agent_to_policy = {} + # Number of "inference" steps taken in the environment. + # Regardless of the number of agents involved in each of these steps. + self.count = 0 + + @override(_SampleCollector) + def add_init_obs(self, episode_id: EpisodeID, agent_id: AgentID, + env_id: EnvID, policy_id: PolicyID, + obs: TensorType) -> None: + # Make sure our mappings are up to date. + if agent_id not in self.agent_to_policy: + self.agent_to_policy[agent_id] = policy_id + else: + assert self.agent_to_policy[agent_id] == policy_id + + # Add initial obs to Trajectory. + self.rollout_sample_collectors[policy_id].add_init_obs( + episode_id, agent_id, env_id, chunk_num=0, init_obs=obs) + + @override(_SampleCollector) + def add_action_reward_next_obs(self, episode_id: EpisodeID, + agent_id: AgentID, env_id: EnvID, + policy_id: PolicyID, agent_done: bool, + values: Dict[str, TensorType]) -> None: + assert policy_id in self.rollout_sample_collectors + + # Make sure our mappings are up to date. + if agent_id not in self.agent_to_policy: + self.agent_to_policy[agent_id] = policy_id + else: + assert self.agent_to_policy[agent_id] == policy_id + + # Include the current agent id for multi-agent algorithms. + if agent_id != _DUMMY_AGENT_ID: + values["agent_id"] = agent_id + + # Add action/reward/next-obs (and other data) to Trajectory. + self.rollout_sample_collectors[policy_id].add_action_reward_next_obs( + episode_id, agent_id, env_id, agent_done, values) + + @override(_SampleCollector) + def total_env_steps(self) -> int: + return sum(a.timesteps_since_last_reset + for a in self.rollout_sample_collectors.values()) + + def total(self): + # TODO: (sven) deprecate; use `self.total_env_steps`, instead. + # Sampler is currently still using `total()`. + return self.total_env_steps() + + @override(_SampleCollector) + def get_inference_input_dict(self, policy_id: PolicyID) -> \ + Dict[str, TensorType]: + policy = self.policy_map[policy_id] + view_reqs = policy.model.inference_view_requirements + return self.rollout_sample_collectors[ + policy_id].get_inference_input_dict(view_reqs) + + @override(_SampleCollector) + def has_non_postprocessed_data(self) -> bool: + return self.total_env_steps() > 0 + + @override(_SampleCollector) + def postprocess_trajectories_so_far( + self, episode: Optional[MultiAgentEpisode] = None) -> None: + # Loop through each per-policy collector and create a view (for each + # agent as SampleBatch) from its buffers for post-processing + all_agent_batches = {} + for pid, rc in self.rollout_sample_collectors.items(): + policy = self.policy_map[pid] + view_reqs = policy.training_view_requirements + agent_batches = rc.get_postprocessing_sample_batches( + episode, view_reqs) + + for agent_key, batch in agent_batches.items(): + other_batches = None + if len(agent_batches) > 1: + other_batches = agent_batches.copy() + del other_batches[agent_key] + + agent_batches[agent_key] = policy.postprocess_trajectory( + batch, other_batches, episode) + # Call the Policy's Exploration's postprocess method. + if getattr(policy, "exploration", None) is not None: + agent_batches[ + agent_key] = policy.exploration.postprocess_trajectory( + policy, agent_batches[agent_key], + getattr(policy, "_sess", None)) + + # Add new columns' data to buffers. + for col in agent_batches[agent_key].new_columns: + data = agent_batches[agent_key].data[col] + rc._build_buffers({col: data[0]}) + timesteps = data.shape[0] + rc.buffers[col][rc.shift_before:rc.shift_before + + timesteps, rc.agent_key_to_slot[ + agent_key]] = data + + all_agent_batches.update(agent_batches) + + if log_once("after_post"): + logger.info("Trajectory fragment after postprocess_trajectory():" + "\n\n{}\n".format(summarize(all_agent_batches))) + + # Append into policy batches and reset + from ray.rllib.evaluation.rollout_worker import get_global_worker + for agent_key, batch in sorted(all_agent_batches.items()): + self.callbacks.on_postprocess_trajectory( + worker=get_global_worker(), + episode=episode, + agent_id=agent_key[0], + policy_id=self.agent_to_policy[agent_key[0]], + policies=self.policy_map, + postprocessed_batch=batch, + original_batches=None) # TODO: (sven) do we really need this? + + @override(_SampleCollector) + def check_missing_dones(self, episode_id: EpisodeID) -> None: + for pid, rc in self.rollout_sample_collectors.items(): + for agent_key in rc.agent_key_to_slot.keys(): + # Only check for given episode and only for last chunk + # (all previous chunks for that agent in the episode are + # non-terminal). + if (agent_key[1] == episode_id + and rc.agent_key_to_chunk_num[agent_key[:2]] == + agent_key[2]): + t = rc.agent_key_to_timestep[agent_key] - 1 + b = rc.agent_key_to_slot[agent_key] + if not rc.buffers["dones"][t][b]: + raise ValueError( + "Episode {} terminated for all agents, but we " + "still don't have a last observation for " + "agent {} (policy {}). ".format(agent_key[0], pid) + + "Please ensure that you include the last " + "observations of all live agents when setting " + "'__all__' done to True. Alternatively, set " + "no_done_at_end=True to allow this.") + + @override(_SampleCollector) + def get_multi_agent_batch_and_reset(self): + self.postprocess_trajectories_so_far() + policy_batches = {} + for pid, rc in self.rollout_sample_collectors.items(): + policy = self.policy_map[pid] + view_reqs = policy.training_view_requirements + policy_batches[pid] = rc.get_train_sample_batch_and_reset( + view_reqs) + + ma_batch = MultiAgentBatch.wrap_as_needed(policy_batches, self.count) + # Reset our across-all-agents env step count. + self.count = 0 + return ma_batch diff --git a/rllib/evaluation/per_policy_sample_collector.py b/rllib/evaluation/per_policy_sample_collector.py new file mode 100644 index 000000000..0834c0fac --- /dev/null +++ b/rllib/evaluation/per_policy_sample_collector.py @@ -0,0 +1,487 @@ +import logging +import numpy as np +from typing import Dict, Optional + +from ray.rllib.evaluation.episode import MultiAgentEpisode +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.view_requirement import ViewRequirement +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.types import AgentID, EnvID, EpisodeID, TensorType + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + +logger = logging.getLogger(__name__) + + +class _PerPolicySampleCollector: + """A class for efficiently collecting samples for a single (fixed) policy. + + Can be used by a _MultiAgentSampleCollector for its different policies. + """ + + def __init__(self, + num_agents: Optional[int] = None, + num_timesteps: Optional[int] = None, + time_major: bool = True, + shift_before: int = 0, + shift_after: int = 0): + """Initializes a _PerPolicySampleCollector object. + + Args: + num_agents (int): The max number of agent slots to pre-allocate + in the buffer. + num_timesteps (int): The max number of timesteps to pre-allocate + in the buffer. + time_major (Optional[bool]): Whether to preallocate buffers and + collect samples in time-major fashion (TxBx...). + shift_before (int): The additional number of time slots to + pre-allocate at the beginning of a time window (for possible + underlying data column shifts, e.g. PREV_ACTIONS). + shift_after (int): The additional number of time slots to + pre-allocate at the end of a time window (for possible + underlying data column shifts, e.g. NEXT_OBS). + """ + + self.num_agents = num_agents or 100 + self.num_timesteps = num_timesteps + self.time_major = time_major + # `shift_before must at least be 1 for the init obs timestep. + self.shift_before = max(shift_before, 1) + self.shift_after = shift_after + + # The offset on the agent dim to start the next SampleBatch build from. + self.sample_batch_offset = 0 + + # The actual underlying data-buffers. + self.buffers = {} + self.postprocessed_agents = [False] * self.num_agents + + # Next agent-slot to be used by a new agent/env combination. + self.agent_slot_cursor = 0 + # Maps agent/episode ID/chunk-num to an agent slot. + self.agent_key_to_slot = {} + # Maps agent/episode ID to the last chunk-num. + self.agent_key_to_chunk_num = {} + # Maps agent slot number to agent keys. + self.slot_to_agent_key = [None] * self.num_agents + # Maps agent/episode ID/chunk-num to a time step cursor. + self.agent_key_to_timestep = {} + + # Total timesteps taken in the env over all agents since last reset. + self.timesteps_since_last_reset = 0 + + # Indices (T,B) to pick from the buffers for the next forward pass. + self.forward_pass_indices = [[], []] + self.forward_pass_size = 0 + # Maps index from the forward pass batch to (agent_id, episode_id, + # env_id) tuple. + self.forward_pass_index_to_agent_info = {} + self.agent_key_to_forward_pass_index = {} + + def add_init_obs(self, episode_id: EpisodeID, agent_id: AgentID, + env_id: EnvID, chunk_num: int, + init_obs: TensorType) -> None: + """Adds a single initial observation (after env.reset()) to the buffer. + + Args: + episode_id (EpisodeID): Unique ID for the episode we are adding the + initial observation for. + agent_id (AgentID): Unique ID for the agent we are adding the + initial observation for. + env_id (EnvID): The env ID to which `init_obs` belongs. + chunk_num (int): The time-chunk number (0-based). Some episodes + may last for longer than self.num_timesteps and therefore + have to be chopped into chunks. + init_obs (TensorType): Initial observation (after env.reset()). + """ + agent_key = (agent_id, episode_id, chunk_num) + agent_slot = self.agent_slot_cursor + self.agent_key_to_slot[agent_key] = agent_slot + self.agent_key_to_chunk_num[agent_key[:2]] = chunk_num + self.slot_to_agent_key[agent_slot] = agent_key + self._next_agent_slot() + + if SampleBatch.OBS not in self.buffers: + self._build_buffers(single_row={SampleBatch.OBS: init_obs}) + if self.time_major: + self.buffers[SampleBatch.OBS][self.shift_before-1, agent_slot] = \ + init_obs + else: + self.buffers[SampleBatch.OBS][agent_slot, self.shift_before-1] = \ + init_obs + self.agent_key_to_timestep[agent_key] = self.shift_before + + self._add_to_next_inference_call(agent_key, env_id, agent_slot, + self.shift_before - 1) + + def add_action_reward_next_obs( + self, episode_id: EpisodeID, agent_id: AgentID, env_id: EnvID, + agent_done: bool, values: Dict[str, TensorType]) -> None: + """Add the given dictionary (row) of values to this batch. + + Args: + episode_id (EpisodeID): Unique ID for the episode we are adding the + values for. + agent_id (AgentID): Unique ID for the agent we are adding the + values for. + env_id (EnvID): The env ID to which the given data belongs. + agent_done (bool): Whether next obs should not be used for an + upcoming inference call. Default: False = next-obs should be + used for upcoming inference. + values (Dict[str, TensorType]): Data dict (interpreted as a single + row) to be added to buffer. Must contain keys: + SampleBatch.ACTIONS, REWARDS, DONES, and NEXT_OBS. + """ + assert (SampleBatch.ACTIONS in values and SampleBatch.REWARDS in values + and SampleBatch.NEXT_OBS in values + and SampleBatch.DONES in values) + + assert SampleBatch.OBS not in values + values[SampleBatch.OBS] = values[SampleBatch.NEXT_OBS] + del values[SampleBatch.NEXT_OBS] + + chunk_num = self.agent_key_to_chunk_num[(agent_id, episode_id)] + agent_key = (agent_id, episode_id, chunk_num) + agent_slot = self.agent_key_to_slot[agent_key] + ts = self.agent_key_to_timestep[agent_key] + for k, v in values.items(): + if k not in self.buffers: + self._build_buffers(single_row=values) + if self.time_major: + self.buffers[k][ts, agent_slot] = v + else: + self.buffers[k][agent_slot, ts] = v + self.agent_key_to_timestep[agent_key] += 1 + + # Time-axis is "full" -> Cut-over to new chunk (only if not DONE). + if self.agent_key_to_timestep[ + agent_key] - self.shift_before == self.num_timesteps and \ + not values[SampleBatch.DONES]: + self._new_chunk_from(agent_slot, agent_key, + self.agent_key_to_timestep[agent_key]) + + self.timesteps_since_last_reset += 1 + + if not agent_done: + self._add_to_next_inference_call(agent_key, env_id, agent_slot, ts) + + def get_inference_input_dict(self, view_reqs: Dict[str, ViewRequirement] + ) -> Dict[str, TensorType]: + """Returns an input_dict for an (inference) forward pass. + + The input_dict can then be used for action computations inside a + Policy via `Policy.compute_actions_from_input_dict()`. + + Args: + view_reqs (Dict[str, ViewRequirement]): The view requirements + dict to use. + + Returns: + Dict[str, TensorType]: The input_dict to be passed into the ModelV2 + for inference/training. + + Examples: + >>> obs, r, done, info = env.step(action) + >>> collector.add_action_reward_next_obs(12345, 0, "pol0", { + ... "action": action, "obs": obs, "reward": r, "done": done + ... }) + >>> input_dict = collector.get_inference_input_dict(policy.model) + >>> action = policy.compute_actions_from_input_dict(input_dict) + >>> # repeat + """ + input_dict = {} + for view_col, view_req in view_reqs.items(): + # Create the batch of data from the different buffers. + data_col = view_req.data_col or view_col + if data_col not in self.buffers: + self._build_buffers({data_col: view_req.space.sample()}) + + indices = self.forward_pass_indices + if self.time_major: + input_dict[view_col] = self.buffers[data_col][indices] + else: + if isinstance(view_req.shift, (list, tuple)): + time_indices = \ + np.array(view_req.shift) + np.array(indices[0]) + input_dict[view_col] = self.buffers[data_col][indices[1], + time_indices] + else: + input_dict[view_col] = \ + self.buffers[data_col][indices[1], indices[0]] + + self._reset_inference_call() + + return input_dict + + def get_postprocessing_sample_batches( + self, + episode: MultiAgentEpisode, + view_reqs: Dict[str, ViewRequirement]) -> \ + Dict[AgentID, SampleBatch]: + """Returns a SampleBatch object ready for postprocessing. + + Args: + episode (MultiAgentEpisode): The MultiAgentEpisode object to + get the to-be-postprocessed SampleBatches for. + view_reqs (Dict[str, ViewRequirement]): The view requirements dict + to use for creating the SampleBatch from our buffers. + + Returns: + Dict[AgentID, SampleBatch]: The sample batch objects to be passed + to `Policy.postprocess_trajectory()`. + """ + # Loop through all agents and create a SampleBatch + # (as "view"; no copying). + + # Construct the SampleBatch-dict. + sample_batch_data = {} + + range_ = self.agent_slot_cursor - self.sample_batch_offset + if range_ < 0: + range_ = self.num_agents + range_ + for i in range(range_): + agent_slot = self.sample_batch_offset + i + if agent_slot >= self.num_agents: + agent_slot = agent_slot % self.num_agents + # Do not postprocess the same slot twice. + if self.postprocessed_agents[agent_slot]: + continue + agent_key = self.slot_to_agent_key[agent_slot] + # Skip other episodes (if episode provided). + if episode and agent_key[1] != episode.episode_id: + continue + end = self.agent_key_to_timestep[agent_key] + # Do not build any empty SampleBatches. + if end == self.shift_before: + continue + self.postprocessed_agents[agent_slot] = True + + assert agent_key not in sample_batch_data + sample_batch_data[agent_key] = {} + batch = sample_batch_data[agent_key] + + for view_col, view_req in view_reqs.items(): + # Skip columns that will only get added through postprocessing + # (these may not even exist yet). + if view_req.created_during_postprocessing: + continue + + data_col = view_req.data_col or view_col + shift = view_req.shift + if data_col == SampleBatch.OBS: + shift -= 1 + + batch[view_col] = self.buffers[data_col][ + self.shift_before + shift:end + shift, agent_slot] + + batches = {} + for agent_key, data in sample_batch_data.items(): + batches[agent_key] = SampleBatch(data) + return batches + + def get_train_sample_batch_and_reset(self, view_reqs) -> SampleBatch: + """Returns the accumulated sample batche for this policy. + + This is usually called to collect samples for policy training. + + Returns: + SampleBatch: Returns the accumulated sample batch for this + policy. + """ + seq_lens = [ + self.agent_key_to_timestep[k] - self.shift_before + for k in self.slot_to_agent_key if k is not None + ] + first_zero_len = len(seq_lens) + if seq_lens[-1] == 0: + first_zero_len = seq_lens.index(0) + # Assert that all zeros lie at the end of the seq_lens array. + try: + assert all(seq_lens[i] == 0 + for i in range(first_zero_len, len(seq_lens))) + except AssertionError as e: + print() + raise e + + t_start = self.shift_before + t_end = t_start + self.num_timesteps + + # The agent_slot cursor that points to the newest agent-slot that + # actually already has at least 1 timestep of data (thus it excludes + # just-rolled over chunks (which only have the initial obs in them)). + valid_agent_cursor = \ + (self.agent_slot_cursor - (len(seq_lens) - first_zero_len)) % \ + self.num_agents + + # Construct the view dict. + view = {} + for view_col, view_req in view_reqs.items(): + data_col = view_req.data_col or view_col + assert data_col in self.buffers + # For OBS, indices must be shifted by -1. + extra_shift = 0 if data_col != SampleBatch.OBS else -1 + # If agent_slot has been rolled-over to beginning, we have to copy + # here. + if valid_agent_cursor < self.sample_batch_offset: + time_slice = self.buffers[data_col][t_start + extra_shift: + t_end + extra_shift] + one_ = time_slice[:, self.sample_batch_offset:] + two_ = time_slice[:, :valid_agent_cursor] + if torch and isinstance(time_slice, torch.Tensor): + view[view_col] = torch.cat([one_, two_], dim=1) + else: + view[view_col] = np.concatenate([one_, two_], axis=1) + else: + view[view_col] = \ + self.buffers[data_col][ + t_start + extra_shift:t_end + extra_shift, + self.sample_batch_offset:valid_agent_cursor] + + # Copy all still ongoing trajectories to new agent slots + # (including the ones that just started (are seq_len=0)). + new_chunk_args = [] + for i, seq_len in enumerate(seq_lens): + if seq_len < self.num_timesteps: + agent_slot = self.sample_batch_offset + i + if agent_slot >= self.num_agents: + agent_slot = agent_slot % self.num_agents + if not self.buffers[SampleBatch. + DONES][seq_len - 1 + + self.shift_before][agent_slot]: + agent_key = self.slot_to_agent_key[agent_slot] + new_chunk_args.append( + (agent_slot, agent_key, + self.agent_key_to_timestep[agent_key])) + # Cut out all 0 seq-lens. + seq_lens = seq_lens[:first_zero_len] + batch = SampleBatch( + view, _seq_lens=np.array(seq_lens), _time_major=True) + + # Reset everything for new data. + self.postprocessed_agents = [False] * self.num_agents + self.agent_key_to_slot.clear() + self.agent_key_to_chunk_num.clear() + self.slot_to_agent_key = [None] * self.num_agents + self.agent_key_to_timestep.clear() + self.timesteps_since_last_reset = 0 + self.forward_pass_size = 0 + self.sample_batch_offset = self.agent_slot_cursor + + for args in new_chunk_args: + self._new_chunk_from(*args) + + return batch + + def _build_buffers(self, single_row: Dict[str, TensorType]) -> None: + """Builds the internal data buffers based on a single given row. + + Args: + single_row (Dict[str, TensorType]): A single datarow with one or + more columns (str as key, np.ndarray|tensor as data). + """ + time_size = self.num_timesteps + self.shift_before + self.shift_after + for col, data in single_row.items(): + if col in self.buffers: + continue + base_shape = (time_size, self.num_agents) if self.time_major else \ + (self.num_agents, time_size) + # Python primitive -> np.array. + if isinstance(data, (int, float, bool)): + t_ = type(data) + dtype = np.float32 if t_ == float else \ + np.int32 if type(data) == int else np.bool_ + self.buffers[col] = np.zeros(shape=base_shape, dtype=dtype) + # np.ndarray, torch.Tensor, or tf.Tensor. + else: + shape = base_shape + data.shape + dtype = data.dtype + if torch and isinstance(data, torch.Tensor): + self.buffers[col] = torch.zeros( + *shape, dtype=dtype, device=data.device) + elif tf and isinstance(data, tf.Tensor): + self.buffers[col] = tf.zeros(shape=shape, dtype=dtype) + else: + self.buffers[col] = np.zeros(shape=shape, dtype=dtype) + + def _next_agent_slot(self): + """Starts a new agent slot at the end of the agent-axis. + + Also makes sure, the new slot is not taken yet. + """ + self.agent_slot_cursor += 1 + if self.agent_slot_cursor >= self.num_agents: + self.agent_slot_cursor = 0 + # Just make sure, there is space in our buffer. + assert self.slot_to_agent_key[self.agent_slot_cursor] is None + + def _new_chunk_from(self, agent_slot, agent_key, timestep): + """Creates a new time-window (chunk) given an agent. + + The agent may already have an unfinished episode going on (in a + previous chunk). The end of that previous chunk will be copied to the + beginning of the new one for proper data-shift handling (e.g. + PREV_ACTIONS/REWARDS). + + Args: + agent_slot (int): The agent to start a new chunk for (from an + ongoing episode (chunk)). + agent_key (Tuple[AgentID, EpisodeID, int]): The internal key to + identify an active agent in some episode. + timestep (int): The timestep in the old chunk being continued. + """ + new_agent_slot = self.agent_slot_cursor + # Increase chunk num by 1. + new_agent_key = agent_key[:2] + (agent_key[2] + 1, ) + # Copy relevant timesteps at end of old chunk into new one. + if self.time_major: + for k in self.buffers.keys(): + self.buffers[k][0:self.shift_before, new_agent_slot] = \ + self.buffers[k][ + timestep - self.shift_before:timestep, agent_slot] + else: + for k in self.buffers.keys(): + self.buffers[k][new_agent_slot, 0:self.shift_before] = \ + self.buffers[k][ + agent_slot, timestep - self.shift_before:timestep] + + self.agent_key_to_slot[new_agent_key] = new_agent_slot + self.agent_key_to_chunk_num[new_agent_key[:2]] = new_agent_key[2] + self.slot_to_agent_key[new_agent_slot] = new_agent_key + self._next_agent_slot() + self.agent_key_to_timestep[new_agent_key] = self.shift_before + + def _add_to_next_inference_call(self, agent_key, env_id, agent_slot, + timestep): + """Registers given T and B (agent_slot) for get_inference_input_dict. + + Calling `get_inference_input_dict` will produce an input_dict (for + Policy.compute_actions_from_input_dict) with all registered agent/time + indices and then automatically reset the registry. + + Args: + agent_key (Tuple[AgentID, EpisodeID, int]): The internal key to + identify an active agent in some episode. + env_id (EnvID): The env ID of the given agent. + agent_slot (int): The agent_slot to register (B axis). + timestep (int): The timestep to register (T axis). + """ + idx = self.forward_pass_size + self.forward_pass_index_to_agent_info[idx] = (agent_key[0], + agent_key[1], env_id) + self.agent_key_to_forward_pass_index[agent_key[:2]] = idx + if self.forward_pass_size == 0: + self.forward_pass_indices[0].clear() + self.forward_pass_indices[1].clear() + self.forward_pass_indices[0].append(timestep) + self.forward_pass_indices[1].append(agent_slot) + self.forward_pass_size += 1 + + def _reset_inference_call(self): + """Resets indices for the next inference call. + + After calling this, new calls to `add_init_obs()` and + `add_action_reward_next_obs()` will count for the next input_dict + returned by `get_inference_input_dict()`. + """ + self.forward_pass_size = 0 diff --git a/rllib/evaluation/rollout_worker.py b/rllib/evaluation/rollout_worker.py index 5c2c24bf6..f4bf4ccd6 100644 --- a/rllib/evaluation/rollout_worker.py +++ b/rllib/evaluation/rollout_worker.py @@ -6,7 +6,7 @@ import pickle import platform import os from typing import Callable, Any, List, Dict, Tuple, Union, Optional, \ - TYPE_CHECKING, TypeVar + TYPE_CHECKING, Type, TypeVar import ray from ray.rllib.env.atari_wrappers import wrap_deepmind, is_atari @@ -36,15 +36,15 @@ from ray.rllib.utils.filter import get_filter, Filter from ray.rllib.utils.framework import try_import_tf, try_import_torch from ray.rllib.utils.sgd import do_minibatch_sgd from ray.rllib.utils.tf_run_builder import TFRunBuilder -from ray.rllib.utils.typing import EnvType, AgentID, PolicyID, EnvConfigDict, \ - ModelConfigDict, TrainerConfigDict, SampleBatchType, ModelWeights, \ - ModelGradients, MultiAgentPolicyConfigDict +from ray.rllib.utils.typing import AgentID, EnvConfigDict, EnvType, \ + ModelConfigDict, ModelGradients, ModelWeights, \ + MultiAgentPolicyConfigDict, PartialTrainerConfigDict, PolicyID, \ + SampleBatchType, TrainerConfigDict from ray.util.debug import log_once, disable_log_once_globally, \ enable_periodic_logging from ray.util.iter import ParallelIteratorWorker if TYPE_CHECKING: - from ray.rllib.agents.callbacks import DefaultCallbacks from ray.rllib.evaluation.observation_function import ObservationFunction # Generic type var for foreach_* methods. @@ -129,63 +129,67 @@ class RolloutWorker(ParallelIteratorWorker): resources=resources)(cls) @DeveloperAPI - def __init__(self, - env_creator: Callable[[EnvContext], EnvType], - policy: type, - policy_mapping_fn: Callable[[AgentID], PolicyID] = None, - policies_to_train: List[PolicyID] = None, - tf_session_creator: Callable[[], Any] = None, - rollout_fragment_length: int = 100, - batch_mode: str = "truncate_episodes", - episode_horizon: int = None, - preprocessor_pref: str = "deepmind", - sample_async: bool = False, - compress_observations: bool = False, - num_envs: int = 1, - observation_fn: "ObservationFunction" = None, - observation_filter: str = "NoFilter", - clip_rewards: bool = None, - clip_actions: bool = True, - env_config: EnvConfigDict = None, - model_config: ModelConfigDict = None, - policy_config: TrainerConfigDict = None, - worker_index: int = 0, - num_workers: int = 0, - monitor_path: str = None, - log_dir: str = None, - log_level: str = None, - callbacks: "DefaultCallbacks" = None, - input_creator: Callable[[ - IOContext - ], InputReader] = lambda ioctx: ioctx.default_sampler_input(), - input_evaluation: List[str] = frozenset([]), - output_creator: Callable[ - [IOContext], OutputWriter] = lambda ioctx: NoopOutput(), - remote_worker_envs: bool = False, - remote_env_batch_wait_ms: int = 0, - soft_horizon: bool = False, - no_done_at_end: bool = False, - seed: int = None, - extra_python_environs: dict = None, - fake_sampler: bool = False): + def __init__( + self, + env_creator: Callable[[EnvContext], EnvType], + policy: Union[type, Dict[str, Tuple[Optional[ + type], gym.Space, gym.Space, PartialTrainerConfigDict]]], + policy_mapping_fn: Callable[[AgentID], PolicyID] = None, + policies_to_train: Optional[List[PolicyID]] = None, + tf_session_creator: Optional[Callable[[], "tf1.Session"]] = None, + rollout_fragment_length: int = 100, + batch_mode: str = "truncate_episodes", + episode_horizon: int = None, + preprocessor_pref: str = "deepmind", + sample_async: bool = False, + compress_observations: bool = False, + num_envs: int = 1, + observation_fn: "ObservationFunction" = None, + observation_filter: str = "NoFilter", + clip_rewards: bool = None, + clip_actions: bool = True, + env_config: EnvConfigDict = None, + model_config: ModelConfigDict = None, + policy_config: TrainerConfigDict = None, + worker_index: int = 0, + num_workers: int = 0, + monitor_path: str = None, + log_dir: str = None, + log_level: str = None, + callbacks: Type["DefaultCallbacks"] = None, + input_creator: Callable[[ + IOContext + ], InputReader] = lambda ioctx: ioctx.default_sampler_input(), + input_evaluation: List[str] = frozenset([]), + output_creator: Callable[ + [IOContext], OutputWriter] = lambda ioctx: NoopOutput(), + remote_worker_envs: bool = False, + remote_env_batch_wait_ms: int = 0, + soft_horizon: bool = False, + no_done_at_end: bool = False, + seed: int = None, + extra_python_environs: dict = None, + fake_sampler: bool = False): """Initialize a rollout worker. - Arguments: - env_creator (func): Function that returns a gym.Env given an - EnvContext wrapped configuration. - policy (class|dict): Either a class implementing - Policy, or a dictionary of policy id strings to - (Policy, obs_space, action_space, config) tuples. If a - dict is specified, then we are in multi-agent mode and a - policy_mapping_fn should also be set. - policy_mapping_fn (func): A function that maps agent ids to - policy ids in multi-agent mode. This function will be called - each time a new agent appears in an episode, to bind that agent - to a policy for the duration of the episode. - policies_to_train (list): Optional list of policies to train, - or None for all policies. - tf_session_creator (func): A function that returns a TF session. - This is optional and only useful with TFPolicy. + Args: + env_creator (Callable[[EnvContext], EnvType]): Function that + returns a gym.Env given an EnvContext wrapped configuration. + policy (Union[type, Dict[str, Tuple[Optional[type], gym.Space, + gym.Space, PartialTrainerConfigDict]]]): Either a Policy class + or a dict of policy id strings to + (Policy (None for default), obs_space, action_space, + config)-tuples. If a dict is specified, then we are in + multi-agent mode and a policy_mapping_fn should also be set. + policy_mapping_fn (Callable[[AgentID], PolicyID]): A function that + maps agent ids to policy ids in multi-agent mode. This function + will be called each time a new agent appears in an episode, to + bind that agent to a policy for the duration of the episode. + policies_to_train (Optional[List[PolicyID]]): Optional list of + policies to train, or None for all policies. + tf_session_creator (Optional[Callable[[], tf1.Session]]): A + function that returns a TF session. This is optional and only + useful with TFPolicy. rollout_fragment_length (int): The target number of env transitions to include in each sample batch returned from this worker. batch_mode (str): One of the following batch modes: @@ -221,10 +225,11 @@ class RolloutWorker(ParallelIteratorWorker): only. clip_actions (bool): Whether to clip action values to the range specified by the policy action space. - env_config (dict): Config to pass to the env creator. - model_config (dict): Config to use when creating the policy model. - policy_config (dict): Config to pass to the policy. In the - multi-agent case, this config will be merged with the + env_config (EnvConfigDict): Config to pass to the env creator. + model_config (ModelConfigDict): Config to use when creating the + policy model. + policy_config (TrainerConfigDict): Config to pass to the policy. + In the multi-agent case, this config will be merged with the per-policy configs specified by `policy`. worker_index (int): For remote workers, this should be set to a non-zero and unique value. This index is passed to created envs @@ -236,17 +241,19 @@ class RolloutWorker(ParallelIteratorWorker): log_dir (str): Directory where logs can be placed. log_level (str): Set the root log level on creation. callbacks (DefaultCallbacks): Custom training callbacks. - input_creator (func): Function that returns an InputReader object - for loading previous generated experiences. - input_evaluation (list): How to evaluate the policy performance. - This only makes sense to set when the input is reading offline - data. The possible values include: + input_creator (Callable[[IOContext], InputReader]): Function that + returns an InputReader object for loading previous generated + experiences. + input_evaluation (List[str]): How to evaluate the policy + performance. This only makes sense to set when the input is + reading offline data. The possible values include: - "is": the step-wise importance sampling estimator. - "wis": the weighted step-wise is estimator. - "simulation": run the environment in the background, but use this data for evaluation only and never for learning. - output_creator (func): Function that returns an OutputWriter object - for saving generated experiences. + output_creator (Callable[[IOContext], OutputWriter]): Function that + returns an OutputWriter object for saving generated + experiences. remote_worker_envs (bool): If using num_envs > 1, whether to create those new envs in remote processes instead of in the current process. This adds overheads, but can make sense if your envs diff --git a/rllib/evaluation/sample_batch_builder.py b/rllib/evaluation/sample_batch_builder.py index 15d471240..ec7164883 100644 --- a/rllib/evaluation/sample_batch_builder.py +++ b/rllib/evaluation/sample_batch_builder.py @@ -44,7 +44,7 @@ class SampleBatchBuilder: self.count = 0 @PublicAPI - def add_values(self, **values: Dict[str, Any]) -> None: + def add_values(self, **values: Any) -> None: """Add the given dictionary (row) of values to this batch.""" for k, v in values.items(): @@ -138,7 +138,7 @@ class MultiAgentSampleBatchBuilder: @DeveloperAPI def add_values(self, agent_id: AgentID, policy_id: AgentID, - **values: Dict[str, Any]) -> None: + **values: Any) -> None: """Add the given dictionary (row) of values to this batch. Arguments: diff --git a/rllib/evaluation/sample_collector.py b/rllib/evaluation/sample_collector.py index f532ba25a..babf65e9d 100644 --- a/rllib/evaluation/sample_collector.py +++ b/rllib/evaluation/sample_collector.py @@ -3,7 +3,6 @@ import logging from typing import Dict, Optional from ray.rllib.evaluation.episode import MultiAgentEpisode -from ray.rllib.models.modelv2 import ModelV2 from ray.rllib.utils.typing import AgentID, EpisodeID, PolicyID, \ TensorType @@ -101,15 +100,15 @@ class _SampleCollector(metaclass=ABCMeta): raise NotImplementedError @abstractmethod - def get_inference_input_dict(self, model: ModelV2) -> \ + def get_inference_input_dict(self, policy_id: PolicyID) -> \ Dict[str, TensorType]: - """Returns input_dict for an inference forward pass from our data. + """Returns an input_dict for an (inference) forward pass from our data. - The input_dict can then be used for action computations. + The input_dict can then be used for action computations inside a + Policy via `Policy.compute_actions_from_input_dict()`. Args: - model (ModelV2): The ModelV2 object for which to generate the view - (input_dict) from `data`. + policy_id (PolicyID): The Policy ID to get the input dict for. Returns: Dict[str, TensorType]: The input_dict to be passed into the ModelV2 @@ -155,23 +154,32 @@ class _SampleCollector(metaclass=ABCMeta): """ raise NotImplementedError + @abstractmethod + def check_missing_dones(self, episode_id: EpisodeID) -> None: + """Checks whether given episode is properly terminated with done=True. + + This applies to all agents in the episode. + + Args: + episode_id (EpisodeID): The episode ID to check for proper + termination. + + Raises: + ValueError: If `episode` has no done=True at the end. + """ + raise NotImplementedError + @abstractmethod def get_multi_agent_batch_and_reset(self): """Returns the accumulated sample batches for each policy. Any unprocessed rows will be first postprocessed with a policy - postprocessor. The internal state of this builder will be reset. - - Args: - episode (Optional[MultiAgentEpisode]): The Episode object that - holds this MultiAgentBatchBuilder object or None. + postprocessor. The internal state of this builder will be reset to + start the next batch. + This is usually called to collect samples for policy training. Returns: MultiAgentBatch: Returns the accumulated sample batches for each policy inside one MultiAgentBatch object. """ raise NotImplementedError - - @abstractmethod - def check_missing_dones(self, episode_id: EpisodeID) -> None: - raise NotImplementedError diff --git a/rllib/evaluation/tests/test_trajectories.py b/rllib/evaluation/tests/test_trajectories.py deleted file mode 100644 index c834d3db0..000000000 --- a/rllib/evaluation/tests/test_trajectories.py +++ /dev/null @@ -1,70 +0,0 @@ -from gym.spaces import Box, Discrete -import numpy as np -import unittest - -from ray.rllib.evaluation.trajectory import Trajectory - - -class TestTrajectories(unittest.TestCase): - """Tests Trajectory classes.""" - - def test_trajectory(self): - """Tests the Trajectory class.""" - - buffer_size = 5 - - # Small trajecory object for testing purposes. - trajectory = Trajectory(buffer_size=buffer_size) - self.assertEqual(trajectory.cursor, 0) - self.assertEqual(trajectory.timestep, 0) - self.assertEqual(trajectory.sample_batch_offset, 0) - assert not trajectory.buffers - observation_space = Box(-1.0, 1.0, shape=(3, )) - action_space = Discrete(2) - trajectory.add_init_obs( - env_id=0, - agent_id="agent", - policy_id="policy", - init_obs=observation_space.sample()) - self.assertEqual(trajectory.cursor, 0) - self.assertEqual(trajectory.initial_obs.shape, observation_space.shape) - - # Fill up the buffer and make it extend if it hits the limit. - cur_buffer_size = buffer_size - for i in range(buffer_size + 1): - trajectory.add_action_reward_next_obs( - env_id=0, - agent_id="agent", - policy_id="policy", - values=dict( - t=i, - actions=action_space.sample(), - rewards=1.0, - dones=i == buffer_size, - new_obs=observation_space.sample(), - action_logp=-0.5, - action_dist_inputs=np.array([[0.5, 0.5]]), - )) - self.assertEqual(trajectory.cursor, i + 1) - self.assertEqual(trajectory.timestep, i + 1) - self.assertEqual(trajectory.sample_batch_offset, 0) - if i == buffer_size - 1: - cur_buffer_size *= 2 - self.assertEqual( - len(trajectory.buffers["new_obs"]), cur_buffer_size) - self.assertEqual( - len(trajectory.buffers["rewards"]), cur_buffer_size) - - # Create a SampleBatch from the Trajectory and reset it. - batch = trajectory.get_sample_batch_and_reset() - self.assertEqual(batch.count, buffer_size + 1) - # Make sure, Trajectory was reset properly. - self.assertEqual(trajectory.cursor, buffer_size + 1) - self.assertEqual(trajectory.timestep, 0) - self.assertEqual(trajectory.sample_batch_offset, buffer_size + 1) - - -if __name__ == "__main__": - import pytest - import sys - sys.exit(pytest.main(["-v", __file__])) diff --git a/rllib/policy/view_requirement.py b/rllib/policy/view_requirement.py index 9ecaf7b70..687d76024 100644 --- a/rllib/policy/view_requirement.py +++ b/rllib/policy/view_requirement.py @@ -1,5 +1,5 @@ import gym -from typing import Optional +from typing import List, Optional, Union from ray.rllib.utils.framework import try_import_torch @@ -21,7 +21,7 @@ class ViewRequirement: Examples: >>> # The default ViewRequirement for a Model is: - >>> req = [ModelV2].inference_view_requirements() + >>> req = [ModelV2].inference_view_requirements >>> print(req) {"obs": ViewRequirement(shift=0)} """ @@ -29,7 +29,8 @@ class ViewRequirement: def __init__(self, data_col: Optional[str] = None, space: gym.Space = None, - shift: int = 0): + shift: Union[int, List[int]] = 0, + created_during_postprocessing: bool = False): """Initializes a ViewRequirement object. Args: @@ -39,15 +40,18 @@ class ViewRequirement: space (gym.Space): The gym Space used in case we need to pad data in inaccessible areas of the trajectory (t<0 or t>H). Default: Simple box space, e.g. rewards. - shift (Union[List[int], int]): Single shift value of list of + shift (Union[int, List[int]]): Single shift value of list of shift values to use relative to the underlying `data_col`. Example: For a view column "prev_actions", you can set `data_col="actions"` and `shift=-1`. Example: For a view column "obs" in an Atari framestacking fashion, you can set `data_col="obs"` and `shift=[-3, -2, -1, 0]`. + created_during_postprocessing (bool): Whether this column only gets + created during postprocessing. """ self.data_col = data_col self.space = space or gym.spaces.Box( float("-inf"), float("inf"), shape=()) self.shift = shift + self.created_during_postprocessing = created_during_postprocessing diff --git a/rllib/rollout.py b/rllib/rollout.py index c036ec35c..3dae24290 100755 --- a/rllib/rollout.py +++ b/rllib/rollout.py @@ -3,13 +3,13 @@ import argparse import collections import copy +import gym import json import os from pathlib import Path import pickle import shelve -import gym import ray from ray.rllib.env import MultiAgentEnv from ray.rllib.env.base_env import _DUMMY_AGENT_ID