evolutionary-policy-optimization 0.0.27__tar.gz → 0.0.29__tar.gz

{evolutionary_policy_optimization-0.0.27 → evolutionary_policy_optimization-0.0.29}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.27
+Version: 0.0.29
 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.27 → evolutionary_policy_optimization-0.0.29}/evolutionary_policy_optimization/epo.py

@@ -311,6 +311,7 @@ class LatentGenePool(Module):
         frac_tournaments = 0.25,         # fraction of genes to participate in tournament - the lower the value, the more chance a less fit gene could be selected
         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)
         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
@@ -348,12 +349,16 @@ class LatentGenePool(Module):
         self.num_natural_selected = int(frac_natural_selected * latents_per_island)
 
         self.num_tournament_participants = int(frac_tournaments * self.num_natural_selected)
+
         self.crossover_random  = crossover_random
 
         self.mutation_strength = mutation_strength
         self.num_elites = int(frac_elitism * latents_per_island)
         self.has_elites = self.num_elites > 0
 
+        latents_without_elites = num_latents - self.num_elites
+        self.num_migrate = int(frac_migrate * latents_without_elites)
+
         if not exists(should_run_genetic_algorithm):
             should_run_genetic_algorithm = ShouldRunGeneticAlgorithm(gamma = default_should_run_ga_gamma)
 
@@ -365,7 +370,6 @@ class LatentGenePool(Module):
         beta0 = 2.,           # exploitation factor, moving fireflies of low light intensity to high
         gamma = 1.,           # controls light intensity decay over distance - setting this to zero will make firefly equivalent to vanilla PSO
         alpha = 0.1,          # exploration factor
-        alpha_decay = 0.995,  # exploration decay each step
         inplace = True,
     ):
         islands = self.num_islands
@@ -411,8 +415,11 @@ class LatentGenePool(Module):
     def genetic_algorithm_step(
         self,
         fitness, # Float['p'],
-        inplace = True
+        inplace = True,
+        migrate = False # trigger a migration in the setting of multiple islands, the loop outside will need to have some `migrate_every` hyperparameter
     ):
+        device = self.latents.device
+
         """
         i - islands
         p - population
@@ -462,7 +469,7 @@ class LatentGenePool(Module):
 
         # 2. for finding pairs of parents to replete gene pool, we will go with the popular tournament strategy
 
-        rand_tournament_gene_ids = torch.randn((islands, pop_size_per_island - self.num_natural_selected, tournament_participants)).argsort(dim = -1)
+        rand_tournament_gene_ids = torch.randn((islands, pop_size_per_island - self.num_natural_selected, tournament_participants), device = device).argsort(dim = -1)
         rand_tournament_gene_ids_for_gather = rearrange(rand_tournament_gene_ids, 'i p t -> i (p t)')
 
         participant_fitness = fitness.gather(1, rand_tournament_gene_ids_for_gather)
@@ -494,6 +501,20 @@ class LatentGenePool(Module):
 
         genes = mutation(genes, mutation_strength = self.mutation_strength)
 
+        # 6. maybe migration
+
+        if migrate:
+            assert self.num_islands > 1
+            randperm = torch.randn(genes.shape[:-1], device = device).argsort(dim = -1)
+
+            migrate_mask = randperm < self.num_migrate
+
+            nonmigrants = rearrange(genes[~migrate_mask], '(i p) g -> i p g', i = islands)
+            migrants = rearrange(genes[migrate_mask], '(i p) g -> i p g', i = islands)
+            migrants = torch.roll(migrants, 1, dims = 0)
+
+            genes = cat((nonmigrants, migrants), dim = 1)
+
         # add back the elites
 
         if self.has_elites:

{evolutionary_policy_optimization-0.0.27 → evolutionary_policy_optimization-0.0.29}/pyproject.toml

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

{evolutionary_policy_optimization-0.0.27 → evolutionary_policy_optimization-0.0.29}/tests/test_epo.py

@@ -37,7 +37,9 @@ def test_readme(
 
     fitness = torch.randn(128)
 
-    latent_pool.genetic_algorithm_step(fitness) # update once
+    latent_pool.genetic_algorithm_step(fitness, migrate = num_islands > 1) # update once
+
+    latent_pool.firefly_step(fitness)
 
 @pytest.mark.parametrize('latent_ids', (2, (2, 4)))
 def test_create_agent(