[RLlib] Implementation of MultiAgentEpisodeReplayBuffer for the new API stack.

rllib/BUILD

@@ -1524,6 +1524,13 @@ py_test(
     srcs = ["utils/replay_buffers/tests/test_episode_replay_buffer.py"]
 )
 
+oy_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"],

rllib/env/multi_agent_episode.py

@@ -181,7 +181,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.