major redesign ob observations and entittes

algorithms/marl/base_ac.py

@@ -1,5 +1,5 @@
 import torch
-from typing import Union, List
+from typing import Union, List, Dict
 import numpy as np
 from torch.distributions import Categorical
 from algorithms.marl.memory import MARLActorCriticMemory
@@ -74,7 +74,7 @@ class BaseActorCritic:
         actions = [Categorical(logits=logits).sample().item() for logits in out[nms.LOGITS]]
         return actions
 
-    def init_hidden(self) -> dict[ListOrTensor]:
+    def init_hidden(self) -> Dict[str, ListOrTensor]:
         pass
 
     def forward(self,
@@ -82,7 +82,7 @@ class BaseActorCritic:
                 actions:       ListOrTensor,
                 hidden_actor:  ListOrTensor,
                 hidden_critic: ListOrTensor
-                ) -> dict[ListOrTensor]:
+                ) -> Dict[str, ListOrTensor]:
         pass
 
     @torch.no_grad()

algorithms/marl/iac.py

@@ -39,7 +39,7 @@ class LoopIAC(BaseActorCritic):
     def forward(self, observations, actions, hidden_actor, hidden_critic):
         outputs = [
             net(
-                self._as_torch(observations[ag_i]).unsqueeze(0).unsqueeze(0),  # agents x time
+                self._as_torch(observations[ag_i]).unsqueeze(0).unsqueeze(0),  # agent x time
                 self._as_torch(actions[ag_i]).unsqueeze(0),
                 hidden_actor[ag_i],
                 hidden_critic[ag_i]

algorithms/marl/mappo.py

@@ -46,7 +46,7 @@ class LoopMAPPO(LoopSNAC):
 
         # monte carlo returns
         mc_returns = self.monte_carlo_returns(batch[nms.REWARD], batch[nms.DONE], gamma)
-        mc_returns = (mc_returns - mc_returns.mean()) / (mc_returns.std() + 1e-8) #todo: norm across agents ok?
+        mc_returns = (mc_returns - mc_returns.mean()) / (mc_returns.std() + 1e-8) #todo: norm across agent ok?
         advantages =  mc_returns - out[nms.CRITIC][:, :-1]
 
         # policy loss

algorithms/marl/memory.py

@@ -120,7 +120,7 @@ class MARLActorCriticMemory(object):
 
     def __getattr__(self, attr):
         all_attrs = [getattr(mem, attr) for mem in self.memories]
-        return torch.cat(all_attrs, 0)  # agents x time ...
+        return torch.cat(all_attrs, 0)  # agent x time ...
 
     def chunk_dataloader(self, chunk_len, k):
         datasets = [ExperienceChunks(mem, chunk_len, k) for mem in self.memories]