evolutionary-policy-optimization 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from collections import namedtuple
+
 import torch
 from torch import nn, cat
 import torch.nn.functional as F
@@ -9,6 +11,8 @@ from torch.nn import Linear, Module, ModuleList
 
 from einops import rearrange, repeat
 
+from assoc_scan import AssocScan
+
 # helpers
 
 def exists(v):
@@ -56,18 +60,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 +84,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(
@@ -183,6 +220,100 @@ class MLP(Module):
 
         return x
 
+# actor, critic, and agent (actor + critic)
+# eventually, should just create a separate repo and aggregate all the MLP related architectures
+
+class Actor(Module):
+    def __init__(
+        self,
+        dim_in,
+        num_actions,
+        dim_hiddens: tuple[int, ...],
+        dim_latent = 0,
+    ):
+        super().__init__()
+
+        assert len(dim_hiddens) >= 2
+        dim_first, *_, dim_last = dim_hiddens
+
+        self.init_layer = nn.Sequential(
+            nn.Linear(dim_in, dim_first),
+            nn.SiLU()
+        )
+
+        self.mlp = MLP(dims = dim_hiddens, dim_latent = dim_latent)
+
+        self.to_out = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim_last, num_actions),
+        )
+
+    def forward(
+        self,
+        state,
+        latent
+    ):
+
+        hidden = self.init_layer(state)
+
+        hidden = self.mlp(state, latent)
+
+        return self.to_out(hidden)
+
+class Critic(Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_hiddens: tuple[int, ...],
+        dim_latent = 0,
+    ):
+        super().__init__()
+
+        assert len(dim_hiddens) >= 2
+        dim_first, *_, dim_last = dim_hiddens
+
+        self.init_layer = nn.Sequential(
+            nn.Linear(dim_in, dim_first),
+            nn.SiLU()
+        )
+
+        self.mlp = MLP(dims = dim_hiddens, dim_latent = dim_latent)
+
+        self.to_out = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim_last, 1),
+            Rearrange('... 1 -> ...')
+        )
+
+    def forward(
+        self,
+        state,
+        latent
+    ):
+
+        hidden = self.init_layer(state)
+
+        hidden = self.mlp(state, latent)
+
+        return self.to_out(hidden)
+
+class Agent(Module):
+    def __init__(
+        self,
+        actor: Actor,
+        critic: Critic,
+    ):
+        super().__init__()
+
+        self.actor = actor
+        self.critic = critic
+
+    def forward(
+        self,
+        memories: list[Memory]
+    ):
+        raise NotImplementedError
+
 # classes
 
 class LatentGenePool(Module):
@@ -333,3 +464,35 @@ class LatentGenePool(Module):
             latent = latent,
             **kwargs
         )
+
+# EPO - which is just PPO with natural selection of a population of latent variables conditioning the agent
+# the tricky part is that the latent ids for each episode / trajectory needs to be tracked
+
+Memory = namedtuple('Memory', [
+    'state',
+    'latent_gene_id',
+    'action',
+    'log_prob',
+    'reward',
+    'values',
+    'done'
+])
+
+class EPO(Module):
+
+    def __init__(
+        self,
+        agent: Agent,
+        latent_gene_pool: LatentGenePool
+    ):
+        super().__init__()
+
+        self.agent = agent
+        self.latent_gene_pool = latent_gene_pool
+
+    def forward(
+        self,
+        env
+    ) -> list[Memory]:
+
+        raise NotImplementedError

{evolutionary_policy_optimization-0.0.3.dist-info → evolutionary_policy_optimization-0.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.3
+Version: 0.0.5
 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.5.dist-info/RECORD

@@ -0,0 +1,7 @@
+evolutionary_policy_optimization/__init__.py,sha256=Nu-_OMA8abe4AiW9Gw6MvbZH0JZpMHMqjeXmkC9-7UU,81
+evolutionary_policy_optimization/epo.py,sha256=lDhMV535MhUw1di7D7RM-Rr_J6aiuLqV-puh4EaNCd8,13455
+evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
+evolutionary_policy_optimization-0.0.5.dist-info/METADATA,sha256=uzkB4DrpzLLxbMEeiTID4CDxDxmEX1pO9fabwryDQcY,4098
+evolutionary_policy_optimization-0.0.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.5.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.5.dist-info/RECORD,,

evolutionary_policy_optimization-0.0.3.dist-info/RECORD

@@ -1,7 +0,0 @@
-evolutionary_policy_optimization/__init__.py,sha256=Nu-_OMA8abe4AiW9Gw6MvbZH0JZpMHMqjeXmkC9-7UU,81
-evolutionary_policy_optimization/epo.py,sha256=gB9_ZPtVjjFilZEeXX1ZWh67ZSd-OBt6Xs6WyfjRSxI,9967
-evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
-evolutionary_policy_optimization-0.0.3.dist-info/METADATA,sha256=qlQ1NxiA-AVWiNlpDcMI9hVqBg1g5DY9AvISpnrKkV0,4072
-evolutionary_policy_optimization-0.0.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.3.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.3.dist-info/RECORD,,