evolutionary-policy-optimization 0.1.1__tar.gz → 0.1.4__tar.gz

{evolutionary_policy_optimization-0.1.1 → evolutionary_policy_optimization-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.1
+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
@@ -34,7 +34,7 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
-Requires-Dist: accelerate
+Requires-Dist: accelerate>=1.6.0
 Requires-Dist: adam-atan2-pytorch
 Requires-Dist: assoc-scan>=0.0.2
 Requires-Dist: einops>=0.8.1
@@ -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.1 → evolutionary_policy_optimization-0.1.4}/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.1 → evolutionary_policy_optimization-0.1.4}/evolutionary_policy_optimization/epo.py

@@ -40,6 +40,8 @@ from ema_pytorch import EMA
 
 from tqdm import tqdm
 
+from accelerate import Accelerator
+
 # helpers
 
 def exists(v):
@@ -60,13 +62,24 @@ 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 maybe(fn):
+
+    @wraps(fn)
+    def decorated(inp, *args, **kwargs):
+        if not exists(inp):
+            return None
+
+        return fn(inp, *args, **kwargs)
+
+    return decorated
+
 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))
+        args, kwargs = tree_map(lambda t: t.cpu().numpy() if is_tensor(t) else t, (args, kwargs))
 
         out = fn(*args, **kwargs)
 
@@ -75,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):
@@ -721,10 +744,22 @@ class Agent(Module):
         actor_optim_kwargs: dict = dict(),
         critic_optim_kwargs: dict = dict(),
         latent_optim_kwargs: dict = dict(),
-        get_fitness_scores: Callable[..., Tensor] = get_fitness_scores
+        get_fitness_scores: Callable[..., Tensor] = get_fitness_scores,
+        wrap_with_accelerate: bool = True,
+        accelerate_kwargs: dict = dict(),
     ):
         super().__init__()
 
+        # hf accelerate
+
+        self.wrap_with_accelerate = wrap_with_accelerate
+
+        if wrap_with_accelerate:
+            accelerate = Accelerator(**accelerate_kwargs)
+            self.accelerate = accelerate
+
+        # actor, critic, and their shared latent gene pool
+
         self.actor = actor
 
         self.critic = critic
@@ -768,12 +803,48 @@ 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))
+        # wrap with accelerate
+
+        self.unwrap_model = identity if not wrap_with_accelerate else self.accelerate.unwrap_model
+
+        dummy = tensor(0)
+
+        if wrap_with_accelerate:
+            (
+                self.actor,
+                self.critic,
+                self.latent_gene_pool,
+                self.actor_optim,
+                self.critic_optim,
+                self.latent_optim,
+            ) = tuple(
+                maybe(self.accelerate.prepare)(m) for m in (
+                    self.actor,
+                    self.critic,
+                    self.latent_gene_pool,
+                    self.actor_optim,
+                    self.critic_optim,
+                    self.latent_optim,
+                )
+            )
+
+            if exists(self.critic_ema):
+                self.critic_ema.to(self.accelerate.device)
+
+            dummy = dummy.to(self.accelerate.device)
+
+        # device tracking
+
+        self.register_buffer('dummy', dummy)
 
     @property
     def device(self):
         return self.dummy.device
 
+    @property
+    def unwrapped_latent_gene_pool(self):
+        return self.unwrap_model(self.latent_gene_pool)
+
     def save(self, path, overwrite = False):
         path = Path(path)
 
@@ -814,19 +885,23 @@ 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,
         latent_id = None,
         latent = None,
         sample = False,
-        temperature = 1.
+        temperature = 1.,
+        use_unwrapped_model = False
     ):
+        maybe_unwrap = identity if not use_unwrapped_model else self.unwrap_model
 
         if not exists(latent) and exists(latent_id):
-            latent = self.latent_gene_pool(latent_id = latent_id)
+            latent = maybe_unwrap(self.latent_gene_pool)(latent_id = latent_id)
 
-        logits = self.actor(state, latent)
+        print(self.device, state.device, next(self.actor.parameters()).device)
+        logits = maybe_unwrap(self.actor)(state, latent)
 
         if not sample:
             return logits
@@ -837,18 +912,22 @@ class Agent(Module):
 
         return actions, log_probs
 
+    @move_input_tensors_to_device
     def get_critic_values(
         self,
         state,
         latent_id = None,
         latent = None,
-        use_ema_if_available = False
+        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):
-            latent = self.latent_gene_pool(latent_id = latent_id)
+            latent = maybe_unwrap(self.latent_gene_pool)(latent_id = latent_id)
 
