evolutionary-policy-optimization 0.0.43__py3-none-any.whl → 0.0.45__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -626,7 +626,9 @@ class Agent(Module):
         actor_lr = 1e-4,
         critic_lr = 1e-4,
         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,
@@ -647,7 +649,9 @@ class Agent(Module):
         self.actor = actor
 
         self.critic = critic
-        self.critic_ema = EMA(critic, beta = critic_ema_beta, include_online_model = False, **ema_kwargs)
+
+        self.use_critic_ema = use_critic_ema
+        self.critic_ema = EMA(critic, beta = critic_ema_beta, include_online_model = False, **ema_kwargs) if use_critic_ema else None
 
         self.num_latents = latent_gene_pool.num_latents
         self.latent_gene_pool = latent_gene_pool
@@ -659,7 +663,11 @@ class Agent(Module):
         self.actor_loss = partial(actor_loss, **actor_loss_kwargs)
         self.calc_gae = partial(calc_generalized_advantage_estimate, **calc_gae_kwargs)
 
+        # learning hparams
+
         self.batch_size = batch_size
+        self.max_grad_norm = max_grad_norm
+        self.has_grad_clip = exists(max_grad_norm)
 
         # optimizers
 
@@ -676,7 +684,7 @@ class Agent(Module):
         pkg = dict(
             actor = self.actor.state_dict(),
             critic = self.critic.state_dict(),
-            critic_ema = self.critic_ema.state_dict(),
+            critic_ema = self.critic_ema.state_dict() if self.use_critic_ema else None,
             latents = self.latent_gene_pool.state_dict(),
             actor_optim = self.actor_optim.state_dict(),
             critic_optim = self.critic_optim.state_dict(),
@@ -695,7 +703,9 @@ class Agent(Module):
         self.actor.load_state_dict(pkg['actor'])
 
         self.critic.load_state_dict(pkg['critic'])
-        self.critic_ema.load_state_dict(pkg['critic_ema'])
+
+        if self.use_critic_ema:
+            self.critic_ema.load_state_dict(pkg['critic_ema'])
 
         self.latent_gene_pool.load_state_dict(pkg['latents'])
 
@@ -733,14 +743,20 @@ class Agent(Module):
         self,
         state,
         latent_id = None,
-        latent = None
+        latent = None,
+        use_ema_if_available = False
     ):
         assert exists(latent_id) or exists(latent)
 
         if not exists(latent):
             latent = self.latent_gene_pool(latent_id = latent_id)
 
-        return self.critic(state, latent)
+        critic_forward = self.critic
+
+        if use_ema_if_available and self.use_critic_ema:
+            critic_forward = self.critic_ema
+
+        return critic_forward(state, latent)
 
     def update_latent_gene_pool_(
         self,
@@ -810,6 +826,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()
 
@@ -823,9 +842,17 @@ 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()
 
+                # maybe ema update critic
+
+                if self.use_critic_ema:
+                    self.critic_ema.update()
+
                 # maybe update latents, if not frozen
 
                 if not self.latent_gene_pool.frozen_latents:
@@ -875,6 +902,7 @@ def create_agent(
     actor_num_actions,
     actor_dim_hiddens: int | tuple[int, ...],
     critic_dim_hiddens: int | tuple[int, ...],
+    use_critic_ema = True,
     latent_gene_pool_kwargs: dict = dict(),
     actor_kwargs: dict = dict(),
     critic_kwargs: dict = dict(),
@@ -901,7 +929,14 @@ def create_agent(
         **critic_kwargs
     )
 
-    return Agent(actor = actor, critic = critic, latent_gene_pool = latent_gene_pool)
+    agent = Agent(
+        actor = actor,
+        critic = critic,
+        latent_gene_pool = latent_gene_pool,
+        use_critic_ema = use_critic_ema
+    )
+
+    return agent
 
 # 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
@@ -978,7 +1013,7 @@ class EPO(Module):
 
                     # values
 
-                    value = temp_batch_dim(self.agent.get_critic_values)(state, latent = latent)
+                    value = temp_batch_dim(self.agent.get_critic_values)(state, latent = latent, use_ema_if_available = True)
 
                     # get the next state, action, and reward
 

{evolutionary_policy_optimization-0.0.43.dist-info → evolutionary_policy_optimization-0.0.45.dist-info}/METADATA

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

@@ -1,8 +1,8 @@
 evolutionary_policy_optimization/__init__.py,sha256=0q0aBuFgWi06MLMD8FiHzBYQ3_W4LYWrwmCtF3u5H2A,201
-evolutionary_policy_optimization/epo.py,sha256=Yf-iw1gqmAUEVzg6_PwYy-q4005eroZKUYGxNgwCsKk,30440
+evolutionary_policy_optimization/epo.py,sha256=NA-d7pWDJyQYULDIVB25lnbpTbwMyxc1U8RU8XGTNts,31500
 evolutionary_policy_optimization/experimental.py,sha256=9FrJGviLESlYysHI3i83efT9g2ZB9ha4u3K9HXN98_w,1100
 evolutionary_policy_optimization/mock_env.py,sha256=6AIc4mwL_C6JkAxwESJgCLxXHMzCAu2FcffVg3HkSm0,920
-evolutionary_policy_optimization-0.0.43.dist-info/METADATA,sha256=pMVLppijepjmI1A9wVVhdX2IXo4BNPsOozpMAAsS6Lo,6213
-evolutionary_policy_optimization-0.0.43.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.43.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.43.dist-info/RECORD,,
+evolutionary_policy_optimization-0.0.45.dist-info/METADATA,sha256=3jXsZBoltrWQJk2Yd6zu1KmCcl9AEuhxES_mX8E1lAk,6213
+evolutionary_policy_optimization-0.0.45.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.45.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.45.dist-info/RECORD,,