x-evolution 0.1.5__tar.gz → 0.1.7__tar.gz

{x_evolution-0.1.5 → x_evolution-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.5
+Version: 0.1.7
 Summary: x-evolution
 Project-URL: Homepage, https://pypi.org/project/x-evolution/
 Project-URL: Repository, https://github.com/lucidrains/x-evolution
@@ -38,7 +38,7 @@ Requires-Dist: accelerate
 Requires-Dist: beartype
 Requires-Dist: einops>=0.8.0
 Requires-Dist: torch>=2.4
-Requires-Dist: x-mlps-pytorch>=0.1.30
+Requires-Dist: x-mlps-pytorch>=0.1.31
 Requires-Dist: x-transformers>=2.11.23
 Provides-Extra: examples
 Requires-Dist: gymnasium[box2d]>=1.0.0; extra == 'examples'
@@ -78,7 +78,7 @@ evo_strat = EvoStrategy(
     model,
     environment = lambda model: torch.randint(0, 100, ()), # environment is just a function that takes in the individual model (with unique noise) and outputs a scalar (the fitness) the measure you are selecting for
     noise_population_size = 30, # increase this for better gradient estimates
-    noise_scale = 1e-3,       # the scale of the perturbation noise
+    noise_scale = 1e-2,         # the scale of the perturbation noise, also the initial noise scale (sigma) if `learned_noise_scale` = True
     num_generations = 100,    # number of generations / training steps
     learning_rate = 1e-3,     # scale on update derived by fitness and perturb noises
     params_to_optimize = None # defaults to all parameters, but can be [str {param name}] or [Parameter]

{x_evolution-0.1.5 → x_evolution-0.1.7}/README.md

@@ -29,7 +29,7 @@ evo_strat = EvoStrategy(
     model,
     environment = lambda model: torch.randint(0, 100, ()), # environment is just a function that takes in the individual model (with unique noise) and outputs a scalar (the fitness) the measure you are selecting for
     noise_population_size = 30, # increase this for better gradient estimates
-    noise_scale = 1e-3,       # the scale of the perturbation noise
+    noise_scale = 1e-2,         # the scale of the perturbation noise, also the initial noise scale (sigma) if `learned_noise_scale` = True
     num_generations = 100,    # number of generations / training steps
     learning_rate = 1e-3,     # scale on update derived by fitness and perturb noises
     params_to_optimize = None # defaults to all parameters, but can be [str {param name}] or [Parameter]

{x_evolution-0.1.5 → x_evolution-0.1.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-evolution"
-version = "0.1.5"
+version = "0.1.7"
 description = "x-evolution"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -27,7 +27,7 @@ dependencies = [
     "beartype",
     "einops>=0.8.0",
     "torch>=2.4",
-    "x-mlps-pytorch>=0.1.30",
+    "x-mlps-pytorch>=0.1.31",
     "x-transformers>=2.11.23"
 ]
 

{x_evolution-0.1.5 → x_evolution-0.1.7}/tests/test_evolution.py

@@ -11,12 +11,16 @@ model = MLP(8, 16, 4)
 @param('noise_low_rank', (None, 1))
 @param('mirror_sampling', (False, True))
 @param('multi_models', (False, True))
+@param('learned_sigma', (False, True))
+@param('use_sigma_optimizer', (False, True))
 def test_evo_strat(
     params_to_optimize,
     use_optimizer,
     noise_low_rank,
     mirror_sampling,
-    multi_models
+    multi_models,
+    learned_sigma,
+    use_sigma_optimizer
 ):
     from random import randrange
 
@@ -31,10 +35,12 @@ def test_evo_strat(
         to_optim,
         environment = lambda model: float(randrange(100)),
         num_generations = 2,
+        learned_noise_scale = learned_sigma,
         params_to_optimize = params_to_optimize,
         use_optimizer = use_optimizer,
         noise_low_rank = noise_low_rank,
-        mirror_sampling = mirror_sampling
+        mirror_sampling = mirror_sampling,
+        use_sigma_optimizer = use_sigma_optimizer
     )
 
     evo_strat('evolve', 1)

{x_evolution-0.1.5 → x_evolution-0.1.7}/train_lunar.py

@@ -19,6 +19,7 @@ class LunarEnvironment(Module):
         video_folder = './recordings',
         render_every_eps = 500,
         max_steps = 500,
+        repeats = 1
     ):
         super().__init__()
 
@@ -36,36 +37,40 @@ class LunarEnvironment(Module):
 
         self.env = env
         self.max_steps = max_steps
+        self.repeats = repeats
 
     def forward(self, model):
 
         device = next(model.parameters()).device
 
         seed = torch.randint(0, int(1e6), ())
-        state, _ = self.env.reset(seed = seed.item())
 
-        step = 0
         cum_reward = 0.
 
-        while step < self.max_steps:
+        for _ in range(self.repeats):
+            state, _ = self.env.reset(seed = seed.item())
 
-            state = torch.from_numpy(state).to(device)
+            step = 0
 
-            action_logits = model(state)
+            while step < self.max_steps:
 
-            action = F.gumbel_softmax(action_logits, hard = True).argmax(dim = -1)
+                state = torch.from_numpy(state).to(device)
 
-            next_state, reward, truncated, terminated, *_ = self.env.step(action.item())
+                action_logits = model(state)
 
-            cum_reward += float(reward)
-            step += 1
+                action = F.gumbel_softmax(action_logits, hard = True).argmax(dim = -1)
 
-            state = next_state
+                next_state, reward, truncated, terminated, *_ = self.env.step(action.item())
 
-            if truncated or terminated:
-                break
+                cum_reward += float(reward)
+                step += 1
 
-        return cum_reward
+                state = next_state
+
+                if truncated or terminated:
+                    break
+
+        return cum_reward / self.repeats
 
 # evo strategy
 
@@ -84,12 +89,15 @@ actor = ResidualNormedMLP(
 
 evo_strat = EvoStrategy(
     actor,
-    environment = LunarEnvironment(),
+    environment = LunarEnvironment(repeats = 2),
     num_generations = 50_000,
     noise_population_size = 60,
     noise_low_rank = 2,
     noise_scale = 1e-1,
-    learning_rate = 1e-3
+    learned_noise_scale = True,
+    use_sigma_optimizer = True,
+    learning_rate = 5e-2,
+    noise_scale_learning_rate = 1e-2
 )
 
 evo_strat()

{x_evolution-0.1.5 → x_evolution-0.1.7}/train_mnist.py

@@ -20,7 +20,7 @@ model = nn.Sequential(
     ResidualNormedMLP(dim_in = 784, dim = 512, depth = 8, residual_every = 2, dim_out = 10)
 ).half()
 
-batch_size = 128
+batch_size = 256
 
 # data
 
@@ -54,7 +54,7 @@ evo_strat = EvoStrategy(
     noise_scale = 1e-2,
     noise_low_rank = 2,
     num_generations = 10_000,
-    learning_rate = 1e-5
+    learning_rate = 1e-3
 )
 
 evo_strat()

{x_evolution-0.1.5 → x_evolution-0.1.7}/train_xor.py

@@ -12,6 +12,8 @@ model = nn.Sequential(
     nn.Linear(16, 2)
 ).half()
 
+batch_size = 128
+
 # fitness as inverse of loss
 
 from x_evolution import EvoStrategy
@@ -19,7 +21,7 @@ from x_evolution import EvoStrategy
 def loss_xor(model):
     device = next(model.parameters()).device
 
-    data = torch.randint(0, 2, (32, 2))
+    data = torch.randint(0, 2, (batch_size, 2))
     labels = data[:, 0] ^ data[:, 1]
 
     data, labels = tuple(t.to(device) for t in (data, labels))
@@ -36,10 +38,14 @@ evo_strat = EvoStrategy(
     model,
     environment = loss_xor,
     noise_population_size = 100,
-    noise_scale = 1e-1,
     noise_low_rank = 1,
     num_generations = 100_000,
-    learning_rate = 1e-4
+    learning_rate = 1e-2,
+    noise_scale = 1e-1,
+    learned_noise_scale = True,
+    noise_scale_learning_rate = 5e-3,
+    sigma_optimizer_klass = torch.optim.Adam,
+    sigma_optimizer_kwargs = dict(betas = (0.9, 0.99))
 )
 
 evo_strat()

{x_evolution-0.1.5 → x_evolution-0.1.7}/x_evolution/x_evolution.py

@@ -7,7 +7,7 @@ from functools import partial
 
 import torch
 from torch import tensor, Tensor, is_tensor, arange, randint
-from torch.nn import Module, ModuleList, Parameter
+from torch.nn import Module, ModuleList, Parameter, ParameterList
 from torch.optim import SGD
 
 import torch.nn.functional as F
@@ -56,14 +56,20 @@ class EvoStrategy(Module):
         num_generations,
         noise_population_size = 30,
         learning_rate = 1e-3,
-        noise_scale = 1e-3,                 # the noise scaling during rollouts with environment, todo - figure out right value and make sure it can also be customized per parameter name through a dict
         mirror_sampling = True,
         params_to_optimize: list[str] | Module | list[Module] | list[Parameter] | None = None,
         noise_low_rank: int | None = None,
         rollout_fixed_seed = True,
+        noise_scale = 1e-2,  # the noise scaling during rollouts with environment, todo - figure out right value and make sure it can also be customized per parameter name through a dict
+        learned_noise_scale = False,
+        noise_scale_learning_rate = 5e-3,
+        noise_scale_clamp_range: tuple[float, float] | None = None,
         use_optimizer = True,
         optimizer_klass = partial(SGD, nesterov = True, momentum = 0.1, weight_decay = 1e-2),
         optimizer_kwargs: dict = dict(),
+        use_sigma_optimizer = True,
+        sigma_optimizer_klass = partial(SGD, nesterov = True, momentum = 0.1),
+        sigma_optimizer_kwargs: dict = dict(),
         transform_fitness: Callable = identity,
         fitness_to_weighted_factor: Callable[[Tensor], Tensor] = normalize,
         checkpoint_every = None,            # saving every number of generations
@@ -141,9 +147,24 @@ class EvoStrategy(Module):
 
         self.mirror_sampling = mirror_sampling # mirror / antithetical sampling - reducing variance by doing positive + negative of noise and subtracting
 
+        self.learning_rate = learning_rate
+
+        # noise scale, which can be fixed or learned
+
         self.noise_scale = noise_scale
 
-        self.learning_rate = learning_rate
+        self.learned_noise_scale = learned_noise_scale
+
+        if learned_noise_scale:
+            self.sigmas = ParameterList([Parameter(torch.ones(()) * noise_scale) for _ in range(len(named_parameters_dict))])
+            self.param_name_to_sigma_index = {name: i for i, name in enumerate(named_parameters_dict.keys())}
+
+            self.sigma_clamp_ = identity
+
+            if exists(noise_scale_clamp_range):
+                self.sigma_clamp_ = lambda t: t.clamp_(*noise_scale_clamp_range)
+
+        self.noise_scale_learning_rate = noise_scale_learning_rate
 
         # rolling out with a fixed seed
 
@@ -157,6 +178,11 @@ class EvoStrategy(Module):
             optim_params = [named_parameters_dict[name] for name in params_to_optimize]
             self.optimizer = optimizer_klass(optim_params, lr = learning_rate, **optimizer_kwargs)
 
+        self.use_sigma_optimizer = use_sigma_optimizer
+
+        if learned_noise_scale and use_sigma_optimizer:
+            self.sigma_optimizer = sigma_optimizer_klass(self.sigmas, lr = noise_scale_learning_rate, **sigma_optimizer_kwargs)
+
         # rejecting the fitnesses for a certain generation if this function is true
 
         self.reject_generation_fitnesses_if = reject_generation_fitnesses_if
@@ -175,6 +201,13 @@ class EvoStrategy(Module):
     def print(self, *args, **kwargs):
         return self.accelerate.print(*args, **kwargs)
 
+    def _get_noise_scale(self, param_name):
+        if not self.learned_noise_scale:
+            return self.noise_scale
+
+        sigma_index = self.param_name_to_sigma_index[param_name]
+        return self.sigmas[sigma_index]
+
     @torch.inference_mode()
     def evolve_(
         self,
@@ -208,21 +241,30 @@ class EvoStrategy(Module):
         else:
             weights = fitnesses
 
-        weights /= self.noise_scale * (2. if self.mirror_sampling else 1.)
+        weights /= self.noise_population_size * (2. if self.mirror_sampling else 1.)
 
         if not use_optimizer:
             weights *= self.learning_rate # some learning rate that subsumes another constant
 
+        # maybe learned sigma
+
+        if self.learned_noise_scale:
+            grad_sigmas = [torch.zeros_like(sigma) for sigma in self.sigmas]
+
         # update one seed at a time for enabling evolutionary strategy for large models
 
         for individual_seed, weight in zip(seeds_for_population.tolist(), weights.tolist()):
 
             individual_param_seeds = with_seed(individual_seed)(torch.randint)(0, MAX_SEED_VALUE, (self.num_params,))
 
+            # setup noise config
+
             noise_config = dict(zip(self.param_names_to_optimize, individual_param_seeds.tolist()))
 
             noise_config_with_weight = dict()
 
+            # loop through each noise
+
             for param_name, seed in noise_config.items():
 
                 shape = self.param_shapes[param_name]
@@ -231,12 +273,31 @@ class EvoStrategy(Module):
                 # reproduce the noise from the seed
 
                 noise = with_seed(seed)(model.create_noise_fn)(shape, dtype = dtype)
+                noise = noise.to(self.device)
+                # scale the weight
+
+                noise_scale = self._get_noise_scale(param_name)
+
+                scaled_weight = weight / noise_scale
 
                 # set the noise weight
 
-                noise_config_with_weight[param_name] = (noise, weight)
+                noise_config_with_weight[param_name] = (noise, scaled_weight)
+
+                # maybe learned sigma
 
-            # now update
+                if self.learned_noise_scale:
+                    one_grad_sigma = weight * ((noise ** 2 - 1) / noise_scale)
+
+                    sigma_index = self.param_name_to_sigma_index[param_name]
+                    sigma = self.sigmas[sigma_index]
+
+                    if sigma.numel() == 1:
+                        one_grad_sigma = one_grad_sigma.mean()
+
+                    grad_sigmas[sigma_index].add_(one_grad_sigma)
+
+            # now update params for one seed
 
             model.add_noise_(noise_config_with_weight, negate = use_optimizer, add_to_grad = use_optimizer)
 
@@ -244,6 +305,21 @@ class EvoStrategy(Module):
             self.optimizer.step()
             self.optimizer.zero_grad()
 
+        if self.learned_noise_scale:
+
+            if self.use_sigma_optimizer:
+                for sigma, grad_sigma in zip(self.sigmas, grad_sigmas):
+                    sigma.grad = -grad_sigma
+
+                self.sigma_optimizer.step()
+                self.sigma_optimizer.zero_grad()
+            else:
+                for sigma, grad_sigma in zip(self.sigmas, grad_sigmas):
+                    sigma.add_(grad_sigma * self.noise_scale_learning_rate)
+
+            for sigma in self.sigmas:
+                self.sigma_clamp_(sigma)
+
     def checkpoint(self, filename = 'evolved.model'):
 
         if self.accelerate.is_main_process:
@@ -262,6 +338,11 @@ class EvoStrategy(Module):
 
         model = self.noisable_model.to(self.device)
 
+        # maybe sigmas
+
+        if self.learned_noise_scale:
+             self.sigmas = self.sigmas.to(self.device)
+
         # get world size, rank, and determine if distributed
 
         rank = self.accelerate.process_index
@@ -333,11 +414,20 @@ class EvoStrategy(Module):
                 individual_param_seeds = with_seed(individual_seed)(randint)(0, MAX_SEED_VALUE, (self.num_params,))
 
                 noise_config = dict(zip(self.param_names_to_optimize, individual_param_seeds.tolist()))
-                noise_config = {param_name: (seed, self.noise_scale) for param_name, seed in noise_config.items()}
+
+                # determine noise scale, which can be fixed or learned
+
+                noise_config_with_scale = dict()
+
+                for param_name, seed in noise_config.items():
+
+                    noise_scale = self._get_noise_scale(param_name)
+
+                    noise_config_with_scale[param_name] = (seed, noise_scale)
 
                 # maybe roll out with a fixed seed
 
-                with model.temp_add_noise_(noise_config):
+                with model.temp_add_noise_(noise_config_with_scale):
                     fitness = with_seed(maybe_rollout_seed)(rollout_for_fitness)()
 
                 if not self.mirror_sampling:
@@ -346,7 +436,7 @@ class EvoStrategy(Module):
 
                 # handle mirror sampling
 
-                with model.temp_add_noise_(noise_config, negate = True):
+                with model.temp_add_noise_(noise_config_with_scale, negate = True):
                     fitness_mirrored = with_seed(maybe_rollout_seed)(rollout_for_fitness)()
 
                 fitnesses.append([fitness, fitness_mirrored])
@@ -360,6 +450,11 @@ class EvoStrategy(Module):
             if is_distributed:
                 fitnesses = self.accelerate.gather(fitnesses)
 
+            # remove padding
+
+            fitnesses = fitnesses[:pop_size]
+            seeds_for_population = seeds_for_population[:pop_size]
+
             # validate fitnesses
 
             if exists(self.reject_generation_fitnesses_if) and self.reject_generation_fitnesses_if(fitnesses):
@@ -368,14 +463,17 @@ class EvoStrategy(Module):
 
             # pass fitnesses to evolve function
 
-            self.evolve_(
-                fitnesses[:pop_size],
-                seeds_for_population[:pop_size]
-            )
+            self.evolve_(fitnesses, seeds_for_population)
 
             # log
 
-            self.print(f'[{generation}] average fitness: {fitnesses.mean():.3f} | fitness std: {fitnesses.std():.3f}')
+            msg = f'[{generation}] average fitness: {fitnesses.mean():.3f} | fitness std: {fitnesses.std():.3f}'
+
+            if self.learned_noise_scale:
+                avg_sigma = torch.stack(list(self.sigmas)).mean().item()
+                msg += f' | avg sigma: {avg_sigma:.3f}'
+
+            self.print(msg)
 
             # maybe checkpoint