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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.14
+Version: 0.1.16
 Summary: x-evolution
 Project-URL: Homepage, https://pypi.org/project/x-evolution/
 Project-URL: Repository, https://github.com/lucidrains/x-evolution

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

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

{x_evolution-0.1.14 → x_evolution-0.1.16}/tests/test_evolution.py

@@ -45,3 +45,5 @@ def test_evo_strat(
 
     evo_strat('evolve', 1)
     evo_strat('more.evolve', 1)
+
+    fitnesses = evo_strat('more.evolve', 2, rollback_model_at_end = True)

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

@@ -80,6 +80,8 @@ from x_mlps_pytorch.residual_normed_mlp import ResidualNormedMLP
 
 actor = ResidualNormedMLP(dim_in = 8, dim = 24, depth = 2, residual_every = 1, dim_out = 4)
 
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
 evo_strat = EvoStrategy(
     actor,
     environment = LunarEnvironment(repeats = 2),
@@ -91,7 +93,10 @@ evo_strat = EvoStrategy(
     learned_noise_scale = True,
     use_sigma_optimizer = True,
     learning_rate = 1e-3,
-    noise_scale_learning_rate = 1e-4
+    noise_scale_learning_rate = 1e-4,
+    use_scheduler = True,
+    scheduler_klass = CosineAnnealingLR,
+    scheduler_kwargs = dict(T_max = 50_000)
 )
 
 evo_strat()

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

@@ -366,11 +366,17 @@ class EvoStrategy(Module):
         self,
         filename = 'evolved.model',
         num_generations = None,
-        disable_distributed = False
+        disable_distributed = False,
+        rollback_model_at_end = False
     ):
 
         model = self.noisable_model.to(self.device)
 
+        # maybe save model for rolling back (for meta-evo)
+
+        if rollback_model_at_end:
+            self.checkpoint('initial.model')
+
         # maybe sigmas
 
         if self.learned_noise_scale:
@@ -515,7 +521,7 @@ class EvoStrategy(Module):
             if self.learned_noise_scale:
                 packed_sigma, _ = pack(list(self.sigmas), '*')
                 avg_sigma = packed_sigma.mean().item()
-                msg += f' | avg sigma: {avg_sigma:.3f}'
+                msg += f' | average sigma: {avg_sigma:.3f}'
 
             self.print(msg)
 
@@ -533,8 +539,17 @@ class EvoStrategy(Module):
 
         self.print('evolution complete')
 
+        # final checkpoint
+
         self.checkpoint(f'{filename}.final.{generation}')
 
+        # maybe rollback
+
+        if rollback_model_at_end:
+            orig_state_dict = torch.load(str(self.checkpoint_folder / 'initial.model.pt'), weights_only = True)
+
+            self.model.load_state_dict(orig_state_dict)
+
         # return fitnesses across generations
         # for meta-evolutionary (nesting EvoStrategy within the environment of another and optimizing some meta-network)