evolutionary-policy-optimization 0.0.4__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
@@ -85,9 +87,9 @@ def critic_loss(
 # generalized advantage estimate
 
 def calc_generalized_advantage_estimate(
-    rewards: Float['g n'],
-    values: Float['g n+1'],
-    masks: Bool['n'],
+    rewards, # Float[g n]
+    values,  # Float[g n+1]
+    masks,   # Bool[n]
     gamma = 0.99,
     lam = 0.95,
     use_accelerated = None
@@ -218,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):
@@ -368,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.4.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.4
+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

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.4.dist-info/RECORD

@@ -1,7 +0,0 @@
-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,,