evolutionary-policy-optimization 0.1.18__tar.gz → 0.1.19__tar.gz

{evolutionary_policy_optimization-0.1.18 → evolutionary_policy_optimization-0.1.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.18
+Version: 0.1.19
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization

{evolutionary_policy_optimization-0.1.18 → evolutionary_policy_optimization-0.1.19}/evolutionary_policy_optimization/epo.py

@@ -1068,6 +1068,12 @@ class Agent(Module):
     def unwrapped_latent_gene_pool(self):
         return self.unwrap_model(self.latent_gene_pool)
 
+    def log(self, **data_kwargs):
+        if not self.wrap_with_accelerate:
+            return
+
+        self.accelerate.log(data_kwargs, step = self.step)
+
     def save(self, path, overwrite = False):
         path = Path(path)
 
@@ -1283,6 +1289,14 @@ class Agent(Module):
                 self.critic_optim.step()
                 self.critic_optim.zero_grad()
 
+                # log actor critic loss
+
+                self.log(
+                    actor_loss = actor_loss.item(),
+                    critic_loss = critic_loss.item(),
+                    fitness_scores = fitness_scores
+                )
+
                 # maybe ema update critic
 
                 if self.use_critic_ema:
@@ -1307,6 +1321,11 @@ class Agent(Module):
                 self.latent_optim.step()
                 self.latent_optim.zero_grad()
 
+                if self.has_diversity_loss:
+                    self.log(
+                        diversity_loss = diversity_loss.item()
+                    )
+
         # apply evolution
 
         if self.has_latent_genes:

{evolutionary_policy_optimization-0.1.18 → evolutionary_policy_optimization-0.1.19}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.1.18"
+version = "0.1.19"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }