### Installing evosax with Examples

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Installs the `evosax` library from PyPi, including the necessary dependencies for running examples. This command ensures all required components for the tutorials are available.

```python
%pip install -U "evosax[examples]"
```

--------------------------------

### Installing JAX with CUDA

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Installs the JAX library with CUDA support using pip. This is a prerequisite for running `evosax` on NVIDIA GPUs, ensuring accelerated computations.

```python
%pip install -U "jax[cuda]"
```

--------------------------------

### Installing EvoSax with Examples

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Installs the EvoSax library from PyPi, including additional dependencies required for running the provided examples. This command ensures all necessary components for the reinforcement learning demonstrations are available.

```Python
%pip install -U "evosax[examples]"
```

--------------------------------

### Installing evosax with Examples (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/04_restart_es.ipynb

Installs the `evosax` library from PyPi, including additional dependencies required for running the examples. This is a necessary step to use the evolution strategy algorithms.

```Python
%pip install -U "evosax[examples]"
```

--------------------------------

### Installing EvoSax with Examples for Python

Source: https://github.com/roberttlange/evosax/blob/main/examples/07_persistent_es.ipynb

Installs the EvoSax library, including its example dependencies. This command uses pip to upgrade the existing EvoSax installation or install it if not present.

```Python
%pip install -U "evosax[examples]"
```

--------------------------------

### Installing Evosax with Examples (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Installs the `evosax` library, including its example dependencies. `evosax` is a JAX-based library for evolutionary algorithms.

```python
%pip install -U "evosax[examples]"
```

--------------------------------

### Importing Core Libraries

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Imports essential libraries: `jax` for numerical computation, `jax.numpy` for NumPy-like array operations on JAX arrays, and `matplotlib.pyplot` for plotting results. These are fundamental for setting up and visualizing the optimization process.

```python
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
```

--------------------------------

### Installing `evosax` for Development

Source: https://github.com/roberttlange/evosax/blob/main/CONTRIBUTING.md

This snippet clones the `evosax` repository, navigates into its directory, and installs it in editable development mode, including development dependencies. This setup allows for immediate testing of code modifications.

```bash
git clone https://github.com/RobertTLange/evosax
cd evosax
pip install -e ".[dev]"
```

--------------------------------

### Initializing JAX Random Key

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Initializes a JAX random key with a fixed seed. This key is crucial for reproducible stochastic operations within JAX, such as sampling or splitting for subkeys in evolutionary algorithms.

```python
seed = 0
key = jax.random.key(seed)
```

--------------------------------

### Executing One Generation of Ask-Eval-Tell Cycle

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Demonstrates a single iteration of the 'ask-eval-tell' cycle, which is the core of `evosax`'s API. It involves asking the ES for candidate solutions, evaluating their fitness using the problem function, and telling the ES the results to update its internal state.

```python
key, subkey = jax.random.split(key)
key_ask, key_eval, key_tell = jax.random.split(subkey, 3)

# Generate a set of candidate solutions to evaluate
population, state = es.ask(key_ask, state, params)

# Evaluate the fitness of the population
fitness, problem_state, info = problem.eval(key_eval, population, problem_state)

# Update the evolution strategy
state, metrics = es.tell(key_tell, population, fitness, state, params)
```

--------------------------------

### Initializing CMA-ES Evolution Strategy

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Sets up the CMA-ES (Covariance Matrix Adaptation Evolution Strategy) algorithm. It defines the number of generations and population size, instantiates the `CMA_ES` solver with a dummy solution, and initializes its internal state and parameters for the optimization run.

```python
from evosax.algorithms import CMA_ES as ES

num_generations = 64
population_size = 16

# Instantiate evolution strategy
key, subkey = jax.random.split(key)
solution = problem.sample(subkey)
es = ES(
    population_size=population_size,
    solution=solution,  # requires a dummy solution
)

# Use default parameters
params = es.default_params

# Initialize evolution strategy
key, subkey = jax.random.split(key)
state = es.init(subkey, solution, params)
```

--------------------------------

### Inspecting CMA-ES Default Parameters

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Displays the default parameters configured for the CMA-ES algorithm. Users can inspect and modify these parameters as needed to fine-tune the optimization process for specific problems.

```python
params  # You can have a look at the parameters and change those if desired
```

--------------------------------

### Running CMA-ES Optimization Loop

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Executes the full optimization process for a specified number of generations using the CMA-ES algorithm. It iteratively performs the 'ask-eval-tell' cycle, logging performance metrics at each step to track the optimization progress.

```python
key, subkey = jax.random.split(key)
state = es.init(subkey, solution, params)

metrics_log = []
for i in range(num_generations):
    key, subkey = jax.random.split(key)
    key_ask, key_eval, key_tell = jax.random.split(subkey, 3)

    population, state = es.ask(key_ask, state, params)

    fitness, problem_state, info = problem.eval(key_eval, population, problem_state)

    state, metrics = es.tell(key_tell, population, fitness, state, params)

    # Log metrics
    metrics_log.append(metrics)
```

--------------------------------

