evolutionary-policy-optimization 0.1.14__py3-none-any.whl → 0.1.16__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -76,6 +76,7 @@ def maybe(fn):
 def interface_torch_numpy(fn, device):
     # for a given function, move all inputs from torch tensor to numpy, and all outputs from numpy to torch tensor
 
+    @maybe
     def to_torch_tensor(t):
         if isinstance(t, (np.ndarray, np.float64)):
             t = from_numpy(np.array(t))
@@ -287,6 +288,54 @@ class PowerLawDist(Module):
 
         return self.values[sampled]
 
+# FiLM for latent to mlp conditioning
+
+class FiLM(Module):
+    def __init__(self, dim, dim_out):
+        super().__init__()
+        self.to_gamma = nn.Linear(dim, dim_out, bias = False)
+        self.to_beta = nn.Linear(dim, dim_out, bias = False)
+
+        nn.init.zeros_(self.to_gamma.weight)
+        nn.init.zeros_(self.to_beta.weight)
+
+    def forward(self, x, cond):
+        gamma, beta = self.to_gamma(cond), self.to_beta(cond)
+
+        return x * (gamma + 1.) + beta
+
+# layer integrated memory
+
+class DynamicLIMe(Module):
+    def __init__(
+        self,
+        dim,
+        num_layers
+    ):
+        super().__init__()
+        self.num_layers = num_layers
+
+        self.to_weights = nn.Sequential(
+            nn.RMSNorm(dim),
+            nn.Linear(dim, num_layers),
+            nn.Softmax(dim = -1)
+        )
+
+    def forward(
+        self,
+        x,
+        hiddens
+    ):
+
+        if not is_tensor(hiddens):
+            hiddens = stack(hiddens)
+
+        assert hiddens.shape[0] == self.num_layers, f'expected hiddens to have {self.num_layers} layers but received {tuple(hiddens.shape)} instead (first dimension must be layers)'
+
+        weights = self.to_weights(x)
+
+        return einsum(hiddens, weights, 'l b d, b l -> b d')
+
 # simple MLP networks, but with latent variables
 # the latent variables are the "genes" with the rest of the network as the scaffold for "gene expression" - as suggested in the paper
 
@@ -306,7 +355,7 @@ class MLP(Module):
         self.needs_latent = dim_latent > 0
 
         self.encode_latent = nn.Sequential(
-            Linear(dim_latent, dim),
+            Linear(dim_latent, dim * 2),
             nn.SiLU()
         ) if self.needs_latent else None
 
@@ -316,15 +365,28 @@ class MLP(Module):
 
         layers = []
 
-        for _ in range(depth):
+        for ind in range(depth):
+            is_first = ind == 0
+
+            film = None
+
+            if self.needs_latent:
+                film = FiLM(dim * 2, dim)
+
+            lime = DynamicLIMe(dim, num_layers = ind + 1) if not is_first else None
+
             layer = nn.Sequential(
-                nn.LayerNorm(dim, bias = False),
+                nn.RMSNorm(dim),
                 nn.Linear(dim, dim_hidden),
                 nn.SiLU(),
                 nn.Linear(dim_hidden, dim),
             )
 
-            layers.append(layer)
+            layers.append(ModuleList([
+                lime,
+                film,
+                layer
+            ]))
 
         # modules across layers
 
@@ -350,14 +412,23 @@ class MLP(Module):
 
             assert latent.shape[0] == x.shape[0], f'received state with batch size {x.shape[0]} but latent ids received had batch size {latent_id.shape[0]}'
 
-            x = x * latent
-
         # layers
 
-        for ind, layer in enumerate(self.layers, start = 1):
-            is_last = ind == len(self.layers)
+        prev_layer_inputs = [x]
+
+        for lime, film, layer in self.layers:
+
+            layer_inp = x
+
+            if exists(lime):
+                layer_inp = lime(x, prev_layer_inputs)
+
+            if exists(film):
+                layer_inp = film(layer_inp, latent)
+
+            x = layer(layer_inp) + x
 
