evolutionary-policy-optimization 0.0.66__py3-none-any.whl → 0.0.68__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -55,6 +55,9 @@ def xnor(x, y):
 def divisible_by(num, den):
     return (num % den) == 0
 
+def to_device(inp, device):
+    return tree_map(lambda t: t.to(device) if is_tensor(t) else t, inp)
+
 # tensor helpers
 
 def l2norm(t):
@@ -181,6 +184,8 @@ class MLP(Module):
         dim_latent = 0,
     ):
         super().__init__()
+        dim_latent = default(dim_latent, 0)
+
         assert len(dims) >= 2, 'must have at least two dimensions'
 
         # add the latent to the first dim
@@ -376,6 +381,7 @@ class LatentGenePool(Module):
         init_latent_fn: Callable | None = None
     ):
         super().__init__()
+        assert num_latents > 1
 
         maybe_l2norm = l2norm if l2norm_latent else identity
 
@@ -670,7 +676,7 @@ class Agent(Module):
         self,
         actor: Actor,
         critic: Critic,
-        latent_gene_pool: LatentGenePool,
+        latent_gene_pool: LatentGenePool | None,
         optim_klass = AdoptAtan2,
         actor_lr = 1e-4,
         critic_lr = 1e-4,
@@ -705,10 +711,14 @@ class Agent(Module):
         self.use_critic_ema = use_critic_ema
         self.critic_ema = EMA(critic, beta = critic_ema_beta, include_online_model = False, **ema_kwargs) if use_critic_ema else None
 
-        self.num_latents = latent_gene_pool.num_latents
         self.latent_gene_pool = latent_gene_pool
+        self.num_latents = latent_gene_pool.num_latents if exists(latent_gene_pool) else 1
+        self.has_latent_genes = exists(latent_gene_pool)
+
+        assert actor.dim_latent == critic.dim_latent
 
-        assert actor.dim_latent == critic.dim_latent == latent_gene_pool.dim_latent
+        if self.has_latent_genes:
+            assert latent_gene_pool.dim_latent == actor.dim_latent
 
         # gae function
 
@@ -730,13 +740,19 @@ class Agent(Module):
         self.actor_optim = optim_klass(actor.parameters(), lr = actor_lr, **actor_optim_kwargs)
         self.critic_optim = optim_klass(critic.parameters(), lr = critic_lr, **critic_optim_kwargs)
 
-        self.latent_optim = optim_klass(latent_gene_pool.parameters(), lr = latent_lr, **latent_optim_kwargs) if not latent_gene_pool.frozen_latents else None
+        self.latent_optim = optim_klass(latent_gene_pool.parameters(), lr = latent_lr, **latent_optim_kwargs) if exists(latent_gene_pool) and not latent_gene_pool.frozen_latents else None
 
         # promotes latents to be farther apart for diversity maintenance
 
         self.has_diversity_loss = diversity_aux_loss_weight > 0.
         self.diversity_aux_loss_weight = diversity_aux_loss_weight
 
+        self.register_buffer('dummy', tensor(0))
+
+    @property
+    def device(self):
+        return self.dummy.device
+
     def save(self, path, overwrite = False):
         path = Path(path)
 
@@ -746,7 +762,7 @@ class Agent(Module):
             actor = self.actor.state_dict(),
             critic = self.critic.state_dict(),
             critic_ema = self.critic_ema.state_dict() if self.use_critic_ema else None,
-            latents = self.latent_gene_pool.state_dict(),
+            latents = self.latent_gene_pool.state_dict() if self.has_latent_genes else None,
             actor_optim = self.actor_optim.state_dict(),
             critic_optim = self.critic_optim.state_dict(),
             latent_optim = self.latent_optim.state_dict() if exists(self.latent_optim) else None
@@ -768,7 +784,8 @@ class Agent(Module):
         if self.use_critic_ema:
             self.critic_ema.load_state_dict(pkg['critic_ema'])
 
-        self.latent_gene_pool.load_state_dict(pkg['latents'])
+        if exists(pkg.get('latents', None)):
+            self.latent_gene_pool.load_state_dict(pkg['latents'])
 
         self.actor_optim.load_state_dict(pkg['actor_optim'])
         self.critic_optim.load_state_dict(pkg['critic_optim'])
@@ -784,9 +801,8 @@ class Agent(Module):
         sample = False,
         temperature = 1.
     ):
-        assert exists(latent_id) or exists(latent)
 
-        if not exists(latent):
+        if not exists(latent) and exists(latent_id):
             latent = self.latent_gene_pool(latent_id = latent_id)
 
         logits = self.actor(state, latent)
@@ -807,9 +823,8 @@ class Agent(Module):
         latent = None,
         use_ema_if_available = False
     ):
-        assert exists(latent_id) or exists(latent)
 
-        if not exists(latent):
+        if not exists(latent) and exists(latent_id):
             latent = self.latent_gene_pool(latent_id = latent_id)
 
         critic_forward = self.critic
@@ -823,13 +838,19 @@ class Agent(Module):
         self,
         fitnesses
     ):