### Executing One Generation with Custom Metrics

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Runs a single generation of the 'ask-eval-tell' cycle using the `Sep_CMA_ES` algorithm configured with a custom metrics function. This demonstrates how the custom function is integrated into the optimization loop to log additional state information.

```python
key, subkey = jax.random.split(key)
state = es.init(subkey, solution, params)

key, subkey = jax.random.split(key)
key_ask, key_eval, key_tell = jax.random.split(subkey, 3)

population, state = es.ask(key_ask, state, params)

fitness, problem_state, info = problem.eval(key_eval, population, problem_state)

state, metrics = es.tell(key_tell, population, fitness, state, params)
```

--------------------------------

### Defining a Black Box Optimization Problem (Rosenbrock)

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Instantiates a `BBOBProblem` from `evosax.problems` to define a classic black-box optimization benchmark, specifically the 2-dimensional Rosenbrock function. It sets up the problem's dimensions, optimal solution, and initializes its state for evaluation.

```python
from evosax.problems import BBOBProblem as Problem

num_dims = 2
fn_name = "rosenbrock"

problem = Problem(
    fn_name=fn_name,
    num_dims=num_dims,
    x_opt=2.5 * jnp.ones(num_dims),
    f_opt=0.0,
    sample_rotations=False,
    seed=seed,
)

key, subkey = jax.random.split(key)
problem_state = problem.init(subkey)
```

--------------------------------

### Initializing OpenAI-ES with Custom Fitness Shaping

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Configures the OpenAI-ES algorithm with a custom fitness shaping function. It combines `add_weight_decay` and `standardize_fitness_shaping_fn` to apply both weight decay and z-score standardization to the fitness values, influencing how the algorithm updates its state.

```python
from evosax.algorithms import Open_ES as ES
from evosax.core.fitness_shaping import add_weight_decay, standardize_fitness_shaping_fn

# Add weight decay and change the default fitness shaping function to standardize (z-score)
fitness_shaping_fn = add_weight_decay(standardize_fitness_shaping_fn, weight_decay=0.01)

num_generations = 64
population_size = 16

key, subkey = jax.random.split(key)
problem_state = problem.init(subkey)

# Instantiate evolution strategy
key, subkey = jax.random.split(key)
solution = problem.sample(subkey)
es = ES(
    population_size=population_size,
    solution=solution,  # requires a dummy solution
    fitness_shaping_fn=fitness_shaping_fn,
)

# Use default parameters
params = es.default_params
```

--------------------------------

### Maximizing Objective by Inverting Fitness

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Demonstrates how to convert a minimization problem (default in `evosax`) into a maximization problem. By simply negating the fitness value before passing it to the `tell` method, the algorithm will effectively maximize the original objective.

```python
fitness, info = problem.eval(key_eval, population)  # Fitness to maximize
state, metrics = es.tell(key_tell, population, -fitness, state, params)  # add minus sign to maximize
```

--------------------------------

### Defining Custom Metrics for Sep-CMA-ES

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Defines a custom metrics function to extend the default metrics logged by `evosax` algorithms. This example specifically adds the diagonal matrix `D` from the `Sep_CMA_ES` state to the logged metrics, providing more detailed insights into the algorithm's internal workings.

```python
from evosax.algorithms import Sep_CMA_ES as ES
from evosax.algorithms.distribution_based.base import metrics_fn

def custom_metrics_fn(
    key,
    population,
    fitness,
    state,
    params,
):
    metrics = metrics_fn(key, population, fitness, state, params)
    return metrics | {
        "D": state.D,
    }

key, subkey = jax.random.split(key)
problem_state = problem.init(subkey)

es = ES(
    population_size=population_size,
    solution=solution,  # requires a dummy solution
    metrics_fn=custom_metrics_fn,
)

# Use default parameters
params = es.default_params
```

--------------------------------

### Running OpenAI-ES Optimization Loop with Custom Shaping

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Executes the optimization process using OpenAI-ES with the previously defined custom fitness shaping function. It iterates through generations, applying the 'ask-eval-tell' cycle and logging metrics, demonstrating the effect of custom fitness shaping on the algorithm's behavior.

```python
key, subkey = jax.random.split(key)
state = es.init(subkey, solution, params)

metrics_log = []
for i in range(num_generations):
    key, subkey = jax.random.split(key)
    key_ask, key_eval, key_tell = jax.random.split(subkey, 3)

    population, state = es.ask(key_ask, state, params)

    fitness, problem_state, info = problem.eval(key_eval, population, problem_state)

    state, metrics = es.tell(key_tell, population, fitness, state, params)

    # Log metrics
    metrics_log.append(metrics)
```

--------------------------------

### Installing Evosax via pip

Source: https://github.com/roberttlange/evosax/blob/main/README.md

This command installs the evosax library from PyPi using pip, making it available for use in Python projects. It requires Python 3.10 or later and a working JAX installation.

```bash
pip install evosax
```

--------------------------------

### Plotting Best Fitness Over Generations (OpenAI-ES)

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Generates a plot illustrating the best fitness achieved per generation for the OpenAI-ES optimization run with custom fitness shaping. This visualization helps assess the convergence and performance of the algorithm under the specified fitness shaping.