-        critic_forward = self.critic
+        critic_forward = maybe_unwrap(self.critic)
 
         if use_ema_if_available and self.use_critic_ema:
             critic_forward = self.critic_ema
@@ -922,6 +1001,9 @@ class Agent(Module):
 
         dataloader = DataLoader(dataset, batch_size = self.batch_size, shuffle = True)
 
+        if self.wrap_with_accelerate:
+            dataloader = self.accelerate.prepare(dataloader)
+
         # updating actor and critic
 
         self.actor.train()
@@ -955,7 +1037,7 @@ class Agent(Module):
                 actor_loss.backward()
 
                 if exists(self.has_grad_clip):
-                    nn.utils.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
+                    self.accelerate.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
 
                 self.actor_optim.step()
                 self.actor_optim.zero_grad()
@@ -971,7 +1053,7 @@ class Agent(Module):
                 critic_loss.backward()
 
                 if exists(self.has_grad_clip):
-                    nn.utils.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm)
+                    self.accelerate.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm)
 
                 self.critic_optim.step()
                 self.critic_optim.zero_grad()
@@ -994,6 +1076,9 @@ class Agent(Module):
 
                     (diversity_loss * self.diversity_aux_loss_weight).backward()
 
+                if exists(self.has_grad_clip):
+                    self.accelerate.clip_grad_norm_(self.latent_gene_pool.parameters(), self.max_grad_norm)
+
                 self.latent_optim.step()
                 self.latent_optim.zero_grad()
 
@@ -1128,7 +1213,7 @@ class EPO(Module):
         self.max_episode_length = max_episode_length
         self.fix_environ_across_latents = fix_environ_across_latents
 
-        self.register_buffer('dummy', tensor(0))
+        self.register_buffer('dummy', tensor(0, device = agent.device))
 
     @property
     def device(self):
@@ -1215,7 +1300,7 @@ class EPO(Module):
 
             # get latent from pool
 
-            latent = self.agent.latent_gene_pool(latent_id = latent_id) if self.agent.has_latent_genes else None
+            latent = self.agent.unwrapped_latent_gene_pool(latent_id = latent_id) if self.agent.has_latent_genes else None
 
             # until maximum episode length
 
@@ -1225,11 +1310,11 @@ class EPO(Module):
 
                 # sample action
 
-                action, log_prob = temp_batch_dim(self.agent.get_actor_actions)(state, latent = latent, sample = True, temperature = self.action_sample_temperature)
+                action, log_prob = temp_batch_dim(self.agent.get_actor_actions)(state, latent = latent, sample = True, temperature = self.action_sample_temperature, use_unwrapped_model = True)
 
                 # values
 
-                value = temp_batch_dim(self.agent.get_critic_values)(state, latent = latent, use_ema_if_available = True)
+                value = temp_batch_dim(self.agent.get_critic_values)(state, latent = latent, use_ema_if_available = True, use_unwrapped_model = True)
 
                 # get the next state, action, and reward
 

{evolutionary_policy_optimization-0.1.1 → evolutionary_policy_optimization-0.1.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.1.1"
+version = "0.1.4"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -24,7 +24,7 @@ classifiers = [
   ]
 
 dependencies = [
-    "accelerate",
+    "accelerate>=1.6.0",
     "adam-atan2-pytorch",
     'assoc-scan>=0.0.2',
     'einx>=0.3.0',

{evolutionary_policy_optimization-0.1.1 → evolutionary_policy_optimization-0.1.4}/train_gym.py

@@ -26,18 +26,19 @@ agent = env.to_epo_agent(
     latent_gene_pool_kwargs = dict(
         frac_natural_selected = 0.5,
         frac_tournaments = 0.5
+    ),
+    accelerate_kwargs = dict(
+        cpu = False
     )
 )
 
 epo = EPO(
     agent,
-    episodes_per_latent = 1,
+    episodes_per_latent = 5,
     max_episode_length = 10,
-    action_sample_temperature = 1.
+    action_sample_temperature = 1.,
 )
 
-epo.to('cpu' if not torch.cuda.is_available() else 'cuda')
-
-epo(agent, env, num_learning_cycles = 1)
+epo(agent, env, num_learning_cycles = 5)
 
 agent.save('./agent.pt', overwrite = True)