firs commit for our new MARL algorithms library, contains working implementations of IAC, SNAC and SEAC

algorithms/marl/seac.py

@@ -0,0 +1,55 @@
+import torch
+from torch.distributions import Categorical
+from algorithms.marl.iac import LoopIAC
+from algorithms.marl.memory import MARLActorCriticMemory
+
+
+class LoopSEAC(LoopIAC):
+    def __init__(self, cfg):
+        super(LoopSEAC, self).__init__(cfg)
+
+    def actor_critic(self, tm, networks, gamma, entropy_coef, vf_coef, **kwargs):
+        obs, actions, done, reward = tm.observation, tm.action, tm.done, tm.reward
+        outputs = [net(obs, actions, tm.hidden_actor, tm.hidden_critic) for net in networks]
+
+        with torch.inference_mode(True):
+            true_action_logp = torch.stack([
+                torch.log_softmax(out['logits'][ag_i, :-1], -1)
+                    .gather(index=actions[ag_i, 1:, None], dim=-1)
+                for ag_i, out in enumerate(outputs)
+            ], 0).squeeze()
+
+        losses = []
+
+        for ag_i, out in enumerate(outputs):
+            logits = out['logits'][:, :-1]  # last one only needed for v_{t+1}
+            critic = out['critic']
+
+            entropy_loss = Categorical(logits=logits[ag_i]).entropy().mean()
+            advantages = self.compute_advantages(critic, reward, done, gamma)
+
+            # policy loss
+            log_ap = torch.log_softmax(logits, -1)
+            log_ap = torch.gather(log_ap, dim=-1, index=actions[:, 1:].unsqueeze(-1)).squeeze()
+
+            # importance weights
+            iw = (log_ap - true_action_logp).exp().detach()  # importance_weights
+
+            a2c_loss = (-iw*log_ap * advantages.detach()).mean(-1)
+
+
+            value_loss = (iw*advantages.pow(2)).mean(-1)  # n_agent
+
+            # weighted loss
+            loss = (a2c_loss + vf_coef*value_loss - entropy_coef * entropy_loss).mean()
+            losses.append(loss)
+
+        return losses
+
+    def learn(self, tms: MARLActorCriticMemory, **kwargs):
+        losses = self.actor_critic(tms, self.net, **self.cfg['algorithm'], **kwargs)
+        for ag_i, loss in enumerate(losses):
+            self.optimizer[ag_i].zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(self.net[ag_i].parameters(), 0.5)
+            self.optimizer[ag_i].step()