evolutionary-policy-optimization 0.0.12__py3-none-any.whl → 0.0.15__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -4,16 +4,17 @@ from collections import namedtuple
 
 import torch
 from torch import nn, cat
-import torch.nn.functional as F
-
 import torch.nn.functional as F
 from torch.nn import Linear, Module, ModuleList
+from torch.utils.data import TensorDataset, DataLoader
 
 from einops import rearrange, repeat, einsum
 from einops.layers.torch import Rearrange
 
 from assoc_scan import AssocScan
 
+from adam_atan2_pytorch import AdoptAtan2
+
 # helpers
 
 def exists(v):
@@ -49,42 +50,6 @@ def gather_log_prob(
     log_prob = log_probs.gather(-1, indices)
     return rearrange(log_prob, '... 1 -> ...')
 
-# reinforcement learning related - ppo
-
-def actor_loss(
-    logits,         # Float[b l]
-    old_log_probs,  # Float[b]
-    actions,        # Int[b]
-    advantages,     # Float[b]
-    eps_clip = 0.2,
-    entropy_weight = .01,
-):
-    log_probs = gather_log_prob(logits, actions)
-
-    ratio = (log_probs - old_log_probs).exp()
-
-    # 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
-
-def critic_loss(
-    pred_values,  # Float[b]
-    advantages,   # Float[b]
-    old_values    # Float[b]
-):
-    discounted_values = advantages + old_values
-    return F.mse_loss(pred_values, discounted_values)
-
 # generalized advantage estimate
 
 def calc_generalized_advantage_estimate(
@@ -500,46 +465,21 @@ class LatentGenePool(Module):
             **kwargs
         )
 
-# agent contains the actor, critic, and the latent genetic pool
-
-def create_agent(
-    dim_state,
-    num_latents,
-    dim_latent,
-    actor_num_actions,
-    actor_dim_hiddens: int | tuple[int, ...],
-    critic_dim_hiddens: int | tuple[int, ...],
-    num_latent_sets = 1
-) -> Agent:
-
-    actor = Actor(
-        num_actions = actor_num_actions,
-        dim_state = dim_state,
-        dim_latent = dim_latent,
-        dim_hiddens = actor_dim_hiddens
-    )
-
-    critic = Critic(
-        dim_state = dim_state,
-        dim_latent = dim_latent,
-        dim_hiddens = critic_dim_hiddens
-    )
-
-    latent_gene_pool = LatentGenePool(
-        dim_state = dim_state,
-        num_latents = num_latents,
-        dim_latent = dim_latent,
-        num_latent_sets = num_latent_sets
-    )
-
-    return Agent(actor = actor, critic = critic, latent_gene_pool = latent_gene_pool)
+# agent class
 
 class Agent(Module):
     def __init__(
         self,
         actor: Actor,
         critic: Critic,
-        latent_gene_pool: LatentGenePool
+        latent_gene_pool: LatentGenePool,
+        optim_klass = AdoptAtan2,
+        actor_lr = 1e-4,
+        critic_lr = 1e-4,
+        latent_lr = 1e-5,
+        actor_optim_kwargs: dict = dict(),
+        critic_optim_kwargs: dict = dict(),
+        latent_optim_kwargs: dict = dict(),
     ):
         super().__init__()
 
@@ -548,6 +488,13 @@ class Agent(Module):
 
         self.latent_gene_pool = latent_gene_pool
 
+        # optimizers
+
+        self.actor_optim = optim_klass(actor.parameters(), lr = actor_lr, **actor_optim_kwargs)
+        self.critic_optim = optim_klass(critic.parameters(), lr = critic_lr, **critic_optim_kwargs)
+
+        self.latent_optim = optim_klass(latent_gene_pool.parameters(), lr = latent_lr, **latent_optim_kwargs) if latent_gene_pool.needs_latent_gate else None
+
     def get_actor_actions(
         self,
         state,
@@ -576,6 +523,76 @@ class Agent(Module):
     ):
         raise NotImplementedError
 
+# reinforcement learning related - ppo
+
+def actor_loss(
+    logits,         # Float[b l]
+    old_log_probs,  # Float[b]
+    actions,        # Int[b]
+    advantages,     # Float[b]
+    eps_clip = 0.2,
+    entropy_weight = .01,
+):
+    log_probs = gather_log_prob(logits, actions)
+
+    ratio = (log_probs - old_log_probs).exp()
+
+    # 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
+
+def critic_loss(
+    pred_values,  # Float[b]
+    advantages,   # Float[b]
+    old_values    # Float[b]
+):
+    discounted_values = advantages + old_values
+    return F.mse_loss(pred_values, discounted_values)
+
+# agent contains the actor, critic, and the latent genetic pool
+
+def create_agent(
+    dim_state,
+    num_latents,
+    dim_latent,
+    actor_num_actions,
+    actor_dim_hiddens: int | tuple[int, ...],
+    critic_dim_hiddens: int | tuple[int, ...],
+    num_latent_sets = 1
+) -> Agent:
+
+    actor = Actor(
+        num_actions = actor_num_actions,
+        dim_state = dim_state,
+        dim_latent = dim_latent,
+        dim_hiddens = actor_dim_hiddens
+    )
+
+    critic = Critic(
+        dim_state = dim_state,
+        dim_latent = dim_latent,
+        dim_hiddens = critic_dim_hiddens
+    )
+
+    latent_gene_pool = LatentGenePool(
+        dim_state = dim_state,
+        num_latents = num_latents,
+        dim_latent = dim_latent,
+        num_latent_sets = num_latent_sets
+    )
+
+    return Agent(actor = actor, critic = critic, latent_gene_pool = latent_gene_pool)
+
 # 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
 
@@ -593,13 +610,10 @@ class EPO(Module):
 
     def __init__(
         self,
-        agent: Agent,
-        latent_gene_pool: LatentGenePool
+        agent: Agent
     ):
         super().__init__()
-
         self.agent = agent
-        self.latent_gene_pool = latent_gene_pool
 
     def forward(
         self,

{evolutionary_policy_optimization-0.0.12.dist-info → evolutionary_policy_optimization-0.0.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.12
+Version: 0.0.15
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -95,6 +95,7 @@ value = critic(state, latent)
 
 fitness = torch.randn(128)
 
+latent_pool.genetic_algorithm_step(fitness) # update latent genes with genetic algorithm
 ```
 
 ## Citations

evolutionary_policy_optimization-0.0.15.dist-info/RECORD

@@ -0,0 +1,7 @@
+evolutionary_policy_optimization/__init__.py,sha256=Qavcia0n13jjaWIS_LPW7QrxSLT_BBeKujCjF9kQjbA,133
+evolutionary_policy_optimization/epo.py,sha256=f_e-TkJRFF1VHG3psJDgLGNIzlEvDSjX0nOsbLaOBrw,17543
+evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
+evolutionary_policy_optimization-0.0.15.dist-info/METADATA,sha256=dGjFvEBt10Ac9PhsYSnluSbw_BixLTA_pEekW7TPF3U,4446
+evolutionary_policy_optimization-0.0.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.15.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.15.dist-info/RECORD,,

evolutionary_policy_optimization-0.0.12.dist-info/RECORD

@@ -1,7 +0,0 @@
-evolutionary_policy_optimization/__init__.py,sha256=Qavcia0n13jjaWIS_LPW7QrxSLT_BBeKujCjF9kQjbA,133
-evolutionary_policy_optimization/epo.py,sha256=dX7UANZk4RI_rwoSPi8GnZDoF9H1EdVT6Z7WA30cZ3Q,16934
-evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
-evolutionary_policy_optimization-0.0.12.dist-info/METADATA,sha256=Uyee4UOgNs04nrJlU6m7drQFTUuka3s8nLhblHxfiEg,4357
-evolutionary_policy_optimization-0.0.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.12.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.12.dist-info/RECORD,,