evolutionary-policy-optimization 0.0.10__py3-none-any.whl → 0.0.11__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -163,7 +163,6 @@ class MLP(Module):
         self,
         dims: tuple[int, ...],
         dim_latent = 0,
-        num_latent_sets = 1
     ):
         super().__init__()
         assert len(dims) >= 2, 'must have at least two dimensions'
@@ -173,24 +172,15 @@ class MLP(Module):
         first_dim, *rest_dims = dims
         dims = (first_dim + dim_latent, *rest_dims)
 
-        assert num_latent_sets >= 1
-
         self.dim_latent = dim_latent
-        self.num_latent_sets = num_latent_sets
 
         self.needs_latent = dim_latent > 0
-        self.needs_latent_gate = num_latent_sets > 1
 
         self.encode_latent = nn.Sequential(
             Linear(dim_latent, dim_latent),
             nn.SiLU()
         ) if self.needs_latent else None
 
-        self.to_latent_gate = nn.Sequential(
-            Linear(first_dim, num_latent_sets),
-            nn.Softmax(dim = -1)
-        ) if self.needs_latent_gate else None
-
         # pairs of dimension
 
         dim_pairs = tuple(zip(dims[:-1], dims[1:]))
@@ -210,15 +200,6 @@ class MLP(Module):
 
         assert xnor(self.needs_latent, exists(latent))
 
-        if exists(latent) and self.needs_latent_gate:
-            # an improvisation where set of genes with controlled expression by environment
-
-            gates = self.to_latent_gate(x)
-            latent = einsum(latent, gates, 'n g, b n -> b g')
-        else:
-            assert latent.shape[0] == 1
-            latent = latent[0]
-
         if exists(latent):
             # start with naive concatenative conditioning
             # but will also offer some alternatives once a spark is seen (film, adaptive linear from stylegan, etc)
@@ -248,7 +229,7 @@ class MLP(Module):
 class Actor(Module):
     def __init__(
         self,
-        dim_in,
+        dim_state,
         num_actions,
         dim_hiddens: tuple[int, ...],
         dim_latent = 0,
@@ -259,7 +240,7 @@ class Actor(Module):
         dim_first, *_, dim_last = dim_hiddens
 
         self.init_layer = nn.Sequential(
-            nn.Linear(dim_in, dim_first),
+            nn.Linear(dim_state, dim_first),
             nn.SiLU()
         )
 
@@ -285,7 +266,7 @@ class Actor(Module):
 class Critic(Module):
     def __init__(
         self,
-        dim_in,
+        dim_state,
         dim_hiddens: tuple[int, ...],
         dim_latent = 0,
     ):
@@ -295,7 +276,7 @@ class Critic(Module):
         dim_first, *_, dim_last = dim_hiddens
 
         self.init_layer = nn.Sequential(
-            nn.Linear(dim_in, dim_first),
+            nn.Linear(dim_state, dim_first),
             nn.SiLU()
         )
 
@@ -346,6 +327,7 @@ class LatentGenePool(Module):
         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
+        dim_state = None,
         crossover_random = True,         # random interp from parent1 to parent2 for crossover, set to `False` for averaging (0.5 constant value)
         l2norm_latent = False,           # whether to enforce latents on hypersphere,
         frac_tournaments = 0.25,         # fraction of genes to participate in tournament - the lower the value, the more chance a less fit gene could be selected
@@ -365,11 +347,23 @@ class LatentGenePool(Module):
             latents = maybe_l2norm(latents, dim = -1)
 
         self.num_latents = num_latents
-        self.num_latent_sets = num_latent_sets
+        self.needs_latent_gate = num_latent_sets > 1
         self.latents = nn.Parameter(latents, requires_grad = False)
 
         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 0. < frac_tournaments < 1.
@@ -471,6 +465,7 @@ class LatentGenePool(Module):
     def forward(
         self,
         *args,
+        state: Tensor | None = None,
         latent_id: int | None = None,
         net: Module | None = None,
         **kwargs,
@@ -487,6 +482,15 @@ 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'
+
+            gates = self.to_latent_gate(state)
+            latent = einsum(latent, gates, 'n g, b n -> b g')
+        else:
+            assert latent.shape[0] == 1
+            latent = latent[0]
+
         if not exists(net):
             return latent
 

{evolutionary_policy_optimization-0.0.10.dist-info → evolutionary_policy_optimization-0.0.11.dist-info}/METADATA

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

@@ -0,0 +1,7 @@
+evolutionary_policy_optimization/__init__.py,sha256=A07bhbBI_p-GlSTkI15pioQ1XgtJ0V4tBN6v3vs2nuU,115
+evolutionary_policy_optimization/epo.py,sha256=JGow9ofx7IgFy7QNL0dL0K_SCL_bVkBUznMG8aSGM9Q,15591
+evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
+evolutionary_policy_optimization-0.0.11.dist-info/METADATA,sha256=fkouRBZU5nrPgHt0eT5izSHdOiYGAg67N5Gn3t039mQ,4357
+evolutionary_policy_optimization-0.0.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.11.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.11.dist-info/RECORD,,

evolutionary_policy_optimization-0.0.10.dist-info/RECORD

@@ -1,7 +0,0 @@
-evolutionary_policy_optimization/__init__.py,sha256=A07bhbBI_p-GlSTkI15pioQ1XgtJ0V4tBN6v3vs2nuU,115
-evolutionary_policy_optimization/epo.py,sha256=66GOQq8_s5kmQI7G-2Z0J_0g4E5QarjQPJfWEP7mmKg,15442
-evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
-evolutionary_policy_optimization-0.0.10.dist-info/METADATA,sha256=bD3fw2Zw1IxhfkCvzjsRhODyL_XIC5ZsvNQqFbZXNc4,4357
-evolutionary_policy_optimization-0.0.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.10.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.10.dist-info/RECORD,,