evolutionary-policy-optimization 0.1.5__tar.gz → 0.1.7__tar.gz

{evolutionary_policy_optimization-0.1.5 → evolutionary_policy_optimization-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.5
+Version: 0.1.7
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -185,4 +185,15 @@ agent.load('./agent.pt')
 }
 ```
 
+```bibtex
+@article{Banerjee2022BoostingEI,
+    title   = {Boosting Exploration in Actor-Critic Algorithms by Incentivizing Plausible Novel States},
+    author  = {Chayan Banerjee and Zhiyong Chen and Nasimul Noman},
+    journal = {2023 62nd IEEE Conference on Decision and Control (CDC)},
+    year    = {2022},
+    pages   = {7009-7014},
+    url     = {https://api.semanticscholar.org/CorpusID:252682944}
+}
+```
+
 *Evolution is cleverer than you are.* - Leslie Orgel

{evolutionary_policy_optimization-0.1.5 → evolutionary_policy_optimization-0.1.7}/README.md

@@ -132,4 +132,15 @@ agent.load('./agent.pt')
 }
 ```
 
+```bibtex
+@article{Banerjee2022BoostingEI,
+    title   = {Boosting Exploration in Actor-Critic Algorithms by Incentivizing Plausible Novel States},
+    author  = {Chayan Banerjee and Zhiyong Chen and Nasimul Noman},
+    journal = {2023 62nd IEEE Conference on Decision and Control (CDC)},
+    year    = {2022},
+    pages   = {7009-7014},
+    url     = {https://api.semanticscholar.org/CorpusID:252682944}
+}
+```
+
 *Evolution is cleverer than you are.* - Leslie Orgel

{evolutionary_policy_optimization-0.1.5 → evolutionary_policy_optimization-0.1.7}/evolutionary_policy_optimization/epo.py

@@ -364,11 +364,37 @@ class Critic(Module):
             hl_gauss_loss = hl_gauss_loss_kwargs
         )
 
+        self.use_regression = use_regression
+
+        hl_gauss_loss = self.to_pred.hl_gauss_loss
+
+        self.maybe_bins_to_value = hl_gauss_loss if not use_regression else identity
+        self.loss_fn = hl_gauss_loss if not use_regression else F.mse_loss
+
+    def forward_for_loss(
+        self,
+        state,
+        latent,
+        old_values,
+        target,
+        eps_clip = 0.4
+    ):
+        logits = self.forward(state, latent, return_logits = True)
+
+        value = self.maybe_bins_to_value(logits)
+
+        clipped_value = old_values + (value - old_values).clamp(1. - eps_clip, 1. + eps_clip)
+
+        loss = self.loss_fn(logits, target, reduction = 'none')
+        clipped_loss = self.loss_fn(clipped_value, target, reduction = 'none')
+
+        return torch.max(loss, clipped_loss).mean()
+
     def forward(
         self,
         state,
         latent,
-        target = None
+        return_logits = False
     ):
 
         hidden = self.init_layer(state)
@@ -377,7 +403,8 @@ class Critic(Module):
 
         hidden = self.final_act(hidden)
 
-        return self.to_pred(hidden, target = target)
+        pred_kwargs = dict(return_logits = return_logits) if not self.use_regression else dict()
+        return self.to_pred(hidden, **pred_kwargs)
 
 # criteria for running genetic algorithm
 
@@ -740,6 +767,9 @@ class Agent(Module):
             entropy_weight = .01,
             norm_advantages = True
         ),
+        critic_loss_kwargs: dict = dict(
+            eps_clip = 0.4
+        ),
         ema_kwargs: dict = dict(),
         actor_optim_kwargs: dict = dict(),
         critic_optim_kwargs: dict = dict(),
@@ -778,9 +808,13 @@ class Agent(Module):
 
         # gae function
 
-        self.actor_loss = partial(actor_loss, **actor_loss_kwargs)
         self.calc_gae = partial(calc_generalized_advantage_estimate, **calc_gae_kwargs)
 
+        # actor critic loss related
+
+        self.actor_loss = partial(actor_loss, **actor_loss_kwargs)
+        self.critic_loss_kwargs = critic_loss_kwargs
+
         # fitness score related
 
         self.get_fitness_scores = get_fitness_scores
@@ -809,7 +843,11 @@ class Agent(Module):
 
         dummy = tensor(0)
 
+        self.clip_grad_norm_ = nn.utils.clip_grad_norm_
+
         if wrap_with_accelerate:
+            self.clip_grad_norm_ = self.accelerate.clip_grad_norm_
+
             (
                 self.actor,
                 self.critic,
@@ -1036,23 +1074,25 @@ class Agent(Module):
                 actor_loss.backward()
 
                 if exists(self.has_grad_clip):
-                    self.accelerate.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
+                    self.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
 
                 self.actor_optim.step()
                 self.actor_optim.zero_grad()
 
                 # learn critic with maybe classification loss
 
-                critic_loss = self.critic(
+                critic_loss = self.critic.forward_for_loss(
                     states,
                     latents,
-                    target = advantages + old_values
+                    old_values = old_values,
+                    target = advantages + old_values,
+                    **self.critic_loss_kwargs
                 )
 
                 critic_loss.backward()
 
                 if exists(self.has_grad_clip):
-                    self.accelerate.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm)
+                    self.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm)
 
                 self.critic_optim.step()
                 self.critic_optim.zero_grad()
@@ -1076,7 +1116,7 @@ class Agent(Module):
                     (diversity_loss * self.diversity_aux_loss_weight).backward()
 
                 if exists(self.has_grad_clip):
-                    self.accelerate.clip_grad_norm_(self.latent_gene_pool.parameters(), self.max_grad_norm)
+                    self.clip_grad_norm_(self.latent_gene_pool.parameters(), self.max_grad_norm)
 
                 self.latent_optim.step()
                 self.latent_optim.zero_grad()

{evolutionary_policy_optimization-0.1.5 → evolutionary_policy_optimization-0.1.7}/pyproject.toml

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

{evolutionary_policy_optimization-0.1.5 → evolutionary_policy_optimization-0.1.7}/tests/test_epo.py

@@ -75,13 +75,15 @@ def test_create_agent(
 
 @pytest.mark.parametrize('frozen_latents', (False, True))
 @pytest.mark.parametrize('use_critic_ema', (False, True))
+@pytest.mark.parametrize('critic_use_regression', (False, True))
 @pytest.mark.parametrize('num_latents', (1, 8))
 @pytest.mark.parametrize('diversity_aux_loss_weight', (0., 1e-3))
 def test_e2e_with_mock_env(
     frozen_latents,
     use_critic_ema,
     num_latents,
-    diversity_aux_loss_weight
+    diversity_aux_loss_weight,
+    critic_use_regression
 ):
     from evolutionary_policy_optimization import create_agent, EPO, Env
 
@@ -94,11 +96,15 @@ def test_e2e_with_mock_env(
         critic_dim_hiddens = (256, 128, 64),
         use_critic_ema = use_critic_ema,
         diversity_aux_loss_weight = diversity_aux_loss_weight,
+        critic_kwargs = dict(
+            use_regression = critic_use_regression
+        ),
         latent_gene_pool_kwargs = dict(
             frozen_latents = frozen_latents,
             frac_natural_selected = 0.75,
             frac_tournaments = 0.9
-        )
+        ),
+        wrap_with_accelerate = False
     )
 
     epo = EPO(