evolutionary-policy-optimization 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -79,7 +79,7 @@ def interface_torch_numpy(fn, device):
     @wraps(fn)
     def decorated_fn(*args, **kwargs):
 
-        args, kwargs = tree_map(lambda t: t.cpu().numpy() if isinstance(t, Tensor) else t, (args, kwargs))
+        args, kwargs = tree_map(lambda t: t.cpu().numpy() if is_tensor(t) else t, (args, kwargs))
 
         out = fn(*args, **kwargs)
 
@@ -88,6 +88,16 @@ def interface_torch_numpy(fn, device):
 
     return decorated_fn
 
+def move_input_tensors_to_device(fn):
+
+    @wraps(fn)
+    def decorated_fn(self, *args, **kwargs):
+        args, kwargs = tree_map(lambda t: t.to(self.device) if is_tensor(t) else t, (args, kwargs))
+
+        return fn(self, *args, **kwargs)
+
+    return decorated_fn
+
 # tensor helpers
 
 def l2norm(t):
@@ -797,6 +807,8 @@ class Agent(Module):
 
         self.unwrap_model = identity if not wrap_with_accelerate else self.accelerate.unwrap_model
 
+        dummy = tensor(0)
+
         if wrap_with_accelerate:
             (
                 self.actor,
@@ -816,11 +828,14 @@ class Agent(Module):
                 )
             )
 
-        # device tracking
+            if exists(self.critic_ema):
+                self.critic_ema.to(self.accelerate.device)
 
-        self.register_buffer('dummy', tensor(0, device = self.accelerate.device))
+            dummy = dummy.to(self.accelerate.device)
 
-        self.critic_ema.to(self.accelerate.device)
+        # device tracking
+
+        self.register_buffer('dummy', dummy)
 
     @property
     def device(self):
@@ -870,6 +885,7 @@ class Agent(Module):
         if exists(pkg.get('latent_optim', None)):
             self.latent_optim.load_state_dict(pkg['latent_optim'])
 
+    @move_input_tensors_to_device
     def get_actor_actions(
         self,
         state,
@@ -884,6 +900,7 @@ class Agent(Module):
         if not exists(latent) and exists(latent_id):
             latent = maybe_unwrap(self.latent_gene_pool)(latent_id = latent_id)
 
+        print(self.device, state.device, next(self.actor.parameters()).device)
         logits = maybe_unwrap(self.actor)(state, latent)
 
         if not sample:
@@ -895,6 +912,7 @@ class Agent(Module):
 
         return actions, log_probs
 
+    @move_input_tensors_to_device
     def get_critic_values(
         self,
         state,
@@ -903,6 +921,7 @@ class Agent(Module):
         use_ema_if_available = False,
         use_unwrapped_model = False
     ):
+
         maybe_unwrap = identity if not use_unwrapped_model else self.unwrap_model
 
         if not exists(latent) and exists(latent_id):

{evolutionary_policy_optimization-0.1.2.dist-info → evolutionary_policy_optimization-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.2
+Version: 0.1.4
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -53,7 +53,7 @@ Description-Content-Type: text/markdown
 
 <img width="450px" alt="fig1" src="https://github.com/user-attachments/assets/33bef569-e786-4f09-bdee-56bad7ea9e6d" />
 
-## Evolutionary Policy Optimization (wip)
+## Evolutionary Policy Optimization
 
 Pytorch implementation of [Evolutionary Policy Optimization](https://web3.arxiv.org/abs/2503.19037), from Wang et al. of the Robotics Institute at Carnegie Mellon University
 

{evolutionary_policy_optimization-0.1.2.dist-info → evolutionary_policy_optimization-0.1.4.dist-info}/RECORD

@@ -1,10 +1,10 @@
 evolutionary_policy_optimization/__init__.py,sha256=NyiYDYU7DlpmOTM7xiBQET3r1WwX0ebrgMCBLSQrW3c,288
 evolutionary_policy_optimization/distributed.py,sha256=7KgZdeS_wxBHo_du9XZFB1Cu318J-Bp66Xdr6Log_20,2423
 evolutionary_policy_optimization/env_wrappers.py,sha256=bDL06o9_b1iW6k3fw2xifnOnYlzs643tdW6Yv2gsIdw,803
-evolutionary_policy_optimization/epo.py,sha256=VgiPvaG-Lib1JZRS3oNV7BSv-19YSWKK2yCMtqDqP-M,41682
+evolutionary_policy_optimization/epo.py,sha256=3PEmFE032fQz7mAfIoJzwmysbU6lnjCCngEMc9Jf46M,42180
 evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
 evolutionary_policy_optimization/mock_env.py,sha256=TLyyRm6tOD0Kdn9QqJJQriaSnsR-YmNQHo4OohmZFG4,1410
-evolutionary_policy_optimization-0.1.2.dist-info/METADATA,sha256=kXkidoqrBxyMtQvh7GHr1wh1pexy2WDSDWN5laQSPhI,6336
-evolutionary_policy_optimization-0.1.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.1.2.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.1.2.dist-info/RECORD,,
+evolutionary_policy_optimization-0.1.4.dist-info/METADATA,sha256=81Q8wK2N7zlOsSGuqcXrftpd1TIYkrplSPBuldMj8uU,6330
+evolutionary_policy_optimization-0.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.1.4.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.1.4.dist-info/RECORD,,