```python
# Extract the best fitness values across generations
generations = [metrics["generation_counter"] for metrics in metrics_log]
best_fitness = [metrics["best_fitness"] for metrics in metrics_log]

plt.figure(figsize=(10, 5))
plt.plot(generations, best_fitness, label="Best Fitness", marker="o", markersize=3)

plt.title("Best fitness over generations")
plt.xlabel("Generation")
plt.ylabel("Fitness")
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
```

--------------------------------

### Inspecting Logged Custom Metrics Keys

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Prints the keys available in the `metrics` dictionary after a generation run with a custom metrics function. This verifies that the custom metric, such as 'D' in this case, has been successfully added to the logged data.

```python
print(metrics.keys())
```

--------------------------------

### Initializing Open_ES Algorithm with Schedules

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Initializes the Open_ES algorithm from EvoSax, configuring it with a larger population size and exponential decay schedules for both the learning rate and standard deviation. This setup is tailored for more complex environments like Brax Ant.

```Python
from evosax.algorithms import Open_ES as ES

num_generations = 512
lr_schedule = optax.exponential_decay(
    init_value=0.01,
    transition_steps=num_generations,
    decay_rate=0.1,
)
std_schedule = optax.exponential_decay(
    init_value=0.05,
    transition_steps=num_generations,
    decay_rate=0.2,
)
es = ES(
    population_size=256,
    solution=solution,
    optimizer=optax.adam(learning_rate=lr_schedule),
    std_schedule=std_schedule,
)

params = es.default_params
```

--------------------------------

### Installing JAX with CUDA Support

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Installs the JAX library with CUDA support, enabling GPU acceleration for numerical computations. This is a prerequisite for running EvoSax efficiently on NVIDIA GPUs.

```Python
%pip install -U "jax[cuda]"
```

--------------------------------

### Plotting Best Fitness Over Generations (CMA-ES)

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Visualizes the best fitness achieved across generations during the CMA-ES optimization run. It extracts the generation counter and best fitness values from the logged metrics and plots them to show the convergence of the algorithm.

```python
# Extract the best fitness values across generations
generations = [metrics["generation_counter"] for metrics in metrics_log]
best_fitness = [metrics["best_fitness"] for metrics in metrics_log]

plt.figure(figsize=(10, 5))
plt.plot(generations, best_fitness, label="Best Fitness", marker="o", markersize=3)

plt.title("Best fitness over generations")
plt.xlabel("Generation")
plt.ylabel("Fitness")
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
```

--------------------------------

### Installing JAX with CUDA for Python

Source: https://github.com/roberttlange/evosax/blob/main/examples/07_persistent_es.ipynb

Installs the JAX library with CUDA support, which is required for GPU acceleration. This command uses pip to upgrade the existing JAX installation or install it if not present.

```Python
%pip install -U "jax[cuda]"
```

--------------------------------

### Implementing pmap-compatible SNES with EvoSax in Python

Source: https://github.com/roberttlange/evosax/blob/main/CHANGELOG.md

This snippet demonstrates how to use JAX's `pmap` for device-parallel evaluation of an EvoSax strategy, specifically SNES, on a BBOB problem. It initializes multiple JAX devices, replicates strategy parameters, and performs parallel `ask` and `tell` operations for population-based optimization. The example showcases setting up the environment, initializing the strategy and problem, and executing a single step of the evolutionary loop across multiple devices.

```Python
# set number of cpu devices for jax pmap
import os

num_devices = 4
os.environ["XLA_FLAGS"] = f"--xla_force_host_platform_device_count={num_devices}"

import jax
import jax.numpy as jnp
from flax import jax_utils

print(jax.devices())

from evosax.problems import BBOBFitness
from evosax.v2 import SNES

fn_name = "Sphere"
num_dims = 2
popsize = 64
rng = jax.random.PRNGKey(0)
problem = BBOBFitness(fn_name, num_dims=num_dims, n_devices=num_devices)

strategy = SNES(
    popsize=popsize,
    num_dims=num_dims,
    sigma_init=0.1,
    n_devices=num_devices,
    maximize=False,
)

params = strategy.default_params.replace(init_min=-3.0, init_max=3.0)
params = jax_utils.replicate(params)

init_rng = jnp.tile(rng[None], (num_devices, 1))
state = jax.pmap(strategy.initialize)(init_rng, params)
print("Mean pre-update:", state.mean)  # (num_devices, num_dims)

rng, rng_a, rng_e = jax.random.split(rng, 3)
ask_rng = jax.random.split(rng_a, num_devices)
x, state = jax.pmap(strategy.ask, axis_name="device")(ask_rng, state, params)

print("Population shape:", x.shape)  # (num_devices, popsize/num_devices, num_dims)

fitness = problem.rollout(rng_e, x)
print("Fitness shape:", fitness.shape)  # (num_devices, popsize/num_devices)

state = jax.pmap(strategy.tell, axis_name="device")(x, fitness, state, params)
print("Mean post-update:", state.mean)  # (num_devices, num_dims)
```

