evolutionary-policy-optimization 0.0.25__py3-none-any.whl → 0.0.27__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -303,7 +303,6 @@ class LatentGenePool(Module):
         self,
         num_latents,                     # same as gene pool size
         dim_latent,                      # gene dimension
-        num_latent_sets = 1,             # allow for sets of latents / gene per individual, expression of a set controlled by the environment
         num_islands = 1,                 # add the island strategy, which has been effectively used in a few recent works
         dim_state = None,
         frozen_latents = True,
@@ -320,29 +319,17 @@ class LatentGenePool(Module):
 
         maybe_l2norm = l2norm if l2norm_latent else identity
 
-        latents = torch.randn(num_latents, num_latent_sets, dim_latent)
+        latents = torch.randn(num_latents, dim_latent)
 
         if l2norm_latent:
             latents = maybe_l2norm(latents, dim = -1)
 
         self.num_latents = num_latents
-        self.needs_latent_gate = num_latent_sets > 1
+        self.frozen_latents = frozen_latents
         self.latents = nn.Parameter(latents, requires_grad = not frozen_latents)
 
         self.maybe_l2norm = maybe_l2norm
 
-        # gene expression as a function of environment
-
-        self.num_latent_sets = num_latent_sets
-
-        if self.needs_latent_gate:
-            assert exists(dim_state), '`dim_state` must be passed in if using gated gene expression'
-
-        self.to_latent_gate = nn.Sequential(
-            Linear(dim_state, num_latent_sets),
-            nn.Softmax(dim = -1)
-        ) if self.needs_latent_gate else None
-
         # some derived values
 
         assert num_islands >= 1
@@ -410,6 +397,10 @@ class LatentGenePool(Module):
 
         fireflies = rearrange(fireflies, 'i p ... -> (i p) ...')
 
+        # maybe fireflies on hypersphere
+
+        fireflies = self.maybe_l2norm(fireflies)
+
         if not inplace:
             return fireflies
 
@@ -456,7 +447,7 @@ class LatentGenePool(Module):
 
             return genes
 
-        genes = rearrange(genes, '(i p) n g -> i p n g', i = islands)
+        genes = rearrange(genes, '(i p) ... -> i p ...', i = islands)
 
         orig_genes = genes
 
@@ -465,7 +456,7 @@ class LatentGenePool(Module):
 
         sorted_indices = fitness.sort(dim = -1).indices
         natural_selected_indices = sorted_indices[..., -self.num_natural_selected:]
-        natural_select_gene_indices = repeat(natural_selected_indices, '... -> ... n g', n = genes.shape[-2], g = genes.shape[-1])
+        natural_select_gene_indices = repeat(natural_selected_indices, '... -> ... g', g = genes.shape[-1])
 
         genes, fitness = genes.gather(1, natural_select_gene_indices), fitness.gather(1, natural_selected_indices)
 
@@ -480,7 +471,7 @@ class LatentGenePool(Module):
         parent_indices_at_tournament = participant_fitness.topk(2, dim = -1).indices
         parent_gene_ids = rand_tournament_gene_ids.gather(-1, parent_indices_at_tournament)
 
-        parent_gene_ids_for_gather = repeat(parent_gene_ids, 'i p parents -> i (p parents) n g', n = genes.shape[-2], g = genes.shape[-1])
+        parent_gene_ids_for_gather = repeat(parent_gene_ids, 'i p parents -> i (p parents) g', g = genes.shape[-1])
 
         parents = genes.gather(1, parent_gene_ids_for_gather)
         parents = rearrange(parents, 'i (p parents) ... -> i p parents ...', parents = 2)
@@ -551,22 +542,6 @@ class LatentGenePool(Module):
 
         latent = self.latents[latent_id]
 
-        if self.needs_latent_gate:
-            assert exists(state), 'state must be passed in if greater than number of 1 latent set'
-
-            if not fetching_multiple_latents:
-                latent = repeat(latent, '... -> b ...', b = state.shape[0])
-
-            assert latent.shape[0] == state.shape[0]
-
-            gates = self.to_latent_gate(state)
-            latent = einsum(latent, gates, 'b n g, b n -> b g')
-
-        elif fetching_multiple_latents:
-            latent = latent[:, 0]
-        else:
-            latent = latent[0]
-
         latent = self.maybe_l2norm(latent)
 
         if not exists(net):
@@ -608,7 +583,7 @@ class Agent(Module):
         self.actor_optim = optim_klass(actor.parameters(), lr = actor_lr, **actor_optim_kwargs)
         self.critic_optim = optim_klass(critic.parameters(), lr = critic_lr, **critic_optim_kwargs)
 
-        self.latent_optim = optim_klass(latent_gene_pool.parameters(), lr = latent_lr, **latent_optim_kwargs) if latent_gene_pool.needs_latent_gate else None
+        self.latent_optim = optim_klass(latent_gene_pool.parameters(), lr = latent_lr, **latent_optim_kwargs) if not latent_gene_pool.frozen_latents else None
 
     def get_actor_actions(
         self,
@@ -683,7 +658,6 @@ def create_agent(
     actor_num_actions,
     actor_dim_hiddens: int | tuple[int, ...],
     critic_dim_hiddens: int | tuple[int, ...],
-    num_latent_sets = 1
 ) -> Agent:
 
     actor = Actor(
@@ -703,7 +677,6 @@ def create_agent(
         dim_state = dim_state,
         num_latents = num_latents,
         dim_latent = dim_latent,
-        num_latent_sets = num_latent_sets
     )
 
     return Agent(actor = actor, critic = critic, latent_gene_pool = latent_gene_pool)

{evolutionary_policy_optimization-0.0.25.dist-info → evolutionary_policy_optimization-0.0.27.dist-info}/METADATA

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

@@ -0,0 +1,7 @@
+evolutionary_policy_optimization/__init__.py,sha256=Qavcia0n13jjaWIS_LPW7QrxSLT_BBeKujCjF9kQjbA,133
+evolutionary_policy_optimization/epo.py,sha256=UCCwYK-b20X-5Cq-pah1NTeHFc_35b4xZ3y0aSR8aaI,20783
+evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
+evolutionary_policy_optimization-0.0.27.dist-info/METADATA,sha256=pJ2kQD5YtKDSUp1TCO_hsrRMh6FCMm8dyu6WrpVHiQk,4958
+evolutionary_policy_optimization-0.0.27.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.27.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.27.dist-info/RECORD,,

evolutionary_policy_optimization-0.0.25.dist-info/RECORD

@@ -1,7 +0,0 @@
-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,,