evolutionary-policy-optimization 0.0.70__tar.gz → 0.0.72__tar.gz

{evolutionary_policy_optimization-0.0.70 → evolutionary_policy_optimization-0.0.72}/PKG-INFO

@@ -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.70 → evolutionary_policy_optimization-0.0.72}/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-0.0.70 → evolutionary_policy_optimization-0.0.72}/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 → evolutionary_policy_optimization-0.0.72}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.0.70"
+version = "0.0.72"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }