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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.2.19
+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,33 +355,6 @@ That's it
 }
 ```
 
-```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},
-@misc{kedia2026simtoolrealobjectcentricpolicyzeroshot,
-    title   = {SimToolReal: An Object-Centric Policy for Zero-Shot Dexterous Tool Manipulation}, 
-    author  = {Kushal Kedia and Tyler Ga Wei Lum and Jeannette Bohg and C. Karen Liu},
-    year    = {2026},
-    eprint  = {2602.16863},
-    archivePrefix = {arXiv},
-    primaryClass = {cs.RO},
-    url     = {https://arxiv.org/abs/2602.16863},
-}
-```
-
 ```bibtex
 @article{Eysenbach2018DiversityIA,
     title   = {Diversity is All You Need: Learning Skills without a Reward Function},
@@ -358,7 +363,7 @@ That's it
     year    = {2018},
     volume  = {abs/1802.06070},
     url     = {https://arxiv.org/abs/1802.06070}
-    url     = {https://arxiv.org/abs/2602.16863}, 
+    url     = {https://arxiv.org/abs/2602.16863},
 ```
 
 ```bibtex
@@ -376,13 +381,13 @@ That's it
 
 ```bibtex
 @misc{kedia2026simtoolrealobjectcentricpolicyzeroshot,
-    title   = {SimToolReal: An Object-Centric Policy for Zero-Shot Dexterous Tool Manipulation}, 
+    title   = {SimToolReal: An Object-Centric Policy for Zero-Shot Dexterous Tool Manipulation},
     author  = {Kushal Kedia and Tyler Ga Wei Lum and Jeannette Bohg and C. Karen Liu},
     year    = {2026},
     eprint  = {2602.16863},
     archivePrefix = {arXiv},
     primaryClass = {cs.RO},
-    url     = {https://arxiv.org/abs/2602.16863}, 
+    url     = {https://arxiv.org/abs/2602.16863},
 }
 ```
 

{evolutionary_policy_optimization-0.2.19 → 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,33 +293,6 @@ That's it
 }
 ```
 
-```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},
-@misc{kedia2026simtoolrealobjectcentricpolicyzeroshot,
-    title   = {SimToolReal: An Object-Centric Policy for Zero-Shot Dexterous Tool Manipulation}, 
-    author  = {Kushal Kedia and Tyler Ga Wei Lum and Jeannette Bohg and C. Karen Liu},
-    year    = {2026},
-    eprint  = {2602.16863},
-    archivePrefix = {arXiv},
-    primaryClass = {cs.RO},
-    url     = {https://arxiv.org/abs/2602.16863},
-}
-```
-
 ```bibtex
 @article{Eysenbach2018DiversityIA,
     title   = {Diversity is All You Need: Learning Skills without a Reward Function},
@@ -300,7 +301,7 @@ That's it
     year    = {2018},
     volume  = {abs/1802.06070},
     url     = {https://arxiv.org/abs/1802.06070}
-    url     = {https://arxiv.org/abs/2602.16863}, 
+    url     = {https://arxiv.org/abs/2602.16863},
 ```
 
 ```bibtex
@@ -318,13 +319,13 @@ That's it
 
 ```bibtex
 @misc{kedia2026simtoolrealobjectcentricpolicyzeroshot,
-    title   = {SimToolReal: An Object-Centric Policy for Zero-Shot Dexterous Tool Manipulation}, 
+    title   = {SimToolReal: An Object-Centric Policy for Zero-Shot Dexterous Tool Manipulation},
     author  = {Kushal Kedia and Tyler Ga Wei Lum and Jeannette Bohg and C. Karen Liu},
     year    = {2026},
     eprint  = {2602.16863},
     archivePrefix = {arXiv},
     primaryClass = {cs.RO},
-    url     = {https://arxiv.org/abs/2602.16863}, 
+    url     = {https://arxiv.org/abs/2602.16863},
 }
 ```
 

{evolutionary_policy_optimization-0.2.19 → 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.19 → 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.19 → evolutionary_policy_optimization-0.2.20}/pyproject.toml

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