Reduce KL in the multi-agent implementation

Toni-SM · Jul 16, 2024 · 946a544 · 946a544
1 parent 0af1edc
commit 946a544
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/skrl/multi_agents/torch/ippo/ippo.py b/skrl/multi_agents/torch/ippo/ippo.py
@@ -467,7 +467,12 @@ def compute_gae(rewards: torch.Tensor,
                 # update learning rate
                 if self._learning_rate_scheduler[uid]:
                     if isinstance(self.schedulers[uid], KLAdaptiveLR):
-                        self.schedulers[uid].step(torch.tensor(kl_divergences, device=self.device).mean())
+                        kl = torch.tensor(kl_divergences, device=self.device).mean()
+                        # reduce (collect from all workers/processes) KL in distributed runs
+                        if config.torch.is_distributed:
+                            torch.distributed.all_reduce(kl, op=torch.distributed.ReduceOp.SUM)
+                            kl /= config.torch.world_size
+                        self.schedulers[uid].step(kl.item())
                     else:
                         self.schedulers[uid].step()
 

diff --git a/skrl/multi_agents/torch/mappo/mappo.py b/skrl/multi_agents/torch/mappo/mappo.py
@@ -487,7 +487,12 @@ def compute_gae(rewards: torch.Tensor,
                 # update learning rate
                 if self._learning_rate_scheduler[uid]:
                     if isinstance(self.schedulers[uid], KLAdaptiveLR):
-                        self.schedulers[uid].step(torch.tensor(kl_divergences, device=self.device).mean())
+                        kl = torch.tensor(kl_divergences, device=self.device).mean()
+                        # reduce (collect from all workers/processes) KL in distributed runs
+                        if config.torch.is_distributed:
+                            torch.distributed.all_reduce(kl, op=torch.distributed.ReduceOp.SUM)
+                            kl /= config.torch.world_size
+                        self.schedulers[uid].step(kl.item())
                     else:
                         self.schedulers[uid].step()