evolutionary-policy-optimization 0.0.58__py3-none-any.whl → 0.0.61__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -127,15 +127,9 @@ def calc_generalized_advantage_estimate(
     delta = rewards + gamma * values_next * masks - values
     gates = gamma * lam * masks
 
-    gates, delta = gates[..., :, None], delta[..., :, None]
-
     scan = AssocScan(reverse = True, use_accelerated = use_accelerated)
 
-    gae = scan(gates, delta)
-
-    gae = gae[..., :, 0]
-
-    return gae
+    return scan(gates, delta)    
 
 # evolution related functions
 
@@ -845,6 +839,8 @@ class Agent(Module):
         if is_distributed():
             memories = map(partial(all_gather_variable_dim, dim = 0), memories)
 
+            fitness_scores = all_gather_variable_dim(fitness_scores, dim = 0)
+
         (
             episode_ids,
             states,
@@ -958,7 +954,10 @@ def actor_loss(
     advantages,     # Float[b]
     eps_clip = 0.2,
     entropy_weight = .01,
+    eps = 1e-5
 ):
+    batch = logits.shape[0]
+
     log_probs = gather_log_prob(logits, actions)
 
     ratio = (log_probs - old_log_probs).exp()
@@ -967,6 +966,8 @@ def actor_loss(
 
     clipped_ratio = ratio.clamp(min = 1. - eps_clip, max = 1. + eps_clip)
 
+    advantages = F.layer_norm(advantages, (batch,), eps = eps)
+
     actor_loss = -torch.min(clipped_ratio * advantages, ratio * advantages)
 
     # add entropy loss for exploration

{evolutionary_policy_optimization-0.0.58.dist-info → evolutionary_policy_optimization-0.0.61.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.58
+Version: 0.0.61
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -35,7 +35,7 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: adam-atan2-pytorch
-Requires-Dist: assoc-scan
+Requires-Dist: assoc-scan>=0.0.2
 Requires-Dist: einops>=0.8.1
 Requires-Dist: einx>=0.3.0
 Requires-Dist: ema-pytorch>=0.7.7

{evolutionary_policy_optimization-0.0.58.dist-info → evolutionary_policy_optimization-0.0.61.dist-info}/RECORD

@@ -1,9 +1,9 @@
 evolutionary_policy_optimization/__init__.py,sha256=0q0aBuFgWi06MLMD8FiHzBYQ3_W4LYWrwmCtF3u5H2A,201
 evolutionary_policy_optimization/distributed.py,sha256=7KgZdeS_wxBHo_du9XZFB1Cu318J-Bp66Xdr6Log_20,2423
-evolutionary_policy_optimization/epo.py,sha256=2PlLBrhX5_kkV4NOhxC5VB_plR-_jw3krrmMSjWAiCY,35288
+evolutionary_policy_optimization/epo.py,sha256=kFT49rJdcmaDehfpx3YyhYhvAcp7S-gRWDkS2y20Q2Y,35377
 evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
 evolutionary_policy_optimization/mock_env.py,sha256=202KJ5g57wQvOzhGYzgHfBa7Y2do5uuDvl5kFg5o73g,934
-evolutionary_policy_optimization-0.0.58.dist-info/METADATA,sha256=ni9pyD5CS2QbBmwEUhCYt_RAxz3YnnOv3IY3ywNagSs,6213
-evolutionary_policy_optimization-0.0.58.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.58.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.58.dist-info/RECORD,,
+evolutionary_policy_optimization-0.0.61.dist-info/METADATA,sha256=3IbcY9kg71P6lTNxZaRBw3IYfDjcK4uTJJaFRD0Skwg,6220
+evolutionary_policy_optimization-0.0.61.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.61.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.61.dist-info/RECORD,,