--------------------------------

### Installing JAX with CUDA Support (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/04_restart_es.ipynb

Installs the JAX library with CUDA support, a prerequisite for running `evosax` on NVIDIA GPUs. This command is typically run in an environment like Google Colab or Jupyter.

```Python
%pip install -U "jax[cuda]"
```

--------------------------------

### Accessing Custom Logged Metric 'D'

Source: https://github.com/roberttlange/evosax/blob/main/examples/00_getting_started.ipynb

Retrieves and displays the value of the custom metric 'D' from the `metrics` dictionary. This confirms that the custom metrics function correctly captured and stored the desired internal state information from the `Sep_CMA_ES` algorithm.

```python
metrics["D"]  # Custom data logged in the metrics
```

--------------------------------

### Installing JAX with CUDA Support (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Installs the JAX library with CUDA support, enabling GPU acceleration for numerical computations. This is a prerequisite for running `evosax` efficiently on NVIDIA GPUs.

```python
%pip install -U "jax[cuda]"
```

--------------------------------

### Defining CIFAR-10 CNN Problem (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Defines a Convolutional Neural Network (CNN) architecture for the CIFAR-10 dataset and initializes the `TorchVisionProblem` instance. This setup is similar to MNIST but with adjusted strides and MLP layer sizes to accommodate the larger CIFAR-10 images and complexity. It also initializes the problem state and samples an initial solution structure.

```python
from evosax.problems import CNN, TorchVisionProblem as Problem, identity_output_fn

network = CNN(
    num_filters=[8, 16],
    kernel_sizes=[(5, 5), (5, 5)],
    strides=[(2, 2), (2, 2)],
    mlp_layer_sizes=[128, 10],
    output_fn=identity_output_fn,
)

problem = Problem(
    task_name="CIFAR10",
    network=network,
    batch_size=1024,
)

key, subkey = jax.random.split(key)
problem_state = problem.init(key)

key, subkey = jax.random.split(key)
solution = problem.sample(subkey)
```

--------------------------------

### Importing Core Libraries

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Imports essential libraries: `jax` for JIT compilation and automatic differentiation, `matplotlib.pyplot` for plotting results, and `optax` for optimization algorithms used by EvoSax.

```Python
import jax
import matplotlib.pyplot as plt
import optax
```

--------------------------------

### Importing Libraries for ES and JAX Operations in Python

Source: https://github.com/roberttlange/evosax/blob/main/examples/07_persistent_es.ipynb

Imports necessary modules: `partial` for function currying, `jax` and `jax.numpy` for numerical operations, `optax` for optimization schedules, and `NoiseReuseES`, `PersistentES` from `evosax.algorithms` for evolution strategies.

```Python
from functools import partial

import jax
import jax.numpy as jnp
import optax
from evosax.algorithms import NoiseReuseES, PersistentES
```

--------------------------------

### Setting Up CartPole Environment with MLP Policy (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/04_restart_es.ipynb

Configures the CartPole-v1 environment using `evosax`'s `GymnaxProblem`. It defines a Multi-Layer Perceptron (MLP) policy with specific layer sizes and a categorical output function, then initializes the problem state and samples an initial solution.

```Python
from evosax.problems import GymnaxProblem as Problem
from evosax.problems.networks import MLP, categorical_output_fn

policy = MLP(
    layer_sizes=(64, 64, 2),
    output_fn=categorical_output_fn,
)

problem = Problem(
    env_name="CartPole-v1",
    policy=policy,
    episode_length=200,
    num_rollouts=16,
    use_normalize_obs=True,
)

key, subkey = jax.random.split(key)
problem_state = problem.init(key)

key, subkey = jax.random.split(key)
solution = problem.sample(subkey)
```

--------------------------------

### Implementing Evolution Strategy with CMA-ES in JAX

Source: https://github.com/roberttlange/evosax/blob/main/README.md

This snippet demonstrates the basic 'ask'-'eval'-'tell' API usage of evosax for an Evolution Strategy, specifically using CMA-ES. It shows how to instantiate a strategy, initialize its state with a JAX random key, and then iterate through generations to generate candidate solutions ('ask'), evaluate their fitness, and update the strategy's state ('tell'). The 'dummy_solution', 'num_generations', and 'fitness' evaluation are placeholders that need to be defined by the user.

```python
import jax
from evosax.algorithms import CMA_ES


# Instantiate the search strategy
es = CMA_ES(population_size=32, solution=dummy_solution)
params = es.default_params

# Initialize state
key = jax.random.key(0)
state = es.init(key, params)

# Ask-Eval-Tell loop
for i in range(num_generations):
    key, key_ask, key_eval = jax.random.split(key, 3)

    # Generate a set of candidate solutions to evaluate
    population, state = es.ask(key_ask, state, params)

    # Evaluate the fitness of the population
    fitness = ...

    # Update the evolution strategy
    state = es.tell(population, fitness, state, params)

# Get best solution
state.best_solution, state.best_fitness
```

--------------------------------

