evolutionary-policy-optimization 0.0.57__py3-none-any.whl → 0.0.60__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
 
@@ -856,10 +850,12 @@ class Agent(Module):
             dones
         ) = memories
 
+        masks = 1. - dones.float()
+
         advantages = self.calc_gae(
             rewards[:-1],
             values,
-            dones[:-1],
+            masks[:-1],
         )
 
         valid_episode = episode_ids >= 0
@@ -956,7 +952,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()
@@ -965,6 +964,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.57.dist-info → evolutionary_policy_optimization-0.0.60.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.57
+Version: 0.0.60
 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.57.dist-info → evolutionary_policy_optimization-0.0.60.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=qPj5kRsISY1I6WjCc-ejpuiwOSxtPsSdMABmchXJ3s0,35252
+evolutionary_policy_optimization/epo.py,sha256=pD4j_oP7Cg8vSVQE34oJMcXcYN4oOi2_TtOAI5YbZDQ,35298
 evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
 evolutionary_policy_optimization/mock_env.py,sha256=202KJ5g57wQvOzhGYzgHfBa7Y2do5uuDvl5kFg5o73g,934
-evolutionary_policy_optimization-0.0.57.dist-info/METADATA,sha256=WBHRK98s_lzWbqG4ouq620ayykPF9SHUz3HdvsRUywc,6213
-evolutionary_policy_optimization-0.0.57.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.57.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.57.dist-info/RECORD,,
+evolutionary_policy_optimization-0.0.60.dist-info/METADATA,sha256=vfiTPTi00-ZPqXgeVn4yEzCcD8p2k61aDmp5YX99Uww,6220
+evolutionary_policy_optimization-0.0.60.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.60.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.60.dist-info/RECORD,,