From 7aa4e8b52dc0834450b3a36a0e2f1c1218e5c3f7 Mon Sep 17 00:00:00 2001 From: Matteo Bettini Date: Fri, 22 Sep 2023 11:20:40 +0100 Subject: [PATCH] [Fix] RewardSum transform Signed-off-by: Matteo Bettini --- benchmarl/experiment/experiment.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/benchmarl/experiment/experiment.py b/benchmarl/experiment/experiment.py index 6e6cc596..8e24d8b5 100644 --- 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)