+        if not self.has_latent_genes:
+            return
+
         return self.latent_gene_pool.genetic_algorithm_step(fitnesses)
 
     def forward(
         self,
         memories_and_cumulative_rewards: MemoriesAndCumulativeRewards,
         epochs = 2
+
     ):
+        memories_and_cumulative_rewards = to_device(memories_and_cumulative_rewards, self.device)
+
         memories, rewards_per_latent_episode = memories_and_cumulative_rewards
 
         # stack memories
@@ -893,11 +914,14 @@ class Agent(Module):
                 old_values
             ) in dataloader:
 
-                latents = self.latent_gene_pool(latent_id = latent_gene_ids)
+                if self.has_latent_genes:
+                    latents = self.latent_gene_pool(latent_id = latent_gene_ids)
 
-                orig_latents = latents
-                latents = latents.detach()
-                latents.requires_grad_()
+                    orig_latents = latents
+                    latents = latents.detach()
+                    latents.requires_grad_()
+                else:
+                    latents = None
 
                 # learn actor
 
@@ -936,7 +960,7 @@ class Agent(Module):
 
                 # maybe update latents, if not frozen
 
-                if self.latent_gene_pool.frozen_latents:
+                if not self.has_latent_genes or self.latent_gene_pool.frozen_latents:
                     continue
 
                 orig_latents.backward(latents.grad)
@@ -952,7 +976,9 @@ class Agent(Module):
 
         # apply evolution
 
-        self.latent_gene_pool.genetic_algorithm_step(fitness_scores)
+        if self.has_latent_genes:
+
+            self.latent_gene_pool.genetic_algorithm_step(fitness_scores)
 
 # reinforcement learning related - ppo
 
@@ -1005,11 +1031,16 @@ def create_agent(
     **kwargs
 ) -> Agent:
 
+    has_latent_genes = num_latents > 1
+
+    if not has_latent_genes:
+        dim_latent = None
+
     latent_gene_pool = LatentGenePool(
         num_latents = num_latents,
         dim_latent = dim_latent,
         **latent_gene_pool_kwargs
-    )
+    ) if has_latent_genes else None
 
     actor = Actor(
         num_actions = actor_num_actions,
@@ -1069,7 +1100,7 @@ class EPO(Module):
         self.agent = agent
         self.action_sample_temperature = action_sample_temperature
 
-        self.num_latents = agent.latent_gene_pool.num_latents
+        self.num_latents = agent.latent_gene_pool.num_latents if agent.has_latent_genes else 1
         self.episodes_per_latent = episodes_per_latent
         self.max_episode_length = max_episode_length
         self.fix_environ_across_latents = fix_environ_across_latents
@@ -1159,7 +1190,7 @@ class EPO(Module):
 
             # get latent from pool
 
-            latent = self.agent.latent_gene_pool(latent_id = latent_id)
+            latent = self.agent.latent_gene_pool(latent_id = latent_id) if self.agent.has_latent_genes else None
 
             # until maximum episode length
 
@@ -1198,6 +1229,8 @@ class EPO(Module):
                     terminated
                 )
 
+                memory = Memory(*tuple(t.cpu() for t in memory))
+
                 memories.append(memory)
 
                 time += 1
@@ -1209,7 +1242,7 @@ class EPO(Module):
 
                 memory_for_gae = memory._replace(
                     episode_id = invalid_episode,
-                    value = next_value,
+                    value = next_value.cpu(),
                     done = tensor(True)
                 )
 

{evolutionary_policy_optimization-0.0.66.dist-info → evolutionary_policy_optimization-0.0.68.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.66
+Version: 0.0.68
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization

{evolutionary_policy_optimization-0.0.66.dist-info → evolutionary_policy_optimization-0.0.68.dist-info}/RECORD

@@ -1,9 +1,9 @@
 evolutionary_policy_optimization/__init__.py,sha256=0q0aBuFgWi06MLMD8FiHzBYQ3_W4LYWrwmCtF3u5H2A,201
 evolutionary_policy_optimization/distributed.py,sha256=7KgZdeS_wxBHo_du9XZFB1Cu318J-Bp66Xdr6Log_20,2423
-evolutionary_policy_optimization/epo.py,sha256=-ULtP9-EJDQv5TlWwkPwPiGsU7bKnD2qfAYvyDK8GIU,36912
+evolutionary_policy_optimization/epo.py,sha256=xhE_kHas54xGsgOese9SQEvyK7NKZqEuK3AiVhm0y7Q,38047
 evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
 evolutionary_policy_optimization/mock_env.py,sha256=Bv9ONFRbma8wpjUurc9aCk19A6ceiWitRnS3nwrIR64,1339
-evolutionary_policy_optimization-0.0.66.dist-info/METADATA,sha256=YsfcSX3Vf5ALrePot6Eq4NStsVPzNDJ3Py6sMgh7lOE,6220
-evolutionary_policy_optimization-0.0.66.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.66.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.66.dist-info/RECORD,,
+evolutionary_policy_optimization-0.0.68.dist-info/METADATA,sha256=hOOKOrrPQtQmK3zN1z5nkGJEoaQLyXUzs9ArsEKn1DE,6220
+evolutionary_policy_optimization-0.0.68.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.68.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.68.dist-info/RECORD,,