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

evolutionary_policy_optimization/epo.py

@@ -76,6 +76,14 @@ 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
 
+    def to_torch_tensor(t):
+        if isinstance(t, (np.ndarray, np.float64)):
+            t = from_numpy(np.array(t))
+        elif isinstance(t, (float, int, bool)):
+            t = tensor(t)
+
+        return t.to(device)
+
     @wraps(fn)
     def decorated_fn(*args, **kwargs):
 
@@ -83,7 +91,7 @@ def interface_torch_numpy(fn, device):
 
         out = fn(*args, **kwargs)
 
-        out = tree_map(lambda t: from_numpy(np.array(t)).to(device) if isinstance(t, (np.ndarray, np.float64)) else t, out)
+        out = tree_map(to_torch_tensor, out)
         return out
 
     return decorated_fn
@@ -279,43 +287,87 @@ 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
 
 class MLP(Module):
     def __init__(
         self,
-        dims: tuple[int, ...],
+        dim,
+        depth,
         dim_latent = 0,
+        expansion_factor = 2.
     ):
         super().__init__()
         dim_latent = default(dim_latent, 0)
 
-        assert len(dims) >= 2, 'must have at least two dimensions'
-
-        # add the latent to the first dim
-
-        first_dim, *rest_dims = dims
-        dims = (first_dim + dim_latent, *rest_dims)
-
         self.dim_latent = dim_latent
 
         self.needs_latent = dim_latent > 0
 
         self.encode_latent = nn.Sequential(
-            Linear(dim_latent, dim_latent),
+            Linear(dim_latent, dim),
             nn.SiLU()
         ) if self.needs_latent else None
 
-        # pairs of dimension
+        dim_hidden = int(dim * expansion_factor)
 
-        dim_pairs = tuple(zip(dims[:-1], dims[1:]))
+        # layers
 
-        # modules across layers
+        layers = []
+
+        for ind in range(depth):
+            is_first = ind == 0
 
-        layers = ModuleList([Linear(dim_in, dim_out) for dim_in, dim_out in dim_pairs])
+            lime = DynamicLIMe(dim, num_layers = ind + 1) if not is_first else None
+
+            layer = nn.Sequential(
+                nn.RMSNorm(dim),
+                nn.Linear(dim, dim_hidden),
+                nn.SiLU(),
+                nn.Linear(dim_hidden, dim),
+            )
 
-        self.layers = layers
+            layers.append(ModuleList([
+                lime,
+                layer
+            ]))
+
+        # modules across layers
+
+        self.layers = ModuleList(layers)
 
     def forward(
         self,
@@ -337,17 +389,22 @@ 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 = cat((x, latent), dim = -1)
+            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, layer in self.layers:
 
-            x = layer(x)
+            layer_inp = x
 
-            if not is_last:
-                x = F.silu(x)
+            if exists(lime):
+                layer_inp = lime(x, prev_layer_inputs)
+
+            x = layer(layer_inp) + x
+
+            prev_layer_inputs.append(x)
 
         return x
 
@@ -359,26 +416,24 @@ class Actor(Module):
         self,
         dim_state,
         num_actions,
-        dim_hiddens: tuple[int, ...],
+        dim,
+        mlp_depth,
         dim_latent = 0,
     ):
         super().__init__()
 
-        assert len(dim_hiddens) >= 2
-        dim_first, *_, dim_last = dim_hiddens
-
         self.dim_latent = dim_latent
 
         self.init_layer = nn.Sequential(
-            nn.Linear(dim_state, dim_first),
+            nn.Linear(dim_state, dim),
             nn.SiLU()
         )
 
