[RLlib] Prioritized version of an episode-based replay buffer.

rllib/algorithms/algorithm_config.py

@@ -3409,7 +3409,8 @@ def get_marl_module_spec(
                         "is passed in nor in the default module spec used in "
                         "the algorithm."
                     )
-
+            # TODO (sven): Find a good way to pack module specific parameters from
+            # the algorithms into the `model_config_dict`.
             if module_spec.observation_space is None:
                 module_spec.observation_space = policy_spec.observation_space
             if module_spec.action_space is None:

rllib/utils/replay_buffers/__init__.py

@@ -10,6 +10,9 @@
     MultiAgentReplayBuffer,
     ReplayMode,
 )
+from ray.rllib.utils.replay_buffers.prioritized_episode_replay_buffer import (
+    PrioritizedEpisodeReplayBuffer,
+)
 from ray.rllib.utils.replay_buffers.prioritized_replay_buffer import (
     PrioritizedReplayBuffer,
 )
@@ -23,6 +26,7 @@
     "MultiAgentMixInReplayBuffer",
     "MultiAgentPrioritizedReplayBuffer",
     "MultiAgentReplayBuffer",
+    "PrioritizedEpisodeReplayBuffer",
     "PrioritizedReplayBuffer",
     "ReplayMode",
     "ReplayBuffer",

rllib/utils/replay_buffers/episode_replay_buffer.py

@@ -126,14 +126,16 @@ def add(self, episodes: Union[List["SingleAgentEpisode"], "SingleAgentEpisode"])
                 eps_idx = self.episode_id_to_index[eps.id_]
                 existing_eps = self.episodes[eps_idx - self._num_episodes_evicted]
                 old_len = len(existing_eps)
-                self._indices.extend([(eps_idx, old_len + i) for i in range(len(eps))])
+                self._indices.extend(
+                    [(eps_idx, old_len + i, None) for i in range(len(eps))]
+                )
                 existing_eps.concat_episode(eps)
             # New episode. Add to end of our episodes deque.
             else:
                 self.episodes.append(eps)
                 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))])
+                self._indices.extend([(eps_idx, i, None) for i in range(len(eps))])
 
             # Eject old records from front of deque (only if we have more than 1 episode
             # in the buffer).