from IPython.display import clear_output, Image

!pip install evojax

clear_output()

import os
import numpy as np
import jax
import jax.numpy as jnp

from evojax.task.cartpole import CartPoleSwingUp
from evojax.policy.mlp import MLPPolicy
from evojax.algo import PGPE
from evojax import Trainer
from evojax.util import create_logger

# Let's create a directory to save logs and models.
log_dir = './log'
logger = create_logger(name='EvoJAX', log_dir=log_dir)
logger.info('Welcome to the tutorial on Neuroevolution algorithm creation!')

logger.info('Jax backend: {}'.format(jax.local_devices()))
!nvidia-smi --query-gpu=name --format=csv,noheader

"""
Explanation: <a href="https://colab.research.google.com/github/google/evojax/blob/main/examples/notebooks/TutorialAlgorithmImplementation.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
Tutorial: Creating Neuroevolution Algoritms
Pre-requisite
Before we start, we need to install EvoJAX and import some libraries.
Note In our paper, we ran the experiments on NVIDIA V100 GPU(s). Your results can be different from ours.
End of explanation
"""


seed = 42  # Wish me luck!

# We use the classic cart-pole swing up as our tasks, see
# https://github.com/google/evojax/tree/main/evojax/task for more example tasks.
# The test flag provides the opportunity for a user to
# 1. Return different signals as rewards. For example, in our MNIST example,
#    we use negative cross-entropy loss as the reward in training tasks, and the
#    classification accuracy as the reward in test tasks.
# 2. Perform reward shaping. It is common for RL practitioners to modify the
#    rewards during training so that the agent learns more efficiently. But this
#    modification should not be allowed in tests for fair evaluations.
hard = False
train_task = CartPoleSwingUp(harder=hard, test=False)
test_task = CartPoleSwingUp(harder=hard, test=True)

# We use a feedforward network as our policy.
# By default, MLPPolicy uses "tanh" as its activation function for the output.
policy = MLPPolicy(
    input_dim=train_task.obs_shape[0],
    hidden_dims=[64, 64],
    output_dim=train_task.act_shape[0],
    logger=logger,
)

# We use PGPE as our evolution algorithm.
# If you want to know more about the algorithm, please take a look at the paper:
# https://people.idsia.ch/~juergen/nn2010.pdf 
solver = PGPE(
    pop_size=64,
    param_size=policy.num_params,
    optimizer='adam',
    center_learning_rate=0.05,
    seed=seed,
)

# Now that we have all the three components instantiated, we can create a
# trainer and start the training process.
trainer = Trainer(
    policy=policy,
    solver=solver,
    train_task=train_task,
    test_task=test_task,
    max_iter=600,
    log_interval=100,
    test_interval=200,
    n_repeats=5,
    n_evaluations=128,
    seed=seed,
    log_dir=log_dir,
    logger=logger,
)
_ = trainer.run()

# Let's visualize the learned policy.

def render(task, algo, policy):
    """Render the learned policy."""

    task_reset_fn = jax.jit(test_task.reset)
    policy_reset_fn = jax.jit(policy.reset)
    step_fn = jax.jit(test_task.step)
    act_fn = jax.jit(policy.get_actions)

    params = algo.best_params[None, :]
    task_s = task_reset_fn(jax.random.PRNGKey(seed=seed)[None, :])
    policy_s = policy_reset_fn(task_s)

    images = [CartPoleSwingUp.render(task_s, 0)]
    done = False
    step = 0
    reward = 0
    while not done:
        act, policy_s = act_fn(task_s, params, policy_s)
        task_s, r, d = step_fn(task_s, act)
        step += 1
        reward = reward + r
        done = bool(d[0])
        if step % 3 == 0:
            images.append(CartPoleSwingUp.render(task_s, 0))
    print('reward={}'.format(reward))
    return images


imgs = render(test_task, solver, policy)
gif_file = os.path.join(log_dir, 'cartpole.gif')
imgs[0].save(
    gif_file, save_all=True, append_images=imgs[1:], duration=40, loop=0)
Image(open(os.path.join(log_dir, 'cartpole.gif'),'rb').read())

