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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.14
+Version: 0.1.15
 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.15}/pyproject.toml

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

{x_evolution-0.1.14 → x_evolution-0.1.15}/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.15}/x_evolution/x_evolution.py

@@ -515,7 +515,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)