Compute the multi-agend reward as the sum of each agent reward in jax

skrl/multi_agents/jax/base.py

@@ -369,7 +369,7 @@ def record_transition(self,
         :type timesteps: int
         """
         if self.write_interval > 0:
-            _rewards = next(iter(rewards.values()))
+            _rewards = sum(rewards.values())
             _terminated = next(iter(terminated.values()))
             _truncated = next(iter(truncated.values()))