# Assignment 2 - Implement your agent

Welcome to Course 4, Programming Assignment 2! We have learned about reinforcement learning algorithms for prediction and control in previous courses and extended those algorithms to large state spaces using function approximation. One example of this was in assignment 2 of course 3 where we implemented semi-gradient TD for prediction and used a neural network as the function approximator. In this notebook, we will build a reinforcement learning agent for control, again using a neural network for function approximation. This combination of neural network function approximators and reinforcement learning algorithms, often referred to as Deep RL, is an active area of research and has led to many impressive results (e. g., AlphaGo: https://deepmind.com/research/case-studies/alphago-the-story-so-far).

**In this assignment, you will:**
  1. Extend the neural network code from assignment 2 of course 3 to output action-values instead of state-values.
  2. Write up the Adam algorithm for neural network optimization.
  3. Understand experience replay buffers.
  4. Implement Softmax action-selection.
  5. Build an Expected Sarsa agent by putting all the pieces together.
  6. Solve Lunar Lander with your agent.

## Packages
- [numpy](www.numpy.org) : Fundamental package for scientific computing with Python.
- [matplotlib](http://matplotlib.org) : Library for plotting graphs in Python.
- [RL-Glue](http://www.jmlr.org/papers/v10/tanner09a.html), BaseEnvironment, BaseAgent : Library and abstract classes to inherit from  for reinforcement learning experiments.
- [LunarLanderEnvironment](https://gym.openai.com/envs/LunarLander-v2/) : An RLGlue environment that wraps a LundarLander environment implementation from OpenAI Gym.
- [collections.deque](https://docs.python.org/3/library/collections.html#collections.deque): a double-ended queue implementation. We use deque to implement the experience replay buffer.
- [copy.deepcopy](https://docs.python.org/3/library/copy.html#copy.deepcopy): As objects are not passed by value in python, we often need to make copies of mutable objects. copy.deepcopy allows us to make a new object with the same contents as another object. (Take a look at this link if you are interested to learn more: https://robertheaton.com/2014/02/09/pythons-pass-by-object-reference-as-explained-by-philip-k-dick/)
- [tqdm](https://github.com/tqdm/tqdm) : A package to display progress bar when running experiments
- [os](https://docs.python.org/3/library/os.html): Package used to interface with the operating system. Here we use it for creating a results folder when it does not exist.
- [shutil](https://docs.python.org/3/library/shutil.html): Package used to operate on files and folders. Here we use it for creating a zip file of the results folder.
- plot_script: Used for plotting learning curves using matplotlib.

```python
# Do not modify this cell!

# Import necessary libraries
# DO NOT IMPORT OTHER LIBRARIES - This will break the autograder.
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

from rl_glue import RLGlue
from environment import BaseEnvironment

from lunar_lander import LunarLanderEnvironment

from agent import BaseAgent

from collections import deque

from copy import deepcopy

from tqdm import tqdm

import os

import shutil

from plot_script import plot_result
```

## Section 1: Action-Value Network
This section includes the function approximator that we use in our agent, a neural network. In Course 3 Assignment 2, we used a neural network as the function approximator for a policy evaluation problem. In this assignment, we will use a neural network for approximating the action-value function in a control problem. The main difference between approximating a state-value function and an action-value function using a neural network is that in the former the output layer only includes one unit whereas in the latter the output layer includes as many units as the number of actions.

In the cell below, you will specify the architecture of the action-value neural network. More specifically, you will specify `self.layer_size` in the `__init__()` function.

We have already provided `get_action_values()` and `get_TD_update()` methods. The former computes the action-value function by doing a forward pass and the latter computes the gradient of the action-value function with respect to the weights times the TD error. These `get_action_values()` and `get_TD_update()` methods are similar to the `get_value()` and `get_gradient()` methods that you implemented in Course 3 Assignment 2. The main difference is that in this notebook, they are designed to be applied to batches of states instead of one state. You will later use these functions for implementing the agent.

```python
# Work Required: Yes. Fill in the code for layer_sizes in __init__ (~1 Line).
# Also go through the rest of the code to ensure your understanding is correct.
class ActionValueNetwork:
    # Work Required: Yes. Fill in the layer_sizes member variable (~1 Line).
    def __init__(self, network_config):
        self.state_dim = network_config.get("state_dim")
        self.num_hidden_units = network_config.get("num_hidden_units")
        self.num_actions = network_config.get("num_actions")

        self.rand_generator = np.random.RandomState(network_config.get("seed"))

        # Specify self.layer_size which shows the number of nodes in each layer
        ### START CODE HERE (~1 Line)
        self.layer_sizes = [self.state_dim, self.num_hidden_units, self.num_actions]
        ### END CODE HERE

        # Initialize the weights of the neural network
        # self.weights is an array of dictionaries with each dictionary corresponding to
        # the weights from one layer to the next. Each dictionary includes W and b
        self.weights = [dict() for i in range(0, len(self.layer_sizes) - 1)]
        for i in range(0, len(self.layer_sizes) - 1):
            self.weights[i]['W'] = self.init_saxe(self.layer_sizes[i], self.layer_sizes[i + 1])
            self.weights[i]['b'] = np.zeros((1, self.layer_sizes[i + 1]))

    # Work Required: No.
    def get_action_values(self, s):
        """
        Args:
            s (Numpy array): The state.
        Returns:
            The action-values (Numpy array) calculated using the network's weights.
        """

        W0, b0 = self.weights[0]['W'], self.weights[0]['b']
        psi = np.dot(s, W0) + b0
        x = np.maximum(psi, 0)

        W1, b1 = self.weights[1]['W'], self.weights[1]['b']
        q_vals = np.dot(x, W1) + b1

        return q_vals

    # Work Required: No.
    def get_TD_update(self, s, delta_mat):
        """
        Args:
            s (Numpy array): The state.
            delta_mat (Numpy array): A 2D array of shape (batch_size, num_actions). Each row of delta_mat
            correspond to one state in the batch. Each row has only one non-zero element
            which is the TD-error corresponding to the action taken.
        Returns:
            The TD update (Array of dictionaries with gradient times TD errors) for the network's weights
        """

        W0, b0 = self.weights[0]['W'], self.weights[0]['b']
        W1, b1 = self.weights[1]['W'], self.weights[1]['b']

        psi = np.dot(s, W0) + b0
        x = np.maximum(psi, 0)
        dx = (psi > 0).astype(float)

        # td_update has the same structure as self.weights, that is an array of dictionaries.
        # td_update[0]["W"], td_update[0]["b"], td_update[1]["W"], and td_update[1]["b"] have the same shape as
        # self.weights[0]["W"], self.weights[0]["b"], self.weights[1]["W"], and self.weights[1]["b"] respectively
        td_update = [dict() for i in range(len(self.weights))]

        v = delta_mat
        td_update[1]['W'] = np.dot(x.T, v) * 1. / s.shape[0]
        td_update[1]['b'] = np.sum(v, axis=0, keepdims=True) * 1. / s.shape[0]

        v = np.dot(v, W1.T) * dx
        td_update[0]['W'] = np.dot(s.T, v) * 1. / s.shape[0]
        td_update[0]['b'] = np.sum(v, axis=0, keepdims=True) * 1. / s.shape[0]

        return td_update

    # Work Required: No. You may wish to read the relevant paper for more information on this weight initialization
    # (Exact solutions to the nonlinear dynamics of learning in deep linear neural networks by Saxe, A et al., 2013)
    def init_saxe(self, rows, cols):
        """
        Args:
            rows (int): number of input units for layer.
            cols (int): number of output units for layer.
        Returns:
            NumPy Array consisting of weights for the layer based on the initialization in Saxe et al.
        """
        tensor = self.rand_generator.normal(0, 1, (rows, cols))
        if rows < cols:
            tensor = tensor.T
        tensor, r = np.linalg.qr(tensor)
        d = np.diag(r, 0)
        ph = np.sign(d)
        tensor *= ph

        if rows < cols:
            tensor = tensor.T
        return tensor

    # Work Required: No.
    def get_weights(self):
        """
        Returns:
            A copy of the current weights of this network.
        """
        return deepcopy(self.weights)

    # Work Required: No.
    def set_weights(self, weights):
        """
        Args:
            weights (list of dictionaries): Consists of weights that this network will set as its own weights.
        """
        self.weights = deepcopy(weights)
```

Run the cell below to test your implementation of the `__init__()` function for ActionValueNetwork:

```python
# Do not modify this cell!

## Test Code for ActionValueNetwork __init__() ##

# NOTE: The test below is limited in scope. Additional tests are used in the autograder, so it is recommended
# to test your implementations more carefully for correctness.

network_config = {
    "state_dim": 5,
    "num_hidden_units": 20,
    "num_actions": 3
}

test_network = ActionValueNetwork(network_config)
print("layer_sizes:", test_network.layer_sizes)
assert(np.allclose(test_network.layer_sizes, np.array([5, 20, 3])))

print("Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)")
```

    layer_sizes: [5, 20, 3]
    Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)

**Expected output:**

    layer_sizes: [ 5 20  3]

## Section 2: Adam Optimizer

In this assignment, you will use the Adam algorithm for updating the weights of your action-value network. As you may remember from Course 3 Assignment 2, the Adam algorithm is a more advanced variant of stochastic gradient descent (SGD). The Adam algorithm improves the SGD update with two concepts: adaptive vector stepsizes and momentum. It keeps running estimates of the mean and second moment of the updates, denoted by $\mathbf{m}$ and $\mathbf{v}$ respectively:
$$\mathbf{m_t} = \beta_m \mathbf{m_{t-1}} + (1 - \beta_m)g_t \\
\mathbf{v_t} = \beta_v \mathbf{v_{t-1}} + (1 - \beta_v)g^2_t
$$

Here, $\beta_m$ and $\beta_v$ are fixed parameters controlling the linear combinations above and $g_t$ is the update at time $t$ (generally the gradients, but here the TD error times the gradients).

Given that $\mathbf{m}$ and $\mathbf{v}$ are initialized to zero, they are biased toward zero. To get unbiased estimates of the mean and second moment, Adam defines $\mathbf{\hat{m}}$ and $\mathbf{\hat{v}}$ as:
$$ \mathbf{\hat{m}_t} = \frac{\mathbf{m_t}}{1 - \beta_m^t} \\
\mathbf{\hat{v}_t} = \frac{\mathbf{v_t}}{1 - \beta_v^t}
$$

The weights are then updated as follows:
$$ \mathbf{w_t} = \mathbf{w_{t-1}} + \frac{\alpha}{\sqrt{\mathbf{\hat{v}_t}}+\epsilon} \mathbf{\hat{m}_t}
$$

Here, $\alpha$ is the step size parameter and $\epsilon$ is another small parameter to keep the denominator from being zero.

In the cell below, you will implement the `__init__()` and `update_weights()` methods for the Adam algorithm. In `__init__()`, you will initialize `self.m` and `self.v`. In `update_weights()`, you will compute new weights given the input weights and an update $g$ (here `td_errors_times_gradients`) according to the equations above.

```python
### Work Required: Yes. Fill in code in __init__ and update_weights (~9-11 Lines).
class Adam():
    # Work Required: Yes. Fill in the initialization for self.m and self.v (~4 Lines).
    def __init__(self, layer_sizes,
                 optimizer_info):
        self.layer_sizes = layer_sizes

        # Specify Adam algorithm's hyper parameters
        self.step_size = optimizer_info.get("step_size")
        self.beta_m = optimizer_info.get("beta_m")
        self.beta_v = optimizer_info.get("beta_v")
        self.epsilon = optimizer_info.get("epsilon")

        # Initialize Adam algorithm's m and v
        self.m = [dict() for i in range(1, len(self.layer_sizes))]
        self.v = [dict() for i in range(1, len(self.layer_sizes))]

        for i in range(0, len(self.layer_sizes) - 1):
            ### START CODE HERE (~4 Lines)
            # Hint: The initialization for m and v should look very much like the initializations of the weights
            # except for the fact that initialization here is to zeroes (see description above.)
            self.m[i]["W"] = np.zeros((self.layer_sizes[i], self.layer_sizes[i+1]))
            self.m[i]["b"] = np.zeros((1, self.layer_sizes[i+1]))
            self.v[i]["W"] = np.zeros((self.layer_sizes[i], self.layer_sizes[i+1]))
            self.v[i]["b"] = np.zeros((1, self.layer_sizes[i+1]))
            ### END CODE HERE

        # Notice that to calculate m_hat and v_hat, we use powers of beta_m and beta_v to
        # the time step t. We can calculate these powers using an incremental product. At initialization then,
        # beta_m_product and beta_v_product should be ...? (Note that timesteps start at 1 and if we were to
        # start from 0, the denominator would be 0.)
        self.beta_m_product = self.beta_m
        self.beta_v_product = self.beta_v

    # Work Required: Yes. Fill in the weight updates (~5-7 lines).
    def update_weights(self, weights, td_errors_times_gradients):
        """
        Args:
            weights (Array of dictionaries): The weights of the neural network.
            td_errors_times_gradients (Array of dictionaries): The gradient of the
            action-values with respect to the network's weights times the TD-error
        Returns:
            The updated weights (Array of dictionaries).
        """
        for i in range(len(weights)):
            for param in weights[i].keys():
                ### START CODE HERE (~5-7 Lines)
                # Hint: Follow the equations above. First, you should update m and v and then compute
                # m_hat and v_hat. Finally, compute how much the weights should be incremented by.
                self.m[i][param] = self.beta_m * self.m[i][param] + (1 - self.beta_m) * td_errors_times_gradients[i][param]
                self.v[i][param] = self.beta_v * self.v[i][param] + (1 - self.beta_v) * td_errors_times_gradients[i][param] ** 2
                m_hat = self.m[i][param] / (1 - self.beta_m_product)
                v_hat = self.v[i][param] / (1 - self.beta_v_product)
                weight_update = self.step_size * m_hat / (np.sqrt(v_hat) + self.epsilon)
                ### END CODE HERE

                weights[i][param] = weights[i][param] + weight_update
        # Notice that to calculate m_hat and v_hat, we use powers of beta_m and beta_v to
        ### update self.beta_m_product and self.beta_v_product
        self.beta_m_product *= self.beta_m
        self.beta_v_product *= self.beta_v

        return weights
```

Run the following code to test your implementation of the `__init__()` function:

```python
# Do not modify this cell!

## Test Code for Adam __init__() ##

# NOTE: The test below is limited in scope. Additional tests are used in the autograder, so it is recommended
# to test your implementations more carefully for correctness.

network_config = {"state_dim": 5,
                  "num_hidden_units": 2,
                  "num_actions": 3
                 }

optimizer_info = {"step_size": 0.1,
                  "beta_m": 0.99,
                  "beta_v": 0.999,
                  "epsilon": 0.0001
                 }

network = ActionValueNetwork(network_config)
test_adam = Adam(network.layer_sizes, optimizer_info)

print("m[0][\"W\"] shape: {}".format(test_adam.m[0]["W"].shape))
print("m[0][\"b\"] shape: {}".format(test_adam.m[0]["b"].shape))
print("m[1][\"W\"] shape: {}".format(test_adam.m[1]["W"].shape))
print("m[1][\"b\"] shape: {}".format(test_adam.m[1]["b"].shape), "\n")

assert(np.allclose(test_adam.m[0]["W"].shape, np.array([5, 2])))
assert(np.allclose(test_adam.m[0]["b"].shape, np.array([1, 2])))
assert(np.allclose(test_adam.m[1]["W"].shape, np.array([2, 3])))
assert(np.allclose(test_adam.m[1]["b"].shape, np.array([1, 3])))

print("v[0][\"W\"] shape: {}".format(test_adam.v[0]["W"].shape))
print("v[0][\"b\"] shape: {}".format(test_adam.v[0]["b"].shape))
print("v[1][\"W\"] shape: {}".format(test_adam.v[1]["W"].shape))
print("v[1][\"b\"] shape: {}".format(test_adam.v[1]["b"].shape), "\n")

assert(np.allclose(test_adam.v[0]["W"].shape, np.array([5, 2])))
assert(np.allclose(test_adam.v[0]["b"].shape, np.array([1, 2])))
assert(np.allclose(test_adam.v[1]["W"].shape, np.array([2, 3])))
assert(np.allclose(test_adam.v[1]["b"].shape, np.array([1, 3])))

assert(np.all(test_adam.m[0]["W"]==0))
assert(np.all(test_adam.m[0]["b"]==0))
assert(np.all(test_adam.m[1]["W"]==0))
assert(np.all(test_adam.m[1]["b"]==0))

assert(np.all(test_adam.v[0]["W"]==0))
assert(np.all(test_adam.v[0]["b"]==0))
assert(np.all(test_adam.v[1]["W"]==0))
assert(np.all(test_adam.v[1]["b"]==0))

print("Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)")
```

    m[0]["W"] shape: (5, 2)
    m[0]["b"] shape: (1, 2)
    m[1]["W"] shape: (2, 3)
    m[1]["b"] shape: (1, 3)

    v[0]["W"] shape: (5, 2)
    v[0]["b"] shape: (1, 2)
    v[1]["W"] shape: (2, 3)
    v[1]["b"] shape: (1, 3)

    Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)

**Expected output:**

    m[0]["W"] shape: (5, 2)
    m[0]["b"] shape: (1, 2)
    m[1]["W"] shape: (2, 3)
    m[1]["b"] shape: (1, 3)

    v[0]["W"] shape: (5, 2)
    v[0]["b"] shape: (1, 2)
    v[1]["W"] shape: (2, 3)
    v[1]["b"] shape: (1, 3)

Run the following code to test your implementation of the `update_weights()` function:

```python
# Do not modify this cell!

## Test Code for Adam update_weights() ##

network_config = {"state_dim": 5,
                  "num_hidden_units": 2,
                  "num_actions": 3
                 }

optimizer_info = {"step_size": 0.1,
                  "beta_m": 0.99,
                  "beta_v": 0.999,
                  "epsilon": 0.0001
                 }

network = ActionValueNetwork(network_config)
test_adam = Adam(network.layer_sizes, optimizer_info)

rand_generator = np.random.RandomState(0)

# Initialize m and v
test_adam.m[0]["W"] = rand_generator.normal(0, 1, (5, 2))
test_adam.m[0]["b"] = rand_generator.normal(0, 1, (1, 2))
test_adam.m[1]["W"] = rand_generator.normal(0, 1, (2, 3))
test_adam.m[1]["b"] = rand_generator.normal(0, 1, (1, 3))

test_adam.v[0]["W"] = np.abs(rand_generator.normal(0, 1, (5, 2)))
test_adam.v[0]["b"] = np.abs(rand_generator.normal(0, 1, (1, 2)))
test_adam.v[1]["W"] = np.abs(rand_generator.normal(0, 1, (2, 3)))
test_adam.v[1]["b"] = np.abs(rand_generator.normal(0, 1, (1, 3)))

# Specify weights
weights = [dict() for i in range(1, len(test_adam.layer_sizes))]
weights[0]["W"] = rand_generator.normal(0, 1, (5, 2))
weights[0]["b"] = rand_generator.normal(0, 1, (1, 2))
weights[1]["W"] = rand_generator.normal(0, 1, (2, 3))
weights[1]["b"] = rand_generator.normal(0, 1, (1, 3))

# Specify g
g = [dict() for i in range(1, len(test_adam.layer_sizes))]
g[0]["W"] = rand_generator.normal(0, 1, (5, 2))
g[0]["b"] = rand_generator.normal(0, 1, (1, 2))
g[1]["W"] = rand_generator.normal(0, 1, (2, 3))
g[1]["b"] = rand_generator.normal(0, 1, (1, 3))

# Update weights
updated_weights = test_adam.update_weights(weights, g)

# updated weights asserts
updated_weights_answer = np.load("asserts/update_weights.npz")

print("updated_weights[0][\"W\"]\n", updated_weights[0]["W"], "\n")
print("updated_weights[0][\"b\"]\n", updated_weights[0]["b"], "\n")
print("updated_weights[1][\"W\"]\n", updated_weights[1]["W"], "\n")
print("updated_weights[1][\"b\"]\n", updated_weights[1]["b"], "\n")

assert(np.allclose(updated_weights[0]["W"], updated_weights_answer["W0"]))
assert(np.allclose(updated_weights[0]["b"], updated_weights_answer["b0"]))
assert(np.allclose(updated_weights[1]["W"], updated_weights_answer["W1"]))
assert(np.allclose(updated_weights[1]["b"], updated_weights_answer["b1"]))

print("Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)")
```

    updated_weights[0]["W"]
     [[-1.03112528  2.08618453]
     [-0.15531623  0.02412129]
     [-0.76656476 -0.65405898]
     [-0.92569612 -0.24916335]
     [-0.92180119  0.72137957]]

    updated_weights[0]["b"]
     [[-0.44392532 -0.69588495]]

    updated_weights[1]["W"]
     [[ 0.13962892  0.48820826  0.41311548]
     [ 0.3958054  -0.20738072 -0.47172585]]

    updated_weights[1]["b"]
     [[-0.48917533 -0.61934122 -1.48771198]]

    Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)

**Expected output:**

    updated_weights[0]["W"]
     [[-1.03112528  2.08618453]
     [-0.15531623  0.02412129]
     [-0.76656476 -0.65405898]
     [-0.92569612 -0.24916335]
     [-0.92180119  0.72137957]]

    updated_weights[0]["b"]
     [[-0.44392532 -0.69588495]]

    updated_weights[1]["W"]
     [[ 0.13962892  0.48820826  0.41311548]
     [ 0.3958054  -0.20738072 -0.47172585]]

    updated_weights[1]["b"]
     [[-0.48917533 -0.61934122 -1.48771198]]

## Section 3: Experience Replay Buffers

In Course 3, you implemented agents that update value functions once for each sample. We can use a more efficient approach for updating value functions. You have seen an example of an efficient approach in Course 2 when implementing Dyna. The idea behind Dyna is to learn a model using sampled experience, obtain simulated experience from the model, and improve the value function using the simulated experience.

Experience replay is a simple method that can get some of the advantages of Dyna by saving a buffer of experience and using the data stored in the buffer as a model. This view of prior data as a model works because the data represents actual transitions from the underlying MDP. Furthermore, as a side note, this kind of model that is not learned and simply a collection of experience can be called non-parametric as it can be ever-growing as opposed to a parametric model where the transitions are learned to be represented with a fixed set of parameters or weights.

We have provided the implementation of the experience replay buffer in the cell below. ReplayBuffer includes two main functions: `append()` and `sample()`. `append()` adds an experience transition to the buffer as an array that includes the state, action, reward, terminal flag (indicating termination of the episode), and next_state. `sample()` gets a batch of experiences from the buffer with size `minibatch_size`.

You will use the `append()` and `sample()` functions when implementing the agent.

```python
# Do not modify this cell!

# Work Required: No. However, do go through the code to ensure your understanding is correct.

class ReplayBuffer:
    def __init__(self, size, minibatch_size, seed):
        """
        Args:
            size (integer): The size of the replay buffer.
            minibatch_size (integer): The sample size.
            seed (integer): The seed for the random number generator.
        """
        self.buffer = []
        self.minibatch_size = minibatch_size
        self.rand_generator = np.random.RandomState(seed)
        self.max_size = size

    def append(self, state, action, reward, terminal, next_state):
        """
        Args:
            state (Numpy array): The state.
            action (integer): The action.
            reward (float): The reward.
            terminal (integer): 1 if the next state is a terminal state and 0 otherwise.
            next_state (Numpy array): The next state.
        """
        if len(self.buffer) == self.max_size:
            del self.buffer[0]
        self.buffer.append([state, action, reward, terminal, next_state])

    def sample(self):
        """
        Returns:
            A list of transition tuples including state, action, reward, terinal, and next_state
        """
        idxs = self.rand_generator.choice(np.arange(len(self.buffer)), size=self.minibatch_size)
        return [self.buffer[idx] for idx in idxs]

    def size(self):
        return len(self.buffer)
```

## Section 4: Softmax Policy

In this assignment, you will use a softmax policy. One advantage of a softmax policy is that it explores according to the action-values, meaning that an action with a moderate value has a higher chance of getting selected compared to an action with a lower value. Contrast this with an $\epsilon$-greedy policy which does not consider the individual action values when choosing an exploratory action in a state and instead chooses randomly when doing so.

The probability of selecting each action according to the softmax policy is shown below:
$$Pr{(A_t=a | S_t=s)} \hspace{0.1cm} \dot{=} \hspace{0.1cm} \frac{e^{Q(s, a)/\tau}}{\sum_{b \in A}e^{Q(s, b)/\tau}}$$
where $\tau$ is the temperature parameter which controls how much the agent focuses on the highest valued actions. The smaller the temperature, the more the agent selects the greedy action. Conversely, when the temperature is high, the agent selects among actions more uniformly random.

Given that a softmax policy exponentiates action values, if those values are large, exponentiating them could get very large. To implement the softmax policy in a numerically stable way, we often subtract the maximum action-value from the action-values. If we do so, the probability of selecting each action looks as follows:

$$Pr{(A_t=a | S_t=s)} \hspace{0.1cm} \dot{=} \hspace{0.1cm} \frac{e^{Q(s, a)/\tau - max_{c}Q(s, c)/\tau}}{\sum_{b \in A}e^{Q(s, b)/\tau - max_{c}Q(s, c)/\tau}}$$

In the cell below, you will implement the `softmax()` function. In order to do so, you could break the above computation into smaller steps:
- compute the preference, $H(a)$, for taking each action by dividing the action-values by the temperature parameter $\tau$,
- subtract the maximum preference across the actions from the preferences to avoid overflow, and,
- compute the probability of taking each action.

```python
def softmax(action_values, tau=1.0):
    """
    Args:
        action_values (Numpy array): A 2D array of shape (batch_size, num_actions).
                       The action-values computed by an action-value network.
        tau (float): The temperature parameter scalar.
    Returns:
        A 2D array of shape (batch_size, num_actions). Where each column is a probability distribution over
        the actions representing the policy.
    """
    ### START CODE HERE (~2 Lines)
    # Compute the preferences by dividing the action-values by the temperature parameter tau
    preferences = action_values / tau
    # Compute the maximum preference across the actions
    max_preference = np.max(preferences, axis=1)
    ### END CODE HERE

    # Reshape max_preference array which has shape [Batch,] to [Batch, 1]. This allows NumPy broadcasting
    # when subtracting the maximum preference from the preference of each action.
    reshaped_max_preference = max_preference.reshape((-1, 1))

    ### START CODE HERE (~2 Lines)
    # Compute the numerator, i.e., the exponential of the preference - the max preference.
    exp_preferences = np.exp(preferences - reshaped_max_preference)
    # Compute the denominator, i.e., the sum over the numerator along the actions axis.
    sum_of_exp_preferences = np.sum(exp_preferences, axis=1)
    ### END CODE HERE

    # Reshape sum_of_exp_preferences array which has shape [Batch,] to [Batch, 1] to  allow for NumPy broadcasting
    # when dividing the numerator by the denominator.
    reshaped_sum_of_exp_preferences = sum_of_exp_preferences.reshape((-1, 1))

    ### START CODE HERE (~1 Lines)
    # Compute the action probabilities according to the equation in the previous cell.
    action_probs = exp_preferences / reshaped_sum_of_exp_preferences
    ### END CODE HERE

    # squeeze() removes any singleton dimensions. It is used here because this function is used in the
    # agent policy when selecting an action (for which the batch dimension is 1.) As np.random.choice is used in
    # the agent policy and it expects 1D arrays, we need to remove this singleton batch dimension.
    action_probs = action_probs.squeeze()
    return action_probs
```

Run the cell below to test your implementation of the `softmax()` function:

```python
# Do not modify this cell!

## Test Code for softmax() ##

# NOTE: The test below is limited in scope. Additional tests are used in the autograder, so it is recommended
# to test your implementations more carefully for correctness.

rand_generator = np.random.RandomState(0)
action_values = rand_generator.normal(0, 1, (2, 4))
tau = 0.5

action_probs = softmax(action_values, tau)
print("action_probs", action_probs)

assert(np.allclose(action_probs, np.array([
    [0.25849645, 0.01689625, 0.05374514, 0.67086216],
    [0.84699852, 0.00286345, 0.13520063, 0.01493741]
])))

print("Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)")
```

    action_probs [[0.25849645 0.01689625 0.05374514 0.67086216]
     [0.84699852 0.00286345 0.13520063 0.01493741]]
    Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)

**Expected output:**

    action_probs [[0.25849645 0.01689625 0.05374514 0.67086216]
     [0.84699852 0.00286345 0.13520063 0.01493741]]

## Section 5: Putting the pieces together

In this section, you will combine components from the previous sections to write up an RL-Glue Agent. The main component that you will implement is the action-value network updates with experience sampled from the experience replay buffer.

At time $t$, we have an action-value function represented as a neural network, say $Q_t$. We want to update our action-value function and get a new one we can use at the next timestep. We will get this $Q_{t+1}$ using multiple replay steps that each result in an intermediate action-value function $Q_{t+1}^{i}$ where $i$ indexes which replay step we are at.

In each replay step, we sample a batch of experiences from the replay buffer and compute a minibatch Expected-SARSA update. Across these N replay steps, we will use the current "un-updated" action-value network at time $t$, $Q_t$, for computing the action-values of the next-states. This contrasts using the most recent action-values from the last replay step $Q_{t+1}^{i}$. We make this choice to have targets that are stable across replay steps. Here is the pseudocode for performing the updates:

$$
\begin{align}
& Q_t \leftarrow \text{action-value network at timestep t (current action-value network)}\\
& \text{Initialize } Q_{t+1}^1 \leftarrow Q_t\\
& \text{For } i \text{ in } [1, ..., N] \text{ (i.e. N} \text{  replay steps)}:\\
& \hspace{1cm} s, a, r, t, s'
\leftarrow \text{Sample batch of experiences from experience replay buffer} \\
& \hspace{1cm} \text{Do Expected Sarsa update with } Q_t: Q_{t+1}^{i+1}(s, a) \leftarrow Q_{t+1}^{i}(s, a) + \alpha \cdot \left[r + \gamma \left(\sum_{b} \pi(b | s') Q_t(s', b)\right) - Q_{t+1}^{i}(s, a)\right]\\
& \hspace{1.5cm} \text{ making sure to add the } \gamma \left(\sum_{b} \pi(b | s') Q_t(s', b)\right) \text{ for non-terminal transitions only.} \\
& \text{After N replay steps, we set } Q_{t+1}^{N} \text{ as } Q_{t+1} \text{ and have a new } Q_{t+1} \text{for time step } t + 1 \text{ that we will fix in the next set of updates. }
\end{align}
$$

As you can see in the pseudocode, after sampling a batch of experiences, we do many computations. The basic idea however is that we are looking to compute a form of a TD error. In order to so, we can take the following steps:
- compute the action-values for the next states using the action-value network $Q_{t}$,
- compute the policy $\pi(b | s')$ induced by the action-values $Q_{t}$ (using the softmax function you implemented before),
- compute the Expected sarsa targets $r + \gamma \left(\sum_{b} \pi(b | s') Q_t(s', b)\right)$,
- compute the action-values for the current states using the latest $Q_{t + 1}$, and,
- compute the TD-errors with the Expected Sarsa targets.

For the third step above, you can start by computing $\pi(b | s') Q_t(s', b)$ followed by summation to get $\hat{v}_\pi(s') = \left(\sum_{b} \pi(b | s') Q_t(s', b)\right)$. $\hat{v}_\pi(s')$ is an estimate of the value of the next state. Note for terminal next states, $\hat{v}_\pi(s') = 0$. Finally, we add the rewards to the discount times $\hat{v}_\pi(s')$.

You will implement these steps in the `get_td_error()` function below which given a batch of experiences (including states, next_states, actions, rewards, terminals), fixed action-value network (current_q), and action-value network (network), computes the TD error in the form of a 1D array of size batch_size.

```python
### Work Required: Yes. Fill in code in get_td_error (~9 Lines).
def get_td_error(states, next_states, actions, rewards, discount, terminals, network, current_q, tau):
    """
    Args:
        states (Numpy array): The batch of states with the shape (batch_size, state_dim).
        next_states (Numpy array): The batch of next states with the shape (batch_size, state_dim).
        actions (Numpy array): The batch of actions with the shape (batch_size,).
        rewards (Numpy array): The batch of rewards with the shape (batch_size,).
        discount (float): The discount factor.
        terminals (Numpy array): The batch of terminals with the shape (batch_size,).
        network (ActionValueNetwork): The latest state of the network that is getting replay updates.
        current_q (ActionValueNetwork): The fixed network used for computing the targets,
                                        and particularly, the action-values at the next-states.
    Returns:
        The TD errors (Numpy array) for actions taken, of shape (batch_size,)
    """

    # Note: Here network is the latest state of the network that is getting replay updates. In other words,
    # the network represents Q_{t+1}^{i} whereas current_q represents Q_t, the fixed network used for computing the
    # targets, and particularly, the action-values at the next-states.

    # Compute action values at next states using current_q network
    # Note that q_next_mat is a 2D array of shape (batch_size, num_actions)

    ### START CODE HERE (~1 Line)
    q_next_mat = current_q.get_action_values(next_states)
    ### END CODE HERE

    # Compute policy at next state by passing the action-values in q_next_mat to softmax()
    # Note that probs_mat is a 2D array of shape (batch_size, num_actions)

    ### START CODE HERE (~1 Line)
    probs_mat = softmax(q_next_mat, tau)
    ### END CODE HERE

    # Compute the estimate of the next state value, v_next_vec.
    # Hint: sum the action-values for the next_states weighted by the policy, probs_mat. Then, multiply by
    # (1 - terminals) to make sure v_next_vec is zero for terminal next states.
    # Note that v_next_vec is a 1D array of shape (batch_size,)

    ### START CODE HERE (~3 Lines)
    v_next_vec = np.sum(q_next_mat * probs_mat, axis=1) * (1-terminals)
    ### END CODE HERE

    # Compute Expected Sarsa target
    # Note that target_vec is a 1D array of shape (batch_size,)

    ### START CODE HERE (~1 Line)
    target_vec = rewards + discount * v_next_vec
    ### END CODE HERE

    # Compute action values at the current states for all actions using network
    # Note that q_mat is a 2D array of shape (batch_size, num_actions)

    ### START CODE HERE (~1 Line)
    q_mat = network.get_action_values(states)
    ### END CODE HERE

    # Batch Indices is an array from 0 to the batch size - 1.
    batch_indices = np.arange(q_mat.shape[0])

    # Compute q_vec by selecting q(s, a) from q_mat for taken actions
    # Use batch_indices as the index for the first dimension of q_mat
    # Note that q_vec is a 1D array of shape (batch_size)

    ### START CODE HERE (~1 Line)
    q_vec = q_mat[batch_indices, actions]
    ### END CODE HERE

    # Compute TD errors for actions taken
    # Note that delta_vec is a 1D array of shape (batch_size)

    ### START CODE HERE (~1 Line)
    delta_vec = target_vec - q_vec
    ### END CODE HERE

    return delta_vec
```

Run the following code to test your implementation of the `get_td_error()` function:

```python
# Do not modify this cell!

## Test Code for get_td_error() ##

# NOTE: The test below is limited in scope. Additional tests are used in the autograder, so it is recommended
# to test your implementations more carefully for correctness.

data = np.load("asserts/get_td_error_1.npz", allow_pickle=True)

states = data["states"]
next_states = data["next_states"]
actions = data["actions"]
rewards = data["rewards"]
discount = data["discount"]
terminals = data["terminals"]
tau = 0.001

network_config = {"state_dim": 8,
                  "num_hidden_units": 512,
                  "num_actions": 4
                  }

network = ActionValueNetwork(network_config)
network.set_weights(data["network_weights"])

current_q = ActionValueNetwork(network_config)
current_q.set_weights(data["current_q_weights"])

delta_vec = get_td_error(states, next_states, actions, rewards, discount, terminals, network, current_q, tau)
answer_delta_vec = data["delta_vec"]

assert(np.allclose(delta_vec, answer_delta_vec))
print("Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)")
```

    Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)

Now that you implemented the `get_td_error()` function, you can use it to implement the `optimize_network()` function. In this function, you will:
- get the TD-errors vector from `get_td_error()`,
- make the TD-errors into a matrix using zeroes for actions not taken in the transitions,
- pass the TD-errors matrix to the `get_TD_update()` function of network to calculate the gradients times TD errors, and,
- perform an ADAM optimizer step.

```python
### Work Required: Yes. Fill in code in optimize_network (~2 Lines).
def optimize_network(experiences, discount, optimizer, network, current_q, tau):
    """
    Args:
        experiences (Numpy array): The batch of experiences including the states, actions,
                                   rewards, terminals, and next_states.
        discount (float): The discount factor.
        network (ActionValueNetwork): The latest state of the network that is getting replay updates.
        current_q (ActionValueNetwork): The fixed network used for computing the targets,
                                        and particularly, the action-values at the next-states.
    """

    # Get states, action, rewards, terminals, and next_states from experiences
    states, actions, rewards, terminals, next_states = map(list, zip(*experiences))
    states = np.concatenate(states)
    next_states = np.concatenate(next_states)
    rewards = np.array(rewards)
    terminals = np.array(terminals)
    batch_size = states.shape[0]

    # Compute TD error using the get_td_error function
    # Note that q_vec is a 1D array of shape (batch_size)
    delta_vec = get_td_error(states, next_states, actions, rewards, discount, terminals, network, current_q, tau)

    # Batch Indices is an array from 0 to the batch_size - 1.
    batch_indices = np.arange(batch_size)

    # Make a td error matrix of shape (batch_size, num_actions)
    # delta_mat has non-zero value only for actions taken
    delta_mat = np.zeros((batch_size, network.num_actions))
    delta_mat[batch_indices, actions] = delta_vec

    # Pass delta_mat to compute the TD errors times the gradients of the network's weights from back-propagation

    ### START CODE HERE
    td_update = network.get_TD_update(states,delta_mat)
    ### END CODE HERE

    # Pass network.get_weights and the td_update to the optimizer to get updated weights
    ### START CODE HERE
    weights = optimizer.update_weights(network.get_weights(), td_update)
    ### END CODE HERE

    network.set_weights(weights)
```

Run the following code to test your implementation of the `optimize_network()` function:

```python
# Do not modify this cell!

## Test Code for optimize_network() ##

# NOTE: The test below is limited in scope. Additional tests are used in the autograder, so it is recommended
# to test your implementations more carefully for correctness.

input_data = np.load("asserts/optimize_network_input_1.npz", allow_pickle=True)

experiences = list(input_data["experiences"])
discount = input_data["discount"]
tau = 0.001

network_config = {"state_dim": 8,
                  "num_hidden_units": 512,
                  "num_actions": 4
                  }

network = ActionValueNetwork(network_config)
network.set_weights(input_data["network_weights"])

current_q = ActionValueNetwork(network_config)
current_q.set_weights(input_data["current_q_weights"])

optimizer_config = {'step_size': 3e-5,
                    'beta_m': 0.9,
                    'beta_v': 0.999,
                    'epsilon': 1e-8
                   }
optimizer = Adam(network.layer_sizes, optimizer_config)
optimizer.m = input_data["optimizer_m"]
optimizer.v = input_data["optimizer_v"]
optimizer.beta_m_product = input_data["optimizer_beta_m_product"]
optimizer.beta_v_product = input_data["optimizer_beta_v_product"]

optimize_network(experiences, discount, optimizer, network, current_q, tau)
updated_weights = network.get_weights()

output_data = np.load("asserts/optimize_network_output_1.npz", allow_pickle=True)
answer_updated_weights = output_data["updated_weights"]

assert(np.allclose(updated_weights[0]["W"], answer_updated_weights[0]["W"]))
assert(np.allclose(updated_weights[0]["b"], answer_updated_weights[0]["b"]))
assert(np.allclose(updated_weights[1]["W"], answer_updated_weights[1]["W"]))
assert(np.allclose(updated_weights[1]["b"], answer_updated_weights[1]["b"]))
print("Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)")
```

    Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)

Now that you implemented the `optimize_network()` function, you can implement the agent. In the cell below, you will fill the `agent_step()` and `agent_end()` functions. You should:
- select an action (only in `agent_step()`),
- add transitions (consisting of the state, action, reward, terminal, and next state) to the replay buffer, and,
- update the weights of the neural network by doing multiple replay steps and calling the `optimize_network()` function that you implemented above.

```python
### Work Required: Yes. Fill in code in agent_step and agent_end (~7 Lines).
class Agent(BaseAgent):
    def __init__(self):
        self.name = "expected_sarsa_agent"

    # Work Required: No.
    def agent_init(self, agent_config):
        """Setup for the agent called when the experiment first starts.

        Set parameters needed to setup the agent.

        Assume agent_config dict contains:
        {
            network_config: dictionary,
            optimizer_config: dictionary,
            replay_buffer_size: integer,
            minibatch_sz: integer,
            num_replay_updates_per_step: float
            discount_factor: float,
        }
        """
        self.replay_buffer = ReplayBuffer(agent_config['replay_buffer_size'],
                                          agent_config['minibatch_sz'], agent_config.get("seed"))
        self.network = ActionValueNetwork(agent_config['network_config'])
        self.optimizer = Adam(self.network.layer_sizes, agent_config["optimizer_config"])
        self.num_actions = agent_config['network_config']['num_actions']
        self.num_replay = agent_config['num_replay_updates_per_step']
        self.discount = agent_config['gamma']
        self.tau = agent_config['tau']

        self.rand_generator = np.random.RandomState(agent_config.get("seed"))

        self.last_state = None
        self.last_action = None

        self.sum_rewards = 0
        self.episode_steps = 0

    # Work Required: No.
    def policy(self, state):
        """
        Args:
            state (Numpy array): the state.
        Returns:
            the action.
        """
        action_values = self.network.get_action_values(state)
        probs_batch = softmax(action_values, self.tau)
        action = self.rand_generator.choice(self.num_actions, p=probs_batch.squeeze())
        return action

    # Work Required: No.
    def agent_start(self, state):
        """The first method called when the experiment starts, called after
        the environment starts.
        Args:
            state (Numpy array): the state from the
                environment's evn_start function.
        Returns:
            The first action the agent takes.
        """
        self.sum_rewards = 0
        self.episode_steps = 0
        self.last_state = np.array([state])
        self.last_action = self.policy(self.last_state)
        return self.last_action

    # Work Required: Yes. Fill in the action selection, replay-buffer update,
    # weights update using optimize_network, and updating last_state and last_action (~5 lines).
    def agent_step(self, reward, state):
        """A step taken by the agent.
        Args:
            reward (float): the reward received for taking the last action taken
            state (Numpy array): the state from the
                environment's step based, where the agent ended up after the
                last step
        Returns:
            The action the agent is taking.
        """

        self.sum_rewards += reward
        self.episode_steps += 1

        # Make state an array of shape (1, state_dim) to add a batch dimension and
        # to later match the get_action_values() and get_TD_update() functions
        state = np.array([state])

        # Select action
        ### START CODE HERE (~1 Line)
        action = self.policy(state)
        ### END CODE HERE

        # Append new experience to replay buffer
        # Note: look at the replay_buffer append function for the order of arguments

        ### START CODE HERE (~1 Line)
        self.replay_buffer.append(self.last_state, self.last_action, reward, 0, state)
        ### END CODE HERE

        # Perform replay steps:
        if self.replay_buffer.size() > self.replay_buffer.minibatch_size:
            current_q = deepcopy(self.network)
            for _ in range(self.num_replay):

                # Get sample experiences from the replay buffer
                experiences = self.replay_buffer.sample()

                # Call optimize_network to update the weights of the network (~1 Line)
                ### START CODE HERE
                optimize_network(experiences, self.discount, self.optimizer, self.network, current_q, self.tau)
                ### END CODE HERE

        # Update the last state and last action.
        ### START CODE HERE (~2 Lines)
        self.last_state = state
        self.last_action = action
        ### END CODE HERE

        return action

    # Work Required: Yes. Fill in the replay-buffer update and
    # update of the weights using optimize_network (~2 lines).
    def agent_end(self, reward):
        """Run when the agent terminates.
        Args:
            reward (float): the reward the agent received for entering the
                terminal state.
        """
        self.sum_rewards += reward
        self.episode_steps += 1

        # Set terminal state to an array of zeros
        state = np.zeros_like(self.last_state)

        # Append new experience to replay buffer
        # Note: look at the replay_buffer append function for the order of arguments

        ### START CODE HERE (~1 Line)
        self.replay_buffer.append(self.last_state, self.last_action, reward, 1, state)
        ### END CODE HERE

        # Perform replay steps:
        if self.replay_buffer.size() > self.replay_buffer.minibatch_size:
            current_q = deepcopy(self.network)
            for _ in range(self.num_replay):

                # Get sample experiences from the replay buffer
                experiences = self.replay_buffer.sample()

                # Call optimize_network to update the weights of the network
                ### START CODE HERE (~1 Line)
                optimize_network(experiences, self.discount, self.optimizer, self.network, current_q, self.tau)
                ### END CODE HERE

    def agent_message(self, message):
        if message == "get_sum_reward":
            return self.sum_rewards
        else:
            raise Exception("Unrecognized Message!")
```

Run the following code to test your implementation of the `agent_step()` function:

```python
# Do not modify this cell!

## Test Code for agent_step() ##

# NOTE: The test below is limited in scope. Additional tests are used in the autograder, so it is recommended
# to test your implementations more carefully for correctness.

agent_info = {
             'network_config': {
                 'state_dim': 8,
                 'num_hidden_units': 256,
                 'num_hidden_layers': 1,
                 'num_actions': 4
             },
             'optimizer_config': {
                 'step_size': 3e-5,
                 'beta_m': 0.9,
                 'beta_v': 0.999,
                 'epsilon': 1e-8
             },
             'replay_buffer_size': 32,
             'minibatch_sz': 32,
             'num_replay_updates_per_step': 4,
             'gamma': 0.99,
             'tau': 1000.0,
             'seed': 0}

# Initialize agent
agent = Agent()
agent.agent_init(agent_info)

# load agent network, optimizer, replay_buffer from the agent_input_1.npz file
input_data = np.load("asserts/agent_input_1.npz", allow_pickle=True)
agent.network.set_weights(input_data["network_weights"])
agent.optimizer.m = input_data["optimizer_m"]
agent.optimizer.v = input_data["optimizer_v"]
agent.optimizer.beta_m_product = input_data["optimizer_beta_m_product"]
agent.optimizer.beta_v_product = input_data["optimizer_beta_v_product"]
agent.replay_buffer.rand_generator.seed(int(input_data["replay_buffer_seed"]))
for experience in input_data["replay_buffer"]:
    agent.replay_buffer.buffer.append(experience)

# Perform agent_step multiple times
last_state_array = input_data["last_state_array"]
last_action_array = input_data["last_action_array"]
state_array = input_data["state_array"]
reward_array = input_data["reward_array"]

for i in range(5):
    agent.last_state = last_state_array[i]
    agent.last_action = last_action_array[i]
    state = state_array[i]
    reward = reward_array[i]

    agent.agent_step(reward, state)

    # Load expected values for last_state, last_action, weights, and replay_buffer
    output_data = np.load("asserts/agent_step_output_{}.npz".format(i), allow_pickle=True)
    answer_last_state = output_data["last_state"]
    answer_last_action = output_data["last_action"]
    answer_updated_weights = output_data["updated_weights"]
    answer_replay_buffer = output_data["replay_buffer"]

    # Asserts for last_state and last_action
    assert(np.allclose(answer_last_state, agent.last_state))
    assert(np.allclose(answer_last_action, agent.last_action))

    # Asserts for replay_buffer
    for i in range(answer_replay_buffer.shape[0]):
        for j in range(answer_replay_buffer.shape[1]):
            assert(np.allclose(np.asarray(agent.replay_buffer.buffer)[i, j], answer_replay_buffer[i, j]))

    # Asserts for network.weights
    assert(np.allclose(agent.network.weights[0]["W"], answer_updated_weights[0]["W"]))
    assert(np.allclose(agent.network.weights[0]["b"], answer_updated_weights[0]["b"]))
    assert(np.allclose(agent.network.weights[1]["W"], answer_updated_weights[1]["W"]))
    assert(np.allclose(agent.network.weights[1]["b"], answer_updated_weights[1]["b"]))

print("Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)")

```

    Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)

Run the following code to test your implementation of the `agent_end()` function:

```python
# Do not modify this cell!

## Test Code for agent_end() ##

# NOTE: The test below is limited in scope. Additional tests are used in the autograder, so it is recommended
# to test your implementations more carefully for correctness.

agent_info = {
             'network_config': {
                 'state_dim': 8,
                 'num_hidden_units': 256,
                 'num_hidden_layers': 1,
                 'num_actions': 4
             },
             'optimizer_config': {
                 'step_size': 3e-5,
                 'beta_m': 0.9,
                 'beta_v': 0.999,
                 'epsilon': 1e-8
             },
             'replay_buffer_size': 32,
             'minibatch_sz': 32,
             'num_replay_updates_per_step': 4,
             'gamma': 0.99,
             'tau': 1000,
             'seed': 0
             }

# Initialize agent
agent = Agent()
agent.agent_init(agent_info)

# load agent network, optimizer, replay_buffer from the agent_input_1.npz file
input_data = np.load("asserts/agent_input_1.npz", allow_pickle=True)
agent.network.set_weights(input_data["network_weights"])
agent.optimizer.m = input_data["optimizer_m"]
agent.optimizer.v = input_data["optimizer_v"]
agent.optimizer.beta_m_product = input_data["optimizer_beta_m_product"]
agent.optimizer.beta_v_product = input_data["optimizer_beta_v_product"]
agent.replay_buffer.rand_generator.seed(int(input_data["replay_buffer_seed"]))
for experience in input_data["replay_buffer"]:
    agent.replay_buffer.buffer.append(experience)

# Perform agent_step multiple times
last_state_array = input_data["last_state_array"]
last_action_array = input_data["last_action_array"]
state_array = input_data["state_array"]
reward_array = input_data["reward_array"]

for i in range(5):
    agent.last_state = last_state_array[i]
    agent.last_action = last_action_array[i]
    reward = reward_array[i]

    agent.agent_end(reward)

    # Load expected values for last_state, last_action, weights, and replay_buffer
    output_data = np.load("asserts/agent_end_output_{}.npz".format(i), allow_pickle=True)
    answer_updated_weights = output_data["updated_weights"]
    answer_replay_buffer = output_data["replay_buffer"]

    # Asserts for replay_buffer
    for i in range(answer_replay_buffer.shape[0]):
        for j in range(answer_replay_buffer.shape[1]):
            assert(np.allclose(np.asarray(agent.replay_buffer.buffer)[i, j], answer_replay_buffer[i, j]))

    # Asserts for network.weights
    assert(np.allclose(agent.network.weights[0]["W"], answer_updated_weights[0]["W"]))
    assert(np.allclose(agent.network.weights[0]["b"], answer_updated_weights[0]["b"]))
    assert(np.allclose(agent.network.weights[1]["W"], answer_updated_weights[1]["W"]))
    assert(np.allclose(agent.network.weights[1]["b"], answer_updated_weights[1]["b"]))

print("Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)")
```

    Passed the asserts! (Note: These are however limited in scope, additional testing is encouraged.)

## Section 6: Run Experiment

Now that you implemented the agent, we can use it to run an experiment on the Lunar Lander problem. We will plot the learning curve of the agent to visualize learning progress. To plot the learning curve, we use the sum of rewards in an episode as the performance measure. We have provided for you the experiment/plot code in the cell below which you can go ahead and run. Note that running the cell below has taken approximately 10 minutes in prior testing.

```python
def run_experiment(environment, agent, environment_parameters, agent_parameters, experiment_parameters):

    rl_glue = RLGlue(environment, agent)

    # save sum of reward at the end of each episode
    agent_sum_reward = np.zeros((experiment_parameters["num_runs"],
                                 experiment_parameters["num_episodes"]))

    env_info = {}

    agent_info = agent_parameters

    # one agent setting
    for run in range(1, experiment_parameters["num_runs"]+1):
        agent_info["seed"] = run
        agent_info["network_config"]["seed"] = run
        env_info["seed"] = run

        rl_glue.rl_init(agent_info, env_info)

        for episode in tqdm(range(1, experiment_parameters["num_episodes"]+1)):
            # run episode
            rl_glue.rl_episode(experiment_parameters["timeout"])

            episode_reward = rl_glue.rl_agent_message("get_sum_reward")
            agent_sum_reward[run - 1, episode - 1] = episode_reward
    save_name = "{}".format(rl_glue.agent.name)
    if not os.path.exists('results'):
        os.makedirs('results')
    np.save("results/sum_reward_{}".format(save_name), agent_sum_reward)
    shutil.make_archive('results', 'zip', 'results')

# Run Experiment

# Experiment parameters
experiment_parameters = {
    "num_runs" : 1,
    "num_episodes" : 300,
    # OpenAI Gym environments allow for a timestep limit timeout, causing episodes to end after
    # some number of timesteps. Here we use the default of 1000.
    "timeout" : 1000
}

# Environment parameters
environment_parameters = {}

current_env = LunarLanderEnvironment

# Agent parameters
agent_parameters = {
    'network_config': {
        'state_dim': 8,
        'num_hidden_units': 256,
        'num_actions': 4
    },
    'optimizer_config': {
        'step_size': 1e-3,
        'beta_m': 0.9,
        'beta_v': 0.999,
        'epsilon': 1e-8
    },
    'replay_buffer_size': 50000,
    'minibatch_sz': 8,
    'num_replay_updates_per_step': 4,
    'gamma': 0.99,
    'tau': 0.001
}
current_agent = Agent

# run experiment
run_experiment(current_env, current_agent, environment_parameters, agent_parameters, experiment_parameters)
```

    100%|██████████| 300/300 [11:20<00:00,  2.27s/it]

Run the cell below to see the comparison between the agent that you implemented and a random agent for the one run and 300 episodes. Note that the `plot_result()` function smoothes the learning curve by applying a sliding window on the performance measure.

```python
plot_result(["expected_sarsa_agent", "random_agent"])
```

In the following cell you can visualize the performance of the agent with a correct implementation. As you can see, the agent initially crashes quite quickly (Episode 0). Then, the agent learns to avoid crashing by expending fuel and staying far above the ground. Finally however, it learns to land smoothly within the landing zone demarcated by the two flags (Episode 275).

```python
%%HTML
<div align="middle">
</div>
```

<div align="middle">
</div>

In the learning curve above, you can see that sum of reward over episode has quite a high-variance at the beginning. However, the performance seems to be improving. The experiment that you ran was for 300 episodes and 1 run. To understand how the agent performs in the long run, we provide below the learning curve for the agent trained for 3000 episodes with performance averaged over 30 runs.

You can see that the agent learns a reasonably good policy within 3000 episodes, gaining sum of reward bigger than 200. Note that because of the high-variance in the agent performance, we also smoothed the learning curve.

### Wrapping up!

You have successfully implemented Course 4 Programming Assignment 2.

You have implemented an **Expected Sarsa agent with a neural network and the Adam optimizer** and used it for solving the Lunar Lander problem! You implemented different components of the agent including:

- a neural network for function approximation,
- the Adam algorithm for optimizing the weights of the neural network,
- a Softmax policy,
- the replay steps for updating the action-value function using the experiences sampled from a replay buffer

You tested the agent for a single parameter setting. In the next assignment, you will perform a parameter study on the step-size parameter to gain insight about the effect of step-size on the performance of your agent.

Note: Apart from using the `Submit` button in the notebook, you have to submit an additional zip file containing the 'npy' files that were generated from running the experiment cells. In order to do so:
1. Generate the zip file by running the experiment cells in the notebook. On the top of the notebook, navigate to `File->Open` to open the directory view of this assignment. Select the checkbox next to `results.zip` and click on `Download.` Alternatively, you can download the results folder and run `zip -jr results.zip results/` (The flag 'j' is required by the grader!).
2. Go to the "My submission" tab on the programming assignment and click on "+ Create submission".
3. Click on "PA2 Data-file Grader" and upload your results.zip.

***These account for 25% of the marks, so don't forget to do so!***

4. A complete Reinforcement Learning System(Capstone)/02Implement_your_agent.ipynb


# Quasilinear scalar equation with CentPy in 2d

### Import packages

```python
# Install the centpy package
!pip install centpy
```

    Collecting centpy
      Downloading https://files.pythonhosted.org/packages/92/89/7cbdc92609ea7790eb6444f8a189826582d675f0b7f59ba539159c43c690/centpy-0.1-py3-none-any.whl
    Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from centpy) (1.18.5)
    Installing collected packages: centpy
    Successfully installed centpy-0.1

```python
# Import numpy and centpy for the solution
from numpy import pi, sin, cos, abs, min, max
import centpy
```

```python
# Imports functions from matplotlib and setup for the animation
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import HTML
```

## Equation

We solve the nonlinear scalar conservation law

\begin{equation}
\partial_t u + \partial_x \sin u + \frac{1}{3} \partial_y u^3= 0,
\end{equation}

on the domain $(x,y,t)\in([0,2\pi]\times[0,2\pi]\times[0,6])$ with initial data

\begin{equation}
u(x,y,0) = \sin \left(x+\frac{1}{2}\right) \cos(2x+y)
\end{equation}

and periodic boundary conditions. The solution is computed using a 144 $\times$ 144 mesh and CFL number 0.9.

```python
pars = centpy.Pars2d(
    x_init=0, x_final=2*pi,
    y_init=0.0, y_final=2*pi,
    J=144, K=144,
    t_final=6.0,
    dt_out=0.1,
    cfl=0.9,
    scheme="sd3",
)
```

```python
class Scalar2d(centpy.Equation2d):
    def initial_data(self):
        x = self.xx.T; y = self.yy.T
        return sin(x + 0.5) * cos(2*x + y)

    def boundary_conditions(self, u):
        # x-boundary
        u[0] = u[-4]
        u[1] = u[-3]
        u[-2] = u[2]
        u[-1] = u[3]
        # y-boundary
        u[:, 0] = u[:, -4]
        u[:, 1] = u[:, -3]
        u[:, -2] = u[:, 2]
        u[:, -1] = u[:, 3]

    def flux_x(self, u):
        return sin(u)

    def flux_y(self, u):
        return 1./3 *u**3

    def spectral_radius_x(self, u):
        return abs(cos(u))

    def spectral_radius_y(self, u):
        return u**2
```

## Solution

```python
eqn = Scalar2d(pars)
soln = centpy.Solver2d(eqn)
soln.solve()
```

## Animation

```python
# Animation
j0 = slice(2, -2)

fig = plt.figure()
ax = plt.axes(xlim=(soln.x_init,soln.x_final), ylim=(soln.y_init, soln.y_final))

ax.set_title("Nonlinear scalar")
ax.set_xlabel("x")
ax.set_ylabel("y")

contours=ax.contour(soln.x[j0], soln.y[j0], soln.u_n[0,j0,j0], 8, colors='black')
img=ax.imshow(soln.u_n[0,j0,j0], extent=[0, 6.3, 0, 6.3], origin='lower',
           cmap='ocean', alpha=0.5)

fig.colorbar(img)
def animate(i):
    ax.collections = []
    ax.contour(soln.x[j0], soln.y[j0], soln.u_n[i,j0,j0], 8, colors='black')
    img.set_array(soln.u_n[i,j0,j0])
    img.autoscale()

plt.close()
anim = animation.FuncAnimation(fig, animate, frames=soln.Nt, interval=100, blit=False);
HTML(anim.to_html5_video())
```

```python

```


## Optimising NLTSA functions

```python
import os
tmp = os.getcwd()
os.chdir(tmp.split("jmspack")[0] + "jmspack")
```

```python
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import MinMaxScaler

from jmspack.NLTSA import (
#                             ts_levels,
                           distribution_uniformity,
                           fluctuation_intensity,
#                            complexity_resonance,
                           complexity_resonance_diagram,
#                           cumulative_complexity_peaks,
#                           cumulative_complexity_peaks_plot
                        )

from jmspack.utils import apply_scaling
import jmspack
```

```python
os.chdir(tmp)
```

```python
if "jms_style_sheet" in plt.style.available:
    _ = plt.style.use("jms_style_sheet")
```

```python
print(f"The current version of 'jmspack' used in this notebook is: {jmspack.__version__}")
```

    The current version of 'jmspack' used in this notebook is: 0.0.3

```python
ts_df = pd.read_csv("https://raw.githubusercontent.com/jameshtwose/jmspack/main/datasets/time_series_dataset.csv", index_col=0).loc[0:20, :]
ts_df.head()
```

      lorenz
      rossler
      henon
      Experimental_data_periodic
      Experimental_data_quasi-periodic-2
      Experimental_data_quasi-periodic-3
      Experimental_data_chaotic

      0
      -0.156058
      0.170777
      -1.443430
      3.00425
      2.78453
      3.62505
      3.38143

      1
      -0.071057
      -0.428226
      -0.173122
      2.97957
      2.79728
      3.71557
      3.40631

      2
      0.004560
      -1.014346
      0.937000
      2.95782
      2.81526
      3.81517
      3.44081

      3
      0.072342
      -1.592594
      0.470093
      2.94161
      2.83334
      3.88885
      3.47883

      4
      0.133683
      -2.161230
      1.460110
      2.92549
      2.85939
      3.95306
      3.50761

```python
def fluctuation_intensity_tmp(df: pd.DataFrame,
                          win: int,
                        xmin: int,
                        xmax: int,
                        col_first: int,
                        col_last: int):

    nrow = df.shape[0]
    ncol = col_last - col_first

    ew_data_F = pd.DataFrame(np.zeros((nrow, ncol)))

    newrows = df.iloc[0:1, :].copy()
    newrows.iloc[:, col_first : col_last] = 0.0

    data = df.append(newrows)
    data = data.append(newrows)
    data.reset_index(inplace=True, drop=True)

    s = xmax - xmin
    length_ts = data.shape[0]

    for column in range(col_first, col_last):
        distance = 1
        ts = data.iloc[:, column : column + 1].values

        for i in range(0, length_ts - win - 1):
            y = [0] * (win - 1)
            fluct = [0] * (win - 1)
            k = [0] * (win - 1)
            dist_next = 1

            for j in range(0, win - 1):
                if (ts[i + j + 1] >= ts[i + j]) and (ts[i + j + 1] > ts[i + j + 2]):
                    k[j] = 1
                elif (ts[i + j + 1] <= ts[i + j]) and (ts[i + j + 1] < ts[i + j + 2]):
                    k[j] = 1
                elif (ts[i + j + 1] > ts[i + j]) and (ts[i + j + 1] == ts[i + j + 2]):
                    k[j] = 1
                elif (ts[i + j + 1] < ts[i + j]) and (ts[i + j + 1] == ts[i + j + 2]):
                    k[j] = 1
                elif (ts[i + j + 1] == ts[i + j]) and (ts[i + j + 1] > ts[i + j + 2]):
                    k[j] = 1
                elif (ts[i + j + 1] == ts[i + j]) and (ts[i + j + 1] < ts[i + j + 2]):
                    k[j] = 1
                else:
                    k[j] = 0

            k[win - 2] = 1

            for g in range(0, len(k)):
                if k[g] == 1:
                    y[g] = abs(ts[i + g + 1] - ts[i + g + 1 - dist_next])
                    fluct[g] = y[g] / ((i + g + 2) - (i + g + 2 - dist_next))
                    dist_next = distance
                else:
                    y[g] = 0
                    fluct[g] = 0
                    dist_next += 1

            summation = 0.0
            for num in fluct:
                summation = summation + num / (s * (win - 1))

            ew_data_F.iloc[i + win - 1, column - col_first] = summation

        fluctuation_intensity_df = pd.DataFrame(ew_data_F)

        fluctuation_intensity_df.columns = df.columns.tolist()
        fluctuation_intensity_df.index = df.index.tolist()

    return fluctuation_intensity_df

```

```python
def distribution_uniformity_tmp(df: pd.DataFrame,
                                win: int,
                                xmin: int,
                                xmax: int,
                                col_first: int,
                                col_last: int):

    nrow = len(df) + 2
    ncol = col_last - col_first

    ew_data_D = pd.DataFrame(np.zeros((nrow, ncol)))

    for column in range(col_first, col_last):
        ts = df.iloc[:, column : column + 1].values
        y = np.linspace(xmin, xmax, win)

        for i in range(0, len(ts) - win + 1):
            x = ts[i : i + win]
            x = np.sort(x, axis=None)
            r = 0
            g = 0

            for e in range(0, win - 1):
                for d in range(e + 1, win):
                    for a in range(e, d):
                        for b in range(a + 1, d + 1):
                            h = np.heaviside((y[b] - y[a]) - (x[b] - x[a]), 0)
                            if h == 1:
                                r += (y[b] - y[a]) - (x[b] - x[a])
                                g += y[b] - y[a]

            ew_data_D.iloc[i + win - 1, column - col_first] = 1.0 - r / g

        distribution_uniformity_df = pd.DataFrame(ew_data_D.iloc[0 : len(df), :])

        distribution_uniformity_df.columns = df.columns.tolist()
        distribution_uniformity_df.index = df.index.tolist()

    return distribution_uniformity_df
```

```python
fluctuation_intensity_df = fluctuation_intensity(ts_df.pipe(apply_scaling),
                                                              win=7,
                                                              xmin=0,
                                                              xmax=1,
                                                              col_first=1,
                                                              col_last=7)
_ = complexity_resonance_diagram(fluctuation_intensity_df,
                                plot_title='Fluctuation Intensity Diagram',
                                 cmap_n=2,
                                figsize=(15, 2))
```

```python
fluctuation_intensity_df = fluctuation_intensity_tmp(ts_df.pipe(apply_scaling),
                                                              win=7,
                                                              xmin=0,
                                                              xmax=1,
                                                              col_first=0,
                                                              col_last=7)
_ = complexity_resonance_diagram(fluctuation_intensity_df,
                                plot_title='Fluctuation Intensity Diagram',
                                 labels_n=2,
                                 cmap_n=2,
                                figsize=(15, 2))
```

```python
distribution_uniformity_df = distribution_uniformity(df=ts_df.pipe(apply_scaling),
                                                      win=7,
                                                      xmin=0,
                                                      xmax=1,
                                                      col_first=1,
                                                      col_last=7)

_ = complexity_resonance_diagram(distribution_uniformity_df,
                                plot_title='Distribution Uniformity Diagram',
                                cmap_n=2,
                                figsize=(15, 2))
```

```python
distribution_uniformity_df = distribution_uniformity_tmp(df=ts_df.pipe(apply_scaling),
                                                      win=7,
                                                      xmin=0,
                                                      xmax=1,
                                                      col_first=0,
                                                      col_last=7)

_ = complexity_resonance_diagram(distribution_uniformity_df,
                                plot_title='Distribution Uniformity Diagram',
                                cmap_n=2,
                                figsize=(15, 2))
```

The Fluctuation Intensity is a measure sensitive to the amplitude and frequency of changes which occur in the time series $\mathbf{x} \in \Re^{\mathrm{N}}$ and is given by:

\begin{equation}\label{eq:F}
\mathrm{F}\left(\mathbf{x}, m_{w}\right)= \frac{1}{(x_{\mathrm{M}} - x_{\mathrm{m}}) \, (m_{w}-1)} \sum_{k=1}^{m_{w}-1} \frac{ | x_{k+1} - x_{k} | }{n_{k+1}-n_{k}}, \, \, \mathrm{F} \in \left[0, 1\right]
\end{equation}

where the index $k$ refers to the $\textit{points of return}$, specifically, the number of changes in slope of the time series, whereas $x_{\mathrm{M}}$ and $x_{\mathrm{m}}$ denote the maximum and minimum value of $\mathbf{x}$, respectively.

Conversely, the Distribution Uniformity $\mathrm{D}$ quantifies the irregularities of $\mathbf{x}$ by comparing it with an ideal distribution produced by another time series $\mathbf{y}$ with equal number of time points $\mathrm{N}$, and is defined as:

\begin{equation}\label{eq:D}
\mathrm{D}\left(\mathbf{x},\mathbf{y}, m_{w}\right)=1-\sum_{c=1}^{m_{w}-1} \sum_{d=c+1}^{m_{w}} \sum_{a=c}^{d-1} \sum_{b=a+1}^{d} \frac{\Delta_{b a} \Theta\left(\Delta_{b a}\right)}{y_{b}-y_{a}}, \, \, \mathrm{D} \in \left[0, 1\right]
\end{equation}

with $\Delta_{b a} = \left(y_{b} - x_{b}\right) - \left(y_{a} - x_{a}\right)$ that quantifies the aberration of $\mathbf{x}$ with respect to $\mathbf{y}$, and $\Theta\left(\cdot \right)$ the Heaviside step function. The two outer sums are permutations of all combinations of $c$ and $d$ within the window, whereas the inner sums with index $a$ and $b$ are representing all combinations of positions within the interval given by $c$ and $d$.

```python
ts = ts_df.pipe(apply_scaling).iloc[:, 0]
```

```python
i = 1
```

```python
win=7; xmin=0; xmax=1; col_first=1; col_last=7; df=ts_df
```

```python
constant = 1 / ((xmax - xmin) * (win-1))
constant
```

    0.16666666666666666

```python
constant * sum(abs(ts.iloc[0: 7].values - ts.iloc[0+i: 7+i].values) / 6)
```

    0.010786888632136345

```python
fluctuation_intensity_df
```

      lorenz
      rossler
      henon
      Experimental_data_periodic
      Experimental_data_quasi-periodic-2
      Experimental_data_quasi-periodic-3
      Experimental_data_chaotic

      0
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000

      1
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000

      2
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000

      3
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000

      4
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000

      5
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000

      6
      0.009604
      0.014362
      0.351876
      0.012697
      0.004739
      0.013468
      0.008869

      7
      0.008736
      0.013991
      0.476617
      0.011610
      0.005595
      0.025759
      0.009228

      8
      0.008043
      0.013547
      0.515164
      0.010876
      0.006383
      0.028729
      0.009317

      9
      0.007508
      0.012987
      0.516572
      0.010814
      0.007283
      0.030053
      0.009114

      10
      0.007116
      0.012297
      0.525401
      0.010753
      0.008029
      0.030555
      0.009338

      11
      0.006853
      0.011476
      0.540397
      0.010701
      0.008802
      0.027064
      0.009419

      12
      0.006710
      0.010528
      0.536673
      0.010659
      0.009511
      0.026070
      0.008898

      13
      0.006676
      0.009461
      0.483580
      0.009967
      0.010043
      0.025564
      0.008820

      14
      0.006745
      0.008282
      0.488829
      0.009277
      0.010598
      0.026648
      0.008903

      15
      0.006911
      0.007003
      0.629848
      0.008282
      0.010755
      0.031328
      0.008663

      16
      0.007170
      0.005633
      0.627108
      0.006981
      0.010709
      0.033942
      0.008577

      17
      0.007520
      0.004187
      0.507672
      0.005696
      0.010668
      0.030995
      0.008169

      18
      0.007960
      0.004732
      0.407877
      0.005306
      0.010302
      0.020816
      0.008063

      19
      0.008491
      0.004806
      0.357008
      0.006611
      0.010136
      0.017637
      0.007637

      20
      0.009114
      0.004857
      0.302976
      0.007549
      0.009744
      0.026372
      0.006890

```python
nrow = len(df) + 2
ncol = col_last - col_first

ew_data_D = pd.DataFrame(np.zeros((nrow, ncol)))

for column in range(col_first, col_last):
    ts = df.iloc[:, column : column + 1].values
    y = np.linspace(xmin, xmax, win)

    for i in range(0, len(ts) - win + 1):
        x = ts[i : i + win]
        x = np.sort(x, axis=None)
        r = 0
        g = 0

        for e in range(0, win - 1):
            for d in range(e + 1, win):
                for a in range(e, d):
                    for b in range(a + 1, d + 1):
                        h = np.heaviside((y[b] - y[a]) - (x[b] - x[a]), 0)
                        if h == 1:
                            r += (y[b] - y[a]) - (x[b] - x[a])
                            g += y[b] - y[a]

        #print(r,g)

        ew_data_D.iloc[i + win - 1, column - col_first] = 1.0 - r / g

#     distribution_uniformity_df = pd.DataFrame(ew_data_D.iloc[0 : len(df), :])

#     distribution_uniformity_df.columns = df.columns.tolist()
#     distribution_uniformity_df.index = df.index.tolist()
```

```python
nrow = len(df) + 2
ncol = col_last - col_first

ew_data_D = pd.DataFrame(np.zeros((nrow, ncol)))

# reference series ("ideal distribution")
y = np.linspace(xmin, xmax, win)

for column in range(col_first, col_last):
    ts = df.iloc[:, column: column + 1].values

    for i in range(0, len(ts) - win + 1):
        x = ts[i : i + win]
        x = np.sort(x, axis=None)
        r = 0
        g = 0

        for e in range(0, win - 1):
            for d in range(e + 1, win):
                for a in range(e, d):
                    for b in range(a + 1, d + 1):
                        h = np.heaviside((y[b] - y[a]) - (x[b] - x[a]), 0)
                        if h == 1:
                            r += (y[b] - y[a]) - (x[b] - x[a])
                            g += y[b] - y[a]

        ew_data_D.iloc[i + win - 1, column - col_first] = 1.0 - r / g
```

```python
def mini_du(df: pd.DataFrame,
                                win: int,
                                xmin: int,
                                xmax: int,
                                col_first: int,
                                col_last: int):

    nrow = len(df) + 2
    ncol = col_last - col_first

    ew_data_D = pd.DataFrame(np.zeros((nrow, ncol)))
    # reference series ("ideal distribution")
    y = np.linspace(xmin, xmax, win)

    for i in range(0, len(ts) - win + 1):
        x = ts[i : i + win]
        x = np.sort(x, axis=None)
        r = 0
        g = 0

        for e in range(0, win - 1):
            for d in range(e + 1, win):
                for a in range(e, d):
                    for b in range(a + 1, d + 1):
                        h = np.heaviside((y[b] - y[a]) - (x[b] - x[a]), 0)
                        if h == 1:
                            r += (y[b] - y[a]) - (x[b] - x[a])
                            g += y[b] - y[a]

#         ew_data_D.iloc[i + win - 1, column - col_first] = 1.0 - r / g

        if g==0:
            ew_data_D.iloc[i + win - 1, column - col_first] = 0
        else:
            ew_data_D.iloc[i + win - 1, column - col_first] = 1.0 - r / g

    distribution_uniformity_df = pd.DataFrame(ew_data_D.iloc[0 : len(df), :])

    distribution_uniformity_df.columns = df.columns.tolist()
    distribution_uniformity_df.index = df.index.tolist()

    return distribution_uniformity_df
```

```python
_ = complexity_resonance_diagram(ts_df.pipe(apply_scaling).pipe(mini_du, win=7, xmin=0, xmax=1, col_first=0, col_last=7),
                                plot_title='Distribution Uniformity Diagram',
                                cmap_n=2,
                                figsize=(15, 2))
```

```python
h
```

```python

```


# Polar and Cylindrical Frame of Reference

> Renato Naville Watanabe, Marcos Duarte
> [Laboratory of Biomechanics and Motor Control](http://pesquisa.ufabc.edu.br/bmclab)
> Federal University of ABC, Brazil

Table of Contents

1  Python setup

2  Versors ${\bf\hat{e_R}}$ and ${\bf\hat{e_\theta}}$

2.1  Tridimensional case

3  Time-derivative of the versors ${\bf\hat{e_R}}$ and ${\bf\hat{e_\theta}}$

4  Position, velocity and acceleration

4.1  Position

4.2  Velocity

4.3  Acceleration

5  Important to note

6  Example

6.1  Solving numerically

6.2  Solved symbolically

7  Further reading

8  Video lectures on the Internet

9  Problems

10  Reference

## Python setup

```python
import numpy as np
import sympy as sym
from sympy.plotting import plot_parametric,plot3d_parametric_line
from sympy.vector import CoordSys3D
import matplotlib.pyplot as plt
# from matplotlib import rc
# rc('text', usetex=True)
sym.init_printing()
```

## Versors ${\bf\hat{e_R}}$ and ${\bf\hat{e_\theta}}$

Consider that we have the position vector $\bf\vec{r}$ of a particle, moving in a circular path indicated in the figure below by a dashed line. This vector $\bf\vec{r(t)}$ is described in a fixed reference frame as:

\begin{equation}
{\bf\vec{r}}(t) = x(t){\bf\hat{i}} + y(t){\bf\hat{j}} + z(t){\bf\hat{k}}
\label{eq_1}
\end{equation}

Figure. Position vector of a moving particle moving in a circular path.

Naturally, we could describe all the kinematic variables in the fixed reference frame. But in circular motions, is convenient to define a basis with a vector in the direction of the position vector  $\bf\vec{r}$. So, the vector  $\bf\hat{e_R}$ is defined as:

\begin{equation}
{\bf\hat{e_R}} =  \frac{\bf\vec{r}}{\Vert{\bf\vec{r} }\Vert}
\label{eq_2}
\end{equation}

The second vector of the basis can be obtained by the cross multiplication between $\bf\hat{k}$ and $\bf\hat{e_R}$:

\begin{equation}
{\bf\hat{e_\theta}} = {\bf\hat{k}} \times {\bf\hat{e_R}}
\label{eq_3}
\end{equation}

The third vector of the basis is the conventional ${\bf\hat{k}}$ vector.

Figure. Position vector of a particle moving in a circular path and a time-varying basis>

### Tridimensional case

This basis can be used also for non-circular movements. For a 3D movement, the versor ${\bf\hat{e_R}}$ is obtained by removing the projection of the vector ${\bf\vec{r}}$ onto the versor ${\bf\hat{k}}$:

\begin{equation}
{\bf\hat{e_R}} =  \frac{\bf\vec{r} - ({\bf\vec{r}.{\bf\hat{k}}){\bf\hat{k}}}}{\Vert\bf\vec{r} - ({\bf\vec{r}.{\bf\hat{k}}){\bf\hat{k}}\Vert}}
\label{eq_4}
\end{equation}

Figure. Position vector of a moving particle and a time-varying basis>

Note that if the movement is on a plane, the expression above is equal to ${\bf\hat{e_R}} =  \frac{\bf\vec{r}}{\Vert{\bf\vec{r} }\Vert}$ since the projection of $\bf\vec{r}$ on the $\bf\hat{k}$ versor is zero.

## Time-derivative of the versors ${\bf\hat{e_R}}$ and ${\bf\hat{e_\theta}}$

To obtain the expressions of the velocity and acceleration vectors, it is necessary to obtain the expressions of the time-derivative of the vectors ${\bf\hat{e_R}}$ and ${\bf\hat{e_\theta}}$.

This can be done by noting that:

\begin{align}
{\bf\hat{e_R}} &= \cos(\theta){\bf\hat{i}} + \sin(\theta){\bf\hat{j}}\\
{\bf\hat{e_\theta}} &= -\sin(\theta){\bf\hat{i}} + \cos(\theta){\bf\hat{j}}
\label{eq_5}
\end{align}

Deriving ${\bf\hat{e_R}}$ we obtain:

\begin{equation}
\frac{d{\bf\hat{e_R}}}{dt} = -\sin(\theta)\dot\theta{\bf\hat{i}} + \cos(\theta)\dot\theta{\bf\hat{j}} = \dot{\theta}{\bf\hat{e_\theta}}
\label{eq_6}
\end{equation}

Similarly, we obtain the time-derivative of ${\bf\hat{e_\theta}}$:

\begin{equation}
    \frac{d{\bf\hat{e_\theta}}}{dt} = -\cos(\theta)\dot\theta{\bf\hat{i}} - \sin(\theta)\dot\theta{\bf\hat{j}} = -\dot{\theta}{\bf\hat{e_R}}
\label{eq_7}
\end{equation}

## Position, velocity and acceleration

### Position

The position vector $\bf\vec{r}$, from the definition of $\bf\hat{e_R}$, is:

\begin{equation}
{\bf\vec{r}} = R{\bf\hat{e_R}} + z{\bf\hat{k}}
\label{eq_8}
\end{equation}

where $R = \Vert\bf\vec{r} - ({\bf\vec{r}.{\bf\hat{k}}){\bf\hat{k}}\Vert}$.

### Velocity

The velocity vector $\bf\vec{v}$ is obtained by deriving the vector $\bf\vec{r}$:

\begin{equation}
{\bf\vec{v}} = \frac{d(R{\bf\hat{e_R}})}{dt} + \dot{z}{\bf\hat{k}} = \dot{R}{\bf\hat{e_R}}+R\frac{d\bf\hat{e_R}}{dt}=\dot{R}{\bf\hat{e_R}}+R\dot{\theta}{\bf\hat{e_\theta}}+ \dot{z}{\bf\hat{k}}
\label{eq_9}
\end{equation}

### Acceleration

The acceleration vector $\bf\vec{a}$ is obtained by deriving the velocity vector:

\begin{align}
{\bf\vec{a}} =& \frac{d(\dot{R}{\bf\hat{e_R}}+R\dot{\theta}{\bf\hat{e_\theta}}+\dot{z}{\bf\hat{k}})}{dt}=\\\nonumber
    =&\ddot{R}{\bf\hat{e_R}}+\dot{R}\frac{d\bf\hat{e_R}}{dt} + \dot{R}\dot{\theta}{\bf\hat{e_\theta}} + R\ddot{\theta}{\bf\hat{e_\theta}} + R\dot{\theta}\frac{d{\bf\hat{e_\theta}}}{dt} + \ddot{z}{\bf\hat{k}}=\\\nonumber
    =&\ddot{R}{\bf\hat{e_R}}+\dot{R}\dot{\theta}{\bf\hat{e_\theta}} + \dot{R}\dot{\theta}{\bf\hat{e_\theta}} + R\ddot{\theta}{\bf\hat{e_\theta}} - R\dot{\theta}^2{\bf\hat{e_R}}+ \ddot{z}{\bf\hat{k}} =\\
    =&\ddot{R}{\bf\hat{e_R}}+2\dot{R}\dot{\theta}{\bf\hat{e_\theta}}+ R\ddot{\theta}{\bf\hat{e_\theta}} - {R}\dot{\theta}^2{\bf\hat{e_R}}+ \ddot{z}{\bf\hat{k}} =\\\nonumber
    =&(\ddot{R}-R\dot{\theta}^2){\bf\hat{e_R}}+(2\dot{R}\dot{\theta} + R\ddot{\theta}){\bf\hat{e_\theta}}+ \ddot{z}{\bf\hat{k}}\nonumber
\label{eq_10}
\end{align}

+ The term $\ddot{R}$ is  an acceleration in the radial direction.

+ The term $R\ddot{\theta}$ is an angular acceleration.

+ The term $\ddot{z}$ is an acceleration in the $\bf\hat{k}$ direction.

+ The term $-R\dot{\theta}^2$ is the well known centripetal acceleration.

+ The term $2\dot{R}\dot{\theta}$ is known as Coriolis acceleration. This term may be difficult to understand intuitively. It appears when there is displacement in the radial and angular directions at the same time.

## Important to note

The reader must bear in mind that the use of a different basis to represent the position, velocity or acceleration vectors is only a different representation of the same vector. For example, for the acceleration vector:

\begin{equation}
{\bf\vec{a}} = \ddot{x}{\bf\hat{i}}+ \ddot{y}{\bf\hat{j}} + \ddot{z}{\bf\hat{k}}=(\ddot{R}-R\dot{\theta}^2){\bf\hat{e_R}}+(2\dot{R}\dot{\theta} + R\ddot{\theta}){\bf\hat{e_\theta}} + \ddot{z}{\bf\hat{k}}=\dot{\Vert\bf\vec{v}\Vert}{\bf\hat{e}_t}+{\Vert\bf\vec{v}\Vert}^2\Vert{\bf\vec{C}} \Vert{\bf\hat{e}_n}
\label{eq_11}
\end{equation}

In which the last equality is the acceleration vector represented in the path-coordinate of the particle (see http://nbviewer.jupyter.org/github/BMClab/bmc/blob/master/notebooks/Time-varying%20frames.ipynb).

## Example

Consider a particle following the spiral path described below:

\begin{equation}
{\bf\vec{r}}(t) = (2\sqrt{t}\cos(t)){\bf\hat{i}}+ (2\sqrt{t}\sin(t)){\bf\hat{j}}
\label{eq_12}
\end{equation}

### Solving numerically

```python
t = np.linspace(0.01,10,30).reshape(-1,1) #create a time vector and reshapes it to a column vector
R = 2*np.sqrt(t)
theta = t
rx = R*np.cos(t)
ry = R*np.sin(t)
r = np.hstack((rx, ry)) # creates the position vector by stacking rx and ry horizontally
```

```python
e_r = r/np.linalg.norm(r, axis=1, keepdims=True) # defines e_r vector
e_theta = np.cross([0,0,1],e_r)[:,0:-1] # defines e_theta vector

```

```python
dt = t[1] #defines delta_t
Rdot = np.diff(R, axis=0)/dt #find the R derivative
thetaDot = np.diff(theta, axis=0)/dt #find the angle derivative
v = Rdot*e_r[0:-1,:] +R[0:-1]*thetaDot*e_theta[0:-1,:] # find the linear velocity.
```

```python
Rddot = np.diff(Rdot, axis=0)/dt
thetaddot = np.diff(thetaDot, axis=0)/dt
```

```python
a = ((Rddot - R[1:-1]*thetaDot[0:-1]**2)*e_r[1:-1,:]
     + (2*Rdot[0:-1]*thetaDot[0:-1] + Rdot[0:-1]*thetaddot)*e_theta[1:-1,:])
```

The versors $\bf\hat{e_R}$, $\bf\hat{e_\theta}$ are plotted below at some points of the path.

```python
from matplotlib.patches import FancyArrowPatch
%matplotlib inline
plt.rcParams['figure.figsize'] = (8,8)
fig = plt.figure()

ax = fig.add_axes([0,0,1,1])
plt.plot(r[:,0],r[:,1],'.')
for i in np.arange(len(t)-2):
    vec1 = FancyArrowPatch(r[i,:],r[i,:]+e_r[i,:],mutation_scale=30,color='r', label='e_r')
    vec2 = FancyArrowPatch(r[i,:],r[i,:]+e_theta[i,:],mutation_scale=30,color='g', label='e_theta')
    ax.add_artist(vec1)
    ax.add_artist(vec2)
plt.xlim((-10,10))
plt.ylim((-10,10))
plt.grid()
plt.legend([vec1, vec2],[r'$\vec{e_r}$', r'$\vec{e_{\theta}}$'])
plt.show()
```

The velocity and acceleleration of the particle are plotted below at some points of the path.

```python
from matplotlib.patches import FancyArrowPatch
%matplotlib inline
plt.rcParams['figure.figsize'] = (8,8)
fig = plt.figure()
plt.plot(r[:,0],r[:,1],'.')
ax = fig.add_axes([0,0,1,1])
for i in np.arange(len(t)-2):
    vec1 = FancyArrowPatch(r[i,:],r[i,:]+v[i,:],mutation_scale=10,color='r')
    vec2 = FancyArrowPatch(r[i,:],r[i,:]+a[i,:],mutation_scale=10,color='g')
    ax.add_artist(vec1)
    ax.add_artist(vec2)
plt.xlim((-10,10))
plt.ylim((-10,10))
plt.grid()
plt.legend([vec1, vec2],[r'$\vec{v}$', r'$\vec{a}$'])
plt.show()
```

### Solved symbolically

```python
O = sym.vector.CoordSys3D(' ')
t = sym.symbols('t')
```

```python
r = 2*sym.sqrt(t)*sym.cos(t)*O.i+2*sym.sqrt(t)*sym.sin(t)*O.j
r
```

```python
plot_parametric(r.dot(O.i),r.dot(O.j),(t,0,10))
```

```python
e_r = r - r.dot(O.k)*O.k
e_r = e_r/sym.sqrt(e_r.dot(O.i)**2+e_r.dot(O.j)**2+e_r.dot(O.k)**2)
```

```python
e_r
```

```python
e_theta = O.k.cross(e_r)
e_theta
```

```python
from matplotlib.patches import FancyArrowPatch
plt.rcParams['figure.figsize'] = (8,8)
fig = plt.figure()
ax = fig.add_axes([0, 0, 1, 1])
ax.axis("on")
time = np.linspace(0,10,30)
for instant in time:
    vt = FancyArrowPatch([float(r.dot(O.i).subs(t,instant)),float(r.dot(O.j).subs(t,instant))],
                         [float(r.dot(O.i).subs(t,instant))+float(e_r.dot(O.i).subs(t,instant)),
                          float(r.dot(O.j).subs(t, instant))+float(e_r.dot(O.j).subs(t,instant))],
                         mutation_scale=20,
                         arrowstyle="->",color="r",label='${{e_r}}$')
    vn = FancyArrowPatch([float(r.dot(O.i).subs(t, instant)),float(r.dot(O.j).subs(t,instant))],
                         [float(r.dot(O.i).subs(t, instant))+float(e_theta.dot(O.i).subs(t, instant)),
                          float(r.dot(O.j).subs(t, instant))+float(e_theta.dot(O.j).subs(t, instant))],
                         mutation_scale=20,
                         arrowstyle="->",color="g",label='${{e_{theta}}}$')
    ax.add_artist(vn)
    ax.add_artist(vt)
plt.xlim((-10,10))
plt.ylim((-10,10))
plt.legend(handles=[vt,vn],fontsize=20)
plt.grid()
plt.show()
```

```python
R = 2*sym.sqrt(t)
```

```python
Rdot = sym.diff(R,t)
Rddot = sym.diff(Rdot,t)
Rddot
```

```python
v = Rdot*e_r + R*e_theta
```

```python
v
```

```python
a = (Rddot - R)*e_r + (2*Rdot*1+0)*e_theta
aCor = 2*Rdot*1*e_theta
aCor
```

```python
a
```

```python
from matplotlib.patches import FancyArrowPatch
plt.rcParams['figure.figsize'] = (8,8)
fig = plt.figure()
ax = fig.add_axes([0, 0, 1, 1])
ax.axis("on")
time = np.linspace(0.1,10,30)
for instant in time:
    vt = FancyArrowPatch([float(r.dot(O.i).subs(t,instant)),float(r.dot(O.j).subs(t,instant))],
                         [float(r.dot(O.i).subs(t,instant))+float(v.dot(O.i).subs(t,instant)),
                          float(r.dot(O.j).subs(t, instant))+float(v.dot(O.j).subs(t,instant))],
                         mutation_scale=20,
                         arrowstyle="->",color="r",label='${{v}}$')
    vn = FancyArrowPatch([float(r.dot(O.i).subs(t, instant)),float(r.dot(O.j).subs(t,instant))],
                         [float(r.dot(O.i).subs(t, instant))+float(a.dot(O.i).subs(t, instant)),
                          float(r.dot(O.j).subs(t, instant))+float(a.dot(O.j).subs(t, instant))],
                         mutation_scale=20,
                         arrowstyle="->",color="g",label='${{a}}$')
    vc = FancyArrowPatch([float(r.dot(O.i).subs(t, instant)),float(r.dot(O.j).subs(t,instant))],
                         [float(r.dot(O.i).subs(t, instant))+float(aCor.dot(O.i).subs(t, instant)),
                          float(r.dot(O.j).subs(t, instant))+float(aCor.dot(O.j).subs(t, instant))],
                         mutation_scale=20,
                         arrowstyle="->",color="b",label='${{a_{Cor}}}$')
    ax.add_artist(vn)
    ax.add_artist(vt)
    ax.add_artist(vc)
plt.xlim((-10,10))
plt.ylim((-10,10))
plt.legend(handles=[vt,vn,vc],fontsize=20)
plt.grid()
plt.show()
```

## Further reading

- Read pages 916-931 of the 18th chapter of the [Ruina and Rudra's book] (http://ruina.tam.cornell.edu/Book/index.html) about polar coordinates.

## Video lectures on the Internet

- Khan Academy:
    - [Polar coordinates 1](https://www.youtube.com/watch?v=B5dOy4m6I1E)
    - [Polar coordinates 2](https://www.youtube.com/watch?v=1pQXJuWbz9w)
    - [Polar coordinates 3](https://www.youtube.com/watch?v=roSG9V3zApM)

## Problems

1. Find the polar basis (using the computer) for a projectile motion of a particle following the parametric equation below:

\begin{equation}
\vec{r}(t) = (10t-51){\bf{\hat{i}}} + \left(-\frac{9,81}{2}t^2+50t+100\right){\bf{\hat{j}}}
\label{eq_13}
\end{equation}

2. Problems from 15.1.1 to 15.1.14 from Ruina and Rudra's book,
3. Problems from 18.1.1 to 18.1.8 and 18.1.10 from Ruina and Rudra's book.

## Reference

- Ruina A, Rudra P (2019) [Introduction to Mechanics for Engineers](http://ruina.tam.cornell.edu/Book/index.html). Oxford University Press.