### Configuring CartPole Environment with MLP Policy (Gymnax)

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Sets up the CartPole-v1 environment using `GymnaxProblem` and defines an MLP policy with a categorical output function. It initializes the problem state and samples an initial solution (policy parameters) for the evolutionary algorithm.

```Python
from evosax.problems import GymnaxProblem as Problem
from evosax.problems.networks import MLP, categorical_output_fn

policy = MLP(
    layer_sizes=(64, 64, 2),
    output_fn=categorical_output_fn,
)

problem = Problem(
    env_name="CartPole-v1",
    policy=policy,
    episode_length=200,
    num_rollouts=16,
    use_normalize_obs=True,
)

key, subkey = jax.random.split(key)
problem_state = problem.init(key)

key, subkey = jax.random.split(key)
solution = problem.sample(subkey)
```

--------------------------------

### Initializing SNES Algorithm for Evolution

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Initializes the Separable Natural Evolution Strategies (SNES) algorithm from EvoSax. It configures the population size, provides the initial solution structure, and sets up an Adam optimizer for updating the search distribution.

```Python
from evosax.algorithms import SNES as ES

es = ES(
    population_size=128,
    solution=solution,
    optimizer=optax.adam(learning_rate=0.01),
)

params = es.default_params
```

--------------------------------

### Generating Brax Environment Rollout HTML

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Generates an HTML visualization of the best policy's performance in the Brax Ant environment. It extracts the pipeline states from the evaluation info and uses `brax.io.html.render` to create an interactive simulation.

```Python
from brax.io import html
from IPython.display import HTML

rollout = [
    jax.tree_util.tree_map(lambda x: x[0, 0, t], info["env_states"].pipeline_state)
    for t in range(200)
]

html_content = html.render(
    problem.env.sys.tree_replace({"opt.timestep": problem.env.dt}), rollout
)
HTML(html_content)
```

--------------------------------

### Initializing JAX Random Key

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Initializes a JAX random key with a fixed seed for reproducible random number generation. This key is crucial for various JAX operations, including environment initialization and population sampling.

```Python
seed = 0
key = jax.random.key(seed)
```

--------------------------------

### Configuring Ant Environment with MLP Policy (Brax)

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Sets up the Ant environment using `BraxProblem` and defines an MLP policy with a tanh output function suitable for continuous control. It initializes the problem state and samples an initial solution for the evolutionary algorithm.

```Python
from evosax.problems import BraxProblem as Problem
from evosax.problems.networks import MLP, tanh_output_fn

policy = MLP(
    layer_sizes=(32, 32, 32, 32, 8),
    output_fn=tanh_output_fn,
)

problem = Problem(
    env_name="ant",
    policy=policy,
    episode_length=1000,
    num_rollouts=16,
    use_normalize_obs=True,
)

key, subkey = jax.random.split(key)
problem_state = problem.init(key)

key, subkey = jax.random.split(key)
solution = problem.sample(subkey)
```

--------------------------------

### Importing Core Libraries (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/04_restart_es.ipynb

Imports essential libraries for numerical computation (JAX, JAX NumPy), plotting (Matplotlib), and optimization (Optax). These libraries provide the foundational tools for implementing and visualizing evolution strategies.

```Python
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import optax
```

--------------------------------

### Importing Core Libraries (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Imports `jax` for numerical computation and automatic differentiation, and `optax` for optimization algorithms. These are fundamental libraries for `evosax` and JAX-based machine learning.

```python
import jax
import optax
```

--------------------------------

### Running `evosax` Unit Tests

Source: https://github.com/roberttlange/evosax/blob/main/CONTRIBUTING.md

This command runs all unit tests located in the `tests/` directory using `pytest`. Passing all tests is essential to ensure that new changes do not introduce regressions or break existing functionality.

```bash
pytest tests/
```

--------------------------------

### Saving Brax Visualization HTML to File

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Saves the generated HTML content of the Brax environment rollout to a local file named 'ant_visualization.html'. This allows for offline viewing and sharing of the policy's behavior.

```Python
# Write to file
with open("ant_visualization.html", "w") as f:
    f.write(html_content)

print("Visualization saved to 'ant_visualization.html'")
```

--------------------------------

### Initializing Persistent Evolution Strategy in EvoSax

Source: https://github.com/roberttlange/evosax/blob/main/examples/07_persistent_es.ipynb

Configures and initializes the `PersistentES` strategy from EvoSax. It sets the population size, initial solution, and a constant standard deviation schedule, then overrides default parameters for `T` (total steps) and `K` (unroll steps). The strategy state and inner problem parameters are also initialized.

```Python
population_size = 128

strategy = PersistentES(
    population_size=population_size,
    solution=jnp.zeros(2),
    std_schedule=optax.constant_schedule(0.2),
)
params = strategy.default_params.replace(T=100, K=10)

key, subkey = jax.random.split(key)
state = strategy.init(subkey, jnp.zeros(2), params)

# Initialize inner parameters
xs = jnp.ones((population_size, 2)) * jnp.array([1.0, 1.0])

params
```

--------------------------------

### Citing evosax: BibTeX Entry

Source: https://github.com/roberttlange/evosax/blob/main/README.md

