evolutionary-policy-optimization 0.0.67__py3-none-any.whl → 0.0.68__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):
@@ -744,6 +747,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 +847,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 +1229,8 @@ class EPO(Module):
                     terminated
                 )
 
+                memory = Memory(*tuple(t.cpu() for t in memory))
+
                 memories.append(memory)
 
                 time += 1
@@ -1228,7 +1242,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.68.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.67
+Version: 0.0.68
 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.68.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=xhE_kHas54xGsgOese9SQEvyK7NKZqEuK3AiVhm0y7Q,38047
 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.68.dist-info/METADATA,sha256=hOOKOrrPQtQmK3zN1z5nkGJEoaQLyXUzs9ArsEKn1DE,6220
+evolutionary_policy_optimization-0.0.68.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.68.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.68.dist-info/RECORD,,