evolutionary-policy-optimization 0.1.2__tar.gz → 0.1.5__tar.gz

{evolutionary_policy_optimization-0.1.2 → evolutionary_policy_optimization-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.2
+Version: 0.1.5
 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 → evolutionary_policy_optimization-0.1.5}/README.md

@@ -1,6 +1,6 @@
 <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 → evolutionary_policy_optimization-0.1.5}/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,
@@ -895,6 +911,7 @@ class Agent(Module):
 
         return actions, log_probs
 
+    @move_input_tensors_to_device
     def get_critic_values(
         self,
         state,
@@ -903,6 +920,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 → evolutionary_policy_optimization-0.1.5}/pyproject.toml

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