This BibTeX entry provides the necessary information to cite the evosax paper in academic publications. It includes the author, title, journal, and year of the arXiv preprint, ensuring proper attribution when using the library in research.

```bibtex
@article{evosax2022github,
    author  = {Robert Tjarko Lange},
    title   = {evosax: JAX-based Evolution Strategies},
    journal = {arXiv preprint arXiv:2212.04180},
    year    = {2022},
}
```

--------------------------------

### Running SNES Training Loop with JAX Scan

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Executes the SNES training for a specified number of generations using `jax.lax.scan` for efficient compilation. It initializes the algorithm state and iteratively applies the defined `step` function to evolve the policy.

```Python
num_generations = 64

key, subkey = jax.random.split(key)
state = es.init(subkey, solution, params)

key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, num_generations)
_, metrics = jax.lax.scan(
    step,
    (state, params, problem_state),
    keys,
)
```

--------------------------------

### Initializing OpenAI Evolution Strategy (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/04_restart_es.ipynb

Initializes the OpenAI Evolution Strategy (OpenAI-ES) algorithm with a specified population size, the initial solution, an Adam optimizer for updating the mean, and a constant standard deviation schedule. It then retrieves the default parameters for the algorithm.

```Python
from evosax.algorithms import Open_ES as ES

es = ES(
    population_size=16,
    solution=solution,
    optimizer=optax.adam(learning_rate=0.01),
    std_schedule=optax.constant_schedule(0.1),
)

params = es.default_params
```

--------------------------------

### Displaying Policy Parameter Count (Gymnax)

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Calculates and prints the total number of parameters in the initialized MLP policy. This provides an insight into the complexity of the neural network being evolved.

```Python
print(f"Number of pararmeters: {sum(leaf.size for leaf in jax.tree.leaves(solution))}")
```

--------------------------------

### Displaying Policy Parameter Count (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/04_restart_es.ipynb

Calculates and prints the total number of parameters in the initialized policy network. This provides insight into the complexity of the solution space that the evolution strategy will explore.

```Python
print(f"Number of pararmeters: {sum(leaf.size for leaf in jax.tree.leaves(solution))}")
```

--------------------------------

### Displaying Policy Parameter Count (Brax)

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Calculates and prints the total number of parameters in the initialized MLP policy for the Brax Ant environment. This provides an insight into the complexity of the neural network being evolved.

```Python
print(f"Number of pararmeters: {sum(leaf.size for leaf in jax.tree.leaves(solution))}")
```

--------------------------------

### Committing Changes to Git

Source: https://github.com/roberttlange/evosax/blob/main/CONTRIBUTING.md

These commands stage specified files for commit and then create a new commit with a descriptive message. A well-written commit message is crucial for tracking changes and understanding the purpose of the commit.

```bash
git add file1.py file2.py ...
git commit -m "Your commit message"
```

--------------------------------

### Running Open_ES Training Loop with Periodic Logging

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Executes the Open_ES training for a specified number of generations, using `jax.lax.scan` for efficient computation. It includes periodic logging of the mean fitness to track progress during the long training process for the Brax Ant environment.

```Python
key, subkey = jax.random.split(key)
state = es.init(subkey, solution, params)

fitness_log = []
log_period = 32
for i in range(num_generations // log_period):
    key, subkey = jax.random.split(key)
    keys = jax.random.split(subkey, log_period)
    (state, params, problem_state), metrics = jax.lax.scan(
        step,
        (state, params, problem_state),
        keys,
    )

    mean = es.get_mean(state)

    key, subkey = jax.random.split(key)
    fitness, problem_state, info = problem.eval(
        key, jax.tree.map(lambda x: x[None], mean), problem_state
    )
    print(f"Generation {(i + 1) * log_period:3d} | Mean fitness: {fitness.mean():.2f}")
```

--------------------------------

### Running Lint and Type Checks with Ruff

Source: https://github.com/roberttlange/evosax/blob/main/CONTRIBUTING.md

These commands execute `ruff` to format the code and perform lint and type checks across the entire repository. Ensuring code passes these checks is a prerequisite for contributing to maintain code quality and consistency.

```bash
ruff format .
ruff check .
```

--------------------------------

### Initializing JAX Random Key (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Sets a random seed and initializes a JAX pseudo-random number generator key. This key is crucial for reproducible stochastic operations within JAX, such as network initialization and population sampling.

```python
seed = 0
key = jax.random.key(seed)
```

--------------------------------

### Initializing Noise Reuse Evolution Strategy in EvoSax

Source: https://github.com/roberttlange/evosax/blob/main/examples/07_persistent_es.ipynb

Configures and initializes the `NoiseReuseES` strategy from EvoSax, similar to `PersistentES`. It sets the population size, initial solution, and a constant standard deviation schedule, then overrides default parameters for `T` (total steps) and `K` (unroll steps). The strategy state and inner problem parameters are also initialized.

