evolutionary-policy-optimization 0.0.61__tar.gz → 0.0.62__tar.gz

{evolutionary_policy_optimization-0.0.61 → evolutionary_policy_optimization-0.0.62}/PKG-INFO

@@ -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 → evolutionary_policy_optimization-0.0.62}/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 → evolutionary_policy_optimization-0.0.62}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.0.61"
+version = "0.0.62"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }