evolutionary-policy-optimization 0.0.23__py3-none-any.whl → 0.0.24__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -8,6 +8,7 @@ import torch.nn.functional as F
 from torch.nn import Linear, Module, ModuleList
 from torch.utils.data import TensorDataset, DataLoader
 
+import einx
 from einops import rearrange, repeat, einsum
 from einops.layers.torch import Rearrange
 
@@ -293,7 +294,7 @@ class ShouldRunGeneticAlgorithm(Module):
         # however, this equation does not make much sense to me if fitness increases unbounded
         # just let it be customizable, and offer a variant where mean and variance is over some threshold (could account for skew too)
 
-        return (fitnesses.amax() - fitnesses.amin()) > (self.gamma * torch.median(fitnesses))
+        return (fitnesses.amax(dim = -1) - fitnesses.amin(dim = -1)) > (self.gamma * torch.median(fitnesses, dim = -1).values)
 
 # classes
 
@@ -389,9 +390,6 @@ class LatentGenePool(Module):
         islands = self.num_islands
         tournament_participants = self.num_tournament_participants
 
-        if not self.should_run_genetic_algorithm(fitness):
-            return
-
         assert self.num_latents > 1
 
         genes = self.latents # the latents are the genes
@@ -399,12 +397,25 @@ class LatentGenePool(Module):
         pop_size = genes.shape[0]
         assert pop_size == fitness.shape[0]
 
+        pop_size_per_island = pop_size // islands
+
         # split out the islands
 
-        genes = rearrange(genes, '(i p) n g -> i p n g', i = islands)
         fitness = rearrange(fitness, '(i p) -> i p', i = islands)
 
-        pop_size_per_island = pop_size // islands
+        # from the fitness, decide whether to actually run the genetic algorithm or not
+
+        should_update_per_island = self.should_run_genetic_algorithm(fitness)
+
+        if not should_update_per_island.any():
+            if inplace:
+                return
+
+            return genes
+
+        genes = rearrange(genes, '(i p) n g -> i p n g', i = islands)
+
+        orig_genes = genes
 
         # 1. natural selection is simple in silico
         # you sort the population by the fitness and slice off the least fit end
@@ -456,6 +467,10 @@ class LatentGenePool(Module):
 
         genes = self.maybe_l2norm(genes)
 
+        # account for criteria of whether to actually run GA or not
+
+        genes = einx.where('i, i ..., i ...', should_update_per_island, genes, orig_genes)
+
         # merge island back into pop dimension
 
         genes = rearrange(genes, 'i p ... -> (i p) ...')

{evolutionary_policy_optimization-0.0.23.dist-info → evolutionary_policy_optimization-0.0.24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.23
+Version: 0.0.24
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -36,7 +36,8 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: adam-atan2-pytorch
 Requires-Dist: assoc-scan
-Requires-Dist: einops>=0.8.0
+Requires-Dist: einops>=0.8.1
+Requires-Dist: einx>=0.3.0
 Requires-Dist: hl-gauss-pytorch>=0.1.19
 Requires-Dist: torch>=2.2
 Requires-Dist: tqdm

evolutionary_policy_optimization-0.0.24.dist-info/RECORD

@@ -0,0 +1,7 @@
+evolutionary_policy_optimization/__init__.py,sha256=Qavcia0n13jjaWIS_LPW7QrxSLT_BBeKujCjF9kQjbA,133
+evolutionary_policy_optimization/epo.py,sha256=-kQgrnnOLiCOZ-6EroO057tDx0sS7TQro92cjJhSbZU,20353
+evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
+evolutionary_policy_optimization-0.0.24.dist-info/METADATA,sha256=d3imh1p1-nPpNGhD8cReLdL07_-oHZs3YqJaOEJi1TM,4958
+evolutionary_policy_optimization-0.0.24.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.24.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.24.dist-info/RECORD,,

evolutionary_policy_optimization-0.0.23.dist-info/RECORD

@@ -1,7 +0,0 @@
-evolutionary_policy_optimization/__init__.py,sha256=Qavcia0n13jjaWIS_LPW7QrxSLT_BBeKujCjF9kQjbA,133
-evolutionary_policy_optimization/epo.py,sha256=lMKMGaJdSh71Dnkn6ZvbtVNrCJ00Cv3p_uXJy9D8K90,19908
-evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
-evolutionary_policy_optimization-0.0.23.dist-info/METADATA,sha256=m_EqBjggqqEq09A9KqjA8BQMWECLIvzwbOR8pc-UFrM,4931
-evolutionary_policy_optimization-0.0.23.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.23.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.23.dist-info/RECORD,,