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

evolutionary_policy_optimization/epo.py

@@ -287,6 +287,38 @@ class PowerLawDist(Module):
 
         return self.values[sampled]
 
+# 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.ReLU()
+        )
+
+    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
 
@@ -316,15 +348,22 @@ class MLP(Module):
 
         layers = []
 
-        for _ in range(depth):
+        for ind in range(depth):
+            is_first = ind == 0
+
+            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,
+                layer
+            ]))
 
         # modules across layers
 
@@ -354,10 +393,18 @@ class MLP(Module):
 
         # layers
 
-        for ind, layer in enumerate(self.layers, start = 1):
-            is_last = ind == len(self.layers)
+        prev_layer_inputs = [x]
+
+        for lime, layer in self.layers:
+
+            layer_inp = x
+
+            if exists(lime):
+                layer_inp = lime(x, prev_layer_inputs)
+
+            x = layer(layer_inp) + x
 
-            x = layer(x) + x
+            prev_layer_inputs.append(x)
 
         return x
 
@@ -385,7 +432,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 +473,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.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.14
+Version: 0.1.15
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -226,4 +226,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.15.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=iUxd7gbT1GPGwso4utTaxgtjcxvvGNA8AGGUpSOImTM,47108
 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.15.dist-info/METADATA,sha256=e8ofJe5rpGIyEiMd3mJBU-2VjOfFJ8TpGGv7adSKjRM,7962
+evolutionary_policy_optimization-0.1.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.1.15.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.1.15.dist-info/RECORD,,