evolutionary-policy-optimization 0.0.47__py3-none-any.whl → 0.0.50__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -6,7 +6,7 @@ from collections import namedtuple
 from random import randrange
 
 import torch
-from torch import nn, cat, stack, is_tensor, tensor
+from torch import nn, cat, stack, is_tensor, tensor, Tensor
 import torch.nn.functional as F
 from torch.nn import Linear, Module, ModuleList
 from torch.utils.data import TensorDataset, DataLoader
@@ -357,9 +357,10 @@ class LatentGenePool(Module):
         frac_natural_selected = 0.25,    # number of least fit genes to remove from the pool
         frac_elitism = 0.1,              # frac of population to preserve from being noised
         frac_migrate = 0.1,              # frac of population, excluding elites, that migrate between islands randomly. will use a designated set migration pattern (since for some reason using random it seems to be worse for me)
+        migrate_every = 100,             # how many steps before a migration between islands
         mutation_strength = 1.,          # factor to multiply to gaussian noise as mutation to latents
         should_run_genetic_algorithm: Module | None = None, # eq (3) in paper
-        default_should_run_ga_gamma = 1.5
+        default_should_run_ga_gamma = 1.5,
     ):
         super().__init__()
 
@@ -409,6 +410,22 @@ class LatentGenePool(Module):
 
         self.should_run_genetic_algorithm = should_run_genetic_algorithm
 
+        self.can_migrate = num_islands > 1
+        self.migrate_every = migrate_every
+        self.register_buffer('step', tensor(1))
+
+    def get_distance(self):
+        # returns latent euclidean distance as proxy for diversity
+
+        latents = rearrange(self.latents, '(i p) g -> i p g', i = self.num_islands)
+
+        distance = torch.cdist(latents, latents)
+
+        return distance
+
+    def advance_step_(self):
+        self.step.add_(1)
+
     def firefly_step(
         self,
         fitness,
@@ -460,7 +477,7 @@ class LatentGenePool(Module):
         self,
         fitness, # Float['p'],
         inplace = True,
-        migrate = False # trigger a migration in the setting of multiple islands, the loop outside will need to have some `migrate_every` hyperparameter
+        migrate = None # trigger a migration in the setting of multiple islands, the loop outside will need to have some `migrate_every` hyperparameter
     ):
         device = self.latents.device
 
@@ -547,8 +564,9 @@ class LatentGenePool(Module):
 
         # 6. maybe migration
 
+        migrate = self.can_migrate and default(migrate, divisible_by(self.step.item(), self.migrate_every))
+
         if migrate:
-            assert self.num_islands > 1
             randperm = torch.randn(genes.shape[:-1], device = device).argsort(dim = -1)
 
             migrate_mask = randperm < self.num_migrate
@@ -581,6 +599,8 @@ class LatentGenePool(Module):
 
         self.latents.copy_(genes)
 
+        self.advance_step_()
+
     def forward(
         self,
         *args,
@@ -632,6 +652,7 @@ class Agent(Module):
         actor_lr = 1e-4,
         critic_lr = 1e-4,
         latent_lr = 1e-5,
+        diversity_aux_loss_weight = 0.,
         use_critic_ema = True,
         critic_ema_beta = 0.99,
         max_grad_norm = 0.5,
@@ -687,6 +708,11 @@ class Agent(Module):
 
         self.latent_optim = optim_klass(latent_gene_pool.parameters(), lr = latent_lr, **latent_optim_kwargs) if not latent_gene_pool.frozen_latents else None
 
+        # promotes latents to be farther apart for diversity maintenance
+
+        self.has_diversity_loss = diversity_aux_loss_weight > 0.
+        self.diversity_aux_loss_weight = diversity_aux_loss_weight
+
     def save(self, path, overwrite = False):
         path = Path(path)
 
@@ -868,11 +894,19 @@ class Agent(Module):
 
                 # maybe update latents, if not frozen
 
-                if not self.latent_gene_pool.frozen_latents:
-                    orig_latents.backward(latents.grad)
+                if self.latent_gene_pool.frozen_latents:
+                    continue
+
+                orig_latents.backward(latents.grad)
+
+                if self.has_diversity_loss:
+                    diversity = self.latent_gene_pool.get_distance()
+                    diversity_loss = diversity.mul(-1).exp().mean()
+
+                    (diversity_loss * self.diversity_aux_loss_weight).backward()
 
-                    self.latent_optim.step()
-                    self.latent_optim.zero_grad()
+                self.latent_optim.step()
+                self.latent_optim.zero_grad()
 
         # apply evolution
 
@@ -919,6 +953,7 @@ def create_agent(
     latent_gene_pool_kwargs: dict = dict(),
     actor_kwargs: dict = dict(),
     critic_kwargs: dict = dict(),
+    **kwargs
 ) -> Agent:
 
     latent_gene_pool = LatentGenePool(
@@ -946,7 +981,8 @@ def create_agent(
         actor = actor,
         critic = critic,
         latent_gene_pool = latent_gene_pool,
-        use_critic_ema = use_critic_ema
+        use_critic_ema = use_critic_ema,
+        **kwargs
     )
 
     return agent

{evolutionary_policy_optimization-0.0.47.dist-info → evolutionary_policy_optimization-0.0.50.dist-info}/METADATA

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

@@ -1,8 +1,8 @@
 evolutionary_policy_optimization/__init__.py,sha256=0q0aBuFgWi06MLMD8FiHzBYQ3_W4LYWrwmCtF3u5H2A,201
-evolutionary_policy_optimization/epo.py,sha256=-uRpnD0dKF6h4drVSikm9HnlP2OZ0WYQSWRQcghzd9Y,32242
+evolutionary_policy_optimization/epo.py,sha256=nIZzjz_3WUvHYPwNH9zoKM3fyRJd3MOd6UL7ooSTQV4,33445
 evolutionary_policy_optimization/experimental.py,sha256=9FrJGviLESlYysHI3i83efT9g2ZB9ha4u3K9HXN98_w,1100
 evolutionary_policy_optimization/mock_env.py,sha256=202KJ5g57wQvOzhGYzgHfBa7Y2do5uuDvl5kFg5o73g,934
-evolutionary_policy_optimization-0.0.47.dist-info/METADATA,sha256=oSI5NowsOOlQZ5cPmCs-8kYeG6TmzUybpRZt_6-cFWk,6213
-evolutionary_policy_optimization-0.0.47.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.47.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.47.dist-info/RECORD,,
+evolutionary_policy_optimization-0.0.50.dist-info/METADATA,sha256=iCB34UsXdosJV9q8-qXJpJfI8OdNNXB5o8MhmKy82zY,6213
+evolutionary_policy_optimization-0.0.50.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.50.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.50.dist-info/RECORD,,