[Fix] RewardSum transform

Signed-off-by: Matteo Bettini <[email protected]>
facebookresearch · Sep 22, 2023 · 7aa4e8b · 7aa4e8b
1 parent 2a4051e
commit 7aa4e8b
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/benchmarl/experiment/experiment.py b/benchmarl/experiment/experiment.py
@@ -200,11 +200,12 @@ def _setup_task(self):
         reward_spec = test_env.output_spec["full_reward_spec"]
         transforms = []
         for reward_key in reward_spec.keys(True, True):
-            reward_key = _unravel_key_to_tuple(reward_key)
             transforms.append(
                 RewardSum(
                     in_keys=[reward_key],
-                    out_keys=[reward_key[:-1] + ("episode_reward",)],
+                    out_keys=[
+                        _unravel_key_to_tuple(reward_key)[:-1] + ("episode_reward",)
+                    ],
                 )
             )
         transform = Compose(*transforms)