evolutionary-policy-optimization 0.0.58__tar.gz → 0.0.61__tar.gz

{evolutionary_policy_optimization-0.0.58 → evolutionary_policy_optimization-0.0.61}/PKG-INFO

@@ -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 → evolutionary_policy_optimization-0.0.61}/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 → evolutionary_policy_optimization-0.0.61}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.0.58"
+version = "0.0.61"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -15,7 +15,7 @@ keywords = [
     'genetic algorithms',
     'robotics'
   ]
-classifiers=[
+classifiers = [
     'Development Status :: 4 - Beta',
     'Intended Audience :: Developers',
     'Topic :: Scientific/Engineering :: Artificial Intelligence',
@@ -25,7 +25,7 @@ classifiers=[
 
 dependencies = [
     "adam-atan2-pytorch",
-    'assoc-scan',
+    'assoc-scan>=0.0.2',
     'einx>=0.3.0',
     'einops>=0.8.1',
     'ema-pytorch>=0.7.7',