evolutionary-policy-optimization 0.1.6__tar.gz → 0.1.7__tar.gz

{evolutionary_policy_optimization-0.1.6 → evolutionary_policy_optimization-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.1.6
+Version: 0.1.7
 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.1.6 → evolutionary_policy_optimization-0.1.7}/evolutionary_policy_optimization/epo.py

@@ -369,7 +369,6 @@ class Critic(Module):
         hl_gauss_loss = self.to_pred.hl_gauss_loss
 
         self.maybe_bins_to_value = hl_gauss_loss if not use_regression else identity
-        self.maybe_value_to_bins = hl_gauss_loss.transform_to_logprobs if not use_regression else identity
         self.loss_fn = hl_gauss_loss if not use_regression else F.mse_loss
 
     def forward_for_loss(
@@ -386,7 +385,7 @@ class Critic(Module):
 
         clipped_value = old_values + (value - old_values).clamp(1. - eps_clip, 1. + eps_clip)
 
-        loss = self.loss_fn(value, target, reduction = 'none')
+        loss = self.loss_fn(logits, target, reduction = 'none')
         clipped_loss = self.loss_fn(clipped_value, target, reduction = 'none')
 
         return torch.max(loss, clipped_loss).mean()
@@ -844,7 +843,11 @@ class Agent(Module):
 
         dummy = tensor(0)
 
+        self.clip_grad_norm_ = nn.utils.clip_grad_norm_
+
         if wrap_with_accelerate:
+            self.clip_grad_norm_ = self.accelerate.clip_grad_norm_
+
             (
                 self.actor,
                 self.critic,
@@ -1071,7 +1074,7 @@ class Agent(Module):
                 actor_loss.backward()
 
                 if exists(self.has_grad_clip):
-                    self.accelerate.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
+                    self.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
 
                 self.actor_optim.step()
                 self.actor_optim.zero_grad()
@@ -1089,7 +1092,7 @@ class Agent(Module):
                 critic_loss.backward()
 
                 if exists(self.has_grad_clip):
-                    self.accelerate.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm)
+                    self.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm)
 
                 self.critic_optim.step()
                 self.critic_optim.zero_grad()
@@ -1113,7 +1116,7 @@ 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.clip_grad_norm_(self.latent_gene_pool.parameters(), self.max_grad_norm)
 
                 self.latent_optim.step()
                 self.latent_optim.zero_grad()

{evolutionary_policy_optimization-0.1.6 → evolutionary_policy_optimization-0.1.7}/pyproject.toml

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

{evolutionary_policy_optimization-0.1.6 → evolutionary_policy_optimization-0.1.7}/tests/test_epo.py

@@ -103,7 +103,8 @@ def test_e2e_with_mock_env(
             frozen_latents = frozen_latents,
             frac_natural_selected = 0.75,
             frac_tournaments = 0.9
-        )
+        ),
+        wrap_with_accelerate = False
     )
 
     epo = EPO(