-        self.mlp = MLP(dims = dim_hiddens, dim_latent = dim_latent)
+        self.mlp = MLP(dim = dim, depth = mlp_depth, dim_latent = dim_latent)
 
         self.to_out = nn.Sequential(
-            nn.SiLU(),
-            nn.Linear(dim_last, num_actions),
+            nn.RMSNorm(dim),
+            nn.Linear(dim, num_actions, bias = False),
         )
 
     def forward(
@@ -397,34 +452,31 @@ class Critic(Module):
     def __init__(
         self,
         dim_state,
-        dim_hiddens: tuple[int, ...],
+        dim,
+        mlp_depth,
         dim_latent = 0,
         use_regression = False,
         hl_gauss_loss_kwargs: dict = dict(
-            min_value = -10.,
-            max_value = 10.,
-            num_bins = 25,
-            sigma = 0.5
+            min_value = -100.,
+            max_value = 100.,
+            num_bins = 200
         )
     ):
         super().__init__()
 
-        assert len(dim_hiddens) >= 2
-        dim_first, *_, dim_last = dim_hiddens
-
         self.dim_latent = dim_latent
 
         self.init_layer = nn.Sequential(
-            nn.Linear(dim_state, dim_first),
+            nn.Linear(dim_state, dim),
             nn.SiLU()
         )
 
-        self.mlp = MLP(dims = dim_hiddens, dim_latent = dim_latent)
+        self.mlp = MLP(dim = dim, depth = mlp_depth, dim_latent = dim_latent)
 
-        self.final_act = nn.SiLU()
+        self.final_norm = nn.RMSNorm(dim)
 
         self.to_pred = HLGaussLayer(
-            dim = dim_last,
+            dim = dim,
             use_regression = use_regression,
             hl_gauss_loss = hl_gauss_loss_kwargs
         )
@@ -488,7 +540,7 @@ class Critic(Module):
 
         hidden = self.mlp(hidden, latent)
 
-        hidden = self.final_act(hidden)
+        hidden = self.final_norm(hidden)
 
         pred_kwargs = dict(return_logits = return_logits) if not self.use_regression else dict()
         return self.to_pred(hidden, **pred_kwargs)
@@ -843,16 +895,16 @@ class Agent(Module):
         critic: Critic,
         latent_gene_pool: LatentGenePool | None,
         optim_klass = AdoptAtan2,
-        actor_lr = 1e-4,
-        critic_lr = 1e-4,
+        actor_lr = 8e-4,
+        critic_lr = 8e-4,
         latent_lr = 1e-5,
-        actor_weight_decay = 1e-3,
-        critic_weight_decay = 1e-3,
+        actor_weight_decay = 5e-4,
+        critic_weight_decay = 5e-4,
         diversity_aux_loss_weight = 0.,
         use_critic_ema = True,
-        critic_ema_beta = 0.99,
-        max_grad_norm = 0.5,
-        batch_size = 16,
+        critic_ema_beta = 0.95,
+        max_grad_norm = 1.0,
+        batch_size = 32,
         calc_gae_kwargs: dict = dict(
             use_accelerated = False,
             gamma = 0.99,
@@ -1269,8 +1321,10 @@ def create_agent(
     num_latents,
     dim_latent,
     actor_num_actions,
-    actor_dim_hiddens: int | tuple[int, ...],
-    critic_dim_hiddens: int | tuple[int, ...],
+    actor_dim,
+    actor_mlp_depth,
+    critic_dim,
+    critic_mlp_depth,
     use_critic_ema = True,
     latent_gene_pool_kwargs: dict = dict(),
     actor_kwargs: dict = dict(),
@@ -1293,14 +1347,16 @@ def create_agent(
         num_actions = actor_num_actions,
         dim_state = dim_state,
         dim_latent = dim_latent,
-        dim_hiddens = actor_dim_hiddens,
+        dim = actor_dim,
+        mlp_depth = actor_mlp_depth,
         **actor_kwargs
     )
 
     critic = Critic(
         dim_state = dim_state,
         dim_latent = dim_latent,
-        dim_hiddens = critic_dim_hiddens,
+        dim = critic_dim,
+        mlp_depth = critic_mlp_depth,
         **critic_kwargs
     )
 
@@ -1475,7 +1531,7 @@ class EPO(Module):
                     log_prob,
                     reward,
                     value,
-                    tensor(terminated)
+                    terminated
                 )
 
                 memory = Memory(*tuple(t.cpu() for t in memory))
@@ -1487,7 +1543,7 @@ class EPO(Module):
             if not terminated:
                 # add bootstrap value if truncated
 
-                next_value = temp_batch_dim(self.agent.get_critic_values)(state, latent = latent)
+                next_value = temp_batch_dim(self.agent.get_critic_values)(state, latent = latent, use_ema_if_available = True, use_unwrapped_model = True)
 
                 memory_for_gae = memory._replace(
                     episode_id = invalid_episode,

{evolutionary_policy_optimization-0.1.12.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.12
+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.12.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=aOltJBkZVi2FxXao51zdfYaLynIi8T8v3qO1ex2HElg,46058
+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.12.dist-info/METADATA,sha256=B_eK4c8-dp4rr4K0HxRiNJqY6fK10XmwBbsPm-PE0_k,7625
-evolutionary_policy_optimization-0.1.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.1.12.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.1.12.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,,