evolutionary-policy-optimization 0.0.23__py3-none-any.whl → 0.0.25__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
 
@@ -371,6 +372,49 @@ class LatentGenePool(Module):
 
         self.should_run_genetic_algorithm = should_run_genetic_algorithm
 
+    def firefly_step(
+        self,
+        fitness,
+        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
+        fireflies = self.latents # the latents are the fireflies
+
+        assert fitness.shape[0] == fireflies.shape[0]
+
+        fitness = rearrange(fitness, '(i p) -> i p', i = islands)
+        fireflies = rearrange(fireflies, '(i p) ... -> i p ...', i = islands)
+
+        # fireflies with lower light intensity (high cost) moves towards the higher intensity (lower cost)
+
+        move_mask = einx.less('i x, i y -> i x y', fitness, fitness)
+
+        # get vectors of fireflies to one another
+        # calculate distance and the beta
+
+        delta_positions = einx.subtract('i y ... d, i x ... d -> i x y ... d', fireflies, fireflies)
+
+        distance = delta_positions.norm(dim = -1)
+
+        betas = beta0 * (-gamma * distance ** 2).exp()
+
+        # move the fireflies according to attraction
+
+        fireflies += einsum(move_mask, betas, delta_positions, 'i x y, i x y ..., i x y ... -> i x ...')
+
+        # merge back the islands
+
+        fireflies = rearrange(fireflies, 'i p ... -> (i p) ...')
+
+        if not inplace:
+            return fireflies
+
+        self.latents.copy_(fireflies)
+
     @torch.no_grad()
     # non-gradient optimization, at least, not on the individual level (taken care of by rl component)
     def genetic_algorithm_step(
@@ -389,9 +433,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 +440,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 +510,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.25.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.23
+Version: 0.0.25
 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.25.dist-info/RECORD

@@ -0,0 +1,7 @@
+evolutionary_policy_optimization/__init__.py,sha256=Qavcia0n13jjaWIS_LPW7QrxSLT_BBeKujCjF9kQjbA,133
+evolutionary_policy_optimization/epo.py,sha256=BLwy7PBZOjw6H7MFvMq9CC7Mdm3K8fpzBNH6HbNu6LY,21927
+evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
+evolutionary_policy_optimization-0.0.25.dist-info/METADATA,sha256=p3-_SuLvKs8E0z1l567qA0Pbsv2dOLlrJPX4WYoZaB4,4958
+evolutionary_policy_optimization-0.0.25.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.25.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.25.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,,