Allow for mutation rate schedules for NSGA

PiperOrigin-RevId: 592042863

vizier/_src/algorithms/evolution/nsga2.py

@@ -216,6 +216,9 @@ def __init__(
       ranking_fn: Callable[[np.ndarray], np.ndarray] = _pareto_rank,
       eviction_limit: Optional[int] = None,
       adaptation: Optional[Mutation[Population, Offspring]] = None,
+      adaptation_callable: Optional[
+          Callable[[int], Mutation[Population, Offspring]]
+      ] = None,
       metadata_namespace: str = 'nsga2',
       seed: Optional[int] = None
   ):
@@ -235,6 +238,7 @@ def __init__(
       eviction_limit: Evict a gene that has been alive for this many
         generations.
       adaptation:
+      adaptation_callable:
       metadata_namespace: Metadata namespace to use.
       seed: Random seed.
 
@@ -253,6 +257,9 @@ def __init__(
             ranking_fn=ranking_fn,
             eviction_limit=eviction_limit,
         ),
-        adaptation=adaptation or numpy_populations.LinfMutation(seed=seed),
+        adaptation=adaptation
+        or numpy_populations.LinfMutation(seed=seed, norm=0.001),
         first_survival_after=first_survival_after,
+        adaptation_callable=adaptation_callable,
+        population_size=population_size,
     )

vizier/_src/algorithms/evolution/templates.py

@@ -41,7 +41,7 @@
 """
 
 import abc
-from typing import Generic, Optional, Sequence, TypeVar, Union
+from typing import Callable, Generic, Optional, Sequence, TypeVar, Union
 
 from vizier import algorithms as vza
 from vizier import pyvizier as vz
@@ -128,6 +128,9 @@ def __init__(
       survival: Survival[_PopulationType],
       adaptation: Mutation[_PopulationType, _OffspringsType],
       *,
+      adaptation_callable: Optional[
+          Callable[[int], Mutation[_PopulationType, _OffspringsType]]
+      ] = None,
       initial_population: Optional[_PopulationType] = None,
       first_survival_after: Optional[int] = None,
       population_size: int = 50,
@@ -139,6 +142,7 @@ def __init__(
       sampler:
       survival:
       adaptation:
+      adaptation_callable: Adapation as a function of number of Trials seen.
       initial_population: The initial population to seed the evolution.
       first_survival_after: Apply the survival step after observing this many
         trials. If unset, it defaults to twice the `population_size`.
@@ -149,6 +153,7 @@ def __init__(
     self._converter = converter
     self._sampler = sampler
     self._population_size = population_size
+    self._adaptation_callable = adaptation_callable
     self._first_survival_after = (
         first_survival_after or self._population_size * 2
     )
@@ -168,8 +173,16 @@ def suggest(self,
     count = count or self._population_size
     if self._num_trials_seen < self._first_survival_after:
       return self._converter.to_suggestions(self._sampler.sample(count))
-    return self._converter.to_suggestions(
-        self._adaptation.mutate(self._population, count))
+
+    if self._adaptation_callable is not None:
+      adaptation = self._adaptation_callable(self._num_trials_seen)
+    else:
+      adaptation = self._adaptation
+
+    suggestions = self._converter.to_suggestions(
+        adaptation.mutate(self._population, count)
+    )
+    return suggestions
 
   def update(
       self, completed: vza.CompletedTrials, all_active: vza.ActiveTrials