evolutionary-policy-optimization 0.0.67__py3-none-any.whl → 0.0.69__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -55,6 +55,9 @@ def xnor(x, y):
 def divisible_by(num, den):
     return (num % den) == 0
 
+def to_device(inp, device):
+    return tree_map(lambda t: t.to(device) if is_tensor(t) else t, inp)
+
 # tensor helpers
 
 def l2norm(t):
@@ -678,6 +681,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,
@@ -734,8 +739,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
 
@@ -744,6 +749,12 @@ class Agent(Module):
         self.has_diversity_loss = diversity_aux_loss_weight > 0.
         self.diversity_aux_loss_weight = diversity_aux_loss_weight
 
+        self.register_buffer('dummy', tensor(0))
+
+    @property
+    def device(self):
+        return self.dummy.device
+
     def save(self, path, overwrite = False):
         path = Path(path)
 
@@ -838,7 +849,10 @@ class Agent(Module):
         self,
         memories_and_cumulative_rewards: MemoriesAndCumulativeRewards,
         epochs = 2
+
     ):
+        memories_and_cumulative_rewards = to_device(memories_and_cumulative_rewards, self.device)
+
         memories, rewards_per_latent_episode = memories_and_cumulative_rewards
 
         # stack memories
@@ -1217,6 +1231,8 @@ class EPO(Module):
                     terminated
                 )
 
+                memory = Memory(*tuple(t.cpu() for t in memory))
+
                 memories.append(memory)
 
                 time += 1
@@ -1228,7 +1244,7 @@ class EPO(Module):
 
                 memory_for_gae = memory._replace(
                     episode_id = invalid_episode,
-                    value = next_value,
+                    value = next_value.cpu(),
                     done = tensor(True)
                 )
 

{evolutionary_policy_optimization-0.0.67.dist-info → evolutionary_policy_optimization-0.0.69.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.67
+Version: 0.0.69
 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.67.dist-info → evolutionary_policy_optimization-0.0.69.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=7WOGl22WknudsNLSZ18AWwW7rPt5ITMAtByetUwLp7M,37654
+evolutionary_policy_optimization/epo.py,sha256=e83tghTNXfCW0zhhb4nIjvfbzDvzWRxgTlm3vKJd4rM,38189
 evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
 evolutionary_policy_optimization/mock_env.py,sha256=Bv9ONFRbma8wpjUurc9aCk19A6ceiWitRnS3nwrIR64,1339
-evolutionary_policy_optimization-0.0.67.dist-info/METADATA,sha256=S1biwayyDA4vTOXknMU5KeWtJTFvxePHZZ0OZzuaNms,6220
-evolutionary_policy_optimization-0.0.67.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.67.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.67.dist-info/RECORD,,
+evolutionary_policy_optimization-0.0.69.dist-info/METADATA,sha256=UZEaCY5lfTRMkuyEQs5PLA1AZzSOcsRzXey9kgdd9i0,6220
+evolutionary_policy_optimization-0.0.69.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.69.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.69.dist-info/RECORD,,