evolutionary-policy-optimization 0.2.18__tar.gz → 0.2.20__tar.gz

{evolutionary_policy_optimization-0.2.18 → evolutionary_policy_optimization-0.2.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.2.18
+Version: 0.2.20
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -45,9 +45,13 @@ Requires-Dist: torch>=2.2
 Requires-Dist: tqdm
 Requires-Dist: x-mlps-pytorch>=0.3.0
 Provides-Extra: examples
+Requires-Dist: fire; extra == 'examples'
+Requires-Dist: gymnasium[box2d]>=1.0.0; extra == 'examples'
+Requires-Dist: moviepy; extra == 'examples'
 Requires-Dist: numpy; extra == 'examples'
-Requires-Dist: pufferlib>=2.0.6; extra == 'examples'
+Requires-Dist: swig; extra == 'examples'
 Requires-Dist: tqdm; extra == 'examples'
+Requires-Dist: wandb; extra == 'examples'
 Provides-Extra: experimental
 Requires-Dist: tensordict; extra == 'experimental'
 Requires-Dist: torchvision; extra == 'experimental'
@@ -149,6 +153,34 @@ agent.save('./agent.pt', overwrite = True)
 agent.load('./agent.pt')
 ```
 
+## Train Gym
+
+An example training script for LunarLander-v3 is provided at `train_gym.py`
+
+First install example dependencies
+
+```bash
+$ pip install '.[examples]' # or `uv pip install '.[examples]'`
+```
+
+Single process
+
+```bash
+$ python train_gym.py --cpu
+```
+
+Distributed with 8 processes (one latent per process)
+
+```bash
+$ torchrun --nproc_per_node=8 train_gym.py --cpu
+```
+
+With wandb logging
+
+```bash
+$ torchrun --nproc_per_node=8 train_gym.py --cpu --use_wandb
+```
+
 ## Contributing
 
 At the project root, run
@@ -323,6 +355,30 @@ That's it
 }
 ```
 
+```bibtex
+@article{Eysenbach2018DiversityIA,
+    title   = {Diversity is All You Need: Learning Skills without a Reward Function},
+    author  = {Benjamin Eysenbach and Abhishek Gupta and Julian Ibarz and Sergey Levine},
+    journal = {ArXiv},
+    year    = {2018},
+    volume  = {abs/1802.06070},
+    url     = {https://arxiv.org/abs/1802.06070}
+    url     = {https://arxiv.org/abs/2602.16863},
+```
+
+```bibtex
+@article{Zhu2025Hybrid,
+    author         = {Zhu, Zimo and Yu, Chuanqiang and Wang, Junti},
+    title          = {A Hybrid Genetic Algorithm and Proximal Policy Optimization System for Efficient Multi-Agent Task Allocation},
+    journal        = {Systems},
+    volume         = {13},
+    year           = {2025},
+    number         = {6},
+    article-number = {453},
+    url            = {https://www.mdpi.com/2079-8954/13/6/453},
+    issn           = {2079-8954}
+```
+
 ```bibtex
 @misc{kedia2026simtoolrealobjectcentricpolicyzeroshot,
     title   = {SimToolReal: An Object-Centric Policy for Zero-Shot Dexterous Tool Manipulation},
@@ -335,15 +391,4 @@ That's it
 }
 ```
 
-```bibtex
-@article{Eysenbach2018DiversityIA,
-    title   = {Diversity is All You Need: Learning Skills without a Reward Function},
-    author  = {Benjamin Eysenbach and Abhishek Gupta and Julian Ibarz and Sergey Levine},
-    journal = {ArXiv},
-    year    = {2018},
-    volume  = {abs/1802.06070},
-    url     = {https://arxiv.org/abs/1802.06070}
-}
-```
-
 *Evolution is cleverer than you are.* - Leslie Orgel

{evolutionary_policy_optimization-0.2.18 → evolutionary_policy_optimization-0.2.20}/README.md

@@ -91,6 +91,34 @@ agent.save('./agent.pt', overwrite = True)
 agent.load('./agent.pt')
 ```
 
+## Train Gym
+
+An example training script for LunarLander-v3 is provided at `train_gym.py`
+
+First install example dependencies
+
+```bash
+$ pip install '.[examples]' # or `uv pip install '.[examples]'`
+```
+
+Single process
+
+```bash
+$ python train_gym.py --cpu
+```
+
+Distributed with 8 processes (one latent per process)
+
+```bash
+$ torchrun --nproc_per_node=8 train_gym.py --cpu
+```
+
+With wandb logging
+
+```bash
+$ torchrun --nproc_per_node=8 train_gym.py --cpu --use_wandb
+```
+
 ## Contributing
 
 At the project root, run
@@ -265,6 +293,30 @@ That's it
 }
 ```
 
+```bibtex
+@article{Eysenbach2018DiversityIA,
+    title   = {Diversity is All You Need: Learning Skills without a Reward Function},
+    author  = {Benjamin Eysenbach and Abhishek Gupta and Julian Ibarz and Sergey Levine},
+    journal = {ArXiv},
+    year    = {2018},
+    volume  = {abs/1802.06070},
+    url     = {https://arxiv.org/abs/1802.06070}
+    url     = {https://arxiv.org/abs/2602.16863},
+```
+
+```bibtex
+@article{Zhu2025Hybrid,
+    author         = {Zhu, Zimo and Yu, Chuanqiang and Wang, Junti},
+    title          = {A Hybrid Genetic Algorithm and Proximal Policy Optimization System for Efficient Multi-Agent Task Allocation},
+    journal        = {Systems},
+    volume         = {13},
+    year           = {2025},
+    number         = {6},
+    article-number = {453},
+    url            = {https://www.mdpi.com/2079-8954/13/6/453},
+    issn           = {2079-8954}
+```
+
 ```bibtex
 @misc{kedia2026simtoolrealobjectcentricpolicyzeroshot,
     title   = {SimToolReal: An Object-Centric Policy for Zero-Shot Dexterous Tool Manipulation},
@@ -277,15 +329,4 @@ That's it
 }
 ```
 
-```bibtex
-@article{Eysenbach2018DiversityIA,
-    title   = {Diversity is All You Need: Learning Skills without a Reward Function},
-    author  = {Benjamin Eysenbach and Abhishek Gupta and Julian Ibarz and Sergey Levine},
-    journal = {ArXiv},
-    year    = {2018},
-    volume  = {abs/1802.06070},
-    url     = {https://arxiv.org/abs/1802.06070}
-}
-```
-
 *Evolution is cleverer than you are.* - Leslie Orgel

{evolutionary_policy_optimization-0.2.18 → evolutionary_policy_optimization-0.2.20}/evolutionary_policy_optimization/distributed.py

@@ -86,3 +86,7 @@ def all_gather_variable_dim(t, dim = 0, sizes = None):
     gathered_tensors = gathered_tensors.index_select(dim, indices)
 
     return gathered_tensors, sizes
+
+def all_gather(t, dim = 0):
+    gathered, _ = all_gather_variable_dim(t, dim = dim)
+    return gathered

{evolutionary_policy_optimization-0.2.18 → evolutionary_policy_optimization-0.2.20}/evolutionary_policy_optimization/epo.py

@@ -28,7 +28,7 @@ from evolutionary_policy_optimization.distributed import (
     is_distributed,
     get_world_and_rank,
     maybe_sync_seed,
-    all_gather_variable_dim,
+    all_gather,
     maybe_barrier
 )
 
@@ -1121,6 +1121,7 @@ class Agent(Module):
         get_fitness_scores: Callable[..., Tensor] = get_fitness_scores,
         wrap_with_accelerate: bool = True,
         accelerate_kwargs: dict = dict(),
+        accelerator = None,
     ):
         super().__init__()
 
