evolutionary-policy-optimization 0.0.44__py3-none-any.whl → 0.0.46__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -87,6 +87,11 @@ def temp_batch_dim(fn):
 
     return inner
 
+# fitness related
+
+def get_fitness_scores(cum_rewards, memories):
+    return cum_rewards
+
 # generalized advantage estimate
 
 def calc_generalized_advantage_estimate(
@@ -628,6 +633,7 @@ class Agent(Module):
         latent_lr = 1e-5,
         use_critic_ema = True,
         critic_ema_beta = 0.99,
+        max_grad_norm = 0.5,
         batch_size = 16,
         calc_gae_kwargs: dict = dict(
             use_accelerated = False,
@@ -642,6 +648,7 @@ class Agent(Module):
         actor_optim_kwargs: dict = dict(),
         critic_optim_kwargs: dict = dict(),
         latent_optim_kwargs: dict = dict(),
+        get_fitness_scores: Callable[..., Tensor] = get_fitness_scores
     ):
         super().__init__()
 
@@ -662,7 +669,15 @@ class Agent(Module):
         self.actor_loss = partial(actor_loss, **actor_loss_kwargs)
         self.calc_gae = partial(calc_generalized_advantage_estimate, **calc_gae_kwargs)
 
+        # fitness score related
+
+        self.get_fitness_scores = get_fitness_scores
+
+        # learning hparams
+
         self.batch_size = batch_size
+        self.max_grad_norm = max_grad_norm
+        self.has_grad_clip = exists(max_grad_norm)
 
         # optimizers
 
@@ -761,10 +776,12 @@ class Agent(Module):
 
     def forward(
         self,
-        memories_and_fitness_scores: MemoriesAndFitnessScores,
+        memories_and_cumulative_rewards: MemoriesAndCumulativeRewards,
         epochs = 2
     ):
-        memories, fitness_scores = memories_and_fitness_scores
+        memories, cumulative_rewards = memories_and_cumulative_rewards
+
+        fitness_scores = self.get_fitness_scores(cumulative_rewards, memories)
 
         (
             episode_ids,
@@ -821,6 +838,9 @@ class Agent(Module):
 
                 actor_loss.backward()
 
+                if exists(self.has_grad_clip):
+                    nn.utils.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
+
                 self.actor_optim.step()
                 self.actor_optim.zero_grad()
 
@@ -834,6 +854,9 @@ class Agent(Module):
 
                 critic_loss.backward()
 
+                if exists(self.has_grad_clip):
+                    nn.utils.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm)
+
                 self.critic_optim.step()
                 self.critic_optim.zero_grad()
 
@@ -941,9 +964,9 @@ Memory = namedtuple('Memory', [
     'done'
 ])
 
-MemoriesAndFitnessScores = namedtuple('MemoriesAndFitnessScores', [
+MemoriesAndCumulativeRewards = namedtuple('MemoriesAndCumulativeRewards', [
     'memories',
-    'fitness_scores'
+    'cumulative_rewards'
 ])
 
 class EPO(Module):
@@ -967,7 +990,7 @@ class EPO(Module):
     def forward(
         self,
         env
-    ) -> MemoriesAndFitnessScores:
+    ) -> MemoriesAndCumulativeRewards:
 
         self.agent.eval()
 
@@ -975,7 +998,7 @@ class EPO(Module):
 
         memories: list[Memory] = []
 
-        fitness_scores = torch.zeros((self.num_latents))
+        cumulative_rewards = torch.zeros((self.num_latents))
 
         for episode_id in tqdm(range(self.episodes_per_latent), desc = 'episode'):
 
@@ -1008,9 +1031,9 @@ class EPO(Module):
 
                     state, reward, done = env(action)
 
-                    # update fitness for each gene as cumulative reward received, but make this customizable at some point
+                    # update cumulative rewards per latent, to be used as default fitness score
 
-                    fitness_scores[latent_id] += reward
+                    cumulative_rewards[latent_id] += reward
                     
                     # store memories
 
@@ -1040,7 +1063,7 @@ class EPO(Module):
 
                 memories.append(memory_for_gae)
 
-        return MemoriesAndFitnessScores(
+        return MemoriesAndCumulativeRewards(
             memories = memories,
-            fitness_scores = fitness_scores
+            cumulative_rewards = cumulative_rewards
         )

{evolutionary_policy_optimization-0.0.44.dist-info → evolutionary_policy_optimization-0.0.46.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.44
+Version: 0.0.46
 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.44.dist-info → evolutionary_policy_optimization-0.0.46.dist-info}/RECORD

@@ -1,8 +1,8 @@
 evolutionary_policy_optimization/__init__.py,sha256=0q0aBuFgWi06MLMD8FiHzBYQ3_W4LYWrwmCtF3u5H2A,201
-evolutionary_policy_optimization/epo.py,sha256=8ZyW21nbEokTii7alduJVDfPLhRLsLCjttWni6-BNFE,31072
+evolutionary_policy_optimization/epo.py,sha256=SAhWgRY8uPQEKFg1_nz1mvh8A6S_sHwnDykhd0F5xEI,31853
 evolutionary_policy_optimization/experimental.py,sha256=9FrJGviLESlYysHI3i83efT9g2ZB9ha4u3K9HXN98_w,1100
 evolutionary_policy_optimization/mock_env.py,sha256=6AIc4mwL_C6JkAxwESJgCLxXHMzCAu2FcffVg3HkSm0,920
-evolutionary_policy_optimization-0.0.44.dist-info/METADATA,sha256=8tJisD2xu5Q7udBrmFnfy6oxl_RQrCF9DYrE6yHkP2M,6213
-evolutionary_policy_optimization-0.0.44.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.44.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.44.dist-info/RECORD,,
+evolutionary_policy_optimization-0.0.46.dist-info/METADATA,sha256=xP2kdKo52-X4Z5XXTPpW0M_NFI0spuigeL7fvqFlsRM,6213
+evolutionary_policy_optimization-0.0.46.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.46.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.46.dist-info/RECORD,,