evolutionary-policy-optimization 0.0.68__py3-none-any.whl → 0.0.70__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -416,6 +416,8 @@ class LatentGenePool(Module):
         self.num_natural_selected = int(frac_natural_selected * latents_per_island)
         self.num_tournament_participants = int(frac_tournaments * self.num_natural_selected)
 
+        assert self.num_tournament_participants >= 2
+
         self.crossover_random  = crossover_random
 
         self.mutation_strength = mutation_strength
@@ -681,6 +683,8 @@ class Agent(Module):
         actor_lr = 1e-4,
         critic_lr = 1e-4,
         latent_lr = 1e-5,
+        actor_weight_decay = 1e-3,
+        critic_weight_decay = 1e-3,
         diversity_aux_loss_weight = 0.,
         use_critic_ema = True,
         critic_ema_beta = 0.99,
@@ -737,8 +741,8 @@ class Agent(Module):
 
         # 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.actor_optim = optim_klass(actor.parameters(), lr = actor_lr, weight_decay = actor_weight_decay, **actor_optim_kwargs)
+        self.critic_optim = optim_klass(critic.parameters(), lr = critic_lr, weight_decay = critic_weight_decay, **critic_optim_kwargs)
 
         self.latent_optim = optim_klass(latent_gene_pool.parameters(), lr = latent_lr, **latent_optim_kwargs) if exists(latent_gene_pool) and not latent_gene_pool.frozen_latents else None
 
@@ -843,7 +847,7 @@ class Agent(Module):
 
         return self.latent_gene_pool.genetic_algorithm_step(fitnesses)
 
-    def forward(
+    def learn_from(
         self,
         memories_and_cumulative_rewards: MemoriesAndCumulativeRewards,
         epochs = 2
@@ -851,11 +855,13 @@ class Agent(Module):
     ):
         memories_and_cumulative_rewards = to_device(memories_and_cumulative_rewards, self.device)
 
-        memories, rewards_per_latent_episode = memories_and_cumulative_rewards
+        memories_list, rewards_per_latent_episode = memories_and_cumulative_rewards
 
         # stack memories
 
-        memories = map(stack, zip(*memories))
+        memories = map(stack, zip(*memories_list))
+
+        memories_list.clear()
 
         maybe_barrier()
 
@@ -977,7 +983,6 @@ class Agent(Module):
         # apply evolution
 
         if self.has_latent_genes:
-
             self.latent_gene_pool.genetic_algorithm_step(fitness_scores)
 
 # reinforcement learning related - ppo
@@ -1157,9 +1162,10 @@ class EPO(Module):
             yield latent_id, episode_id, maybe_seed
 
     @torch.no_grad()
-    def forward(
+    def gather_experience_from(
         self,
         env,
+        memories: list[Memory] | None = None,
         fix_environ_across_latents = None
     ) -> MemoriesAndCumulativeRewards:
 
@@ -1169,7 +1175,8 @@ class EPO(Module):
 
         invalid_episode = tensor(-1) # will use `episode_id` value of `-1` for the `next_value`, needed for not discarding last reward for generalized advantage estimate
 
-        memories: list[Memory] = []
+        if not exists(memories):
+            memories = []
 
         rewards_per_latent_episode = torch.zeros((self.num_latents, self.episodes_per_latent))
 
@@ -1252,3 +1259,18 @@ class EPO(Module):
             memories = memories,
             cumulative_rewards = rewards_per_latent_episode
         )
+
+    def forward(
+        self,
+        agent: Agent,
+        env,
+        num_learning_cycles
+    ):
+
+        for _ in tqdm(range(num_learning_cycles), desc = 'learning cycle'):
+
+            memories = self.gather_experience_from(env)
+
+            agent.learn_from(memories)
+
+        print(f'training complete')

{evolutionary_policy_optimization-0.0.68.dist-info → evolutionary_policy_optimization-0.0.70.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.68
+Version: 0.0.70
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -114,11 +114,14 @@ from evolutionary_policy_optimization import (
 
 agent = create_agent(
     dim_state = 512,
-    num_latents = 8,
+    num_latents = 16,
     dim_latent = 32,
     actor_num_actions = 5,
     actor_dim_hiddens = (256, 128),
-    critic_dim_hiddens = (256, 128, 64)
+    critic_dim_hiddens = (256, 128, 64),
+    latent_gene_pool_kwargs = dict(
+        frac_natural_selected = 0.5
+    )
 )
 
 epo = EPO(
@@ -130,9 +133,7 @@ epo = EPO(
 
 env = Env((512,))
 
-memories = epo(env)
-
-agent(memories)
+epo(agent, env, num_learning_cycles = 5)
 
 # saving and loading
 

{evolutionary_policy_optimization-0.0.68.dist-info → evolutionary_policy_optimization-0.0.70.dist-info}/RECORD

@@ -1,9 +1,9 @@
 evolutionary_policy_optimization/__init__.py,sha256=0q0aBuFgWi06MLMD8FiHzBYQ3_W4LYWrwmCtF3u5H2A,201
 evolutionary_policy_optimization/distributed.py,sha256=7KgZdeS_wxBHo_du9XZFB1Cu318J-Bp66Xdr6Log_20,2423
-evolutionary_policy_optimization/epo.py,sha256=xhE_kHas54xGsgOese9SQEvyK7NKZqEuK3AiVhm0y7Q,38047
+evolutionary_policy_optimization/epo.py,sha256=MmsqMwytVqBkb1f2piUygueOfn--Icb817P4bDcfPks,38683
 evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
 evolutionary_policy_optimization/mock_env.py,sha256=Bv9ONFRbma8wpjUurc9aCk19A6ceiWitRnS3nwrIR64,1339
-evolutionary_policy_optimization-0.0.68.dist-info/METADATA,sha256=hOOKOrrPQtQmK3zN1z5nkGJEoaQLyXUzs9ArsEKn1DE,6220
-evolutionary_policy_optimization-0.0.68.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.68.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.68.dist-info/RECORD,,
+evolutionary_policy_optimization-0.0.70.dist-info/METADATA,sha256=fX_hR3dCjKUQ3VZtT-sUQy0qT8sF-nDPyP0QDAGHf60,6304
+evolutionary_policy_optimization-0.0.70.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.70.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.70.dist-info/RECORD,,