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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.0
+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,6 +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>=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.0 → evolutionary_policy_optimization-0.1.2}/evolutionary_policy_optimization/__init__.py

@@ -9,3 +9,7 @@ from evolutionary_policy_optimization.epo import (
 )
 
 from evolutionary_policy_optimization.mock_env import Env
+
+from evolutionary_policy_optimization.env_wrappers import (
+    GymnasiumEnvWrapper
+)

evolutionary_policy_optimization-0.1.2/evolutionary_policy_optimization/env_wrappers.py

@@ -0,0 +1,36 @@
+import torch
+from torch.nn import Module
+
+from evolutionary_policy_optimization.epo import create_agent, Agent
+
+class GymnasiumEnvWrapper(Module):
+    def __init__(
+        self,
+        env
+    ):
+        super().__init__()
+        self.env = env
+
+    def reset(self, *args, **kwargs):
+        return self.env.reset(*args, **kwargs)
+
+    def step(self, *args, **kwargs):
+        return self.env.step(*args, **kwargs)
+
+    def to_agent_hparams(self):
+        return dict(
+            dim_state = self.env.observation_space.shape[0],
+            actor_num_actions = self.env.action_space.n
+        )
+
+    def to_epo_agent(
+        self,
+        *args,
+        **kwargs
+    ) -> Agent:
+
+        return create_agent(
+            *args,
+            **self.to_agent_hparams(),
+            **kwargs
+        )

{evolutionary_policy_optimization-0.1.0 → evolutionary_policy_optimization-0.1.2}/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()
 
@@ -1040,6 +1106,7 @@ def actor_loss(
 # agent contains the actor, critic, and the latent genetic pool
 
 def create_agent(
+    *,
     dim_state,
     num_latents,
     dim_latent,
@@ -1127,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):
@@ -1214,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
 
@@ -1224,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.0 → evolutionary_policy_optimization-0.1.2}/pyproject.toml

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

{evolutionary_policy_optimization-0.1.0 → evolutionary_policy_optimization-0.1.2}/train_gym.py

@@ -1,3 +1,10 @@
+import torch
+
+from evolutionary_policy_optimization import (
+    EPO,
+    GymnasiumEnvWrapper
+)
+
 # gymnasium
 
 import gymnasium as gym
@@ -7,38 +14,31 @@ env = gym.make(
     render_mode = 'rgb_array'
 )
 
-state_dim = env.observation_space.shape[0]
-num_actions = env.action_space.n
+env = GymnasiumEnvWrapper(env)
 
 # epo
 
-import torch
-
-from evolutionary_policy_optimization import (
-    create_agent,
-    EPO,
-    Env
-)
-
-agent = create_agent(
-    dim_state = state_dim,
-    num_latents = 1,
+agent = env.to_epo_agent(
+    num_latents = 8,
     dim_latent = 32,
-    actor_num_actions = num_actions,
     actor_dim_hiddens = (256, 128),
     critic_dim_hiddens = (256, 128, 64),
     latent_gene_pool_kwargs = dict(
-        frac_natural_selected = 0.5
+        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 = 5)
+
+agent.save('./agent.pt', overwrite = True)