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

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)
@@ -919,6 +943,8 @@ class Agent(Module):
             eps_clip = 0.4
         ),
         use_improved_critic_loss = True,
+        shrink_and_perturb_every = None,
+        shrink_and_perturb_kwargs: dict = dict(),
         ema_kwargs: dict = dict(),
         actor_optim_kwargs: dict = dict(),
         critic_optim_kwargs: dict = dict(),
@@ -983,6 +1009,12 @@ class Agent(Module):
 
         self.latent_optim = optim_klass(latent_gene_pool.parameters(), lr = latent_lr, **latent_optim_kwargs) if exists(latent_gene_pool) and not latent_gene_pool.frozen_latents else None
 
+        # shrink and perturb every
+
+        self.should_noise_weights = exists(shrink_and_perturb_every)
+        self.shrink_and_perturb_every = shrink_and_perturb_every
+        self.shrink_and_perturb_ = partial(shrink_and_perturb_, **shrink_and_perturb_kwargs)
+
         # promotes latents to be farther apart for diversity maintenance
 
         self.has_diversity_loss = diversity_aux_loss_weight > 0.
@@ -992,7 +1024,7 @@ class Agent(Module):
 
         self.unwrap_model = identity if not wrap_with_accelerate else self.accelerate.unwrap_model
 
-        dummy = tensor(0)
+        step = tensor(0)
 
         self.clip_grad_norm_ = nn.utils.clip_grad_norm_
 
@@ -1020,15 +1052,15 @@ class Agent(Module):
             if exists(self.critic_ema):
                 self.critic_ema.to(self.accelerate.device)
 
-            dummy = dummy.to(self.accelerate.device)
+            step = step.to(self.accelerate.device)
 
         # device tracking
 
-        self.register_buffer('dummy', dummy)
+        self.register_buffer('step', step)
 
     @property
     def device(self):
-        return self.dummy.device
+        return self.step.device
 
     @property
     def unwrapped_latent_gene_pool(self):
@@ -1278,6 +1310,16 @@ class Agent(Module):
         if self.has_latent_genes:
             self.latent_gene_pool.genetic_algorithm_step(fitness_scores)
 
+        # maybe shrink and perturb
+
+        if self.should_noise_weights and divisible_by(self.step.item(), self.shrink_and_perturb_every):
+            self.shrink_and_perturb_(self.actor)
+            self.shrink_and_perturb_(self.critic)
+
+        # increment step
+
+        self.step.add_(1)
+
 # reinforcement learning related - ppo
 
 def actor_loss(

{evolutionary_policy_optimization-0.1.15.dist-info → evolutionary_policy_optimization-0.1.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.15
+Version: 0.1.17
 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.dist-info → evolutionary_policy_optimization-0.1.17.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=iUxd7gbT1GPGwso4utTaxgtjcxvvGNA8AGGUpSOImTM,47108
+evolutionary_policy_optimization/epo.py,sha256=TW-9l9oRN8XQZxSeG5Glkk4rWuxO9JOjjRJO7hJgHZs,48433
 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.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,,
+evolutionary_policy_optimization-0.1.17.dist-info/METADATA,sha256=-2Lt54pfQkw2_LAwB8pjcvLsc4vSMVQjAt-ZiObCRdw,7979
+evolutionary_policy_optimization-0.1.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.1.17.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.1.17.dist-info/RECORD,,