@@ -1129,8 +1130,7 @@ class Agent(Module):
         self.wrap_with_accelerate = wrap_with_accelerate
 
         if wrap_with_accelerate:
-            accelerate = Accelerator(**accelerate_kwargs)
-            self.accelerate = accelerate
+            self.accelerate = accelerator if exists(accelerator) else Accelerator(**accelerate_kwargs)
 
         # state norm
 
@@ -1239,35 +1239,40 @@ class Agent(Module):
 
         if wrap_with_accelerate:
             self.clip_grad_norm_ = self.accelerate.clip_grad_norm_
+            device = self.accelerate.device
+
+            # device placement for modules without gradient parameters
+
+            for m in (self.state_norm, self.latent_gene_pool):
+                if exists(m):
+                    m.to(device)
+
+            # DDP wrap models with gradient parameters + prepare optimizers
 
             (
-                self.state_norm,
                 self.actor,
                 self.critic,
-                self.latent_gene_pool,
-                self.diversity_discr,
                 self.actor_optim,
                 self.critic_optim,
-                self.latent_optim,
-                self.diversity_discr_optim,
-            ) = tuple(
-                maybe(self.accelerate.prepare)(m) for m in (
-                    self.state_norm,
-                    self.actor,
-                    self.critic,
-                    self.latent_gene_pool,
-                    self.diversity_discr,
-                    self.actor_optim,
-                    self.critic_optim,
-                    self.latent_optim,
-                    self.diversity_discr_optim,
-                )
+            ) = self.accelerate.prepare(
+                self.actor,
+                self.critic,
+                self.actor_optim,
+                self.critic_optim,
             )
 
