[BugFix] Fix meltingpot buffer transforms (#148)

* fixes * empty
facebookresearch · Nov 27, 2024 · 8045155 · 8045155
1 parent 9813807
commit 8045155
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 13 deletions.
diff --git a/benchmarl/environments/common.py b/benchmarl/environments/common.py
@@ -282,13 +282,14 @@ def get_env_transforms(self, env: EnvBase) -> List[Transform]:
         """
         return []
 
-    def get_replay_buffer_transforms(self, env: EnvBase) -> List[Transform]:
+    def get_replay_buffer_transforms(self, env: EnvBase, group: str) -> List[Transform]:
         """
-        Returns a list of :class:`torchrl.envs.Transform` to be applied to the :class:`torchrl.data.ReplayBuffer`.
+        Returns a list of :class:`torchrl.envs.Transform` to be applied to the :class:`torchrl.data.ReplayBuffer`
+        of the specified group.
 
         Args:
             env (EnvBase): An environment created via self.get_env_fun
-
+            group (str): The agent group using the replay buffer
 
         """
         return []

diff --git a/benchmarl/environments/meltingpot/common.py b/benchmarl/environments/meltingpot/common.py
@@ -125,22 +125,16 @@ def get_env_transforms(self, env: EnvBase) -> List[Transform]:
             else []
         )
 
-    def get_replay_buffer_transforms(self, env: EnvBase) -> List[Transform]:
+    def get_replay_buffer_transforms(self, env: EnvBase, group: str) -> List[Transform]:
         return [
             DTypeCastTransform(
                 dtype_in=torch.uint8,
                 dtype_out=torch.float,
                 in_keys=[
                     "RGB",
-                    *[
-                        (group, "observation", "RGB")
-                        for group in self.group_map(env).keys()
-                    ],
+                    (group, "observation", "RGB"),
                     ("next", "RGB"),
-                    *[
-                        ("next", group, "observation", "RGB")
-                        for group in self.group_map(env).keys()
-                    ],
+                    ("next", group, "observation", "RGB"),
                 ],
                 in_keys_inv=[],
             )

diff --git a/benchmarl/experiment/experiment.py b/benchmarl/experiment/experiment.py
@@ -466,7 +466,7 @@ def _setup_algorithm(self):
         self.replay_buffers = {
             group: self.algorithm.get_replay_buffer(
                 group=group,
-                transforms=self.task.get_replay_buffer_transforms(self.test_env),
+                transforms=self.task.get_replay_buffer_transforms(self.test_env, group),
             )
             for group in self.group_map.keys()
         }