import copy import hashlib from collections import defaultdict, deque from typing import Any, Dict, List, Optional, Set, Tuple, Union import numpy as np import scipy from gymnasium.core import ActType, ObsType from ray.rllib.core.columns import Columns from ray.rllib.env.multi_agent_episode import MultiAgentEpisode from ray.rllib.env.single_agent_episode import SingleAgentEpisode from ray.rllib.utils import force_list from ray.rllib.utils.annotations import ( DeveloperAPI, override, ) from ray.rllib.utils.replay_buffers.episode_replay_buffer import EpisodeReplayBuffer from ray.rllib.utils.spaces.space_utils import batch from ray.rllib.utils.typing import AgentID, ModuleID, SampleBatchType @DeveloperAPI class MultiAgentEpisodeReplayBuffer(EpisodeReplayBuffer): """Multi-agent episode replay buffer that stores episodes by their IDs. This class implements a replay buffer as used in "playing Atari with Deep Reinforcement Learning" (Mnih et al., 2013) for multi-agent reinforcement learning, Each "row" (a slot in a deque) in the buffer is occupied by one episode. If an incomplete episode is added to the buffer and then another chunk of that episode is added at a later time, the buffer will automatically concatenate the new fragment to the original episode. This way, episodes can be completed via subsequent `add` calls. Sampling returns a size `B` episode list (number of 'rows'), where each episode holds a tuple tuple of the form `(o_t, a_t, sum(r_t+1:t+n), o_t+n)` where `o_t` is the observation in `t`, `a_t` the action chosen at observation `o_t`, `o_t+n` is the observation `n` timesteps later and `sum(r_t+1:t+n)` is the sum of all rewards collected over the time steps between `t+1` and `t+n`. The `n`-step can be chosen freely when sampling and defaults to `1`. If `n_step` is a tuple it is sampled uniformly across the interval defined by the tuple (for each row in the batch). Each episode contains - in addition to the data tuples presented above - two further elements in its `extra_model_outputs`, namely `n_steps` and `weights`. The former holds the `n_step` used for the sampled timesteps in the episode and the latter the corresponding (importance sampling) weight for the transition. .. testcode:: import gymnasium as gym from ray.rllib.env.multi_agent_episode import MultiAgentEpisode from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole from ray.rllib.utils.replay_buffers import MultiAgentEpisodeReplayBuffer # Create the environment. env = MultiAgentCartPole({"num_agents": 2}) # Set up the loop variables agent_ids = env.agents agent_ids.append("__all__") terminateds = {aid: False for aid in agent_ids} truncateds = {aid: False for aid in agent_ids} num_timesteps = 10000 episodes = [] # Initialize the first episode entries. eps = MultiAgentEpisode() obs, infos = env.reset() eps.add_env_reset(observations=obs, infos=infos) # Sample 10,000 env timesteps. for i in range(num_timesteps): # If terminated we create a new episode. if eps.is_done: episodes.append(eps.to_numpy()) eps = MultiAgentEpisode() terminateds = {aid: False for aid in agent_ids} truncateds = {aid: False for aid in agent_ids} obs, infos = env.reset() eps.add_env_reset(observations=obs, infos=infos) # Sample a random action for all agents that should step in the episode # next. actions = { aid: env.get_action_space(aid).sample() for aid in eps.get_agents_to_act() } obs, rewards, terminateds, truncateds, infos = env.step(actions) eps.add_env_step( obs, actions, rewards, infos, terminateds=terminateds, truncateds=truncateds ) # Add the last (truncated) episode to the list of episodes. if not eps.is_done: episodes.append(eps) # Create the buffer. buffer = MultiAgentEpisodeReplayBuffer() # Add the list of episodes sampled. buffer.add(episodes) # Pull a sample from the buffer using an `n-step` of 3. sample = buffer.sample(num_items=256, gamma=0.95, n_step=3) """ def __init__( self, capacity: int = 10000, *, batch_size_B: int = 16, batch_length_T: int = 1, metrics_num_episodes_for_smoothing: int = 100, **kwargs, ): """Initializes a multi-agent episode replay buffer. Args: capacity: The total number of timesteps to be storable in this buffer. Will start ejecting old episodes once this limit is reached. batch_size_B: The number of episodes returned from `sample()`. batch_length_T: The length of each episode in the episode list returned from `sample()`. """ # Initialize the base episode replay buffer. super().__init__( capacity=capacity, batch_size_B=batch_size_B, batch_length_T=batch_length_T, metrics_num_episodes_for_smoothing=metrics_num_episodes_for_smoothing, **kwargs, ) # Stores indices of module (single-agent) timesteps. Each index is a tuple # of the form: # `(ma_episode_idx, agent_id, timestep)`. # This information is stored for each timestep of an episode and is used in # the `"independent"`` sampling process. The multi-agent episode index amd the # agent ID are used to retrieve the single-agent episode. The timestep is then # needed to retrieve the corresponding timestep data from that single-agent # episode. self._module_to_indices: Dict[ ModuleID, List[Tuple[int, AgentID, int]] ] = defaultdict(list) # Stores the number of single-agent timesteps in the buffer. self._num_agent_timesteps: int = 0 # Stores the number of single-agent timesteps per module. self._num_module_timesteps: Dict[ModuleID, int] = defaultdict(int) # Stores the number of added single-agent timesteps over the # lifetime of the buffer. self._num_agent_timesteps_added: int = 0 # Stores the number of added single-agent timesteps per module # over the lifetime of the buffer. self._num_module_timesteps_added: Dict[ModuleID, int] = defaultdict(int) self._num_module_episodes: Dict[ModuleID, int] = defaultdict(int) # Stores the number of module episodes evicted. Note, this is # important for indexing. self._num_module_episodes_evicted: Dict[ModuleID, int] = defaultdict(int) # Stores hte number of module timesteps sampled. self.sampled_timesteps_per_module: Dict[ModuleID, int] = defaultdict(int) @override(EpisodeReplayBuffer) def add( self, episodes: Union[List["MultiAgentEpisode"], "MultiAgentEpisode"], ) -> None: """Adds episodes to the replay buffer. Note, if the incoming episodes' time steps cause the buffer to overflow, older episodes are evicted. Because episodes usually come in chunks and not complete, this could lead to edge cases (e.g. with very small capacity or very long episode length) where the first part of an episode is evicted while the next part just comes in. To defend against such case, the complete episode is evicted, including the new chunk, unless the episode is the only one in the buffer. In the latter case the buffer will be allowed to overflow in a temporary fashion, i.e. during the next addition of samples to the buffer an attempt is made to fall below capacity again. The user is advised to select a large enough buffer with regard to the maximum expected episode length. Args: episodes: The multi-agent episodes to add to the replay buffer. Can be a single episode or a list of episodes. """ episodes: List["MultiAgentEpisode"] = force_list(episodes) new_episode_ids: Set[str] = {eps.id_ for eps in episodes} total_env_timesteps = sum([eps.env_steps() for eps in episodes]) self._num_timesteps += total_env_timesteps self._num_timesteps_added += total_env_timesteps # Set up some counters for metrics. num_env_steps_added = 0 agent_to_num_steps_added = defaultdict(int) module_to_num_steps_added = defaultdict(int) num_episodes_added = 0 agent_to_num_episodes_added = defaultdict(int) module_to_num_episodes_added = defaultdict(int) num_episodes_evicted = 0 agent_to_num_episodes_evicted = defaultdict(int) module_to_num_episodes_evicted = defaultdict(int) num_env_steps_evicted = 0 agent_to_num_steps_evicted = defaultdict(int) module_to_num_steps_evicted = defaultdict(int) # Evict old episodes. eps_evicted_ids: Set[Union[str, int]] = set() eps_evicted_idxs: Set[int] = set() while ( self._num_timesteps > self.capacity and self._num_remaining_episodes(new_episode_ids, eps_evicted_ids) != 1 ): # Evict episode. evicted_episode = self.episodes.popleft() eps_evicted_ids.add(evicted_episode.id_) eps_evicted_idxs.add(self.episode_id_to_index.pop(evicted_episode.id_)) # If this episode has a new chunk in the new episodes added, # we subtract it again. # TODO (sven, simon): Should we just treat such an episode chunk # as a new episode? if evicted_episode.id_ in new_episode_ids: idx = next( i for i, eps in enumerate(episodes) if eps.id_ == evicted_episode.id_ ) new_eps_to_evict = episodes.pop(idx) self._num_timesteps -= new_eps_to_evict.env_steps() self._num_timesteps_added -= new_eps_to_evict.env_steps() # Remove the timesteps of the evicted episode from the counter. self._num_timesteps -= evicted_episode.env_steps() self._num_agent_timesteps -= evicted_episode.agent_steps() self._num_episodes_evicted += 1 # Increase the counters. num_episodes_evicted += 1 num_env_steps_evicted += evicted_episode.env_steps() for aid, a_eps in evicted_episode.agent_episodes.items(): mid = evicted_episode._agent_to_module_mapping[aid] agent_to_num_episodes_evicted[aid] += 1 module_to_num_episodes_evicted[mid] += 1 agent_to_num_steps_evicted[aid] += a_eps.agent_steps() module_to_num_steps_evicted[mid] += a_eps.agent_steps() # Remove the module timesteps of the evicted episode from the counters. self._evict_module_episodes(evicted_episode) del evicted_episode # Add agent and module steps. for eps in episodes: self._num_agent_timesteps += eps.agent_steps() self._num_agent_timesteps_added += eps.agent_steps() # Update the module counters by the module timesteps. self._update_module_counters(eps) # Remove corresponding indices, if episodes were evicted. if eps_evicted_idxs: # If the episode is not evicted, we keep the index. # Note, each index 2-tuple is of the form (ma_episode_idx, timestep) # and refers to a certain environment timestep in a certain # multi-agent episode. self._indices = [ idx_tuple for idx_tuple in self._indices if idx_tuple[0] not in eps_evicted_idxs ] # Also remove corresponding module indices. for module_id, module_indices in self._module_to_indices.items(): # Each index 3-tuple is of the form # (ma_episode_idx, agent_id, timestep) and refers to a certain # agent timestep in a certain multi-agent episode. self._module_to_indices[module_id] = [ idx_triplet for idx_triplet in module_indices if idx_triplet[0] not in eps_evicted_idxs ] for eps in episodes: eps = copy.deepcopy(eps) # If the episode is part of an already existing episode, concatenate. if eps.id_ in self.episode_id_to_index: eps_idx = self.episode_id_to_index[eps.id_] existing_eps = self.episodes[eps_idx - self._num_episodes_evicted] existing_len = len(existing_eps) self._indices.extend( [ ( eps_idx, existing_len + i, ) for i in range(len(eps)) ] ) # Add new module indices. self._add_new_module_indices(eps, eps_idx, True) # Concatenate the episode chunk. existing_eps.concat_episode(eps) # Otherwise, create a new entry. else: # New episode. self.episodes.append(eps) # Update the counters num_episodes_added += 1 for aid, a_eps in eps.agent_episodes.items(): mid = eps._agent_to_module_mapping[aid] agent_to_num_episodes_added[aid] += 1 module_to_num_episodes_added[mid] += 1 eps_idx = len(self.episodes) - 1 + self._num_episodes_evicted self.episode_id_to_index[eps.id_] = eps_idx self._indices.extend([(eps_idx, i) for i in range(len(eps))]) # Add new module indices. self._add_new_module_indices(eps, eps_idx, False) # Update the step counters. num_env_steps_added += eps.env_steps() for aid, e_eps in eps.agent_episodes.items(): mid = eps._agent_to_module_mapping[aid] agent_to_num_steps_added[aid] += e_eps.agent_steps() module_to_num_steps_added[mid] += e_eps.agent_steps() # Update the adding metrics. self._update_add_metrics( num_episodes_added=num_episodes_added, num_env_steps_added=num_env_steps_added, num_episodes_evicted=num_episodes_evicted, num_env_steps_evicted=num_env_steps_evicted, agent_to_num_episodes_added=agent_to_num_episodes_added, agent_to_num_steps_added=agent_to_num_steps_added, agent_to_num_episodes_evicted=agent_to_num_episodes_evicted, agent_to_num_steps_evicted=agent_to_num_steps_evicted, module_to_num_episodes_added=module_to_num_steps_added, module_to_num_steps_added=module_to_num_episodes_added, module_to_num_episodes_evicted=module_to_num_episodes_evicted, module_to_num_steps_evicted=module_to_num_steps_evicted, ) @override(EpisodeReplayBuffer) def sample( self, num_items: Optional[int] = None, *, batch_size_B: Optional[int] = None, batch_length_T: Optional[int] = None, n_step: Optional[Union[int, Tuple]] = 1, gamma: float = 0.99, include_infos: bool = False, include_extra_model_outputs: bool = False, replay_mode: str = "independent", modules_to_sample: Optional[List[ModuleID]] = None, **kwargs, ) -> Union[List["MultiAgentEpisode"], List["SingleAgentEpisode"]]: """Samples a batch of multi-agent transitions. Multi-agent transitions can be sampled either `"independent"` or `"synchronized"` with the former sampling for each module independent agent steps and the latter sampling agent transitions from the same environment step. The n-step parameter can be either a single integer or a tuple of two integers. In the former case, the n-step is fixed to the given integer and in the latter case, the n-step is sampled uniformly from the given range. Large n-steps could potentially lead to a many retries because not all samples might have a full n-step transition. Sampling returns batches of size B (number of 'rows'), where each row is a tuple of the form `(o_t, a_t, sum(r_t+1:t+n), o_t+n)` where `o_t` is the observation in `t`, `a_t` the action chosen at observation `o_t`, `o_t+n` is the observation `n` timesteps later and `sum(r_t+1:t+n)` is the sum of all rewards collected over the time steps between `t+1` and `t+n`. The n`-step can be chosen freely when sampling and defaults to `1`. If `n_step` is a tuple it is sampled uniformly across the interval defined by the tuple (for each row in the batch). Each batch contains - in addition to the data tuples presented above - two further columns, namely `n_steps` and `weigths`. The former holds the `n_step` used for each row in the batch and the latter a (default) weight of `1.0` for each row in the batch. This weight is used for weighted loss calculations in the training process. Args: num_items: The number of items to sample. If provided, `batch_size_B` should be `None`. batch_size_B: The batch size to sample. If provided, `num_items` should be `None`. batch_length_T: The length of the sampled batch. If not provided, the default batch length is used. This feature is not yet implemented. n_step: The n-step to sample. If the n-step is a tuple, the n-step is sampled uniformly from the given range. If not provided, the default n-step of `1` is used. gamma: The discount factor for the n-step reward calculation. include_infos: Whether to include the infos in the sampled batch. include_extra_model_outputs: Whether to include the extra model outputs in the sampled batch. replay_mode: The replay mode to use for sampling. Either `"independent"` or `"synchronized"`. modules_to_sample: A list of module IDs to sample from. If not provided, transitions for aall modules are sampled. Returns: A dictionary of the form `ModuleID -> SampleBatchType` containing the sampled data for each module or each module in `modules_to_sample`, if provided. """ if num_items is not None: assert batch_size_B is None, ( "Cannot call `sample()` with both `num_items` and `batch_size_B` " "provided! Use either one." ) batch_size_B = num_items # Use our default values if no sizes/lengths provided. batch_size_B = batch_size_B or self.batch_size_B # TODO (simon): Implement trajectory sampling for RNNs. batch_length_T = batch_length_T or self.batch_length_T # Sample for each module independently. if replay_mode == "independent": return self._sample_independent( batch_size_B=batch_size_B, batch_length_T=batch_length_T, n_step=n_step, gamma=gamma, include_infos=include_infos, include_extra_model_outputs=include_extra_model_outputs, modules_to_sample=modules_to_sample, ) else: return self._sample_synchonized( batch_size_B=batch_size_B, batch_length_T=batch_length_T, n_step=n_step, gamma=gamma, include_infos=include_infos, include_extra_model_outputs=include_extra_model_outputs, modules_to_sample=modules_to_sample, ) def get_added_agent_timesteps(self) -> int: """Returns number of agent timesteps that have been added in buffer's lifetime. Note, this could be more than the `get_added_timesteps` returns as an environment timestep could contain multiple agent timesteps (for eaxch agent one). """ return self._num_agent_timesteps_added def get_module_ids(self) -> List[ModuleID]: """Returns a list of module IDs stored in the buffer.""" return list(self._module_to_indices.keys()) def get_num_agent_timesteps(self) -> int: """Returns number of agent timesteps stored in the buffer. Note, this could be more than the `num_timesteps` as an environment timestep could contain multiple agent timesteps (for eaxch agent one). """ return self._num_agent_timesteps @override(EpisodeReplayBuffer) def get_num_episodes(self, module_id: Optional[ModuleID] = None) -> int: """Returns number of episodes stored for a module in the buffer. Note, episodes could be either complete or truncated. Args: module_id: The ID of the module to query. If not provided, the number of episodes for all modules is returned. Returns: The number of episodes stored for the module or all modules. """ return ( self._num_module_episodes[module_id] if module_id else super().get_num_episodes() ) @override(EpisodeReplayBuffer) def get_num_episodes_evicted(self, module_id: Optional[ModuleID] = None) -> int: """Returns number of episodes evicted for a module in the buffer.""" return ( self._num_module_episodes_evicted[module_id] if module_id else super().get_num_episodes_evicted() ) @override(EpisodeReplayBuffer) def get_num_timesteps(self, module_id: Optional[ModuleID] = None) -> int: """Returns number of individual timesteps for a module stored in the buffer. Args: module_id: The ID of the module to query. If not provided, the number of timesteps for all modules are returned. Returns: The number of timesteps stored for the module or all modules. """ return ( self._num_module_timesteps[module_id] if module_id else super().get_num_timesteps() ) @override(EpisodeReplayBuffer) def get_sampled_timesteps(self, module_id: Optional[ModuleID] = None) -> int: """Returns number of timesteps that have been sampled for a module. Args: module_id: The ID of the module to query. If not provided, the number of sampled timesteps for all modules are returned. Returns: The number of timesteps sampled for the module or all modules. """ return ( self.sampled_timesteps_per_module[module_id] if module_id else super().get_sampled_timesteps() ) @override(EpisodeReplayBuffer) def get_added_timesteps(self, module_id: Optional[ModuleID] = None) -> int: """Returns the number of timesteps added in buffer's lifetime for given module. Args: module_id: The ID of the module to query. If not provided, the total number of timesteps ever added. Returns: The number of timesteps added for `module_id` (or all modules if `module_id` is None). """ return ( self._num_module_timesteps_added[module_id] if module_id else super().get_added_timesteps() ) @override(EpisodeReplayBuffer) def get_state(self) -> Dict[str, Any]: """Gets a pickable state of the buffer. This is used for checkpointing the buffer's state. It is specifically helpful, for example, when a trial is paused and resumed later on. The buffer's state can be saved to disk and reloaded when the trial is resumed. Returns: A dict containing all necessary information to restore the buffer's state. """ return super().get_state() | { "_module_to_indices": list(self._module_to_indices.items()), "_num_agent_timesteps": self._num_agent_timesteps, "_num_agent_timesteps_added": self._num_agent_timesteps_added, "_num_module_timesteps": list(self._num_module_timesteps.items()), "_num_module_timesteps_added": list( self._num_module_timesteps_added.items() ), "_num_module_episodes": list(self._num_module_episodes.items()), "_num_module_episodes_evicted": list( self._num_module_episodes_evicted.items() ), "sampled_timesteps_per_module": list( self.sampled_timesteps_per_module.items() ), } @override(EpisodeReplayBuffer) def set_state(self, state) -> None: """Sets the state of a buffer from a previously stored state. See `get_state()` for more information on what is stored in the state. This method is used to restore the buffer's state from a previously stored state. It is specifically helpful, for example, when a trial is paused and resumed later on. The buffer's state can be saved to disk and reloaded when the trial is resumed. Args: state: The state to restore the buffer from. """ # Set the episodes. self._set_episodes(state) # Set the super's state. super().set_state(state) # Now set the remaining attributes. self._module_to_indices = defaultdict(list, dict(state["_module_to_indices"])) self._num_agent_timesteps = state["_num_agent_timesteps"] self._num_agent_timesteps_added = state["_num_agent_timesteps_added"] self._num_module_timesteps = defaultdict( int, dict(state["_num_module_timesteps"]) ) self._num_module_timesteps_added = defaultdict( int, dict(state["_num_module_timesteps_added"]) ) self._num_module_episodes = defaultdict( int, dict(state["_num_module_episodes"]) ) self._num_module_episodes_evicted = defaultdict( int, dict(state["_num_module_episodes_evicted"]) ) self.sampled_timesteps_per_module = defaultdict( list, dict(state["sampled_timesteps_per_module"]) ) def _set_episodes(self, state: Dict[str, Any]) -> None: """Sets the episodes from the state.""" if not self.episodes: self.episodes = deque( [ MultiAgentEpisode.from_state(eps_data) for eps_data in state["episodes"] ] ) def _sample_independent( self, batch_size_B: Optional[int], batch_length_T: Optional[int], n_step: Optional[Union[int, Tuple[int, int]]], gamma: float, include_infos: bool, include_extra_model_outputs: bool, modules_to_sample: Optional[Set[ModuleID]], ) -> List["SingleAgentEpisode"]: """Samples a batch of independent multi-agent transitions.""" actual_n_step = n_step or 1 # Sample the n-step if necessary. random_n_step = isinstance(n_step, (tuple, list)) sampled_episodes = [] # Record the number of samples per module/agent/total. num_env_steps_sampled = 0 agent_to_num_steps_sampled = defaultdict(int) module_to_num_steps_sampled = defaultdict(int) # Record all the env step buffer indices that are contained in the sample. sampled_env_step_idxs = set() agent_to_sampled_env_step_idxs = defaultdict(set) module_to_sampled_env_step_idxs = defaultdict(set) # Record all the episode buffer indices that are contained in the sample. sampled_episode_idxs = set() agent_to_sampled_episode_idxs = defaultdict(set) module_to_sampled_episode_idxs = defaultdict(set) # Record all n-steps that have been used. sampled_n_steps = [] agent_to_sampled_n_steps = defaultdict(list) module_to_sampled_n_steps = defaultdict(list) # Record the number of times a sample needs to be resampled. num_resamples = 0 agent_to_num_resamples = defaultdict(int) module_to_num_resamples = defaultdict(int) # TODO (simon): Ensure that the module has data and if not, skip it. # TODO (sven): Should we then error out or skip? I think the Learner # should handle this case when a module has no train data. modules_to_sample = modules_to_sample or set(self._module_to_indices.keys()) for module_id in modules_to_sample: module_indices = self._module_to_indices[module_id] B = 0 while B < batch_size_B: # Now sample from the single-agent timesteps. index_tuple = module_indices[self.rng.integers(len(module_indices))] # This will be an agent timestep (not env timestep). # TODO (simon, sven): Maybe deprecate sa_episode_idx (_) in the index # quads. Is there any need for it? ma_episode_idx, agent_id, sa_episode_ts = ( index_tuple[0] - self._num_episodes_evicted, index_tuple[1], index_tuple[2], ) # Get the multi-agent episode. ma_episode = self.episodes[ma_episode_idx] # Retrieve the single-agent episode for filtering. sa_episode = ma_episode.agent_episodes[agent_id] # If we use random n-step sampling, draw the n-step for this item. if random_n_step: actual_n_step = int(self.rng.integers(n_step[0], n_step[1])) # If we cannnot make the n-step, we resample. if sa_episode_ts + actual_n_step > len(sa_episode): num_resamples += 1 agent_to_num_resamples[agent_id] += 1 module_to_num_resamples[module_id] += 1 continue # Note, this will be the reward after executing action # `a_(episode_ts)`. For `n_step>1` this will be the discounted sum # of all rewards that were collected over the last n steps. sa_raw_rewards = sa_episode.get_rewards( slice(sa_episode_ts, sa_episode_ts + actual_n_step) ) sa_rewards = scipy.signal.lfilter( [1], [1, -gamma], sa_raw_rewards[::-1], axis=0 )[-1] sampled_sa_episode = SingleAgentEpisode( id_=sa_episode.id_, # Provide the IDs for the learner connector. agent_id=sa_episode.agent_id, module_id=sa_episode.module_id, multi_agent_episode_id=ma_episode.id_, # Ensure that each episode contains a tuple of the form: # (o_t, a_t, sum(r_(t:t+n_step)), o_(t+n_step)) # Two observations (t and t+n). observations=[ sa_episode.get_observations(sa_episode_ts), sa_episode.get_observations(sa_episode_ts + actual_n_step), ], observation_space=sa_episode.observation_space, infos=( [ sa_episode.get_infos(sa_episode_ts), sa_episode.get_infos(sa_episode_ts + actual_n_step), ] if include_infos else None ), actions=[sa_episode.get_actions(sa_episode_ts)], action_space=sa_episode.action_space, rewards=[sa_rewards], # If the sampled single-agent episode is the single-agent episode's # last time step, check, if the single-agent episode is terminated # or truncated. terminated=( sa_episode_ts + actual_n_step >= len(sa_episode) and sa_episode.is_terminated ), truncated=( sa_episode_ts + actual_n_step >= len(sa_episode) and sa_episode.is_truncated ), extra_model_outputs={ "weights": [1.0], "n_step": [actual_n_step], **( { k: [ sa_episode.get_extra_model_outputs(k, sa_episode_ts) ] for k in sa_episode.extra_model_outputs.keys() } if include_extra_model_outputs else {} ), }, # TODO (sven): Support lookback buffers. len_lookback_buffer=0, t_started=sa_episode_ts, ) # Append single-agent episode to the list of sampled episodes. sampled_episodes.append(sampled_sa_episode) # Add the episode indices. sampled_episode_idxs.add(ma_episode_idx) agent_to_sampled_episode_idxs[sa_episode.agent_id].add(sa_episode.id_) module_to_sampled_episode_idxs[module_id].add(sa_episode.id_) # Add the unique step hashes. # Get the corresponding index in the `env_to_agent_t` mapping. # TODO (simon, sven): This has complexity O(n) and could become # expensive when the episode is large. Note, however, that conversion # from list to `numpy.ndarray` is also complexity O(n) and we do this # at many places - also in the `MultiAgentEpisode`s. ma_episode_ts = ma_episode.env_t_to_agent_t[agent_id].data.index( sa_episode_ts ) sampled_env_step_idxs.add( hashlib.sha256( f"{ma_episode.id_}-{ma_episode_ts}".encode() ).hexdigest() ) hashed_agent_step = hashlib.sha256( f"{sa_episode.id_}-{sa_episode_ts}".encode() ).hexdigest() agent_to_sampled_env_step_idxs[agent_id].add(hashed_agent_step) module_to_sampled_env_step_idxs[module_id].add(hashed_agent_step) # Add the actual n-step used in generating this sample. sampled_n_steps.append(actual_n_step) agent_to_sampled_n_steps[agent_id].append(actual_n_step) module_to_sampled_n_steps[module_id].append(actual_n_step) # Increase counter. B += 1 # Increase the per module timesteps counter. self.sampled_timesteps_per_module[module_id] += B # Increase the counter metrics. num_env_steps_sampled += B agent_to_num_steps_sampled[agent_id] += B module_to_num_steps_sampled[module_id] += B # Increase the counter for environment timesteps. self.sampled_timesteps += batch_size_B # Update the sample metrics. num_episodes_per_sample = len(sampled_episode_idxs) num_env_steps_per_sample = len(sampled_env_step_idxs) sampled_n_step = sum(sampled_n_steps) / batch_size_B agent_to_num_episodes_per_sample = { aid: len(l) for aid, l in agent_to_sampled_episode_idxs.items() } module_to_num_episodes_per_sample = { mid: len(l) for mid, l in module_to_sampled_episode_idxs.items() } agent_to_num_steps_per_sample = { aid: len(l) for aid, l in agent_to_sampled_env_step_idxs.items() } module_to_num_steps_per_sample = { mid: len(l) for mid, l in module_to_sampled_env_step_idxs.items() } agent_to_sampled_n_step = { aid: sum(l) / len(l) for aid, l in agent_to_sampled_n_steps.items() } module_to_sampled_n_step = { mid: sum(l) / len(l) for mid, l in module_to_sampled_n_steps.items() } self._update_sample_metrics( num_env_steps_sampled=num_env_steps_sampled, num_episodes_per_sample=num_episodes_per_sample, num_env_steps_per_sample=num_env_steps_per_sample, sampled_n_step=sampled_n_step, num_resamples=num_resamples, agent_to_num_steps_sampled=agent_to_num_steps_sampled, agent_to_num_episodes_per_sample=agent_to_num_episodes_per_sample, agent_to_num_steps_per_sample=agent_to_num_steps_per_sample, agent_to_sampled_n_step=agent_to_sampled_n_step, agent_to_num_resamples=agent_to_num_resamples, module_to_num_steps_sampled=module_to_num_steps_sampled, module_to_num_episodes_per_sample=module_to_num_episodes_per_sample, module_to_num_steps_per_sample=module_to_num_steps_per_sample, module_to_sampled_n_step=module_to_sampled_n_step, module_to_num_resamples=module_to_num_resamples, ) # Return multi-agent dictionary. return sampled_episodes def _sample_synchonized( self, batch_size_B: Optional[int], batch_length_T: Optional[int], n_step: Optional[Union[int, Tuple]], gamma: float, include_infos: bool, include_extra_model_outputs: bool, modules_to_sample: Optional[List[ModuleID]], ) -> SampleBatchType: """Samples a batch of synchronized multi-agent transitions.""" # Sample the n-step if necessary. if isinstance(n_step, tuple): # Use random n-step sampling. random_n_step = True else: actual_n_step = n_step or 1 random_n_step = False # Containers for the sampled data. observations: Dict[ModuleID, List[ObsType]] = defaultdict(list) next_observations: Dict[ModuleID, List[ObsType]] = defaultdict(list) actions: Dict[ModuleID, List[ActType]] = defaultdict(list) rewards: Dict[ModuleID, List[float]] = defaultdict(list) is_terminated: Dict[ModuleID, List[bool]] = defaultdict(list) is_truncated: Dict[ModuleID, List[bool]] = defaultdict(list) weights: Dict[ModuleID, List[float]] = defaultdict(list) n_steps: Dict[ModuleID, List[int]] = defaultdict(list) # If `info` should be included, construct also a container for them. if include_infos: infos: Dict[ModuleID, List[Dict[str, Any]]] = defaultdict(list) # If `extra_model_outputs` should be included, construct a container for them. if include_extra_model_outputs: extra_model_outputs: Dict[ModuleID, List[Dict[str, Any]]] = defaultdict( list ) B = 0 while B < batch_size_B: index_tuple = self._indices[self.rng.integers(len(self._indices))] # This will be an env timestep (not agent timestep) ma_episode_idx, ma_episode_ts = ( index_tuple[0] - self._num_episodes_evicted, index_tuple[1], ) # If we use random n-step sampling, draw the n-step for this item. if random_n_step: actual_n_step = int(self.rng.integers(n_step[0], n_step[1])) # If we are at the end of an episode, continue. # Note, priority sampling got us `o_(t+n)` and we need for the loss # calculation in addition `o_t`. # TODO (simon): Maybe introduce a variable `num_retries` until the # while loop should break when not enough samples have been collected # to make n-step possible. if ma_episode_ts - actual_n_step < 0: continue # Retrieve the multi-agent episode. ma_episode = self.episodes[ma_episode_idx] # Ensure that each row contains a tuple of the form: # (o_t, a_t, sum(r_(t:t+n_step)), o_(t+n_step)) # TODO (simon): Implement version for sequence sampling when using RNNs. eps_observation = ma_episode.get_observations( slice(ma_episode_ts - actual_n_step, ma_episode_ts + 1), return_list=True, ) # Note, `MultiAgentEpisode` stores the action that followed # `o_t` with `o_(t+1)`, therefore, we need the next one. # TODO (simon): This gets the wrong action as long as the getters are not # fixed. eps_actions = ma_episode.get_actions(ma_episode_ts - actual_n_step) # Make sure that at least a single agent should have full transition. # TODO (simon): Filter for the `modules_to_sample`. agents_to_sample = self._agents_with_full_transitions( eps_observation, eps_actions, ) # If not, we resample. if not agents_to_sample: continue # TODO (simon, sven): Do we need to include the common agent rewards? # Note, the reward that is collected by transitioning from `o_t` to # `o_(t+1)` is stored in the next transition in `MultiAgentEpisode`. eps_rewards = ma_episode.get_rewards( slice(ma_episode_ts - actual_n_step, ma_episode_ts), return_list=True, ) # TODO (simon, sven): Do we need to include the common infos? And are # there common extra model outputs? if include_infos: # If infos are included we include the ones from the last timestep # as usually the info contains additional values about the last state. eps_infos = ma_episode.get_infos(ma_episode_ts) if include_extra_model_outputs: # If `extra_model_outputs` are included we include the ones from the # first timestep as usually the `extra_model_outputs` contain additional # values from the forward pass that produced the action at the first # timestep. # Note, we extract them into single row dictionaries similar to the # infos, in a connector we can then extract these into single batch # rows. eps_extra_model_outputs = { k: ma_episode.get_extra_model_outputs( k, ma_episode_ts - actual_n_step ) for k in ma_episode.extra_model_outputs.keys() } # If the sampled time step is the episode's last time step check, if # the episode is terminated or truncated. episode_terminated = False episode_truncated = False if ma_episode_ts == ma_episode.env_t: episode_terminated = ma_episode.is_terminated episode_truncated = ma_episode.is_truncated # TODO (simon): Filter for the `modules_to_sample`. # TODO (sven, simon): We could here also sample for all agents in the # `modules_to_sample` and then adapt the `n_step` for agents that # have not a full transition. for agent_id in agents_to_sample: # Map our agent to the corresponding module we want to # train. module_id = ma_episode._agent_to_module_mapping[agent_id] # Sample only for the modules in `modules_to_sample`. if module_id not in ( modules_to_sample or self._module_to_indices.keys() ): continue # TODO (simon, sven): Here we could skip for modules not # to be sampled in `modules_to_sample`. observations[module_id].append(eps_observation[0][agent_id]) next_observations[module_id].append(eps_observation[-1][agent_id]) # Fill missing rewards with zeros. agent_rewards = [r[agent_id] or 0.0 for r in eps_rewards] rewards[module_id].append( scipy.signal.lfilter([1], [1, -gamma], agent_rewards[::-1], axis=0)[ -1 ] ) # Note, this should exist, as we filtered for agents with full # transitions. actions[module_id].append(eps_actions[agent_id]) if include_infos: infos[module_id].append(eps_infos[agent_id]) if include_extra_model_outputs: extra_model_outputs[module_id].append( { k: eps_extra_model_outputs[agent_id][k] for k in eps_extra_model_outputs[agent_id].keys() } ) # If sampled observation is terminal for the agent. Either MAE # episode is truncated/terminated or SAE episode is truncated/ # terminated at this ts. # TODO (simon, sven): Add method agent_alive(ts) to MAE. # or add slicing to get_terminateds(). agent_ts = ma_episode.env_t_to_agent_t[agent_id][ma_episode_ts] agent_eps = ma_episode.agent_episodes[agent_id] agent_terminated = agent_ts == agent_eps.t and agent_eps.is_terminated agent_truncated = ( agent_ts == agent_eps.t and agent_eps.is_truncated and not agent_eps.is_terminated ) if episode_terminated or agent_terminated: is_terminated[module_id].append(True) is_truncated[module_id].append(False) elif episode_truncated or agent_truncated: is_truncated[module_id].append(True) is_terminated[module_id].append(False) else: is_terminated[module_id].append(False) is_truncated[module_id].append(False) # Increase the per module counter. self.sampled_timesteps_per_module[module_id] += 1 # Increase counter. B += 1 # Increase the counter for environment timesteps. self.sampled_timesteps += batch_size_B # Should be convertible to MultiAgentBatch. ret = { **{ module_id: { Columns.OBS: batch(observations[module_id]), Columns.ACTIONS: batch(actions[module_id]), Columns.REWARDS: np.array(rewards[module_id]), Columns.NEXT_OBS: batch(next_observations[module_id]), Columns.TERMINATEDS: np.array(is_terminated[module_id]), Columns.TRUNCATEDS: np.array(is_truncated[module_id]), "weights": np.array(weights[module_id]), "n_step": np.array(n_steps[module_id]), } for module_id in observations.keys() } } # Return multi-agent dictionary. return ret def _num_remaining_episodes(self, new_eps, evicted_eps): """Calculates the number of remaining episodes. When adding episodes and evicting them in the `add()` method this function calculates iteratively the number of remaining episodes. Args: new_eps: List of new episode IDs. evicted_eps: List of evicted episode IDs. Returns: Number of episodes remaining after evicting the episodes in `evicted_eps` and adding the episode in `new_eps`. """ return len( set(self.episode_id_to_index.keys()).union(set(new_eps)) - set(evicted_eps) ) def _evict_module_episodes(self, ma_episode: MultiAgentEpisode) -> None: """Evicts the module episodes from the buffer adn updates all counters. Args: multi_agent_eps: The multi-agent episode to evict from the buffer. """ # Note we need to take the agent ids from the evicted episode because # different episodes can have different agents and module mappings. for agent_id in ma_episode.agent_episodes: # Retrieve the corresponding module ID and module episode. module_id = ma_episode._agent_to_module_mapping[agent_id] module_eps = ma_episode.agent_episodes[agent_id] # Update all counters. self._num_module_timesteps[module_id] -= module_eps.env_steps() self._num_module_episodes[module_id] -= 1 self._num_module_episodes_evicted[module_id] += 1 def _update_module_counters(self, ma_episode: MultiAgentEpisode) -> None: """Updates the module counters after adding an episode. Args: multi_agent_episode: The multi-agent episode to update the module counters for. """ for agent_id in ma_episode.agent_ids: agent_steps = ma_episode.agent_episodes[agent_id].env_steps() # Only add if the agent has stepped in the episode (chunk). if agent_steps > 0: # Receive the corresponding module ID. module_id = ma_episode.module_for(agent_id) self._num_module_timesteps[module_id] += agent_steps self._num_module_timesteps_added[module_id] += agent_steps # if ma_episode.agent_episodes[agent_id].is_done: # # TODO (simon): Check, if we do not count the same episode # # multiple times. # # Also add to the module episode counter. # self._num_module_episodes[module_id] += 1 def _add_new_module_indices( self, ma_episode: MultiAgentEpisode, episode_idx: int, ma_episode_exists: bool = True, ) -> None: """Adds the module indices for new episode chunks. Args: ma_episode: The multi-agent episode to add the module indices for. episode_idx: The index of the episode in the `self.episodes`. ma_episode_exists: Whether `ma_episode` is already in this buffer (with a predecessor chunk to which we'll concatenate `ma_episode` later). """ existing_ma_episode = None if ma_episode_exists: existing_ma_episode = self.episodes[ self.episode_id_to_index[ma_episode.id_] - self._num_episodes_evicted ] # Note, we iterate through the agent episodes b/c we want to store records # and some agents could not have entered the environment. for agent_id in ma_episode.agent_episodes: # Get the corresponding module id. module_id = ma_episode.module_for(agent_id) # Get the module episode. module_eps = ma_episode.agent_episodes[agent_id] # Is the agent episode already in the buffer's existing `ma_episode`? if ma_episode_exists and agent_id in existing_ma_episode.agent_episodes: existing_sa_eps_len = len(existing_ma_episode.agent_episodes[agent_id]) # Otherwise, it is a new single-agent episode and we increase the counter. else: existing_sa_eps_len = 0 self._num_module_episodes[module_id] += 1 # Add new module indices. self._module_to_indices[module_id].extend( [ ( # Keep the MAE index for sampling episode_idx, agent_id, existing_sa_eps_len + i, ) for i in range(len(module_eps)) ] ) def _agents_with_full_transitions( self, observations: Dict[AgentID, ObsType], actions: Dict[AgentID, ActType] ): """Filters for agents that have full transitions. Args: observations: The observations of the episode. actions: The actions of the episode. Returns: List of agent IDs that have full transitions. """ agents_to_sample = [] for agent_id in observations[0].keys(): # Only if the agent has an action at the first and an observation # at the first and last timestep of the n-step transition, we can sample it. if agent_id in actions and agent_id in observations[-1]: agents_to_sample.append(agent_id) return agents_to_sample