evolutionary-policy-optimization 0.3.1__tar.gz → 0.3.5__tar.gz

{evolutionary_policy_optimization-0.3.1 → evolutionary_policy_optimization-0.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: evolutionary-policy-optimization
-Version: 0.3.1
+Version: 0.3.5
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -33,13 +33,14 @@ Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Requires-Dist: accelerate>=1.6.0
 Requires-Dist: assoc-scan>=0.0.2
 Requires-Dist: einops>=0.8.1
 Requires-Dist: einx>=0.3.0
 Requires-Dist: ema-pytorch>=0.7.7
 Requires-Dist: hl-gauss-pytorch>=0.2.8
+Requires-Dist: mean-conc-beta>=0.0.7
 Requires-Dist: torch>=2.2
 Requires-Dist: tqdm
 Requires-Dist: x-mlps-pytorch>=0.3.0

{evolutionary_policy_optimization-0.3.1 → evolutionary_policy_optimization-0.3.5}/evolutionary_policy_optimization/__init__.py

@@ -1,3 +1,3 @@
 from evolutionary_policy_optimization.env_wrappers import GymnasiumEnvWrapper, rescale_from_to
-from evolutionary_policy_optimization.epo import EPO, MLP, Actor, Agent, BetaActionDistr, CategoricalActionDistr, Critic, LatentGenePool, create_agent
+from evolutionary_policy_optimization.epo import EPO, MLP, Actor, Agent, Beta, CategoricalActionDistr, Critic, LatentGenePool, create_agent
 from evolutionary_policy_optimization.mock_env import Env, VectorEnv

{evolutionary_policy_optimization-0.3.1 → evolutionary_policy_optimization-0.3.5}/evolutionary_policy_optimization/env_wrappers.py

@@ -5,9 +5,8 @@ from torch.nn import Module
 
 from evolutionary_policy_optimization.epo import Agent, create_agent, exists
 
-
-def rescale_from_to(x, from_range = (0., 1.), to_range = (-1., 1.)):
-    # e.g. beta actions on (0, 1) -> the env's action bounds
+def rescale_from_to(x, from_range = (-1., 1.), to_range = (-1., 1.)):
+    # e.g. beta actions on (-1, 1) -> the env's action bounds
 
     from_low, from_high = from_range
     to_low, to_high = to_range
@@ -36,7 +35,7 @@ class GymnasiumEnvWrapper(Module):
         return self.env.reset(*args, **kwargs)
 
     def step(self, actions, *args, **kwargs):
-        # beta lives on (0, 1) - rescale to the env's bounds at the interface
+        # continuous actions are on (-1, 1) - rescale to env bounds at interface
 
         if exists(self.rescale_to):
             actions = rescale_from_to(actions, to_range = self.rescale_to)

{evolutionary_policy_optimization-0.3.1 → evolutionary_policy_optimization-0.3.5}/evolutionary_policy_optimization/epo.py

@@ -15,15 +15,15 @@ import torch
 import torch.distributed as dist
 import torch.nn.functional as F
 from accelerate import Accelerator
-from adam_atan2_pytorch import AdoptAtan2
+from torch.optim import AdamW
 from assoc_scan import AssocScan
 from einops import einsum, rearrange, reduce, repeat
 from einops.layers.torch import Rearrange
 from ema_pytorch import EMA
 from hl_gauss_pytorch import HLGaussLayer
+from mean_conc_beta import Beta
 from torch import Tensor, cat, from_numpy, is_tensor, nn, stack, tensor
-from torch.distributions import Beta, Categorical, Distribution
-
+from torch.distributions import Categorical, Distribution
 from torch.nn import Linear, Module, ModuleList
 from torch.utils._pytree import tree_map
 from torch.utils.data import DataLoader, TensorDataset
@@ -545,7 +545,7 @@ class DiversityDiscr(Module):
 
 # action distributions - the actor always returns a `Distribution`, either
 # categorical (discrete) or beta mean-conc (continuous). beta is defined over
-# (0, 1) - rescale to the env's action range at the interface
+# (-1, 1) - multiplied by some constant usually (e.g. 0.4) for the env action range
 
 class CategoricalActionDistr(Module):
     def forward(self, logits, temperature = 1.):
@@ -554,58 +554,6 @@ class CategoricalActionDistr(Module):
 
         return Categorical(logits = logits)
 
-class BetaActionDistr(Module):
-    def __init__(
-        self,
-        init_conc = 2.,
-        min_conc = 0.,
-        eps = 1e-5
-    ):
-        super().__init__()
-        assert init_conc > min_conc
-
-        self.init_conc = init_conc
-        self.min_conc = min_conc
-        self.eps = eps
-
-        # softplus offset so the concentration at raw_conc = 0 is exactly init_conc
-
-        self.raw_init_conc = math.log(math.expm1(init_conc - min_conc))
-
-    def mean(self, params):
-        # the beta mean is exactly (tanh(raw_mean) + 1) / 2 by construction
-
-        raw_mean, _ = params.unbind(dim = -1)
-        return ((torch.tanh(raw_mean) + 1.) * 0.5).clamp(min = self.eps, max = 1. - self.eps)
-
-    def entropy(self, params, temperature = 1.):
-        # entropy is over the rescaled (2x - 1) in [-1, 1], with log 2 jacobian
-        # adjustment for the differential entropy of an affine transform
-
-        distr = self.forward(params, temperature = temperature)
-        return distr.entropy() + math.log(2.)
-
-    def forward(self, params, temperature = 1.):
-        _, raw_conc = params.unbind(dim = -1)
-
-        mean = self.mean(params)
-
-        conc = F.softplus(raw_conc + self.raw_init_conc) + self.min_conc
-
-        # concentration floor - unimodal (alpha > 1 and beta > 1), mean kept exact
-
-        conc = conc + 1. / torch.minimum(mean, 1. - mean).clamp(min = self.eps)
-
-        # temperature scales the concentration - lower temperature, sharper policy
-
-        if temperature > 0. and temperature != 1.:
-            conc = conc / temperature
-
-        alpha = mean * conc
-        beta = (1. - mean) * conc
-
-        return Beta(alpha, beta)
-
 # self-predictive representations (SPR) over the actor's hidden embedding
 
 class HiddenSpr(Module):
@@ -689,6 +637,7 @@ class Actor(Module):
         state_norm: StateNorm | None = None,
         dim_latent = 0,
         action_is_continuous = False, # continuous control - beta policy
+        beta_kwargs: dict = dict(),
     ):
         super().__init__()
 
