x-evolution 0.1.16__tar.gz → 0.1.19__tar.gz

{x_evolution-0.1.16 → x_evolution-0.1.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.16
+Version: 0.1.19
 Summary: x-evolution
 Project-URL: Homepage, https://pypi.org/project/x-evolution/
 Project-URL: Repository, https://github.com/lucidrains/x-evolution
@@ -153,3 +153,4 @@ $ accelerate launch train.py
 }
 ```
 
+*Nothing makes sense except in the light of evolution* - Theodosius Dobzhansky

{x_evolution-0.1.16 → x_evolution-0.1.19}/README.md

@@ -104,3 +104,4 @@ $ accelerate launch train.py
 }
 ```
 
+*Nothing makes sense except in the light of evolution* - Theodosius Dobzhansky

{x_evolution-0.1.16 → x_evolution-0.1.19}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-evolution"
-version = "0.1.16"
+version = "0.1.19"
 description = "x-evolution"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_evolution-0.1.16 → x_evolution-0.1.19}/train_lunar.py

@@ -25,20 +25,25 @@ class LunarEnvironment(Module):
 
         env = gym.make('LunarLander-v3', render_mode = 'rgb_array')
 
-        rmtree(video_folder, ignore_errors = True)
+        self.env = env
+        self.max_steps = max_steps
+        self.repeats = repeats
+        self.video_folder = video_folder
+        self.render_every_eps = render_every_eps
+
+    def pre_main_callback(self):
+        # the `pre_main_callback` on the environment passed in is called before the start of the evolutionary strategies loop
+
+        rmtree(self.video_folder, ignore_errors = True)
 
-        env = gym.wrappers.RecordVideo(
-            env = env,
-            video_folder = video_folder,
+        self.env = gym.wrappers.RecordVideo(
+            env = self.env,
+            video_folder = self.video_folder,
             name_prefix = 'recording',
-            episode_trigger = lambda eps_num: (eps_num % render_every_eps) == 0,
+            episode_trigger = lambda eps_num: (eps_num % self.render_every_eps) == 0,
             disable_logger = True
         )
 
-        self.env = env
-        self.max_steps = max_steps
-        self.repeats = repeats
-
     def forward(self, model):
 
         device = next(model.parameters()).device

{x_evolution-0.1.16 → x_evolution-0.1.19}/x_evolution/x_evolution.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from functools import partial
 
 import torch
-from torch import tensor, Tensor, is_tensor, arange, randint
+from torch import tensor, Tensor, stack, is_tensor, arange, randint
 from torch.nn import Module, ModuleList, Parameter, ParameterList
 from torch.optim import SGD, Adam, Optimizer
 from torch.optim.lr_scheduler import LRScheduler
@@ -96,20 +96,30 @@ class EvoStrategy(Module):
         reject_generation_fitnesses_if: Callable[[Tensor], bool] | None = None
     ):
         super().__init__()
+        self.verbose = verbose
 
         if not exists(accelerator):
             accelerator = Accelerator(cpu = cpu, **accelerate_kwargs)
 
         self.accelerate = accelerator
 
+        # environment - with optional init
+
+        self.environment = environment
+
+        if accelerator.is_main_process:
+            if hasattr(environment, 'pre_main_callback') and callable(environment.pre_main_callback):
+                self.print('pre_main_callback detected on environment passed in and is invoked')
+                environment.pre_main_callback()
+
+        # take care of model and parameters
+
         if isinstance(model, list):
             model = ModuleList(model)
 
         self.model = model
         self.noisable_model = Noisable(model, low_rank = noise_low_rank)
 
-        self.environment = environment
-
         named_parameters_dict = dict(model.named_parameters())
 
         param_to_name_index = {param: name for name, param in named_parameters_dict.items()}
@@ -217,10 +227,6 @@ class EvoStrategy(Module):
 
         self.reject_generation_fitnesses_if = reject_generation_fitnesses_if
 
-        # verbose
-
-        self.verbose = verbose
-
         # checkpointing
 
         self.checkpoint_every = checkpoint_every
@@ -367,8 +373,10 @@ class EvoStrategy(Module):
         filename = 'evolved.model',
         num_generations = None,
         disable_distributed = False,
-        rollback_model_at_end = False
+        rollback_model_at_end = False,
+        verbose = None
     ):
+        verbose = default(verbose, self.verbose)
 
         model = self.noisable_model.to(self.device)
 
@@ -553,4 +561,4 @@ class EvoStrategy(Module):
         # return fitnesses across generations
         # for meta-evolutionary (nesting EvoStrategy within the environment of another and optimizing some meta-network)
 
-        return fitnesses_across_generations
+        return stack(fitnesses_across_generations)