evolutionary-policy-optimization 0.1.17__py3-none-any.whl → 0.1.19__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -392,6 +392,8 @@ class MLP(Module):
 
         self.layers = ModuleList(layers)
 
+        self.final_lime = DynamicLIMe(dim, depth + 1)
+
     def forward(
         self,
         x,
@@ -430,7 +432,7 @@ class MLP(Module):
 
             prev_layer_inputs.append(x)
 
-        return x
+        return self.final_lime(x, prev_layer_inputs)
 
 # actor, critic, and agent (actor + critic)
 # eventually, should just create a separate repo and aggregate all the MLP related architectures
@@ -1066,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)
 
@@ -1281,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:
@@ -1305,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.17.dist-info → evolutionary_policy_optimization-0.1.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.17
+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.17.dist-info → evolutionary_policy_optimization-0.1.19.dist-info}/RECORD

@@ -1,10 +1,10 @@
 evolutionary_policy_optimization/__init__.py,sha256=NyiYDYU7DlpmOTM7xiBQET3r1WwX0ebrgMCBLSQrW3c,288
 evolutionary_policy_optimization/distributed.py,sha256=7KgZdeS_wxBHo_du9XZFB1Cu318J-Bp66Xdr6Log_20,2423
 evolutionary_policy_optimization/env_wrappers.py,sha256=bDL06o9_b1iW6k3fw2xifnOnYlzs643tdW6Yv2gsIdw,803
-evolutionary_policy_optimization/epo.py,sha256=TW-9l9oRN8XQZxSeG5Glkk4rWuxO9JOjjRJO7hJgHZs,48433
+evolutionary_policy_optimization/epo.py,sha256=JzZdNbzerIMgPg6dlL4eLNJ9_LbmW0xNkgQrgSNoSKA,49084
 evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
 evolutionary_policy_optimization/mock_env.py,sha256=TLyyRm6tOD0Kdn9QqJJQriaSnsR-YmNQHo4OohmZFG4,1410
-evolutionary_policy_optimization-0.1.17.dist-info/METADATA,sha256=-2Lt54pfQkw2_LAwB8pjcvLsc4vSMVQjAt-ZiObCRdw,7979
-evolutionary_policy_optimization-0.1.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.1.17.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.1.17.dist-info/RECORD,,
+evolutionary_policy_optimization-0.1.19.dist-info/METADATA,sha256=Jmbgx_z8dJv1W-FQVRJCI14MW2Tv7wbc4VhP_YN3WNw,7979
+evolutionary_policy_optimization-0.1.19.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.1.19.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.1.19.dist-info/RECORD,,