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

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,6 +62,17 @@ 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
 
@@ -721,10 +734,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 +793,43 @@ 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
+
+        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,
+                )
+            )
+
+        # device tracking
+
+        self.register_buffer('dummy', tensor(0, device = self.accelerate.device))
+
+        self.critic_ema.to(self.accelerate.device)
 
     @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)
 
@@ -820,13 +876,15 @@ class Agent(Module):
         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)
+        logits = maybe_unwrap(self.actor)(state, latent)
 
         if not sample:
             return logits
@@ -842,13 +900,15 @@ class Agent(Module):
         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 +982,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 +1018,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 +1034,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 +1057,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 +1194,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 +1281,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 +1291,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.dist-info → evolutionary_policy_optimization-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.1
+Version: 0.1.2
 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

{evolutionary_policy_optimization-0.1.1.dist-info → evolutionary_policy_optimization-0.1.2.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=hxXOUt4tCmlqN2zjvG6Qfp_icq2ApuMowx3xK0Wm3wY,39498
+evolutionary_policy_optimization/epo.py,sha256=VgiPvaG-Lib1JZRS3oNV7BSv-19YSWKK2yCMtqDqP-M,41682
 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.1.dist-info/METADATA,sha256=ToFlj-YAqNqTZx2W9LZXmhteQBBUCTO00uAFd-p9s48,6329
-evolutionary_policy_optimization-0.1.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.1.1.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.1.1.dist-info/RECORD,,
+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,,