evolutionary-policy-optimization 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -9,6 +9,8 @@ from torch.nn import Linear, Module, ModuleList
 
 from einops import rearrange, repeat
 
+from assoc_scan import AssocScan
+
 # helpers
 
 def exists(v):
@@ -56,18 +58,18 @@ def actor_loss(
 ):
     log_probs = gather_log_prob(logits, actions)
 
-    entropy = calc_entropy(logits)
-
     ratio = (log_probs - old_log_probs).exp()
 
-    clipped_ratio = ratio.clamp(min = 1. - eps_clip, max = 1. + eps_clip)
-
     # classic clipped surrogate loss from ppo
 
+    clipped_ratio = ratio.clamp(min = 1. - eps_clip, max = 1. + eps_clip)
+
     actor_loss = -torch.min(clipped_ratio * advantage, ratio * advantage)
 
     # add entropy loss for exploration
 
+    entropy = calc_entropy(logits)
+
     entropy_aux_loss = -entropy_weight * entropy
 
     return actor_loss + entropy_aux_loss
@@ -80,6 +82,39 @@ def critic_loss(
     discounted_values = advantages + old_values
     return F.mse_loss(pred_values, discounted_values)
 
+# generalized advantage estimate
+
+def calc_generalized_advantage_estimate(
+    rewards: Float['g n'],
+    values: Float['g n+1'],
+    masks: Bool['n'],
+    gamma = 0.99,
+    lam = 0.95,
+    use_accelerated = None
+
+):
+    assert values.shape[-1] == (rewards.shape[-1] + 1)
+
+    use_accelerated = default(use_accelerated, rewards.is_cuda)
+    device = rewards.device
+
+    masks = repeat(masks, 'n -> g n', g = rewards.shape[0])
+
+    values, values_next = values[:, :-1], values[:, 1:]
+
+    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
+
 # evolution related functions
 
 def crossover_latents(
@@ -319,7 +354,7 @@ class LatentGenePool(Module):
 
         # if only 1 latent, assume doing ablation and get lone gene
 
-        if self.num_latents == 1:
+        if not exists(latent_id) and self.num_latents == 1:
             latent_id = 0
 
         assert 0 <= latent_id < self.num_latents

{evolutionary_policy_optimization-0.0.2.dist-info → evolutionary_policy_optimization-0.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.2
+Version: 0.0.4
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -34,6 +34,7 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.8
+Requires-Dist: assoc-scan
 Requires-Dist: einops>=0.8.0
 Requires-Dist: torch>=2.2
 Requires-Dist: tqdm

evolutionary_policy_optimization-0.0.4.dist-info/RECORD

@@ -0,0 +1,7 @@
+evolutionary_policy_optimization/__init__.py,sha256=Nu-_OMA8abe4AiW9Gw6MvbZH0JZpMHMqjeXmkC9-7UU,81
+evolutionary_policy_optimization/epo.py,sha256=jW6wZ_IbTdO05agc9AghDHawLb0rStfOzHKpSh-vEe0,10783
+evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
+evolutionary_policy_optimization-0.0.4.dist-info/METADATA,sha256=ZmVUGRQkqOYs1fAyPXjyvIeyc_mShKVTfRVZsIE_Z1Q,4098
+evolutionary_policy_optimization-0.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.4.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.4.dist-info/RECORD,,

evolutionary_policy_optimization-0.0.2.dist-info/RECORD

@@ -1,7 +0,0 @@
-evolutionary_policy_optimization/__init__.py,sha256=Nu-_OMA8abe4AiW9Gw6MvbZH0JZpMHMqjeXmkC9-7UU,81
-evolutionary_policy_optimization/epo.py,sha256=6IcrAooUY8csFMv4ho5bf6TVAsk1cYyXB5hG3NA-jbA,9941
-evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
-evolutionary_policy_optimization-0.0.2.dist-info/METADATA,sha256=XqR4EZUWWXRjskU3_eHlwNzCeEddjm20un4E04xmNLk,4072
-evolutionary_policy_optimization-0.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.2.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.2.dist-info/RECORD,,