```Python
population_size = 128

strategy = NoiseReuseES(
    population_size=population_size,
    solution=jnp.zeros(2),
    std_schedule=optax.constant_schedule(0.2),
)
params = strategy.default_params.replace(T=100, K=10)

key, subkey = jax.random.split(key)
state = strategy.init(subkey, jnp.zeros(2), params)

# Initialize inner parameters
xs = jnp.ones((population_size, 2)) * jnp.array([1.0, 1.0])

params
```

--------------------------------

### Syncing Local Branch with Upstream Main

Source: https://github.com/roberttlange/evosax/blob/main/CONTRIBUTING.md

These commands fetch the latest changes from the upstream `evosax` repository and then rebase the current branch onto the `upstream/main` branch. This ensures your changes are built on top of the most recent codebase, minimizing merge conflicts.

```bash
git fetch upstream
git rebase upstream/main
```

--------------------------------

### Evaluating Best Evolved Policy (Brax)

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Retrieves the best solution found by the Open_ES algorithm and evaluates its fitness in the Brax environment. This step confirms the performance of the final evolved policy before visualization.

```Python
mean = es.get_mean(state)
mean = es._unravel_solution(state.best_solution)

key, subkey = jax.random.split(key)
fitness, problem_state, info = problem.eval(
    key, jax.tree.map(lambda x: x[None], mean), problem_state
)
fitness[0]
```

--------------------------------

### Initializing JAX Random Key (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/04_restart_es.ipynb

Sets a random seed and initializes a JAX random key. This key is crucial for ensuring reproducibility and for generating random numbers in JAX-based operations, such as problem initialization and population sampling.

```Python
seed = 0
key = jax.random.key(seed)
```

--------------------------------

### Initializing JAX Random Key with Seed in Python

Source: https://github.com/roberttlange/evosax/blob/main/examples/07_persistent_es.ipynb

Sets a fixed seed for reproducibility and initializes a JAX random key. This key is crucial for all random operations within JAX, ensuring consistent results across runs.

```Python
seed = 0
key = jax.random.key(seed)
```

--------------------------------

### Defining Open_ES Training Step Function

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Defines a single training step for the Open_ES algorithm, identical in structure to the SNES step function. It handles population generation, fitness evaluation in the environment, and algorithm state updates based on the results.

```Python
def step(carry, key):
    state, params, problem_state = carry
    key_ask, key_eval, key_tell = jax.random.split(key, 3)

    population, state = es.ask(key_ask, state, params)

    fitness, problem_state, _ = problem.eval(key_eval, population, problem_state)

    state, metrics = es.tell(
        key_tell, population, -fitness, state, params
    )  # Minimize fitness

    return (state, params, problem_state), metrics
```

--------------------------------

### Displaying MNIST Model Parameter Count (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Calculates and prints the total number of parameters in the initialized CNN solution for the MNIST problem. This provides an insight into the model's complexity.

```python
print(f"Number of pararmeters: {sum(leaf.size for leaf in jax.tree.leaves(solution))}")
```

--------------------------------

### Visualizing SNES Fitness on CartPole

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Generates a plot showing the best fitness achieved by the SNES algorithm over generations on the CartPole task. This visualization helps assess the convergence and performance of the evolutionary process.

```Python
plt.figure(figsize=(6, 3))
plt.plot(-metrics["best_fitness"])

plt.title("SNES on CartPole task")
plt.xlabel("Generations")
plt.ylabel("Fitness")

plt.grid(True)
plt.tight_layout()

plt.show()
```

--------------------------------

### Upgrading Evosax to Latest Version

Source: https://github.com/roberttlange/evosax/blob/main/README.md

This command upgrades the evosax library to its latest version directly from the main branch of its GitHub repository. It ensures users have access to the most recent features and bug fixes.

```bash
pip install git+https://github.com/RobertTLange/evosax.git@main
```

--------------------------------

### Initializing ES State for CIFAR-10 (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Initializes the internal state of the Open-ES algorithm for the CIFAR-10 problem. This prepares the algorithm for the optimization loop, setting up the initial mean solution and other necessary parameters.

```python
key, subkey = jax.random.split(key)
state = es.init(subkey, solution, params)
```

--------------------------------

### Initializing ES State for MNIST (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Initializes the internal state of the Open-ES algorithm. This state holds information necessary for the evolutionary process, such as the current mean solution, covariance matrix, and other optimizer-specific parameters, preparing the algorithm for the optimization loop.

```python
key, subkey = jax.random.split(key)
state = es.init(subkey, solution, params)
```

--------------------------------

### Defining SNES Training Step Function

Source: https://github.com/roberttlange/evosax/blob/main/examples/02_rl.ipynb

Defines a single training step for the SNES algorithm, suitable for `jax.lax.scan`. It involves asking for a new population, evaluating their fitness in the environment, and telling the algorithm the results to update its internal state.

```Python
def step(carry, key):
    state, params, problem_state = carry
    key_ask, key_eval, key_tell = jax.random.split(key, 3)

    population, state = es.ask(key_ask, state, params)

    fitness, problem_state, _ = problem.eval(key_eval, population, problem_state)

    state, metrics = es.tell(
        key_tell, population, -fitness, state, params
    )  # Minimize fitness

    return (state, params, problem_state), metrics
```

--------------------------------