+            if exists(self.diversity_discr):
+                self.diversity_discr, self.diversity_discr_optim = self.accelerate.prepare(
+                    self.diversity_discr, self.diversity_discr_optim
+                )
+
+            if exists(self.latent_optim):
+                self.latent_optim = self.accelerate.prepare(self.latent_optim)
+
             if exists(self.critic_ema):
-                self.critic_ema.to(self.accelerate.device)
+                self.critic_ema.to(device)
 
-            step = step.to(self.accelerate.device)
+            step = step.to(device)
 
         # device tracking
 
@@ -1418,9 +1423,8 @@ class Agent(Module):
         maybe_barrier()
 
         if is_distributed():
-            memories = map(partial(all_gather_variable_dim, dim = 0), memories)
-
-            rewards_per_latent_episode = dist.all_reduce(rewards_per_latent_episode)
+            memories = [all_gather(m) for m in memories]
+            dist.all_reduce(rewards_per_latent_episode)
 
         # calculate fitness scores
 
@@ -1500,7 +1504,7 @@ class Agent(Module):
 
                 # learn critic with maybe classification loss
 
-                critic_loss = self.critic.forward_for_loss(
+                critic_loss = self.unwrap_model(self.critic).forward_for_loss(
                     states,
                     latents,
                     old_values = old_values,
@@ -1668,6 +1672,7 @@ def create_agent(
     critic_dim,
     critic_mlp_depth,
     use_critic_ema = True,
+    use_state_norm = False,
     latent_gene_pool_kwargs: dict = dict(),
     actor_kwargs: dict = dict(),
     critic_kwargs: dict = dict(),
@@ -1685,7 +1690,7 @@ def create_agent(
         **latent_gene_pool_kwargs
     ) if has_latent_genes else None
 
-    state_norm = StateNorm(dim = dim_state)
+    state_norm = StateNorm(dim = dim_state) if use_state_norm else None
 
     actor = Actor(
         num_actions = actor_num_actions,

{evolutionary_policy_optimization-0.2.18 → evolutionary_policy_optimization-0.2.20}/evolutionary_policy_optimization/experimental.py

@@ -7,6 +7,7 @@ import torch.nn.functional as F
 from torch.func import vmap, functional_call
 from torch.nn import Module, ParameterList
 
+import einx
 from einops import rearrange, reduce, repeat
 
 def exists(v):
@@ -31,11 +32,12 @@ def shrink_and_perturb_(
 
 def crossover_weights(
     w1, w2,
+    alpha = 0.5,
     shrink_perturb = False,
     shrink_factor = 0.4,
     perturb_factor = 0.1
 ):
-    assert w2.shape == w2.shape
+    assert w1.shape == w2.shape
 
     no_batch = w1.ndim == 2
 
@@ -58,12 +60,17 @@ def crossover_weights(
     u1, s1, v1 = torch.svd(w1)
     u2, s2, v2 = torch.svd(w2)
 
+    sign = einx.dot('b j d, b j d -> b d', v1, v2).sign()
+
+    u2 = einx.multiply('b i d, b d -> b i d', u2, sign)
+    v2 = einx.multiply('b j d, b d -> b j d', v2, sign)
+
     batch_randperm = torch.randn((batch, rank), device = w1.device).argsort(dim = -1)
-    mask = batch_randperm < (rank // 2)
+    mask = batch_randperm < int(rank * alpha)
 
-    u = torch.where(mask[:, None, :], u1, u2)
-    s = torch.where(mask, s1, s2)
-    v = torch.where(mask[:, :, None], v1, v2)
+    u = einx.where('b d, b i d, b i d -> b i d', mask, u1, u2)
+    s = einx.where('b d, b d, b d -> b d', mask, s1, s2)
+    v = einx.where('b d, b j d, b j d -> b j d', mask, v1, v2)
 
     out = u @ torch.diag_embed(s) @ v.mT
 

{evolutionary_policy_optimization-0.2.18 → evolutionary_policy_optimization-0.2.20}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.2.18"
+version = "0.2.20"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -14,7 +14,7 @@ keywords = [
     'reinforcement learning',
     'genetic algorithms',
     'robotics'
-  ]
+]
 classifiers = [
     'Development Status :: 4 - Beta',
     'Intended Audience :: Developers',
@@ -42,9 +42,13 @@ Repository = "https://github.com/lucidrains/evolutionary-policy-optimization"
 
 [project.optional-dependencies]
 examples = [
+    "fire",
+    "gymnasium[box2d]>=1.0.0",
+    "moviepy",
     "numpy",
-    "pufferlib>=2.0.6",
+    "swig",
     "tqdm",
+    "wandb",
 ]
 
 experimental = [