evolutionary-policy-optimization 0.0.61__py3-none-any.whl → 0.0.62__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -104,8 +104,11 @@ def temp_batch_dim(fn):
 
 # fitness related
 
-def get_fitness_scores(cum_rewards, memories):
-    return cum_rewards
+def get_fitness_scores(
+    cum_rewards, # Float['gene episodes']
+    memories
+): # Float['gene']
+    return cum_rewards.sum(dim = -1) # sum all rewards across episodes, but could override this function for normalizing with whatever
 
 # generalized advantage estimate
 
@@ -684,7 +687,8 @@ class Agent(Module):
         ),
         actor_loss_kwargs: dict = dict(
             eps_clip = 0.2,
-            entropy_weight = .01
+            entropy_weight = .01,
+            norm_advantages = True
         ),
         ema_kwargs: dict = dict(),
         actor_optim_kwargs: dict = dict(),
@@ -826,9 +830,7 @@ class Agent(Module):
         memories_and_cumulative_rewards: MemoriesAndCumulativeRewards,
         epochs = 2
     ):
-        memories, cumulative_rewards = memories_and_cumulative_rewards
-
-        fitness_scores = self.get_fitness_scores(cumulative_rewards, memories)
+        memories, rewards_per_latent_episode = memories_and_cumulative_rewards
 
         # stack memories
 
@@ -839,7 +841,13 @@ class Agent(Module):
         if is_distributed():
             memories = map(partial(all_gather_variable_dim, dim = 0), memories)
 
-            fitness_scores = all_gather_variable_dim(fitness_scores, dim = 0)
+            rewards_per_latent_episode = dist.all_reduce(rewards_per_latent_episode)
+
+        # calculate fitness scores
+
+        fitness_scores = self.get_fitness_scores(rewards_per_latent_episode, memories)
+
+        # process memories
 
         (
             episode_ids,
@@ -854,12 +862,16 @@ class Agent(Module):
 
         masks = 1. - dones.float()
 
+        # generalized advantage estimate
+
         advantages = self.calc_gae(
             rewards[:-1],
             values,
             masks[:-1],
         )
 
+        # dataset and dataloader
+
         valid_episode = episode_ids >= 0
 
         dataset = TensorDataset(
@@ -871,6 +883,8 @@ class Agent(Module):
 
         dataloader = DataLoader(dataset, batch_size = self.batch_size, shuffle = True)
 
+        # updating actor and critic
+
         self.actor.train()
         self.critic.train()
 
@@ -954,7 +968,8 @@ def actor_loss(
     advantages,     # Float[b]
     eps_clip = 0.2,
     entropy_weight = .01,
-    eps = 1e-5
+    eps = 1e-5,
+    norm_advantages = True
 ):
     batch = logits.shape[0]
 
@@ -966,7 +981,8 @@ def actor_loss(
 
     clipped_ratio = ratio.clamp(min = 1. - eps_clip, max = 1. + eps_clip)
 
-    advantages = F.layer_norm(advantages, (batch,), eps = eps)
+    if norm_advantages:
+        advantages = F.layer_norm(advantages, (batch,), eps = eps)
 
     actor_loss = -torch.min(clipped_ratio * advantages, ratio * advantages)
 
@@ -1041,7 +1057,7 @@ Memory = namedtuple('Memory', [
 
 MemoriesAndCumulativeRewards = namedtuple('MemoriesAndCumulativeRewards', [
     'memories',
-    'cumulative_rewards'
+    'cumulative_rewards' # Float['latent episodes']
 ])
 
 class EPO(Module):
@@ -1098,7 +1114,7 @@ class EPO(Module):
 
         memories: list[Memory] = []
 
-        cumulative_rewards = torch.zeros((self.num_latents))
+        rewards_per_latent_episode = torch.zeros((self.num_latents, self.episodes_per_latent))
 
         latent_ids_gen = self.latents_for_machine()
 
@@ -1147,7 +1163,7 @@ class EPO(Module):
 
                     # update cumulative rewards per latent, to be used as default fitness score
 
-                    cumulative_rewards[latent_id] += reward
+                    rewards_per_latent_episode[latent_id, episode_id] += reward
                     
                     # store memories
 
@@ -1179,5 +1195,5 @@ class EPO(Module):
 
         return MemoriesAndCumulativeRewards(
             memories = memories,
-            cumulative_rewards = cumulative_rewards
+            cumulative_rewards = rewards_per_latent_episode
         )

{evolutionary_policy_optimization-0.0.61.dist-info → evolutionary_policy_optimization-0.0.62.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.61
+Version: 0.0.62
 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.61.dist-info → evolutionary_policy_optimization-0.0.62.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=kFT49rJdcmaDehfpx3YyhYhvAcp7S-gRWDkS2y20Q2Y,35377
+evolutionary_policy_optimization/epo.py,sha256=lWhHpsfq6vpri6yeDXSTLRMKGPwl0kt3klh0fVaInSs,35921
 evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
 evolutionary_policy_optimization/mock_env.py,sha256=202KJ5g57wQvOzhGYzgHfBa7Y2do5uuDvl5kFg5o73g,934
-evolutionary_policy_optimization-0.0.61.dist-info/METADATA,sha256=3IbcY9kg71P6lTNxZaRBw3IYfDjcK4uTJJaFRD0Skwg,6220
-evolutionary_policy_optimization-0.0.61.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.61.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.61.dist-info/RECORD,,
+evolutionary_policy_optimization-0.0.62.dist-info/METADATA,sha256=oqJyUOXJwHrdf6JCVKPfOmhGJbXgqOmPWN_46l0JtWs,6220
+evolutionary_policy_optimization-0.0.62.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.62.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.62.dist-info/RECORD,,