evolutionary-policy-optimization 0.0.50__tar.gz → 0.0.52__tar.gz

{evolutionary_policy_optimization-0.0.50 → evolutionary_policy_optimization-0.0.52}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.50
+Version: 0.0.52
 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.50 → evolutionary_policy_optimization-0.0.52}/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()
 
@@ -393,7 +396,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
@@ -530,7 +532,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)
@@ -901,7 +905,7 @@ class Agent(Module):
 
                 if self.has_diversity_loss:
                     diversity = self.latent_gene_pool.get_distance()
-                    diversity_loss = diversity.mul(-1).exp().mean()
+                    diversity_loss = (-diversity).tril(-1).exp().mean()
 
                     (diversity_loss * self.diversity_aux_loss_weight).backward()
 

{evolutionary_policy_optimization-0.0.50 → evolutionary_policy_optimization-0.0.52}/pyproject.toml

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

{evolutionary_policy_optimization-0.0.50 → evolutionary_policy_optimization-0.0.52}/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