evolutionary-policy-optimization 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

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.dist-info → evolutionary_policy_optimization-0.1.10.dist-info}/METADATA

@@ -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.dist-info → evolutionary_policy_optimization-0.1.10.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=Ua0o4Xe-Z6gy76-nbB1yKndePGurSwW_otXXrrJWhgc,44835
+evolutionary_policy_optimization/epo.py,sha256=lYK6gfkAfYA5lm2fB0XC1UjvtjmlKCsvQNFrigj4JN0,45669
 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.8.dist-info/METADATA,sha256=tEVMyHVZjknJMQ0QEIVJhMj6QTDYW5Uqcq6nqa7LHpo,7088
-evolutionary_policy_optimization-0.1.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.1.8.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.1.8.dist-info/RECORD,,
+evolutionary_policy_optimization-0.1.10.dist-info/METADATA,sha256=X2DgvM3fkJej6brFyPXSgR5qX-n7Mdg0FaB79AI_3l8,7317
+evolutionary_policy_optimization-0.1.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.1.10.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.1.10.dist-info/RECORD,,