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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.15
+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
     )

{evolutionary_policy_optimization-0.1.15 → 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
     )

{evolutionary_policy_optimization-0.1.15 → 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,22 @@ 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):
@@ -301,7 +318,7 @@ class DynamicLIMe(Module):
         self.to_weights = nn.Sequential(
             nn.RMSNorm(dim),
             nn.Linear(dim, num_layers),
-            nn.ReLU()
+            nn.Softmax(dim = -1)
         )
 
     def forward(
@@ -338,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
 
@@ -351,6 +368,11 @@ class MLP(Module):
         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(
@@ -362,6 +384,7 @@ class MLP(Module):
 
             layers.append(ModuleList([
                 lime,
+                film,
                 layer
             ]))
 
@@ -389,19 +412,20 @@ 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
 
         prev_layer_inputs = [x]
 
-        for lime, layer in self.layers:
+        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
 
             prev_layer_inputs.append(x)

{evolutionary_policy_optimization-0.1.15 → evolutionary_policy_optimization-0.1.16}/pyproject.toml

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

{evolutionary_policy_optimization-0.1.15 → 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.15 → 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.,
 )