Evosax

rex.evo.evo(loss: Loss, solver: EvoSolver, init_state: evx.strategy.EvoState, transform: Transform, max_steps: int = 100, rng: jax.Array = None, verbose: bool = True, logger: LogState = None) -> Tuple[evx.strategy.EvoState, LogState, jax.Array]

Run the Evolutionary Solver (can be jit-compiled).

Parameters:

• loss (Loss) –

  Loss function.

• solver (EvoSolver) –

  Evolutionary Solver.

• init_state (EvoState) –

  Initial state of the Evolutionary Solver.

• transform (Transform) –

  Transform function to go from a normalized set of trainable parameters to the denormalized and extended set of parameters.

• max_steps (int, default: 100 ) –

  Maximum number of steps to run the Evolutionary Solver.

• rng (Array, default: None ) –

  Random number generator.

• verbose (bool, default: True ) –

  Whether to print the progress.

• logger (LogState, default: None ) –

  Logger for the Evolutionary Solver.

Returns:

• final_state ( EvoState ) –

  Final state of the Evolutionary Solver.

• logger ( LogState ) –

  Logger for the Evolutionary Solver.

• losses ( Array ) –

  Losses at each step.

---

rex.evo.EvoSolver

Evolutionary Solver class to manage the evolutionary strategy and its parameters.

Attributes:

• strategy_params (EvoParams) –

  Parameters for the evolutionary strategy.

• strategy (Strategy) –

  Instance of the evolutionary strategy.

• strategy_name (str) –

  Name of the strategy used.

init(u_min: Dict[str, Params], u_max: Dict[str, Params], strategy: str, strategy_kwargs: Dict = None, fitness_kwargs: Dict = None) -> EvoSolver classmethod

Initialize the Evolutionary Solver.

Parameters:

• u_min (Dict[str, Params]) –

  (Normalized) Minimum values for the parameters (pytree).

• u_max (Dict[str, Params]) –

  (Normalized) Maximum values for the parameters (pytree).

• strategy (str) –

  Name of the strategy to use from evosax.Strategies.

• strategy_kwargs (Dict, default: None ) –

  Keyword arguments to pass to the strategy.

• fitness_kwargs (Dict, default: None ) –

  Keyword arguments to pass to the fitness function of the strategy.

Returns:

• EvoSolver –

  EvoSolver instance.

init_state(mean: Dict[str, Params], rng: jax.Array = None) -> EvoState

Initialize the state of the Evolutionary Solver.

Parameters:

• mean (Dict[str, Params]) –

  Normalized mean values for the parameters (pytree).

• rng (Array, default: None ) –

  Random number generator.

Returns:

• EvoState ( EvoState ) –

  The initialized state of the Evolutionary Solver.

init_logger(num_generations: int, top_k: int = 5, maximize: bool = False) -> LogState

Initialize the logger for the Evolutionary Solver.

Parameters:

• num_generations (int) –

  Number of generations to log.

• top_k (int, default: 5 ) –

  Number of top individuals to log.

• maximize (bool, default: False ) –

  Whether the strategy is maximizing or minimizing.

Returns:

• LogState ( LogState ) –

  The initialized log state.

---

rex.evo.EvoState = evx.strategy.EvoState module-attribute

---

rex.evo.LogState

LogState class to manage the logging of evolutionary strategy states.

Attributes:

• state (Dict) –

  The current state of the logger.

• logger (ESLog) –

  The logger instance used for logging.

save(filename: str)

Save the log state to a file.

Parameters:

• filename (str) –

  The name of the file to save to.

load(filename: str) -> LogState

Load the log state from a file.

Parameters:

• filename (str) –

  The name of the file to load from.

Returns:

• LogState ( LogState ) –

  The loaded log state.

plot(title: str, ylims: List[int] = None, fig: plt.Figure = None, ax: plt.Axes = None, no_legend: bool = False) -> Tuple[plt.Figure, plt.Axes]

Plot the log state.

Parameters:

• title (str) –

  The title of the plot.

• ylims (List[int], default: None ) –

  The y-axis limits.

• fig (Figure, default: None ) –

  The figure to plot on.

• ax (Axes, default: None ) –

  The axes to plot on.

• no_legend (bool, default: False ) –

  Whether to omit the legend.

Returns:

• Tuple[Figure, Axes] –

  The plot.