evolutionary-policy-optimization 0.1.14__tar.gz → 0.1.16__tar.gz

{evolutionary_policy_optimization-0.1.14 → evolutionary_policy_optimization-0.1.16}/PKG-INFO

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

@@ -65,8 +65,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
     )
@@ -173,4 +175,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 → evolutionary_policy_optimization-0.1.16}/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 → evolutionary_policy_optimization-0.1.16}/pyproject.toml

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

{evolutionary_policy_optimization-0.1.14 → evolutionary_policy_optimization-0.1.16}/tests/test_epo.py

@@ -26,8 +26,8 @@ def test_readme(
 
     state = torch.randn(2, 512)
 
-    actor = Actor(dim_state = 512, dim_hiddens = (256, 128), num_actions = 4, dim_latent = 32)
-    critic = Critic(dim_state = 512, dim_hiddens = (256, 128, 64), dim_latent = 32)
+    actor = Actor(dim_state = 512, dim = 256, mlp_depth = 2, num_actions = 4, dim_latent = 32)
+    critic = Critic(dim_state = 512, dim = 256, mlp_depth = 4, dim_latent = 32)
 
     latent = latent_pool(latent_id = latent_ids, state = state)
 
@@ -55,8 +55,10 @@ def test_create_agent(
         num_latents = 128,
         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 = 4
     )
 
     state = torch.randn(2, 512)
@@ -98,8 +100,10 @@ def test_e2e_with_mock_env(
         num_latents = num_latents,
         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 = 4,
         use_critic_ema = use_critic_ema,
         diversity_aux_loss_weight = diversity_aux_loss_weight,
         critic_kwargs = dict(

{evolutionary_policy_optimization-0.1.14 → evolutionary_policy_optimization-0.1.16}/train_gym.py

@@ -21,7 +21,7 @@ env = gym.wrappers.RecordVideo(
     env = env,
     video_folder = './recordings',
     name_prefix = 'lunar-video',
-    episode_trigger = lambda eps_num: (eps_num % 250) == 0,
+    episode_trigger = lambda eps_num: (eps_num % (250 * 4)) == 0,
     disable_logger = True
 )
 
@@ -30,12 +30,12 @@ env = GymnasiumEnvWrapper(env)
 # epo
 
 agent = env.to_epo_agent(
-    num_latents = 1,
+    num_latents = 8,
     dim_latent = 32,
     actor_dim = 128,
-    actor_mlp_depth = 2,
+    actor_mlp_depth = 3,
     critic_dim = 256,
-    critic_mlp_depth = 4,
+    critic_mlp_depth = 5,
     latent_gene_pool_kwargs = dict(
         frac_natural_selected = 0.5,
         frac_tournaments = 0.5
@@ -53,7 +53,7 @@ agent = env.to_epo_agent(
 
 epo = EPO(
     agent,
-    episodes_per_latent = 50,
+    episodes_per_latent = 5,
     max_episode_length = 500,
     action_sample_temperature = 1.,
 )