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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: evolutionary-policy-optimization
-Version: 0.3.1
+Version: 0.4.0
 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
@@ -366,6 +367,16 @@ That's it
 }
 ```
 
+```bibtex
+@inproceedings{Sharma2020DynamicsAwareUD,
+    title   = {Dynamics-Aware Unsupervised Discovery of Skills},
+    author  = {Archit Sharma and Shixiang Gu and Sergey Levine and Vikash Kumar and Karol Hausman},
+    booktitle = {International Conference on Learning Representations},
+    year    = {2020},
+    url     = {https://arxiv.org/abs/1907.01657}
+}
+```
+
 ```bibtex
 @article{Zhu2025Hybrid,
     author  = {Zhu, Zimo and Yu, Chuanqiang and Wang, Junti},

{evolutionary_policy_optimization-0.3.1 → evolutionary_policy_optimization-0.4.0}/README.md

@@ -305,6 +305,16 @@ That's it
 }
 ```
 
+```bibtex
+@inproceedings{Sharma2020DynamicsAwareUD,
+    title   = {Dynamics-Aware Unsupervised Discovery of Skills},
+    author  = {Archit Sharma and Shixiang Gu and Sergey Levine and Vikash Kumar and Karol Hausman},
+    booktitle = {International Conference on Learning Representations},
+    year    = {2020},
+    url     = {https://arxiv.org/abs/1907.01657}
+}
+```
+
 ```bibtex
 @article{Zhu2025Hybrid,
     author  = {Zhu, Zimo and Yu, Chuanqiang and Wang, Junti},

{evolutionary_policy_optimization-0.3.1 → evolutionary_policy_optimization-0.4.0}/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, DiversityDiscr, LatentGenePool, create_agent
 from evolutionary_policy_optimization.mock_env import Env, VectorEnv

{evolutionary_policy_optimization-0.3.1 → evolutionary_policy_optimization-0.4.0}/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.4.0}/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
@@ -513,7 +513,7 @@ class MLP(Module):
 
         return self.final_lime(x, prev_layer_inputs)
 
-# discriminator for predicting latent code from states and actions (DIAYN - Eysenbach et al. 2018)
+# discriminator for predicting latent code from states or state transitions (DIAYN / DADS)
 
 class DiversityDiscr(Module):
     def __init__(
@@ -521,15 +521,19 @@ class DiversityDiscr(Module):
         dim_state,
         num_latents,
         dim = 64,
-        depth = 2
+        depth = 2,
+        time_gap = 1
     ):
         super().__init__()
+        self.time_gap = time_gap
+        self.state_transition = time_gap > 0
+
         self.state_proj = nn.Linear(dim_state, dim)
 
         self.net = AttnResidualNormedMLP(
             dim = dim,
             depth = depth,
-            dim_in = dim * 2,
+            dim_in = dim * 2 if self.state_transition else dim,
             dim_out = num_latents
         )
 
@@ -538,14 +542,18 @@ class DiversityDiscr(Module):
             if isinstance(module, (nn.Linear, nn.LayerNorm, nn.RMSNorm)):
                 module.reset_parameters()
 
-    def forward(self, state, next_state):
+    def forward(self, state, next_state = None):
         state_embed = self.state_proj(state)
+
+        if not self.state_transition or not exists(next_state):
+            return self.net(state_embed)
+
         next_state_embed = self.state_proj(next_state)
         return self.net((state_embed, next_state_embed))
 
 # 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 +562,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 +645,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 +677,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 +1036,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 +1188,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,
@@ -1264,6 +1223,7 @@ class Agent(Module):
         critic_optim_kwargs: dict = dict(),
         latent_optim_kwargs: dict = dict(),
         use_diversity_discr = False,
+        diversity_time_gap = 1, # 0 for DIAYN (state only), >= 1 for DADS (state transition with given time gap)
         diversity_discr_kwargs: dict = dict(dim = 64, depth = 2),
         diversity_discr_lr = 3e-4,
         diversity_discr_optim_kwargs: dict = dict(),
@@ -1354,14 +1314,16 @@ class Agent(Module):
         # diversity discriminator
 
         self.use_diversity_discr = use_diversity_discr
+        self.diversity_time_gap = diversity_time_gap
 
         if use_diversity_discr:
-            assert exists(latent_gene_pool), 'latent_gene_pool must be present to use DIAYN'
+            assert exists(latent_gene_pool), 'latent_gene_pool must be present to use DIAYN / DADS'
             dim_state = actor.init_layer[0].in_features
 
             self.diversity_discr = DiversityDiscr(
                 dim_state = dim_state,
                 num_latents = latent_gene_pool.num_latents,
+                time_gap = diversity_time_gap,
                 **diversity_discr_kwargs
             )
 
@@ -1492,6 +1454,7 @@ class Agent(Module):
             critic_ema = self.critic_ema.state_dict() if self.use_critic_ema else None,
             latents = unwrap(self.latent_gene_pool).state_dict() if self.has_latent_genes else None,
             diversity_discr = unwrap(self.diversity_discr).state_dict() if self.use_diversity_discr else None,
+            diversity_time_gap = self.diversity_time_gap,
             has_diversity_discr_warmed_up = self.has_diversity_discr_warmed_up.item(),
             actor_optim = unwrap_optim(self.actor_optim).state_dict(),
             critic_optim = unwrap_optim(self.critic_optim).state_dict(),
@@ -1524,6 +1487,9 @@ class Agent(Module):
         if self.use_diversity_discr and 'diversity_discr' in pkg and exists(pkg['diversity_discr']):
             unwrap(self.diversity_discr).load_state_dict(pkg['diversity_discr'])
 
+        if 'diversity_time_gap' in pkg:
+            self.diversity_time_gap = pkg['diversity_time_gap']
+
         if 'has_diversity_discr_warmed_up' in pkg:
             self.has_diversity_discr_warmed_up.copy_(tensor(pkg['has_diversity_discr_warmed_up']))
 
@@ -1771,7 +1737,8 @@ class Agent(Module):
                 diversity_discr_loss = self.zero
 
                 if self.use_diversity_discr:
-                    diversity_discr_logits = self.diversity_discr(states, next_states)
+                    discr_inp = (states,) if self.diversity_time_gap == 0 else (states, next_states)
+                    diversity_discr_logits = self.diversity_discr(*discr_inp)
                     diversity_discr_loss = F.cross_entropy(diversity_discr_logits, latent_gene_ids)
 
                     diversity_discr_loss.backward()
@@ -1927,6 +1894,7 @@ def create_agent(
     use_critic_ema = True,
     use_state_norm = False,
     action_is_continuous = False, # continuous control - beta policy
+    diversity_time_gap = 1,
     latent_gene_pool_kwargs: dict = dict(),
     actor_kwargs: dict = dict(),
     critic_kwargs: dict = dict(),
@@ -1972,6 +1940,7 @@ def create_agent(
         state_norm = state_norm,
         latent_gene_pool = latent_gene_pool,
         use_critic_ema = use_critic_ema,
+        diversity_time_gap = diversity_time_gap,
         **kwargs
     )
 
@@ -2208,7 +2177,16 @@ class EPO(Module):
 
                 if self.agent.use_diversity_discr and self.agent.has_diversity_discr_warmed_up.item():
                     with torch.no_grad():
-                        logits = self.agent.unwrap_model(self.agent.diversity_discr)(rearrange(state, '... -> 1 ...'), rearrange(next_state, '... -> 1 ...'))
+                        discr = self.agent.unwrap_model(self.agent.diversity_discr)
+                        time_gap = self.agent.diversity_time_gap
+
+                        if time_gap == 0:
+                            logits = discr(rearrange(state, '... -> 1 ...'))
+                        elif time_gap == 1 or len(slot.memories) < (time_gap - 1):
+                            logits = discr(rearrange(state, '... -> 1 ...'), rearrange(next_state, '... -> 1 ...'))
+                        else:
+                            past_state = slot.memories[-(time_gap - 1)].state
+                            logits = discr(rearrange(past_state, '... -> 1 ...'), rearrange(next_state, '... -> 1 ...'))
 
                         latent_log_probs = logits.log_softmax(dim = -1)
                         diversity_reward = latent_log_probs[0, latent_id] + math.log(self.agent.num_latents)
@@ -2233,6 +2211,11 @@ class EPO(Module):
 
                 slot.memories.append(memory)
 
+                if self.agent.use_diversity_discr and self.agent.diversity_time_gap > 1:
+                    time_gap = self.agent.diversity_time_gap
+                    if len(slot.memories) >= time_gap:
+                        slot.memories[-time_gap] = slot.memories[-time_gap]._replace(next_state = next_state)
+
                 rewards_per_latent_episode[latent_id, episode_id] += reward
 
                 # episode is done - bootstrap value if truncated, then refill

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

@@ -1,12 +1,12 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.3.1"
+version = "0.4.0"
 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",