<jupyter_start><jupyter_text>00 - Getting Started with `evosax` - The Ask-Eval-Tell API [Last Update: March 2022][![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/RobertTLange/evosax/blob/main/examples/00_getting_started.ipynb)<jupyter_code>%matplotlib inline
%load_ext autoreload
%autoreload 2
%config InlineBackend.figure_format = 'retina'

!pip install git+https://github.com/RobertTLange/evosax.git@main<jupyter_output><empty_output><jupyter_text>Evolution Strategy Instantiation<jupyter_code>import jax
import jax.numpy as jnp
from evosax import CMA_ES
from evosax.problems import ClassicFitness

# Instantiate the evolution strategy instance
strategy = CMA_ES(num_dims=2, popsize=10)

# Get default hyperparameters (e.g. lrate, etc.)
es_params = strategy.default_params
es_params["init_min"] = -3
es_params["init_max"] = 3

# Initialize the strategy
rng = jax.random.PRNGKey(0)
state = strategy.initialize(rng, es_params)

# Have a look at the hyperparameters (change if desired)
es_params<jupyter_output><empty_output><jupyter_text>Classic Evolution Strategy Benchmarks<jupyter_code># Instantiate helper class for classic evolution strategies benchmarks
evaluator = ClassicFitness("rosenbrock", num_dims=2)<jupyter_output><empty_output><jupyter_text>Given our initialized strategy we are now ready to `ask` for a set of candidate parameters. Afterwards, we evaluate these on the 2D Rosenbrock problem and `tell` them to our strategy. The strategy will update its `state` and we can iterate.<jupyter_code># Ask for a set of candidate solutions to evaluate
x, state = strategy.ask(rng, state, es_params)
# Evaluate the population members
fitness = evaluator.rollout(rng, x)
# Update the evolution strategy
state = strategy.tell(x, fitness, state, es_params)
state<jupyter_output><empty_output><jupyter_text>Running the ES Loop with Logging<jupyter_code>from evosax.utils import ESLog
# Jittable logging helper
num_gens = 50
es_logging = ESLog(num_dims=2, num_generations=num_gens, top_k=3, maximize=False)
log = es_logging.initialize()
state = strategy.initialize(rng, es_params)
for i in range(num_gens):
    rng, rng_ask = jax.random.split(rng)
    # Ask for a set candidates
    x, state = strategy.ask(rng_ask, state, es_params)
    # Evaluate the candidates
    fitness = evaluator.rollout(rng, x)
    # Update the strategy based on fitness
    state = strategy.tell(x, fitness, state, es_params)
    # Update the log with results
    log = es_logging.update(log, x, fitness)
    
es_logging.plot(log, "2D Rosenbrock CMA-ES", ylims=(0, 30))<jupyter_output><empty_output><jupyter_text>Simultaneous PyTree Evaluation with `evosax`'s `ParameterReshaper``evosax` supports the automatic reshaping of proposed flat vectors into Pytrees for smooth neural network fitness evaluations. The transformation is again JAX-composable (`jit`, `vmap`, etc.). Below you find an example for a flax-based multi-layer perceptron:<jupyter_code>from flax import linen as nn


class MLP(nn.Module):
    """Simple ReLU MLP."""

    num_hidden_units: int
    num_hidden_layers: int
    num_output_units: int

    @nn.compact
    def __call__(self, x, rng):
        for l in range(self.num_hidden_layers):
            x = nn.Dense(features=self.num_hidden_units)(x)
            x = nn.relu(x)
        x = nn.Dense(features=self.num_output_units)(x)
        return jax.random.categorical(rng, x)
    

# Instantiate the model callables and get a placeholder pytree
network = MLP(64, 2, 2)
policy_params = network.init(rng, jnp.zeros(4,), rng)["params"]
from evosax.utils import ParameterReshaper

# Instantiate the reshape helper & get total number of parameters to reshape
param_reshaper = ParameterReshaper(policy_params)
param_reshaper.total_params<jupyter_output>ParameterReshaper: More than one device detected. Please make sure that the ES population size divides evenly across the number of devices to pmap/parallelize over.<jupyter_text>Now let's instantiate another evolution strategy (`DE` - Differential Evolution) and generate a set of generation members:<jupyter_code>from evosax import DE
strategy = DE(popsize=100, num_dims=param_reshaper.total_params)
state = strategy.initialize(rng, strategy.default_params)
x, state = strategy.ask(rng, state, strategy.default_params)
x.shape<jupyter_output><empty_output><jupyter_text>As we can see this is simply an array with (population members, parameters) shape. In order to reshape this array into stacked pytrees, we can simply pass it to the reshaper:<jupyter_code>net_params = param_reshaper.reshape(x)
net_params.keys(), net_params['Dense_0']['kernel'].shape<jupyter_output>/cognition/home/RobTLange/anaconda/envs/snippets/lib/python3.8/site-packages/jax/_src/dispatch.py:232: UserWarning: The jitted function reshape_network includes a pmap. Using jit-of-pmap can lead to inefficient data movement, as the outer jit does not preserve sharded data representations and instead collects input and output arrays onto a single device. Consider removing the outer jit unless you know what you're doing. See https://github.com/google/jax/issues/2926.
  warnings.warn(<jupyter_text>If you now want to map over the population member axis, you can do so with the of the `vmap_dict` (more about this later):<jupyter_code># Get dictionary to vectorize/parallelize rollouts with
param_reshaper.vmap_dict<jupyter_output><empty_output><jupyter_text>Fitness Shaping with `evosax`'s `FitnessShaper`By default `evosax` will minimize the objective. If you want to instead maximize it (as you commonly do with RL returns) or want to apply any other common ES fitness shaping, you can use the `FitnessShaper`:<jupyter_code>from evosax import FitnessShaper
fit_shaper = FitnessShaper(centered_rank=True, w_decay=0.01, maximize=True)

x = jnp.array([[1.0], [2.0], [3.0]])
fit = jnp.array([0.0, 1.0, 2.0])
fit_shaper.apply(x, fit)<jupyter_output><empty_output><jupyter_text>ARS on CartPole Task`evosax` also comes with a simple fitness evaluation helper for a JAX-based version of Cartpole. You will have to make use of the `vmap_dict` in order to vectorize the rollouts along the population axis:<jupyter_code>from evosax.problems import GymFitness

evaluator = GymFitness("CartPole-v1", num_env_steps=200, num_rollouts=16)
evaluator.set_apply_fn(param_reshaper.vmap_dict, network.apply)
from evosax import ARS

strategy = ARS(popsize=100,
               num_dims=param_reshaper.total_params,
               elite_ratio=0.1, opt_name="sgd")

es_params = {
        "lrate_init": 0.01,  # Initial learning rate
        "lrate_decay": 0.999, # Multiplicative decay factor
        "lrate_limit": 0.001,  # Smallest possible lrate
        "momentum": 0.0,  # Momentum coefficient
        "sigma_init": 0.02,  # Initial scale of isotropic Gaussian noise
        "sigma_decay": 0.999,  # Multiplicative decay factor
        "sigma_limit": 0.01,  # Smallest possible scale
        "init_min": -0.1,  # Range of parameter archive initialization - Min
        "init_max": 0.1,  # Range of parameter archive initialization - Max
        "clip_min": -10,  # Range of parameter proposals - Min
        "clip_max": 10  # Range of parameter proposals - Max
}
num_generations = 250
num_rollouts = 20
print_every_k_gens = 20

es_logging = ESLog(param_reshaper.total_params,
                   num_generations,
                   top_k=5,
                   maximize=True)
log = es_logging.initialize()
fit_shaper = FitnessShaper(maximize=True)

state = strategy.initialize(rng, es_params)

for gen in range(num_generations):
    rng, rng_init, rng_ask, rng_eval = jax.random.split(rng, 4)
    x, state = strategy.ask(rng_ask, state, es_params)
    reshaped_params = param_reshaper.reshape(x)
    fitness = evaluator.rollout(rng_eval, reshaped_params).mean(axis=1)
    fit_re = fit_shaper.apply(x, fitness)
    state = strategy.tell(x, fit_re, state, es_params)
    log = es_logging.update(log, x, fitness)
    
    if gen % print_every_k_gens == 0:
        print("Generation: ", gen, "Performance: ", -state["best_fitness"])
    #break
        
es_logging.plot(log, "CartPole Augmented Random Search")<jupyter_output>/cognition/home/RobTLange/anaconda/envs/snippets/lib/python3.8/site-packages/jax/_src/dispatch.py:232: UserWarning: The jitted function rollout includes a pmap. Using jit-of-pmap can lead to inefficient data movement, as the outer jit does not preserve sharded data representations and instead collects input and output arrays onto a single device. Consider removing the outer jit unless you know what you're doing. See https://github.com/google/jax/issues/2926.
  warnings.warn(
{'fraction_non_alphanumeric': 0.0925776289968833, 'fraction_numerical': 0.013852014313748124, 'mean_word_length': 4.869918699186992, 'pattern_counts': {'":': 11, '<': 42, '<?xml version=': 0, '>': 42, 'https://': 5, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '130984', 'license': 'Apache-2.0', 'n_tokens_mistral': 2913, 'n_tokens_neox': 2776, 'n_words': 801, 'path': 'examples/00_getting_started.ipynb', 'repo_name': 'mahi97/evosax', 'size': '8663'}