x-evolution 0.1.25__py3-none-any.whl → 0.1.26__py3-none-any.whl

x_evolution/x_evolution.py

@@ -93,11 +93,16 @@ class EvoStrategy(Module):
         verbose = True,
         accelerator: Accelerator | None = None,
         accelerate_kwargs: dict = dict(),
-        reject_generation_fitnesses_if: Callable[[Tensor], bool] | None = None
+        reject_generation_fitnesses_if: Callable[[Tensor], bool] | None = None,
+        vectorized = False,
+        vector_size: int | None = None
     ):
         super().__init__()
         self.verbose = verbose
 
+        self.vectorized = vectorized
+        self.vector_size = vector_size
+
         if not exists(accelerator):
             accelerator = Accelerator(cpu = cpu, **accelerate_kwargs)
 
@@ -475,24 +480,28 @@ class EvoStrategy(Module):
                     fitnesses.append([0., 0.] if self.mirror_sampling else 0.)
                     continue
 
-                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()))
-
-                # determine noise scale, which can be fixed or learned
+                def get_fitness(negate = False):
+                    individual_param_seeds = with_seed(individual_seed.item())(randint)(0, MAX_SEED_VALUE, (self.num_params,))
+                    noise_config = dict(zip(self.param_names_to_optimize, individual_param_seeds.tolist()))
 
-                noise_config_with_scale = dict()
+                    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)
 
-                for param_name, seed in noise_config.items():
+                    with model.temp_add_noise_(noise_config_with_scale, negate = negate):
+                        fitness = with_seed(maybe_rollout_seed)(self.environment)(model)
 
-                    noise_scale = self._get_noise_scale(param_name)
+                    if isinstance(fitness, Tensor) and fitness.numel() > 1:
+                        fitness = fitness.mean().item()
+                    elif isinstance(fitness, Tensor):
+                        fitness = fitness.item()
 
-                    noise_config_with_scale[param_name] = (seed, noise_scale)
+                    return fitness
 
-                # maybe roll out with a fixed seed
+                # evaluate
 
-                with model.temp_add_noise_(noise_config_with_scale):
-                    fitness = with_seed(maybe_rollout_seed)(rollout_for_fitness)()
+                fitness = get_fitness(negate = False)
 
                 if not self.mirror_sampling:
                     fitnesses.append(fitness)
@@ -500,8 +509,7 @@ class EvoStrategy(Module):
 
                 # handle mirror sampling
 
-                with model.temp_add_noise_(noise_config_with_scale, negate = True):
-                    fitness_mirrored = with_seed(maybe_rollout_seed)(rollout_for_fitness)()
+                fitness_mirrored = get_fitness(negate = True)
 
                 fitnesses.append([fitness, fitness_mirrored])
 

{x_evolution-0.1.25.dist-info → x_evolution-0.1.26.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.25
+Version: 0.1.26
 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.31
+Requires-Dist: x-mlps-pytorch>=0.2.0
 Requires-Dist: x-transformers>=2.11.23
 Provides-Extra: examples
 Requires-Dist: gymnasium[box2d]>=1.0.0; extra == 'examples'

x_evolution-0.1.26.dist-info/RECORD

@@ -0,0 +1,6 @@
+x_evolution/__init__.py,sha256=XcwXJgIMPnCWGfGws3-vKgoR_7IfVslJBtiMvmEeSg0,57
+x_evolution/x_evolution.py,sha256=Jln3wpkIQIp7xwa3KMibh0kSuob1NIi7Aj7Miz8RJdY,19491
+x_evolution-0.1.26.dist-info/METADATA,sha256=7zamSGrDtvOUQzpyYCkrt42FmWC4GPYpVSWUAO8a6OA,5853
+x_evolution-0.1.26.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+x_evolution-0.1.26.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+x_evolution-0.1.26.dist-info/RECORD,,

x_evolution-0.1.25.dist-info/RECORD

@@ -1,6 +0,0 @@
-x_evolution/__init__.py,sha256=XcwXJgIMPnCWGfGws3-vKgoR_7IfVslJBtiMvmEeSg0,57
-x_evolution/x_evolution.py,sha256=lvN3ePqD6a5dW1gOv0d1I9yQ4rdv6OuIVvKvXa0yRBM,19126
-x_evolution-0.1.25.dist-info/METADATA,sha256=1fAtssjj_t76rXwLu728z_ohlRZKQDhCm0oOl2eeIxA,5854
-x_evolution-0.1.25.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-x_evolution-0.1.25.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-x_evolution-0.1.25.dist-info/RECORD,,