"""
Explanation: Introduction
EvoJAX has three major components: the task, the policy network and the neuroevolution algorithm. Once these components are implemented and instantiated, we can use a trainer to start the training process. The following code snippet provides an example of how we use EvoJAX.
End of explanation
"""


import cma
from evojax.algo.base import NEAlgorithm


class CMAWrapper(NEAlgorithm):
    """This is a wrapper of CMA-ES."""

    def __init__(self, param_size, pop_size, init_stdev=0.1, seed=0):
        
        self.pop_size = pop_size
        self.params = None
        self._best_params = None

        # We create CMA-ES in a simplest form.
        self.cma = cma.CMAEvolutionStrategy(
            x0=np.zeros(param_size),
            sigma0=init_stdev,
            inopts={
                'popsize': pop_size,
                'seed': seed if seed > 0 else 42,
                'randn': np.random.randn,
            },
        )
        
        # We jit-compile some utility functions.
        self.jnp_array = jax.jit(jnp.array)
        self.jnp_stack = jax.jit(jnp.stack)

    def ask(self):
        self.params = self.cma.ask()
        return self.jnp_stack(self.params)

    def tell(self, fitness):
        # CMA-ES minimizes, so we negate the fitness.
        self.cma.tell(self.params, -np.array(fitness))
        self._best_params = np.array(self.cma.result.xfavorite)

    @property
    def best_params(self):
        return self.jnp_array(self._best_params)

    @best_params.setter
    def best_params(self, params):
        self._best_params = np.array(params)

"""
Explanation: This tutorial walks you through the process of creating a new neuroevolution algoritm.  
To contribute an algorithm implementation to EvoJAX, all you need to do is to implement the NEAlgorithm interface.
The interface is defined as the following and you can see the related Python file here:
```python
class NEAlgorithm(ABC):
    """Interface of all Neuro-evolution algorithms in EvoJAX."""
pop_size: int

@abstractmethod
def ask(self) -&gt; jnp.ndarray:
    """Ask the algorithm for a population of parameters.
    Returns
        A Jax array of shape (population_size, param_size).
    """
    raise NotImplementedError()

@abstractmethod
def tell(self, fitness: Union[np.ndarray, jnp.ndarray]) -&gt; None:
    """Report the fitness of the population to the algorithm.
    Args:
        fitness - The fitness scores array.
    """
    raise NotImplementedError()

@property
def best_params(self) -&gt; jnp.ndarray:
    raise NotImplementedError()

@best_params.setter
def best_params(self, params: Union[np.ndarray, jnp.ndarray]) -&gt; None:
    raise NotImplementedError()

```
Wrap an existing implementation
NEAlgorithm adopts the well-known “ask” and “tell” interfaces, where the former requests the algorithm to generate a population of parameters and the latter reports the parameters' fitness scores so that the algorithm can update its internal states. We think the conventional interface for the neuroevolution algorithms brings familiarity to the developers and thus reduces the required learning effort. Moreover, the interface is also used by many existing algorithms, it is therefore possible for the practitioners to quickly plug in existing algorithms for sanity checks. In the first part of this tutorial, we will create an implementation that wraps CMA-ES. Please take a look at this wonderful tutorial for more information about CMA-ES.
End of explanation
"""


# Instead of PGPE, we use our CMAWrapper now.
solver = CMAWrapper(
    pop_size=64,
    param_size=policy.num_params,
    seed=seed,
)
trainer = Trainer(
    policy=policy,
    solver=solver,
    train_task=train_task,
    test_task=test_task,
    max_iter=600,
    log_interval=100,
    test_interval=200,
    n_repeats=5,
    n_evaluations=128,
    seed=seed,
    log_dir=log_dir,
    logger=logger,
)
_ = trainer.run()

"""
Explanation: Notice that our implementation above is extremely simple, we haven't used many options or functions provided by CMA-ES.
But let's plug in this implementation to our cart-pole earlier example and see how it works.
<font color='red'>Alert</font> Depending on your CPUs, running the following cell may take some time.
End of explanation
"""


from evojax.algo.base import NEAlgorithm


