evolutionary-policy-optimization 0.1.8__tar.gz → 0.1.10__tar.gz

{evolutionary_policy_optimization-0.1.8 → evolutionary_policy_optimization-0.1.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.8
+Version: 0.1.10
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -206,4 +206,13 @@ agent.load('./agent.pt')
 }
 ```
 
+```bibtex
+@article{Lee2024AnalysisClippedCritic
+    title   = {On Analysis of Clipped Critic Loss in Proximal Policy Gradient},
+    author  = {Yongjin Lee, Moonyoung Chung},
+    journal = {Authorea},
+    year    = {2024}
+}
+```
+
 *Evolution is cleverer than you are.* - Leslie Orgel

{evolutionary_policy_optimization-0.1.8 → evolutionary_policy_optimization-0.1.10}/README.md

@@ -153,4 +153,13 @@ agent.load('./agent.pt')
 }
 ```
 
+```bibtex
+@article{Lee2024AnalysisClippedCritic
+    title   = {On Analysis of Clipped Critic Loss in Proximal Policy Gradient},
+    author  = {Yongjin Lee, Moonyoung Chung},
+    journal = {Authorea},
+    year    = {2024}
+}
+```
+
 *Evolution is cleverer than you are.* - Leslie Orgel

{evolutionary_policy_optimization-0.1.8 → evolutionary_policy_optimization-0.1.10}/evolutionary_policy_optimization/epo.py

@@ -424,18 +424,40 @@ class Critic(Module):
         latent,
         old_values,
         target,
-        eps_clip = 0.4
+        eps_clip = 0.4,
+        use_improved = True
     ):
         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_fn = partial(self.loss_fn, reduction = 'none')
 
-        loss = self.loss_fn(logits, target, reduction = 'none')
-        clipped_loss = self.loss_fn(clipped_value, target, reduction = 'none')
+        if use_improved:
+            clipped_target = target.clamp(-eps_clip, eps_clip)
 
-        return torch.max(loss, clipped_loss).mean()
+            old_values_lo = old_values - eps_clip
+            old_values_hi = old_values + eps_clip
+
+            is_between = lambda lo, hi: (lo < value) & (value < hi)
+
+            clipped_loss = loss_fn(logits, clipped_target)
+            loss = loss_fn(logits, target)
+
+            value_loss = torch.where(
+                is_between(target, old_values_lo) | is_between(old_values_hi, target),
+                0.,
+                torch.min(loss, clipped_loss)
+            )
+        else:
+            clipped_value = old_values + (value - old_values).clamp(-eps_clip, eps_clip)
+
+            loss = loss_fn(logits, target)
+            clipped_loss = loss_fn(clipped_value, target)
+
+            value_loss = torch.max(loss, clipped_loss)
+
+        return value_loss.mean()
 
     def forward(
         self,
@@ -826,6 +848,7 @@ class Agent(Module):
         critic_loss_kwargs: dict = dict(
             eps_clip = 0.4
         ),
+        use_improved_critic_loss = True,
         ema_kwargs: dict = dict(),
         actor_optim_kwargs: dict = dict(),
         critic_optim_kwargs: dict = dict(),
@@ -871,6 +894,8 @@ class Agent(Module):
         self.actor_loss = partial(actor_loss, **actor_loss_kwargs)
         self.critic_loss_kwargs = critic_loss_kwargs
 
+        self.use_improved_critic_loss = use_improved_critic_loss
+
         # fitness score related
 
         self.get_fitness_scores = get_fitness_scores
@@ -1142,6 +1167,7 @@ class Agent(Module):
                     latents,
                     old_values = old_values,
                     target = advantages + old_values,
+                    use_improved = self.use_improved_critic_loss,
                     **self.critic_loss_kwargs
                 )
 

{evolutionary_policy_optimization-0.1.8 → evolutionary_policy_optimization-0.1.10}/pyproject.toml

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

{evolutionary_policy_optimization-0.1.8 → evolutionary_policy_optimization-0.1.10}/tests/test_epo.py

@@ -79,6 +79,7 @@ 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('use_improved_critic_loss', (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(
@@ -86,7 +87,8 @@ def test_e2e_with_mock_env(
     use_critic_ema,
     num_latents,
     diversity_aux_loss_weight,
-    critic_use_regression
+    critic_use_regression,
+    use_improved_critic_loss
 ):
     from evolutionary_policy_optimization import create_agent, EPO, Env
 
@@ -102,6 +104,7 @@ def test_e2e_with_mock_env(
         critic_kwargs = dict(
             use_regression = critic_use_regression
         ),
+        use_improved_critic_loss = use_improved_critic_loss,
         latent_gene_pool_kwargs = dict(
             frozen_latents = frozen_latents,
             frac_natural_selected = 0.75,