evolutionary-policy-optimization 0.0.51__tar.gz → 0.0.53__tar.gz

{evolutionary_policy_optimization-0.0.51 → evolutionary_policy_optimization-0.0.53}/PKG-INFO

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

@@ -48,6 +48,9 @@ def divisible_by(num, den):
 def l2norm(t):
     return F.normalize(t, p = 2, dim = -1)
 
+def batch_randperm(shape, device):
+    return torch.randn(shape, device = device).argsort(dim = -1)
+
 def log(t, eps = 1e-20):
     return t.clamp(min = eps).log()
 
@@ -357,10 +360,11 @@ 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,
+        migrate_every = 100,                 # how many steps before a migration between islands
+        apply_genetic_algorithm_every = 2    # how many steps before crossover + mutation happens for genes
     ):
         super().__init__()
 
@@ -393,7 +397,6 @@ class LatentGenePool(Module):
 
         latents_per_island = num_latents // num_islands
         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
@@ -411,7 +414,10 @@ class LatentGenePool(Module):
         self.should_run_genetic_algorithm = should_run_genetic_algorithm
 
         self.can_migrate = num_islands > 1
+
         self.migrate_every = migrate_every
+        self.apply_genetic_algorithm_every = apply_genetic_algorithm_every
+
         self.register_buffer('step', tensor(1))
 
     def get_distance(self):
@@ -481,6 +487,10 @@ class LatentGenePool(Module):
     ):
         device = self.latents.device
 
+        if not divisible_by(self.step.item(), self.apply_genetic_algorithm_every):
+            self.advance_step_()
+            return
+
         """
         i - islands
         p - population
@@ -530,7 +540,9 @@ 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), device = device).argsort(dim = -1)
+        tournament_shape = (islands, pop_size_per_island - self.num_natural_selected, self.num_natural_selected) # (island, num children needed, natural selected population to be bred)
+
+        rand_tournament_gene_ids = batch_randperm(tournament_shape, device)[..., :tournament_participants]
         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)

{evolutionary_policy_optimization-0.0.51 → evolutionary_policy_optimization-0.0.53}/pyproject.toml

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

{evolutionary_policy_optimization-0.0.51 → evolutionary_policy_optimization-0.0.53}/tests/test_epo.py

@@ -78,7 +78,8 @@ def test_create_agent(
 @pytest.mark.parametrize('diversity_aux_loss_weight', (0., 1e-3))
 def test_e2e_with_mock_env(
     frozen_latents,
-    use_critic_ema
+    use_critic_ema,
+    diversity_aux_loss_weight
 ):
     from evolutionary_policy_optimization import create_agent, EPO, Env