evolutionary-policy-optimization 0.0.70__py3-none-any.whl → 0.0.72__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -8,8 +8,10 @@ from functools import partial, wraps
 from collections import namedtuple
 from random import randrange
 
+import numpy as np
+
 import torch
-from torch import nn, cat, stack, is_tensor, tensor, Tensor
+from torch import nn, cat, stack, is_tensor, tensor, from_numpy, Tensor
 import torch.nn.functional as F
 import torch.distributed as dist
 from torch.nn import Linear, Module, ModuleList
@@ -58,6 +60,21 @@ def divisible_by(num, den):
 def to_device(inp, device):
     return tree_map(lambda t: t.to(device) if is_tensor(t) else t, inp)
 
+def interface_torch_numpy(fn, device):
+    # for a given function, move all inputs from torch tensor to numpy, and all outputs from numpy to torch tensor
+
+    @wraps(fn)
+    def decorated_fn(*args, **kwargs):
+
+        args, kwargs = tree_map(lambda t: t.cpu().numpy() if isinstance(t, Tensor) else t, (args, kwargs))
+
+        out = fn(*args, **kwargs)
+
+        out = tree_map(lambda t: from_numpy(t).to(device) if isinstance(t, np.ndarray) else t, out)
+        return out
+
+    return decorated_fn
+
 # tensor helpers
 
 def l2norm(t):
@@ -1178,7 +1195,7 @@ class EPO(Module):
         if not exists(memories):
             memories = []
 
-        rewards_per_latent_episode = torch.zeros((self.num_latents, self.episodes_per_latent))
+        rewards_per_latent_episode = torch.zeros((self.num_latents, self.episodes_per_latent), device = self.device)
 
         rollout_gen = self.rollouts_for_machine(fix_environ_across_latents)
 
@@ -1193,7 +1210,7 @@ class EPO(Module):
             if fix_environ_across_latents:
                 reset_kwargs.update(seed = maybe_seed)
 
-            state = env.reset(**reset_kwargs)
+            state = interface_torch_numpy(env.reset, device = self.device)(**reset_kwargs)
 
             # get latent from pool
 
@@ -1215,7 +1232,7 @@ class EPO(Module):
 
                 # get the next state, action, and reward
 
-                state, reward, truncated, terminated = env(action)
+                state, reward, truncated, terminated = interface_torch_numpy(env.forward, device = self.device)(action)
 
                 done = truncated or terminated
 
@@ -1264,9 +1281,14 @@ class EPO(Module):
         self,
         agent: Agent,
         env,
-        num_learning_cycles
+        num_learning_cycles,
+        seed = None
     ):
 
+        if exists(seed):
+            torch.manual_seed(seed)
+            np.random.seed(seed)
+
         for _ in tqdm(range(num_learning_cycles), desc = 'learning cycle'):
 
             memories = self.gather_experience_from(env)

evolutionary_policy_optimization/mock_env.py

@@ -34,7 +34,7 @@ class Env(Module):
     ):
         state = randn(self.state_shape, device = self.device)
         self.step.zero_()
-        return state
+        return state.numpy()
 
     def forward(
         self,
@@ -51,4 +51,5 @@ class Env(Module):
 
         self.step.add_(1)
 
-        return state, reward, truncated, terminated
+        out = state, reward, truncated, terminated
+        return tuple(t.numpy() for t in out)

{evolutionary_policy_optimization-0.0.70.dist-info → evolutionary_policy_optimization-0.0.72.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.70
+Version: 0.0.72
 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.72.dist-info/RECORD

@@ -0,0 +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=h_BaNLFeU5SfVaOcVeMdEdjfYop6enOPAsESQWZflfA,39449
+evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
+evolutionary_policy_optimization/mock_env.py,sha256=QhPYIazU3uFPKTMZyFw70KhXiDUBr5aU3v1idotfVFI,1391
+evolutionary_policy_optimization-0.0.72.dist-info/METADATA,sha256=kZZupJU1gYm6MEmDYB2D-gIF2NZsAY4qcK9Z5pAT1mQ,6304
+evolutionary_policy_optimization-0.0.72.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.72.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.72.dist-info/RECORD,,

evolutionary_policy_optimization-0.0.70.dist-info/RECORD

@@ -1,9 +0,0 @@
-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=MmsqMwytVqBkb1f2piUygueOfn--Icb817P4bDcfPks,38683
-evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
-evolutionary_policy_optimization/mock_env.py,sha256=Bv9ONFRbma8wpjUurc9aCk19A6ceiWitRnS3nwrIR64,1339
-evolutionary_policy_optimization-0.0.70.dist-info/METADATA,sha256=fX_hR3dCjKUQ3VZtT-sUQy0qT8sF-nDPyP0QDAGHf60,6304
-evolutionary_policy_optimization-0.0.70.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.70.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.70.dist-info/RECORD,,