evolutionary-policy-optimization 0.1.12__tar.gz → 0.1.15__tar.gz

{evolutionary_policy_optimization-0.1.12 → evolutionary_policy_optimization-0.1.15}/PKG-INFO

@@ -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 → evolutionary_policy_optimization-0.1.15}/README.md

@@ -173,4 +173,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 → evolutionary_policy_optimization-0.1.15}/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 → evolutionary_policy_optimization-0.1.15}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.1.12"
+version = "0.1.15"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

evolutionary_policy_optimization-0.1.15/train_gym.py

@@ -0,0 +1,63 @@
+import torch
+
+from evolutionary_policy_optimization import (
+    EPO,
+    GymnasiumEnvWrapper
+)
+
+# gymnasium
+
+from shutil import rmtree
+import gymnasium as gym
+
+env = gym.make(
+    'LunarLander-v3',
+    render_mode = 'rgb_array'
+)
+
+rmtree('./recordings', ignore_errors = True)
+
+env = gym.wrappers.RecordVideo(
+    env = env,
+    video_folder = './recordings',
+    name_prefix = 'lunar-video',
+    episode_trigger = lambda eps_num: (eps_num % 250) == 0,
+    disable_logger = True
+)
+
+env = GymnasiumEnvWrapper(env)
+
+# epo
+
+agent = env.to_epo_agent(
+    num_latents = 1,
+    dim_latent = 32,
+    actor_dim = 128,
+    actor_mlp_depth = 2,
+    critic_dim = 256,
+    critic_mlp_depth = 4,
+    latent_gene_pool_kwargs = dict(
+        frac_natural_selected = 0.5,
+        frac_tournaments = 0.5
+    ),
+    accelerate_kwargs = dict(
+        cpu = False
+    ),
+    actor_optim_kwargs = dict(
+        cautious_factor = 0.1,
+    ),
+    critic_optim_kwargs = dict(
+        cautious_factor = 0.1,
+    ),
+)
+
+epo = EPO(
+    agent,
+    episodes_per_latent = 50,
+    max_episode_length = 500,
+    action_sample_temperature = 1.,
+)
+
+epo(agent, env, num_learning_cycles = 100)
+
+agent.save('./agent.pt', overwrite = True)

evolutionary_policy_optimization-0.1.12/train_gym.py

@@ -1,44 +0,0 @@
-import torch
-
-from evolutionary_policy_optimization import (
-    EPO,
-    GymnasiumEnvWrapper
-)
-
-# gymnasium
-
-import gymnasium as gym
-
-env = gym.make(
-    'LunarLander-v3',
-    render_mode = 'rgb_array'
-)
-
-env = GymnasiumEnvWrapper(env)
-
-# epo
-
-agent = env.to_epo_agent(
-    num_latents = 8,
-    dim_latent = 32,
-    actor_dim_hiddens = (256, 128),
-    critic_dim_hiddens = (256, 128, 64),
-    latent_gene_pool_kwargs = dict(
-        frac_natural_selected = 0.5,
-        frac_tournaments = 0.5
-    ),
-    accelerate_kwargs = dict(
-        cpu = False
-    )
-)
-
-epo = EPO(
-    agent,
-    episodes_per_latent = 5,
-    max_episode_length = 10,
-    action_sample_temperature = 1.,
-)
-
-epo(agent, env, num_learning_cycles = 5)
-
-agent.save('./agent.pt', overwrite = True)