-            x = layer(x) + x
+            prev_layer_inputs.append(x)
 
         return x
 
@@ -385,7 +456,7 @@ class Actor(Module):
         self.mlp = MLP(dim = dim, depth = mlp_depth, dim_latent = dim_latent)
 
         self.to_out = nn.Sequential(
-            nn.LayerNorm(dim, bias = False),
+            nn.RMSNorm(dim),
             nn.Linear(dim, num_actions, bias = False),
         )
 
@@ -426,7 +497,7 @@ class Critic(Module):
 
         self.mlp = MLP(dim = dim, depth = mlp_depth, dim_latent = dim_latent)
 
-        self.final_norm = nn.LayerNorm(dim, bias = False)
+        self.final_norm = nn.RMSNorm(dim)
 
         self.to_pred = HLGaussLayer(
             dim = dim,

{evolutionary_policy_optimization-0.1.14.dist-info → evolutionary_policy_optimization-0.1.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.14
+Version: 0.1.16
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -118,8 +118,10 @@ agent = create_agent(
     num_latents = 16,
     dim_latent = 32,
     actor_num_actions = 5,
-    actor_dim_hiddens = (256, 128),
-    critic_dim_hiddens = (256, 128, 64),
+    actor_dim = 256,
+    actor_mlp_depth = 2,
+    critic_dim = 256,
+    critic_mlp_depth = 3,
     latent_gene_pool_kwargs = dict(
         frac_natural_selected = 0.5
     )
@@ -226,4 +228,13 @@ agent.load('./agent.pt')
 }
 ```
 
+```bibtex
+@inproceedings{Gerasimov2025YouDN,
+    title   = {You Do Not Fully Utilize Transformer's Representation Capacity},
+    author  = {Gleb Gerasimov and Yaroslav Aksenov and Nikita Balagansky and Viacheslav Sinii and Daniil Gavrilov},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:276317819}
+}
+```
+
 *Evolution is cleverer than you are.* - Leslie Orgel

{evolutionary_policy_optimization-0.1.14.dist-info → evolutionary_policy_optimization-0.1.16.dist-info}/RECORD

@@ -1,10 +1,10 @@
 evolutionary_policy_optimization/__init__.py,sha256=NyiYDYU7DlpmOTM7xiBQET3r1WwX0ebrgMCBLSQrW3c,288
 evolutionary_policy_optimization/distributed.py,sha256=7KgZdeS_wxBHo_du9XZFB1Cu318J-Bp66Xdr6Log_20,2423
 evolutionary_policy_optimization/env_wrappers.py,sha256=bDL06o9_b1iW6k3fw2xifnOnYlzs643tdW6Yv2gsIdw,803
-evolutionary_policy_optimization/epo.py,sha256=RTMVCo1joMEKIkqTQLsLgTeOuJVbvkNbX9hcOOL0oCw,46088
+evolutionary_policy_optimization/epo.py,sha256=kM1jKH3Jl0jB7iF7BKToY_xcNt6j5DkooP5k_VRUDDs,47791
 evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
 evolutionary_policy_optimization/mock_env.py,sha256=TLyyRm6tOD0Kdn9QqJJQriaSnsR-YmNQHo4OohmZFG4,1410
-evolutionary_policy_optimization-0.1.14.dist-info/METADATA,sha256=TCj1gzhViuNYP-TnVHCcvWHkTIyqiz-AAO-xhoVahBo,7625
-evolutionary_policy_optimization-0.1.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.1.14.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.1.14.dist-info/RECORD,,
+evolutionary_policy_optimization-0.1.16.dist-info/METADATA,sha256=JrwLijMYTjJdjoTQb-LCYyTuEibgffGUEEPRQEIsOhY,7979
+evolutionary_policy_optimization-0.1.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.1.16.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.1.16.dist-info/RECORD,,