### Displaying CIFAR-10 Model Parameter Count (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Calculates and prints the total number of parameters in the initialized CNN solution for the CIFAR-10 problem. This provides an insight into the model's complexity, which is typically higher for CIFAR-10 than MNIST.

```python
print(f"Number of pararmeters: {sum(leaf.size for leaf in jax.tree.leaves(solution))}")
```

--------------------------------

### Defining MNIST CNN Problem (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Defines a Convolutional Neural Network (CNN) architecture for the MNIST dataset and initializes the `TorchVisionProblem` instance. The `network` specifies layers, filters, kernel sizes, and MLP layers, while the `problem` links the network to the MNIST task, handles data loading, and provides methods for evaluating solutions. It also initializes the problem state and samples an initial solution structure.

```python
from evosax.problems import CNN, TorchVisionProblem as Problem, identity_output_fn

network = CNN(
    num_filters=[8, 16],
    kernel_sizes=[(5, 5), (5, 5)],
    strides=[(1, 1), (1, 1)],
    mlp_layer_sizes=[10],
    output_fn=identity_output_fn,
)

problem = Problem(
    task_name="MNIST",
    network=network,
    batch_size=1024,
)

key, subkey = jax.random.split(key)
problem_state = problem.init(key)

key, subkey = jax.random.split(key)
solution = problem.sample(subkey)
```

--------------------------------

### Running Persistent Evolution Strategy Optimization Loop

Source: https://github.com/roberttlange/evosax/blob/main/examples/07_persistent_es.ipynb

Executes the main optimization loop for `PersistentES` over 5000 generations. In each generation, it asks for a new population, unrolls the inner problem, tells the strategy the fitness, and periodically evaluates and prints the mean fitness of the best solution found.

```Python
for i in range(5000):
    key, key_ask, key_tell = jax.random.split(key, 3)

    if state.inner_step_counter == 0:
        # Reset the inner problem: iteration, parameters
        xs = jnp.ones((population_size, 2)) * jnp.array([1.0, 1.0])

    population, state = strategy.ask(key_ask, state, params)

    # Unroll inner problem for K steps using antithetic perturbations
    fitness, xs = jax.vmap(unroll, in_axes=(0, 0, None, None, None))(
        xs, population, state.inner_step_counter, params.T, params.K
    )

    state, metrics = strategy.tell(key_tell, population, fitness, state, params)

    # Evaluation!
    if i % 500 == 0:
        L, _ = unroll(jnp.array([1.0, 1.0]), state.mean, 0, params.T, params.T)
        print(f"Generation {i:4d} | Mean fitness: {L:.2f}")
```

--------------------------------

### Adding Upstream Remote for `evosax`

Source: https://github.com/roberttlange/evosax/blob/main/CONTRIBUTING.md

This command adds the main `evosax` repository as an 'upstream' remote to your local fork. This is crucial for syncing your local branch with the latest changes from the original repository.

```bash
git remote add upstream https://github.com/RobertTLange/evosax
```

--------------------------------

### Running ES Training Loop for MNIST (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Executes the main training loop for the Open-ES algorithm on the MNIST problem. It uses `jax.lax.scan` to efficiently run multiple `step` iterations, periodically evaluates the mean solution's performance on the test set, and prints the current fitness and accuracy. This loop drives the evolutionary optimization process.

```python
fitness_log = []
log_period = 32
for i in range(num_generations // log_period):
    key, subkey = jax.random.split(key)
    keys = jax.random.split(subkey, log_period)
    (state, params, problem_state), metrics = jax.lax.scan(
        step,
        (state, params, problem_state),
        keys,
    )

    mean = es.get_mean(state)

    key, subkey = jax.random.split(key)
    fitness, problem_state, info = problem.eval_test(
        key, jax.tree.map(lambda x: x[None], mean), problem_state
    )
    print(
        f"Generation {(i + 1) * log_period:03d} | Mean fitness: {fitness.mean():.2f} | Accuracy: {info['accuracy'].mean():.2f}"
    )
```

--------------------------------

### Running ES Training Loop for CIFAR-10 (Python)

Source: https://github.com/roberttlange/evosax/blob/main/examples/03_vision.ipynb

Executes the main training loop for the Open-ES algorithm on the CIFAR-10 problem. Similar to the MNIST loop, it uses `jax.lax.scan` for efficient iteration, periodically evaluates the mean solution's performance on the test set, and prints the current fitness and accuracy.

```python
fitness_log = []
log_period = 32
for i in range(num_generations // log_period):
    key, subkey = jax.random.split(key)
    keys = jax.random.split(subkey, log_period)
    (state, params, problem_state), metrics = jax.lax.scan(
        step,
        (state, params, problem_state),
        keys,
    )

    mean = es.get_mean(state)

    key, subkey = jax.random.split(key)
    fitness, problem_state, info = problem.eval_test(
        key, jax.tree.map(lambda x: x[None], mean), problem_state
    )
    print(
        f"Generation {(i + 1) * log_period:03d} | Mean fitness: {fitness.mean():.2f} | Accuracy: {info['accuracy'].mean():.2f}"
    )
```