Update environment_loop.py

pddub97 · Jan 19, 2023 · 0ceea5b · 0ceea5b
1 parent fd4b4a7
commit 0ceea5b
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/mava/environment_loop.py b/mava/environment_loop.py
@@ -131,9 +131,12 @@ def run_episode(self) -> loggers.LoggingData:
         # For evaluation, this keeps track of the total undiscounted reward
         # for each agent accumulated during the episode.
         rewards: Dict[str, float] = {}
+        rewards_SOC: Dict[str, float] = {}
         episode_returns: Dict[str, float] = {}
         for agent, spec in self._environment.reward_spec().items():
             rewards.update({agent: generate_zeros_from_spec(spec)})
+            rewards_SOC.update({agent: generate_zeros_from_spec(spec)})
+            rewards_load.update({agent: generate_zeros_from_spec(spec)})
             episode_returns.update({agent: generate_zeros_from_spec(spec)})
 
         # Run an episode.