class SimplePGPE(NEAlgorithm):
    """A simplified version of PGPE."""

    def __init__(self, param_size, pop_size,
                 lr_mu=0.05, lr_sigma=0.1, init_stdev=0.1, seed=0):
        
        self.pop_size = pop_size
        assert pop_size % 2 == 0, "pop_size must be a multpile of 2."
        n_directs = pop_size // 2
        self.noises = jnp.zeros(param_size)
        self.params = jnp.zeros(param_size)
        self.mu = jnp.zeros(param_size)
        self.sigma = jnp.ones(param_size) * init_stdev
        self.rand_key = jax.random.PRNGKey(seed=seed)
        
        def ask_fn(key, mu, sigma):
            next_key, sample_key = jax.random.split(key=key, num=2)
            perturbations = jax.random.normal(
                key=sample_key, shape=(n_directs, param_size)) * sigma[None, :]
            params = jnp.vstack([perturbations, -perturbations]) + mu[None, :]
            return params, perturbations, next_key

        self.ask_fn = jax.jit(ask_fn)

        def tell_fn(rewards, mu, sigma, perturbations):            
            fitness = jnp.array(rewards).reshape([2, n_directs])
            
            # To map to the formulae above:
            # (r - b) = (avg_fitness - b) and (theta - mu) = perturbations
            avg_fitness = fitness.mean(axis=0)
            b = jnp.mean(fitness)

            # Update the means.
            grad_mu = (
                (avg_fitness - b)[:, None] * perturbations
            ).mean(axis=0)
            new_mu = mu + lr_mu * grad_mu

            # Update the sigmas.
            # We constrain the change of sigma to prevent numerical errors.
            grad_sigma = (
                (avg_fitness - b)[:, None] *
                (perturbations ** 2 - (sigma ** 2)[None, :]) / sigma[None, :]
            ).mean(axis=0)
            new_sigma = jnp.clip(
                sigma + lr_sigma * grad_sigma, 0.8 * sigma, 1.2 * sigma)
            
            return new_mu, new_sigma

        self.tell_fn = jax.jit(tell_fn)

    def ask(self):
        self.params, self.noises, self.rand_key = self.ask_fn(
            self.rand_key, self.mu, self.sigma)
        return self.params

    def tell(self, fitness):
        self.mu, self.sigma = self.tell_fn(
            fitness, self.mu, self.sigma, self.noises)

    @property
    def best_params(self):
        return self.mu

    @best_params.setter
    def best_params(self, params):
        self.mu = jnp.array(params)

# Let's test our simple PGPE.
solver = SimplePGPE(
    pop_size=64,
    param_size=policy.num_params,
    seed=seed,
)
trainer = Trainer(
    policy=policy,
    solver=solver,
    train_task=train_task,
    test_task=test_task,
    max_iter=1000,
    log_interval=100,
    test_interval=200,
    n_repeats=5,
    n_evaluations=128,
    seed=seed,
    log_dir=log_dir,
    logger=logger,
)
_ = trainer.run()

"""
Explanation: The simple CMA-ES wrapper worked! However, we also notice that the training time increased significantly.
Although the task and the policy networks are accelerated by GPUs, the cma.CMAEvolutionStrategy implementation we used in the code above relies on CPUs, and that is why we see the drop in training speed.  
Nevertheless, being able to wrapper an existing algorithm and plug that in EvoJAX's training pipeline serves as sanity checks and helps debugging when you migrate algorithms to EvoJAX. Next, we will show you how to implement an algorithm in JAX from scratch.
Simple PGPE in JAX
We are going to implement a very simple version of PGPE, users interested in the algorithm can take a look at the paper and also check out some popular implementations (example1, example2).
In a nutshell, PGPE samples the policy network parameters $\theta$ from Gaussian distributions. It maintains the means $\mu$ and the standard deviations $\sigma$ of the Gaussian distributions, and then estimates the gradients of these parameters using the following formulae:
$\Delta \mu_i = \alpha (r - b) (\theta_i - \mu_i)$ ,
$\Delta \sigma_i = \alpha (r - b) \frac{(\theta_i - \mu_i)^2 - {\sigma_i}^2}{\sigma_i}$
where $\alpha$ is the learning rate and $b$ is a baseline from the reward $r$.
The following code snippet provides a sample implementation of PGPE.
Note This simplified version ignores popular tricks such as converting the rewards to ranks, using modern optimizers for parameter update, etc.
End of explanation
"""


