fixed metrics logging for PyMARL integration with lifelong scenario

AIRI-Institute · Oct 29, 2023 · 4a3d06e · 4a3d06e
1 parent 43fbb1b
commit 4a3d06e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pogema/integrations/pymarl.py b/pogema/integrations/pymarl.py
@@ -28,9 +28,9 @@ def step(self, actions):
         info = {}
         done = all(terminated) or all(truncated)
         if done:
-            info.update(CSR=infos[0]['metrics']['CSR'])
-            info.update(ISR=infos[0]['metrics']['ISR'])
-
+            for key, value in infos[0]['metrics'].items():
+                info[key] = value
+            
         return sum(rewards), done, info
 
     def get_obs(self):