evolutionary-policy-optimization 0.4.1__tar.gz → 0.4.2__tar.gz

{evolutionary_policy_optimization-0.4.1 → evolutionary_policy_optimization-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: evolutionary-policy-optimization
-Version: 0.4.1
+Version: 0.4.2
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -44,6 +44,7 @@ Requires-Dist: mean-conc-beta>=0.2.7
 Requires-Dist: torch>=2.2
 Requires-Dist: tqdm
 Requires-Dist: x-mlps-pytorch>=0.3.0
+Requires-Dist: x-ppo>=0.0.16
 Provides-Extra: examples
 Requires-Dist: fire; extra == 'examples'
 Requires-Dist: gymnasium[box2d]>=1.0.0; extra == 'examples'

{evolutionary_policy_optimization-0.4.1 → evolutionary_policy_optimization-0.4.2}/evolutionary_policy_optimization/epo.py

@@ -29,6 +29,7 @@ from torch.utils._pytree import tree_map
 from torch.utils.data import DataLoader, TensorDataset
 from tqdm import tqdm
 from x_mlps_pytorch import AttnResidualNormedMLP
+from x_ppo import hpo_actor_loss, ppo_actor_loss, spo_actor_loss
 
 from evolutionary_policy_optimization.distributed import all_gather, get_world_and_rank, is_distributed, maybe_barrier, maybe_sync_seed
 
@@ -1863,39 +1864,31 @@ def actor_loss(
     advantages,     # Float[b]
     eps_clip = 0.2,
     entropy_weight = .01,
-    eps = 1e-5,
     norm_advantages = True,
-    use_spo = False
+    use_spo = False,
+    use_hpo = False
 ):
-    batch = advantages.shape[0]
-
     log_probs = sum_to_batch(distr.log_prob(actions))
     entropy = sum_to_batch(distr.entropy())
 
-    ratio = (log_probs - old_log_probs).exp()
-
-    if norm_advantages:
-        advantages = F.layer_norm(advantages, (batch,), eps = eps)
-
-    if use_spo:
-        # simple policy optimization - line 14 Algorithm 1 https://arxiv.org/abs/2401.16025v9
-
-        actor_loss = - (
-            ratio * advantages -
-            advantages.abs() / (2 * eps_clip) * (ratio - 1.).square()
-        )
+    if use_hpo:
+        loss_fn = hpo_actor_loss
+    elif use_spo:
+        loss_fn = spo_actor_loss
     else:
-        # classic clipped surrogate loss from ppo
-
-        clipped_ratio = ratio.clamp(min = 1. - eps_clip, max = 1. + eps_clip)
-
-        actor_loss = -torch.min(clipped_ratio * advantages, ratio * advantages)
-
-    # add entropy loss for exploration
+        loss_fn = ppo_actor_loss
+
+    loss = loss_fn(
+        log_probs,
+        old_log_probs,
+        advantages,
+        eps_clip = eps_clip,
+        normalize_advantages = norm_advantages,
+    )
 
-    entropy_aux_loss = -entropy_weight * entropy
+    # entropy bonus
 
-    return (actor_loss + entropy_aux_loss).mean()
+    return (loss - entropy_weight * entropy).mean()
 
 # agent contains the actor, critic, and the latent genetic pool
 

{evolutionary_policy_optimization-0.4.1 → evolutionary_policy_optimization-0.4.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.4.1"
+version = "0.4.2"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -34,6 +34,7 @@ dependencies = [
     'torch>=2.2',
     'tqdm',
     "x-mlps-pytorch>=0.3.0",
+    "x-ppo>=0.0.16",
 ]
 
 [project.urls]