[RLlib] Implementation of MultiAgentEpisodeReplayBuffer for the new…

… API stack. (#44450)
ray-project · Apr 25, 2024 · 97e910b · 97e910b
1 parent 611e7d5
commit 97e910b
Show file tree

Hide file tree

Showing 6 changed files with 1,302 additions and 8 deletions.
diff --git a/rllib/BUILD b/rllib/BUILD
@@ -1526,6 +1526,13 @@ py_test(
     srcs = ["utils/replay_buffers/tests/test_episode_replay_buffer.py"]
 )
 
+py_test(
+    name = "test_multi_agent_episode_replay_buffer",
+    tags = ["team:rllib", "utils"],
+    size = "small",
+    srcs = ["utils/replay_buffers/tests/test_multi_agent_episode_replay_buffer.py"]
+)
+
 py_test(
     name = "test_multi_agent_mixin_replay_buffer",
     tags = ["team:rllib", "utils"],

diff --git a/rllib/env/multi_agent_episode.py b/rllib/env/multi_agent_episode.py
@@ -182,7 +182,7 @@ def __init__(
         # via the `module_for()` API even before the agent has entered the episode
         # (and has its SingleAgentEpisode created), we store all aldeary done mappings
         # in this dict here.
-        self._agent_to_module_mapping: Dict[AgentID, ModuleID] = {}
+        self._agent_to_module_mapping: Dict[AgentID, ModuleID] = agent_module_ids or {}
 
         # Lookback buffer length is not provided. Interpret all provided data as
         # lookback buffer.

diff --git a/rllib/utils/replay_buffers/__init__.py b/rllib/utils/replay_buffers/__init__.py
@@ -3,6 +3,9 @@
 from ray.rllib.utils.replay_buffers.multi_agent_mixin_replay_buffer import (
     MultiAgentMixInReplayBuffer,
 )
+from ray.rllib.utils.replay_buffers.multi_agent_episode_replay_buffer import (
+    MultiAgentEpisodeReplayBuffer,
+)
 from ray.rllib.utils.replay_buffers.multi_agent_prioritized_replay_buffer import (
     MultiAgentPrioritizedReplayBuffer,
 )
@@ -23,6 +26,7 @@
 __all__ = [
     "EpisodeReplayBuffer",
     "FifoReplayBuffer",
+    "MultiAgentEpisodeReplayBuffer",
     "MultiAgentMixInReplayBuffer",
     "MultiAgentPrioritizedReplayBuffer",
     "MultiAgentReplayBuffer",