@@ -720,7 +669,7 @@ class Actor(Module):
                 nn.Linear(dim, num_actions, bias = False),
             )
 
-        self.action_distr = BetaActionDistr() if self.beta_actions else CategoricalActionDistr()
+        self.action_distr = Beta(**beta_kwargs) if self.beta_actions else CategoricalActionDistr()
 
     def latent(
         self,
@@ -1079,6 +1028,8 @@ class LatentGenePool(Module):
         should_update_per_island = self.should_run_genetic_algorithm(fitness)
 
         if not should_update_per_island.any():
+            self.advance_step_()
+
             if inplace:
                 return False, None
 
@@ -1229,13 +1180,13 @@ class Agent(Module):
         actor: Actor,
         critic: Critic,
         latent_gene_pool: LatentGenePool | None,
-        optim_klass = AdoptAtan2,
+        optim_klass = AdamW,
         state_norm: StateNorm | None = None,
         actor_lr = 8e-4,
         critic_lr = 8e-4,
         latent_lr = 1e-5,
-        actor_weight_decay = 5e-4,
-        critic_weight_decay = 5e-4,
+        actor_weight_decay = 0.,
+        critic_weight_decay = 0.,
         diversity_aux_loss_weight = 0.,
         use_critic_ema = True,
         critic_ema_beta = 0.95,

{evolutionary_policy_optimization-0.3.1 → evolutionary_policy_optimization-0.3.5}/pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.3.1"
+version = "0.3.5"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
 ]
 readme = "README.md"
-requires-python = ">= 3.9"
+requires-python = ">= 3.10"
 license = { file = "LICENSE" }
 keywords = [
     'artificial intelligence',
@@ -30,6 +30,7 @@ dependencies = [
     'einops>=0.8.1',
     'ema-pytorch>=0.7.7',
     'hl-gauss-pytorch>=0.2.8',
+    'mean-conc-beta>=0.0.7',
     'torch>=2.2',
     'tqdm',
     "x-mlps-pytorch>=0.3.0",