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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.16
+Version: 0.1.18
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization

{evolutionary_policy_optimization-0.1.16 → evolutionary_policy_optimization-0.1.18}/evolutionary_policy_optimization/epo.py

@@ -392,6 +392,8 @@ class MLP(Module):
 
         self.layers = ModuleList(layers)
 
+        self.final_lime = DynamicLIMe(dim, depth + 1)
+
     def forward(
         self,
         x,
@@ -430,7 +432,7 @@ class MLP(Module):
 
             prev_layer_inputs.append(x)
 
-        return x
+        return self.final_lime(x, prev_layer_inputs)
 
 # actor, critic, and agent (actor + critic)
 # eventually, should just create a separate repo and aggregate all the MLP related architectures
@@ -943,6 +945,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(),
@@ -1007,6 +1011,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.
@@ -1016,7 +1026,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_
 
@@ -1044,15 +1054,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):
@@ -1302,6 +1312,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.16 → evolutionary_policy_optimization-0.1.18}/pyproject.toml

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

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

@@ -85,13 +85,15 @@ def test_create_agent(
 @pytest.mark.parametrize('use_improved_critic_loss', (False, True))
 @pytest.mark.parametrize('num_latents', (1, 8))
 @pytest.mark.parametrize('diversity_aux_loss_weight', (0., 1e-3))
+@pytest.mark.parametrize('shrink_and_perturb_every', (None, 1))
 def test_e2e_with_mock_env(
     frozen_latents,
     use_critic_ema,
     num_latents,
     diversity_aux_loss_weight,
     critic_use_regression,
-    use_improved_critic_loss
+    use_improved_critic_loss,
+    shrink_and_perturb_every
 ):
     from evolutionary_policy_optimization import create_agent, EPO, Env
 
@@ -106,6 +108,7 @@ def test_e2e_with_mock_env(
         critic_mlp_depth = 4,
         use_critic_ema = use_critic_ema,
         diversity_aux_loss_weight = diversity_aux_loss_weight,
+        shrink_and_perturb_every = shrink_and_perturb_every,
         critic_kwargs = dict(
             use_regression = critic_use_regression
         ),
@@ -115,7 +118,7 @@ def test_e2e_with_mock_env(
             frac_natural_selected = 0.75,
             frac_tournaments = 0.9
         ),
-        wrap_with_accelerate = False
+        wrap_with_accelerate = False,
     )
 
     epo = EPO(