# Notebook from KayvanShah1/training-data-analyst
Path: quests/rl/a2c/a2c_on_gcp.ipynb

# Policy Gradients and A2C

In the previous notebook, we learned how to use hyperparameter tuning to help DQN agents balance a pole on a cart. In this notebook, we'll explore two other types of alogrithms: Policy Gradients and A2C.

## Setup

Hypertuning takes some time, and in this case, it can take anywhere between **10 - 30 minutes**. If this hasn't been done already, run the cell below to kick off the training job now. We'll step through what the code is doing while our agents learn._____no_output_____

!pip install tensorflow==2.1 --user_____no_output_____

**Note**: Please ignore any incompatibility warnings and errors.
_____no_output_____

%%bash
BUCKET=<your-bucket-here> # Change to your bucket name
JOB_NAME=pg_on_gcp_$(date -u +%y%m%d_%H%M%S)
REGION='us-central1' # Change to your bucket region
IMAGE_URI=gcr.io/cloud-training-prod-bucket/pg:latest

gcloud ai-platform jobs submit training $JOB_NAME \
    --staging-bucket=gs://$BUCKET \
    --region=$REGION \
    --master-image-uri=$IMAGE_URI \
    --scale-tier=BASIC_GPU \
    --job-dir=gs://$BUCKET/$JOB_NAME \
    --config=templates/hyperparam.yaml_____no_output_____!pip install gym==0.12.5 --user_____no_output_____

**Note**: Restart the kernel if the above libraries needed to be installed
_____no_output_____Thankfully, we can use the same environment for these algorithms as DQN, so this notebook will focus less on the operational work of feeding our agents the data, and more on the theory behind these algorthims. Let's start by loading our libraries and environment._____no_output_____

import gym
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras import backend as K

CLIP_EDGE = 1e-8

def print_state(state, step, reward=None):
    format_string = 'Step {0} - Cart X: {1:.3f}, Cart V: {2:.3f}, Pole A: {3:.3f}, Pole V:{4:.3f}, Reward:{5}'
    print(format_string.format(step, *tuple(state), reward))

env = gym.make('CartPole-v0')_____no_output_____

## The Theory Behind Policy Gradients

Whereas Q-learning attempts to assign each state a value, Policy Gradients tries to find actions directly, increasing or decreaing a chance to take an action depending on how an episode plays out.

To compare, Q-learning has a table that keeps track of the value of each combination of state and action:

|| Meal | Snack | Wait |
|-|-|-|-|
| Hangry | 1 | .5 | -1 |
| Hungry | .5 | 1 | 0 |
| Full | -1 | -.5 | 1.5 |

Instead for Policy Gradients, we can imagine that we have a similar table, but instead of recording the values, we'll keep track of the probability to take the column action given the row state.

|| Meal | Snack | Wait |
|-|-|-|-|
| Hangry | 70% | 20% | 10% |
| Hungry | 30% | 50% | 20% |
| Full | 5% | 15% | 80% |

With Q learning, whenever we take one step in our environment, we can update the value of the old state based on the value of the new state plus any rewards we picked up based on the [Q equation](https://en.wikipedia.org/wiki/Q-learning):

Could we do the same thing if we have a table of probabilities instead values? No, because we don't have a way to calculate the value of each state from our table. Instead, we'll use a different  Temporal Difference Learning strategy.

Q Learning is an evolution of TD(0), and for Policy Gradients, we'll use TD(1). We'll calculate TD(1) accross and entire episode, and use that to indicate whether to increase or decrease the probability correspoding to the action we took. Let's look at a full day of eating.

| Hour | State | Action | Reward |
|-|-|-|-|
|9| Hangry | Wait | -.9 |
|10| Hangry | Meal | 1.2 |
|11| Full | Wait | .5 |
|12| Full | Snack | -.6 |
|13| Full | Wait | 1 |
|14| Full | Wait | .6 |
|15| Full | Wait | .2 |
|16| Hungry | Wait | 0 |
|17| Hungry | Meal | .4 |
|18| Full | Wait| .5 |

We'll work backwards from the last day, using the same discount, or `gamma`, as we did with DQNs. The `total_rewards` variable is equivalent to the value of state prime. Using the [Bellman Equation](https://en.wikipedia.org/wiki/Bellman_equation), everytime we calculate the value of a state, st, we'll set that as the value of state prime for the state before, st-1. _____no_output_____

test_gamma = .5  # Please change me to be between zero and one
episode_rewards = [-.9, 1.2, .5, -.6, 1, .6, .2, 0, .4, .5]

def discount_episode(rewards, gamma):
    discounted_rewards = np.zeros_like(rewards)
    total_rewards = 0
    for t in reversed(range(len(rewards))):
        total_rewards = rewards[t] + total_rewards * gamma
        discounted_rewards[t] = total_rewards
    return discounted_rewards

discount_episode(episode_rewards, test_gamma)_____no_output_____

Wherever our discounted reward is positive, we'll increase the probability corresponding to the action we took. Similarly, wherever our discounted reward is negative, we'll decrease the probabilty.

However, with this strategy, any actions with a positive reward will have it's probability increase, not necessarily the most optimal action. This puts us in a feedback loop, where we're more likely to pick less optimal actions which could further increase their probability. To counter this, we'll divide the size of our increases by the probability to choose the corresponding action, which will slow the growth of popular actions to give other actions a chance.

Here is our update rule for our neural network, where alpha is our learning rate, and pi is our optimal policy, or the probability to take the optimal action, a*, given our current state, s.

Doing some fancy calculus, we can combine the numerator and denominator with a log function. Since it's not clear what the optimal action is, we'll instead use our discounted rewards, or G, to increase or decrease the weights of the respective action the agent took. A full breakdown of the math can be found in [this article by Chris Yoon](https://medium.com/@thechrisyoon/deriving-policy-gradients-and-implementing-reinforce-f887949bd63).

Below is what it looks like in code. `y_true` is the [one-hot encoding](https://en.wikipedia.org/wiki/One-hot) of the action that was taken. `y_pred` is the probabilty to take each action given the state the agent was in._____no_output_____

def custom_loss(y_true, y_pred):
    y_pred_clipped = K.clip(y_pred, CLIP_EDGE, 1-CLIP_EDGE)
    log_likelihood = y_true * K.log(y_pred_clipped)
    return K.sum(-log_likelihood*g)_____no_output_____

We won't have the discounted rewards, or `g`, when our agent is acting in the environment. No problem, we'll have one neural network with two types of pathways. One pathway, `predict`, will be the probability to take an action given an inputed state. It's only used for prediction and is not used for backpropogation. The other pathway, `policy`, will take both a state and a discounted reward, so it can be used for training.

The code in its entirety looks like this. As with Deep Q Networks, the hidden layers of a Policy Gradient can use a CNN if the input state is pixels, but the last layer is typically a [Dense](https://keras.io/layers/core/) layer with a [Softmax](https://en.wikipedia.org/wiki/Softmax_function) activation function to convert the output into probabilities._____no_output_____

def build_networks(
        state_shape, action_size, learning_rate, hidden_neurons):
    """Creates a Policy Gradient Neural Network.

    Creates a two hidden-layer Policy Gradient Neural Network. The loss
    function is altered to be a log-likelihood function weighted
    by the discounted reward, g.

    Args:
        space_shape: a tuple of ints representing the observation space.
        action_size (int): the number of possible actions.
        learning_rate (float): the nueral network's learning rate.
        hidden_neurons (int): the number of neurons to use per hidden
            layer.
    """
    state_input = layers.Input(state_shape, name='frames')
    g = layers.Input((1,), name='g')

    hidden_1 = layers.Dense(hidden_neurons, activation='relu')(state_input)
    hidden_2 = layers.Dense(hidden_neurons, activation='relu')(hidden_1)
    probabilities = layers.Dense(action_size, activation='softmax')(hidden_2)

    def custom_loss(y_true, y_pred):
        y_pred_clipped = K.clip(y_pred, CLIP_EDGE, 1-CLIP_EDGE)
        log_lik = y_true*K.log(y_pred_clipped)
        return K.sum(-log_lik*g)

    policy = models.Model(
        inputs=[state_input, g], outputs=[probabilities])
    optimizer = tf.keras.optimizers.Adam(lr=learning_rate)
    policy.compile(loss=custom_loss, optimizer=optimizer)

    predict = models.Model(inputs=[state_input], outputs=[probabilities])
    return policy, predict_____no_output_____

Let's get a taste of how these networks function. Run the below cell to build our test networks._____no_output_____

space_shape = env.observation_space.shape
action_size = env.action_space.n

# Feel free to play with these
test_learning_rate = .2
test_hidden_neurons = 10

test_policy, test_predict = build_networks(
    space_shape, action_size, test_learning_rate, test_hidden_neurons)_____no_output_____

We can't use the policy network until we build our learning function, but we can feed a state to the predict network so we can see our chances to pick our actions._____no_output_____

state = env.reset()
test_predict.predict(np.expand_dims(state, axis=0))_____no_output_____

Right now, the numbers should be close to `[.5, .5]`, with a little bit of variance due to the randomization of initializing the weights and the cart's starting position. In order to train, we'll need some memories to train on. The memory buffer here is simpler than DQN, as we don't have to worry about random sampling. We'll clear the buffer every time we train as we'll only hold one episode's worth of memory._____no_output_____

class Memory():
    """Sets up a memory replay buffer for Policy Gradient methods.

    Args:
        gamma (float): The "discount rate" used to assess TD(1) values.
    """
    def __init__(self, gamma):
        self.buffer = []
        self.gamma = gamma

    def add(self, experience):
        """Adds an experience into the memory buffer.

        Args:
            experience: a (state, action, reward) tuple.
        """
        self.buffer.append(experience)

    def sample(self):
        """Returns the list of episode experiences and clears the buffer.

        Returns:
            (list): A tuple of lists with structure (
                [states], [actions], [rewards]
            }
        """
        batch = np.array(self.buffer).T.tolist()
        states_mb = np.array(batch[0], dtype=np.float32)
        actions_mb = np.array(batch[1], dtype=np.int8)
        rewards_mb = np.array(batch[2], dtype=np.float32)
        self.buffer = []
        return states_mb, actions_mb, rewards_mb_____no_output_____

Let's make a fake buffer to get a sense of the data we'll be training on. The cell below initializes our memory and runs through one episode of the game by alternating pushing the cart left and right.

Try running it to see the data we'll be using for training._____no_output_____

test_memory = Memory(test_gamma)
actions = [x % 2 for x in range(200)]
state = env.reset()
step = 0
episode_reward = 0
done = False

while not done and step < len(actions):
    action = actions[step]  # In the future, our agents will define this.
    state_prime, reward, done, info = env.step(action)
    episode_reward += reward
    test_memory.add((state, action, reward))
    step += 1
    state = state_prime

test_memory.sample()_____no_output_____

Ok, time to start putting together the agent! Let's start by giving it the ability to act. Here, we don't need to worry about exploration vs exploitation because we already have a random chance to take each of our actions. As the agent learns, it will naturally shift from exploration to exploitation. How conveient!_____no_output_____

class Partial_Agent():
    """Sets up a reinforcement learning agent to play in a game environment."""
    def __init__(self, policy, predict, memory, action_size):
        """Initializes the agent with Policy Gradient networks
            and memory sub-classes.

        Args:
            policy: The policy network created from build_networks().
            predict: The predict network created from build_networks().
            memory: A Memory class object.
            action_size (int): The number of possible actions to take.
        """
        self.policy = policy
        self.predict = predict
        self.action_size = action_size
        self.memory = memory

    def act(self, state):
        """Selects an action for the agent to take given a game state.

        Args:
            state (list of numbers): The state of the environment to act on.

        Returns:
            (int) The index of the action to take.
        """
        # If not acting randomly, take action with highest predicted value.
        state_batch = np.expand_dims(state, axis=0)
        probabilities = self.predict.predict(state_batch)[0]
        action = np.random.choice(self.action_size, p=probabilities)
        return action_____no_output_____

Let's see the act function in action. First, let's build our agent._____no_output_____

test_agent = Partial_Agent(test_policy, test_predict, test_memory, action_size)_____no_output_____

Next, run the below cell a few times to test the `act` method. Is it about a 50/50 chance to push right instead of left?_____no_output_____

action = test_agent.act(state)
print("Push Right" if action else "Push Left")Push Right

Now for the most important part. We need to give our agent a way to learn! To start, we'll [one-hot encode](https://en.wikipedia.org/wiki/One-hot) our actions. Since the output of our network is a probability for each action, we'll have a 1 corresponding to the action that was taken and 0's for the actions we didn't take.

That doesn't give our agent enough information on whether the action that was taken was actually a good idea, so we'll also use our `discount_episode` to calculate the TD(1) value of each step within the episode.

One thing to note, is that CartPole doesn't have any negative rewards, meaning, even if it does terribly, the agent will still think the run went well. To help counter this, we'll take the mean and standard deviation of our discounted rewards, or `discount_mb`, and use that to find the [Standard Score](https://en.wikipedia.org/wiki/Standard_score) for each discounted reward. With this, steps close to dropping the poll will have a negative reward._____no_output_____

def learn(self, print_variables=False):
    """Trains a Policy Gradient policy network based on stored experiences."""
    state_mb, action_mb, reward_mb = self.memory.sample()
    # One hot enocde actions
    actions = np.zeros([len(action_mb), self.action_size])
    actions[np.arange(len(action_mb)), action_mb] = 1
    if print_variables:
        print("action_mb:", action_mb)
        print("actions:", actions)

    # Apply TD(1) and normalize
    discount_mb = discount_episode(reward_mb, self.memory.gamma)
    discount_mb = (discount_mb - np.mean(discount_mb)) / np.std(discount_mb)
    if print_variables:
        print("reward_mb:", reward_mb)
        print("discount_mb:", discount_mb)
    return self.policy.train_on_batch([state_mb, discount_mb], actions)

Partial_Agent.learn = learn
test_agent = Partial_Agent(test_policy, test_predict, test_memory, action_size)_____no_output_____

Try adding in some print statements to the code above to get a sense of how the data is transformed before feeding it into the model, then run the below code to see it in action._____no_output_____Finally, it's time to put it all together. Policy Gradient Networks have less hypertuning parameters than DQNs, but since our custom loss constructs a [TensorFlow Graph](https://www.tensorflow.org/api_docs/python/tf/Graph) under the hood, we'll set up lazy execution by wrapping our traing steps in a default graph.

By changing `test_gamma`, `test_learning_rate`, and `test_hidden_neurons`, can you help the agent reach a score of 200 within 200 episodes? It takes a little bit of thinking and a little bit of luck.

Hover the curser  on this bold text to see a solution to the challenge._____no_output_____

test_gamma = .5
test_learning_rate = .01
test_hidden_neurons = 100

with tf.Graph().as_default():
    test_memory = Memory(test_gamma)
    test_policy, test_predict = build_networks(
    space_shape, action_size, test_learning_rate, test_hidden_neurons)
    test_agent = Partial_Agent(test_policy, test_predict, test_memory, action_size)
    for episode in range(200):
        state = env.reset()
        episode_reward = 0
        done = False

        while not done:
            action = test_agent.act(state)
            state_prime, reward, done, info = env.step(action)
            episode_reward += reward
            test_agent.memory.add((state, action, reward))
            state = state_prime

        test_agent.learn()
        print("Episode", episode, "Score =", episode_reward)WARNING:tensorflow:From /usr/local/lib/python3.5/dist-packages/tensorflow_core/python/ops/resource_variable_ops.py:1630: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.
Episode 0 Score = 11.0
Episode 1 Score = 11.0
Episode 2 Score = 24.0
Episode 3 Score = 27.0
Episode 4 Score = 24.0
Episode 5 Score = 30.0
Episode 6 Score = 24.0
Episode 7 Score = 30.0
Episode 8 Score = 27.0
Episode 9 Score = 33.0
Episode 10 Score = 36.0
Episode 11 Score = 40.0
Episode 12 Score = 34.0
Episode 13 Score = 32.0
Episode 14 Score = 99.0
Episode 15 Score = 61.0
Episode 16 Score = 70.0
Episode 17 Score = 46.0
Episode 18 Score = 90.0
Episode 19 Score = 125.0
Episode 20 Score = 132.0
Episode 21 Score = 154.0
Episode 22 Score = 200.0
Episode 23 Score = 200.0
Episode 24 Score = 200.0
Episode 25 Score = 200.0
Episode 26 Score = 93.0
Episode 27 Score = 188.0
Episode 28 Score = 134.0
Episode 29 Score = 126.0
Episode 30 Score = 157.0
Episode 31 Score = 169.0
Episode 32 Score = 161.0
Episode 33 Score = 128.0
Episode 34 Score = 115.0
Episode 35 Score = 110.0
Episode 36 Score = 167.0
Episode 37 Score = 200.0
Episode 38 Score = 102.0
Episode 39 Score = 109.0
Episode 40 Score = 61.0
Episode 41 Score = 32.0
Episode 42 Score = 14.0
Episode 43 Score = 18.0
Episode 44 Score = 10.0
Episode 45 Score = 11.0
Episode 46 Score = 10.0
Episode 47 Score = 12.0
Episode 48 Score = 17.0
Episode 49 Score = 10.0
Episode 50 Score = 8.0
Episode 51 Score = 9.0
Episode 52 Score = 13.0
Episode 53 Score = 10.0
Episode 54 Score = 9.0
Episode 55 Score = 8.0
Episode 56 Score = 10.0
Episode 57 Score = 9.0
Episode 58 Score = 10.0
Episode 59 Score = 10.0
Episode 60 Score = 10.0
Episode 61 Score = 10.0
Episode 62 Score = 10.0
Episode 63 Score = 9.0
Episode 64 Score = 10.0
Episode 65 Score = 8.0
Episode 66 Score = 10.0
Episode 67 Score = 10.0
Episode 68 Score = 9.0
Episode 69 Score = 9.0
Episode 70 Score = 10.0
Episode 71 Score = 9.0
Episode 72 Score = 9.0
Episode 73 Score = 10.0
Episode 74 Score = 8.0
Episode 75 Score = 11.0
Episode 76 Score = 10.0
Episode 77 Score = 8.0
Episode 78 Score = 10.0
Episode 79 Score = 9.0
Episode 80 Score = 10.0
Episode 81 Score = 9.0
Episode 82 Score = 9.0
Episode 83 Score = 10.0
Episode 84 Score = 8.0
Episode 85 Score = 10.0
Episode 86 Score = 10.0
Episode 87 Score = 9.0
Episode 88 Score = 9.0
Episode 89 Score = 10.0
Episode 90 Score = 9.0
Episode 91 Score = 10.0
Episode 92 Score = 9.0
Episode 93 Score = 9.0
Episode 94 Score = 10.0
Episode 95 Score = 9.0
Episode 96 Score = 11.0
Episode 97 Score = 8.0
Episode 98 Score = 10.0
Episode 99 Score = 8.0
Episode 100 Score = 10.0
Episode 101 Score = 10.0
Episode 102 Score = 9.0
Episode 103 Score = 9.0
Episode 104 Score = 8.0
Episode 105 Score = 10.0
Episode 106 Score = 11.0
Episode 107 Score = 8.0
Episode 108 Score = 10.0
Episode 109 Score = 9.0
Episode 110 Score = 9.0
Episode 111 Score = 9.0
Episode 112 Score = 9.0
Episode 113 Score = 10.0
Episode 114 Score = 11.0
Episode 115 Score = 9.0
Episode 116 Score = 9.0
Episode 117 Score = 8.0
Episode 118 Score = 9.0
Episode 119 Score = 10.0
Episode 120 Score = 10.0
Episode 121 Score = 9.0
Episode 122 Score = 9.0
Episode 123 Score = 9.0
Episode 124 Score = 9.0
Episode 125 Score = 9.0
Episode 126 Score = 8.0
Episode 127 Score = 8.0
Episode 128 Score = 8.0
Episode 129 Score = 8.0
Episode 130 Score = 10.0
Episode 131 Score = 9.0
Episode 132 Score = 9.0
Episode 133 Score = 8.0
Episode 134 Score = 9.0
Episode 135 Score = 10.0
Episode 136 Score = 9.0
Episode 137 Score = 8.0
Episode 138 Score = 8.0
Episode 139 Score = 10.0
Episode 140 Score = 10.0
Episode 141 Score = 10.0
Episode 142 Score = 9.0
Episode 143 Score = 9.0
Episode 144 Score = 10.0
Episode 145 Score = 9.0
Episode 146 Score = 10.0
Episode 147 Score = 9.0
Episode 148 Score = 9.0
Episode 149 Score = 9.0
Episode 150 Score = 10.0
Episode 151 Score = 9.0
Episode 152 Score = 9.0
Episode 153 Score = 8.0
Episode 154 Score = 9.0
Episode 155 Score = 10.0
Episode 156 Score = 9.0
Episode 157 Score = 8.0
Episode 158 Score = 9.0
Episode 159 Score = 9.0
Episode 160 Score = 10.0
Episode 161 Score = 10.0
Episode 162 Score = 9.0
Episode 163 Score = 10.0
Episode 164 Score = 9.0
Episode 165 Score = 8.0
Episode 166 Score = 10.0
Episode 167 Score = 9.0
Episode 168 Score = 10.0
Episode 169 Score = 10.0
Episode 170 Score = 10.0
Episode 171 Score = 8.0
Episode 172 Score = 10.0
Episode 173 Score = 9.0
Episode 174 Score = 9.0
Episode 175 Score = 9.0
Episode 176 Score = 10.0
Episode 177 Score = 9.0
Episode 178 Score = 9.0
Episode 179 Score = 9.0
Episode 180 Score = 9.0
Episode 181 Score = 9.0
Episode 182 Score = 10.0
Episode 183 Score = 9.0
Episode 184 Score = 9.0
Episode 185 Score = 10.0
Episode 186 Score = 10.0
Episode 187 Score = 8.0
Episode 188 Score = 9.0
Episode 189 Score = 9.0
Episode 190 Score = 8.0
Episode 191 Score = 8.0
Episode 192 Score = 10.0
Episode 193 Score = 9.0
Episode 194 Score = 10.0
Episode 195 Score = 10.0
Episode 196 Score = 9.0
Episode 197 Score = 9.0
Episode 198 Score = 9.0
Episode 199 Score = 10.0

# The Theory Behind Actor - Critic

Now that we have the hang of Policy Gradients, let's combine this strategy with Deep Q Agents. We'll have one architecture to rule them all!

Below is the setup for our neural networks. There are plenty of ways to go combining the two strategies. We'll be focusing on one varient called A2C, or Advantage Actor Critic.

Here's the philosophy: We'll use our critic pathway to estimate the value of a state, or V(s). Given a state-action-new state transition, we can use our critic and the Bellman Equation to calculate the discounted value of the new state, or r + γ * V(s').

Like DQNs, this discounted value is the label the critic will train on. While that is happening, we can subtract V(s) and the discounted value of the new state to get the advantage, or A(s,a). In human terms, how much value was the action the agent took? This is what the actor, or the policy gradient portion or our network, will train on.

Too long, didn't read: the critic's job is to learn how to asses the value of a state. The actor's job is to assign probabilities to it's available actions such that it increases its chance to move into a higher valued state.

Below is our new `build_networks` function. Each line has been tagged with whether it comes from Deep Q Networks (`# DQN`), Policy Gradients (`# PG`), or is something new (`# New`)._____no_output_____

def build_networks(state_shape, action_size, learning_rate, critic_weight,
                   hidden_neurons, entropy):
    """Creates Actor Critic Neural Networks.

    Creates a two hidden-layer Policy Gradient Neural Network. The loss
    function is altered to be a log-likelihood function weighted
    by an action's advantage.

    Args:
        space_shape: a tuple of ints representing the observation space.
        action_size (int): the number of possible actions.
        learning_rate (float): the nueral network's learning rate.
        critic_weight (float): how much to weigh the critic's training loss.
        hidden_neurons (int): the number of neurons to use per hidden layer.
        entropy (float): how much to enourage exploration versus exploitation.
    """
    state_input = layers.Input(state_shape, name='frames')
    advantages = layers.Input((1,), name='advantages')  # PG, A instead of G

    # PG
    actor_1 = layers.Dense(hidden_neurons, activation='relu')(state_input)
    actor_2 = layers.Dense(hidden_neurons, activation='relu')(actor_1)
    probabilities = layers.Dense(action_size, activation='softmax')(actor_2)

    # DQN
    critic_1 = layers.Dense(hidden_neurons, activation='relu')(state_input)
    critic_2 = layers.Dense(hidden_neurons, activation='relu')(critic_1)
    values = layers.Dense(1, activation='linear')(critic_2)

    def actor_loss(y_true, y_pred):  # PG
        y_pred_clipped = K.clip(y_pred, CLIP_EDGE, 1-CLIP_EDGE)
        log_lik = y_true*K.log(y_pred_clipped)
        entropy_loss = y_pred * K.log(K.clip(y_pred, CLIP_EDGE, 1-CLIP_EDGE))  # New
        return K.sum(-log_lik * advantages) - (entropy * K.sum(entropy_loss))

    # Train both actor and critic at the same time.
    actor = models.Model(
        inputs=[state_input, advantages], outputs=[probabilities, values])
    actor.compile(
        loss=[actor_loss, 'mean_squared_error'],  # [PG, DQN]
        loss_weights=[1, critic_weight],  # [PG, DQN]
        optimizer=tf.keras.optimizers.Adam(lr=learning_rate))

    critic = models.Model(inputs=[state_input], outputs=[values])
    policy = models.Model(inputs=[state_input], outputs=[probabilities])
    return actor, critic, policy_____no_output_____

The above is one way to go about combining both of the algorithms. Here, we're combining training of both pwathways into on operation. Keras allows for the [training against multiple outputs](https://keras.io/models/model/). They can even have their own loss functions as we have above. When minimizing the loss, Keras will take the weighted sum of all the losses, with the weights provided in `loss_weights`. The `critic_weight` is now another hyperparameter for us to tune.

We could even have completely separate networks for the actor and the critic, and that type of design choice is going to be problem dependent. Having shared nodes and training between the two will be more efficient to train per batch, but more complicated problems could justify keeping the two separate.

The loss function we used here is also slightly different than the one for Policy Gradients. Let's take a look._____no_output_____

def actor_loss(y_true, y_pred):  # PG
    y_pred_clipped = K.clip(y_pred, 1e-8, 1-1e-8)
    log_lik = y_true*K.log(y_pred_clipped)
    entropy_loss = y_pred * K.log(K.clip(y_pred, 1e-8, 1-1e-8))  # New
    return K.sum(-log_lik * advantages) - (entropy * K.sum(entropy_loss))_____no_output_____

We've added a new tool called [entropy](https://arxiv.org/pdf/1912.01557.pdf). We're calculating the [log-likelihood](https://en.wikipedia.org/wiki/Likelihood_function#Log-likelihood) again, but instead of comparing the probabilities of our actions versus the action that was taken, we calculating it for the probabilities of our actions against themselves.

Certainly a mouthful, but the idea is to encourage exploration: if our probability prediction is very confident (or close to 1), our entropy will be close to 0. Similary, if our probability isn't confident at all (or close to 0), our entropy will again be zero. Anywhere inbetween, our entropy will be non-zero. This encourages exploration versus exploitation, as the entropy will discourage overconfident predictions.

Now that the networks are out of the way, let's look at the `Memory`. We could go with Experience Replay, like with DQNs, or we could calculate TD(1) like with Policy Gradients. This time, we'll do something in between. We'll give our memory a `batch_size`. Once there are enough experiences in the buffer, we'll use all the experiences to train and then clear the buffer to start fresh.

In order to speed up training, instead of recording state_prime, we'll record the value of state prime in `state_prime_values` or `next_values`. This will give us enough information to calculate the discounted values and advantages._____no_output_____

class Memory():
    """Sets up a memory replay for actor-critic training.

    Args:
        gamma (float): The "discount rate" used to assess state values.
        batch_size (int): The number of elements to include in the buffer.
    """
    def __init__(self, gamma, batch_size):
        self.buffer = []
        self.gamma = gamma
        self.batch_size = batch_size

    def add(self, experience):
        """Adds an experience into the memory buffer.

        Args:
            experience: (state, action, reward, state_prime_value, done) tuple.
        """
        self.buffer.append(experience)

    def check_full(self):
        return len(self.buffer) >= self.batch_size

    def sample(self):
        """Returns formated experiences and clears the buffer.

        Returns:
            (list): A tuple of lists with structure [
                [states], [actions], [rewards], [state_prime_values], [dones]
            ]
        """
        # Columns have different data types, so numpy array would be awkward.
        batch = np.array(self.buffer).T.tolist()
        states_mb = np.array(batch[0], dtype=np.float32)
        actions_mb = np.array(batch[1], dtype=np.int8)
        rewards_mb = np.array(batch[2], dtype=np.float32)
        dones_mb = np.array(batch[3], dtype=np.int8)
        value_mb = np.squeeze(np.array(batch[4], dtype=np.float32))
        self.buffer = []
        return states_mb, actions_mb, rewards_mb, dones_mb, value_mb_____no_output_____

Ok, time to build out the agent! The `act` method is the exact same as it was for Policy Gradients. Nice! The `learn` method is where things get interesting. We'll find the discounted future state like we did for DQN to train our critic. We'll then subtract the value of the discount state from the value of the current state to find the advantage, which is what the actor will train on._____no_output_____

class Agent():
    """Sets up a reinforcement learning agent to play in a game environment."""
    def __init__(self, actor, critic, policy, memory, action_size):
        """Initializes the agent with DQN and memory sub-classes.

        Args:
            network: A neural network created from deep_q_network().
            memory: A Memory class object.
            epsilon_decay (float): The rate at which to decay random actions.
            action_size (int): The number of possible actions to take.
        """
        self.actor = actor
        self.critic = critic
        self.policy = policy
        self.action_size = action_size
        self.memory = memory

    def act(self, state):
        """Selects an action for the agent to take given a game state.

        Args:
            state (list of numbers): The state of the environment to act on.
            traning (bool): True if the agent is training.

        Returns:
            (int) The index of the action to take.
        """
        # If not acting randomly, take action with highest predicted value.
        state_batch = np.expand_dims(state, axis=0)
        probabilities = self.policy.predict(state_batch)[0]
        action = np.random.choice(self.action_size, p=probabilities)
        return action

    def learn(self, print_variables=False):
        """Trains the Deep Q Network based on stored experiences."""
        gamma = self.memory.gamma
        experiences = self.memory.sample()
        state_mb, action_mb, reward_mb, dones_mb, next_value = experiences

        # One hot enocde actions
        actions = np.zeros([len(action_mb), self.action_size])
        actions[np.arange(len(action_mb)), action_mb] = 1

        #Apply TD(0)
        discount_mb = reward_mb + next_value * gamma * (1 - dones_mb)
        state_values = self.critic.predict([state_mb])
        advantages = discount_mb - np.squeeze(state_values)
        if print_variables:
            print("discount_mb", discount_mb)
            print("next_value", next_value)
            print("state_values", state_values)
            print("advantages", advantages)
        else:
            self.actor.train_on_batch(
                [state_mb, advantages], [actions, discount_mb])_____no_output_____

Run the below cell to initialize an agent, and the cell after that to see the variables used for training. Since it's early, the critic hasn't learned to estimate the values yet, and the advatanges are mostly positive because of it.

Once the crtic has learned how to properly assess states, the actor will start to see negative advantages. Try playing around with the variables to help the agent see this change sooner._____no_output_____

# Change me please.
test_gamma = .9
test_batch_size = 32
test_learning_rate = .02
test_hidden_neurons = 50
test_critic_weight = 0.5
test_entropy = 0.0001

test_memory = Memory(test_gamma, test_batch_size)
test_actor, test_critic, test_policy = build_networks(
    space_shape, action_size,
    test_learning_rate, test_critic_weight,
    test_hidden_neurons, test_entropy)
test_agent = Agent(
    test_actor, test_critic, test_policy, test_memory, action_size)_____no_output_____state = env.reset()
episode_reward = 0
done = False

while not done:
    action = test_agent.act(state)
    state_prime, reward, done, _ = env.step(action)
    episode_reward += reward
    next_value = test_agent.critic.predict([[state_prime]])
    test_agent.memory.add((state, action, reward, done, next_value))
    state = state_prime
test_agent.learn(print_variables=True)discount_mb [1.0772394 1.1522162 1.0811032 1.0144521 1.0888239 1.1633745 1.0974793
 1.0297184 1.1090584 1.0408443 1.1215113 1.2000881 1.133828  1.2161343
 1.1495053 1.0845523 1.0276814 1.       ]
next_value [0.08582156 0.16912904 0.09011465 0.01605788 0.09869321 0.18152726
 0.10831045 0.03302046 0.12117595 0.04538257 0.13501257 0.22232015
 0.14869787 0.24014927 0.16611691 0.09394698 0.0307571  0.12024279]
state_values [[0.00508265]
 [0.08582155]
 [0.16912904]
 [0.09011464]
 [0.01605787]
 [0.09869322]
 [0.1815272 ]
 [0.10831043]
 [0.03302046]
 [0.12117594]
 [0.04538257]
 [0.13501257]
 [0.22232018]
 [0.14869787]
 [0.24014927]
 [0.1661169 ]
 [0.09394696]
 [0.0307571 ]]
advantages [1.0721568  1.0663947  0.9119742  0.92433745 1.0727661  1.0646813
 0.91595215 0.92140794 1.0760379  0.9196684  1.0761287  1.0650756
 0.91150784 1.0674365  0.909356   0.9184354  0.9337344  0.96924293]

Have a set of variables you're happy with? Ok, time to shine! Run the below cell to see how the agent trains._____no_output_____

with tf.Graph().as_default():
    test_memory = Memory(test_gamma, test_batch_size)
    test_actor, test_critic, test_policy = build_networks(
        space_shape, action_size,
        test_learning_rate, test_critic_weight,
        test_hidden_neurons, test_entropy)
    test_agent = Agent(
        test_actor, test_critic, test_policy, test_memory, action_size)
    for episode in range(200):
        state = env.reset()
        episode_reward = 0
        done = False

        while not done:
            action = test_agent.act(state)
            state_prime, reward, done, _ = env.step(action)
            episode_reward += reward
            next_value = test_agent.critic.predict([[state_prime]])
            test_agent.memory.add((state, action, reward, done, next_value))

            #if test_agent.memory.check_full():
            #test_agent.learn(print_variables=True)
            state = state_prime
        test_agent.learn()
        print("Episode", episode, "Score =", episode_reward)Episode 0 Score = 13.0
Episode 1 Score = 13.0
Episode 2 Score = 13.0
Episode 3 Score = 11.0
Episode 4 Score = 10.0
Episode 5 Score = 9.0
Episode 6 Score = 11.0
Episode 7 Score = 9.0
Episode 8 Score = 11.0
Episode 9 Score = 10.0
Episode 10 Score = 10.0
Episode 11 Score = 9.0
Episode 12 Score = 9.0
Episode 13 Score = 9.0
Episode 14 Score = 10.0
Episode 15 Score = 10.0
Episode 16 Score = 9.0
Episode 17 Score = 9.0
Episode 18 Score = 10.0
Episode 19 Score = 10.0
Episode 20 Score = 10.0
Episode 21 Score = 10.0
Episode 22 Score = 9.0
Episode 23 Score = 9.0
Episode 24 Score = 10.0
Episode 25 Score = 10.0
Episode 26 Score = 9.0
Episode 27 Score = 8.0
Episode 28 Score = 9.0
Episode 29 Score = 10.0
Episode 30 Score = 10.0
Episode 31 Score = 10.0
Episode 32 Score = 9.0
Episode 33 Score = 9.0
Episode 34 Score = 10.0
Episode 35 Score = 10.0
Episode 36 Score = 10.0
Episode 37 Score = 10.0
Episode 38 Score = 10.0
Episode 39 Score = 11.0
Episode 40 Score = 9.0
Episode 41 Score = 10.0
Episode 42 Score = 10.0
Episode 43 Score = 9.0
Episode 44 Score = 11.0
Episode 45 Score = 10.0
Episode 46 Score = 9.0
Episode 47 Score = 10.0
Episode 48 Score = 10.0
Episode 49 Score = 9.0
Episode 50 Score = 9.0
Episode 51 Score = 9.0
Episode 52 Score = 9.0
Episode 53 Score = 10.0
Episode 54 Score = 9.0
Episode 55 Score = 9.0
Episode 56 Score = 10.0
Episode 57 Score = 9.0
Episode 58 Score = 10.0
Episode 59 Score = 9.0
Episode 60 Score = 10.0
Episode 61 Score = 9.0
Episode 62 Score = 8.0
Episode 63 Score = 10.0
Episode 64 Score = 11.0
Episode 65 Score = 9.0
Episode 66 Score = 9.0
Episode 67 Score = 10.0
Episode 68 Score = 10.0
Episode 69 Score = 8.0
Episode 70 Score = 10.0
Episode 71 Score = 9.0
Episode 72 Score = 9.0
Episode 73 Score = 9.0
Episode 74 Score = 8.0
Episode 75 Score = 9.0
Episode 76 Score = 8.0
Episode 77 Score = 9.0
Episode 78 Score = 10.0
Episode 79 Score = 9.0
Episode 80 Score = 9.0
Episode 81 Score = 10.0
Episode 82 Score = 11.0
Episode 83 Score = 9.0
Episode 84 Score = 8.0
Episode 85 Score = 9.0
Episode 86 Score = 8.0
Episode 87 Score = 10.0
Episode 88 Score = 10.0
Episode 89 Score = 9.0
Episode 90 Score = 10.0
Episode 91 Score = 10.0
Episode 92 Score = 10.0
Episode 93 Score = 9.0
Episode 94 Score = 10.0
Episode 95 Score = 8.0
Episode 96 Score = 9.0
Episode 97 Score = 10.0
Episode 98 Score = 9.0
Episode 99 Score = 10.0
Episode 100 Score = 9.0
Episode 101 Score = 11.0
Episode 102 Score = 9.0
Episode 103 Score = 9.0
Episode 104 Score = 9.0
Episode 105 Score = 10.0
Episode 106 Score = 9.0
Episode 107 Score = 10.0
Episode 108 Score = 9.0
Episode 109 Score = 9.0
Episode 110 Score = 10.0
Episode 111 Score = 10.0
Episode 112 Score = 9.0
Episode 113 Score = 9.0
Episode 114 Score = 8.0
Episode 115 Score = 9.0
Episode 116 Score = 9.0
Episode 117 Score = 8.0
Episode 118 Score = 10.0
Episode 119 Score = 8.0
Episode 120 Score = 10.0
Episode 121 Score = 10.0
Episode 122 Score = 9.0
Episode 123 Score = 10.0
Episode 124 Score = 8.0
Episode 125 Score = 9.0
Episode 126 Score = 9.0
Episode 127 Score = 10.0
Episode 128 Score = 10.0
Episode 129 Score = 10.0
Episode 130 Score = 10.0
Episode 131 Score = 8.0
Episode 132 Score = 10.0
Episode 133 Score = 9.0
Episode 134 Score = 9.0
Episode 135 Score = 9.0
Episode 136 Score = 11.0
Episode 137 Score = 9.0
Episode 138 Score = 10.0
Episode 139 Score = 10.0
Episode 140 Score = 9.0
Episode 141 Score = 10.0
Episode 142 Score = 8.0
Episode 143 Score = 10.0
Episode 144 Score = 10.0
Episode 145 Score = 8.0
Episode 146 Score = 9.0
Episode 147 Score = 9.0
Episode 148 Score = 10.0
Episode 149 Score = 10.0
Episode 150 Score = 9.0
Episode 151 Score = 10.0
Episode 152 Score = 8.0
Episode 153 Score = 10.0
Episode 154 Score = 10.0
Episode 155 Score = 9.0
Episode 156 Score = 10.0
Episode 157 Score = 9.0
Episode 158 Score = 9.0
Episode 159 Score = 10.0
Episode 160 Score = 9.0
Episode 161 Score = 9.0
Episode 162 Score = 10.0
Episode 163 Score = 10.0
Episode 164 Score = 9.0
Episode 165 Score = 10.0
Episode 166 Score = 10.0
Episode 167 Score = 8.0
Episode 168 Score = 10.0
Episode 169 Score = 10.0
Episode 170 Score = 10.0
Episode 171 Score = 9.0
Episode 172 Score = 9.0
Episode 173 Score = 11.0
Episode 174 Score = 8.0
Episode 175 Score = 10.0
Episode 176 Score = 9.0
Episode 177 Score = 10.0
Episode 178 Score = 11.0
Episode 179 Score = 9.0
Episode 180 Score = 10.0
Episode 181 Score = 9.0
Episode 182 Score = 9.0
Episode 183 Score = 10.0
Episode 184 Score = 10.0
Episode 185 Score = 10.0
Episode 186 Score = 8.0
Episode 187 Score = 9.0
Episode 188 Score = 9.0
Episode 189 Score = 9.0
Episode 190 Score = 9.0
Episode 191 Score = 10.0
Episode 192 Score = 9.0
Episode 193 Score = 9.0
Episode 194 Score = 9.0
Episode 195 Score = 9.0
Episode 196 Score = 8.0
Episode 197 Score = 9.0
Episode 198 Score = 11.0
Episode 199 Score = 9.0

Any luck? No sweat if not! It turns out that by combining the power of both algorithms, we also combined some of their setbacks. For instance, actor-critic can fall into local minimums like Policy Gradients, and has a large number of hyperparameters to tune like DQNs.

Time to check how our agents did [in the cloud](https://console.cloud.google.com/ai-platform/jobs)! Any lucky winners? Find it in [your bucket](https://console.cloud.google.com/storage/browser) to watch a recording of it play._____no_output_____Copyright 2020 Google Inc.
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License._____no_output_____


# Notebook from janeite/course-content
Path: tutorials/Bonus_Autoencoders/student/Bonus_Tutorial3.ipynb

_____no_output_____# Tutorial 3: Autoencoders applications
**Bonus Day: Autoencoders**

**By Neuromatch Academy**

__Content creators:__ Marco Brigham and the [CCNSS](https://www.ccnss.org/) team (2014-2018)

__Content reviewers:__ Itzel Olivos, Karen Schroeder, Karolina Stosio, Kshitij Dwivedi, Spiros Chavlis, Michael Waskom_____no_output_____---
# Tutorial Objectives

## Autoencoder applications

How do autoencoders with rich internal representations perform on the MNIST cognitive task?

How do autoencoders perceive unseen digit classes?

How does ANN image encoding differ from human vision?

We are equipped with tools and techniques to answer these questions, and hopefully, many others you may encounter in your research!

![MNIST cognitive task](https://github.com/mpbrigham/colaboratory-figures/raw/master/nma/autoencoders/mnist_task.png)

In this tutorial, you will:
- Analyze how autoencoders perceive transformed data (added noise, occluded parts, and rotations), and how that evolves with short re-train sessions
- Use autoencoders to visualize unseen digit classes
- Understand visual encoding for fully connected ANN autoencoders_____no_output_____

# @title Video 1: Applications
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  from IPython.display import IFrame
  class BiliVideo(IFrame):
      def __init__(self, id, page=1, width=400, height=300, **kwargs):
          self.id=id
          src = 'https://player.bilibili.com/player.html?bvid={0}&page={1}'.format(id, page)
          super(BiliVideo, self).__init__(src, width, height, **kwargs)

  video = BiliVideo(id="BV12v411q7nS", width=854, height=480, fs=1)
  print('Video available at https://www.bilibili.com/video/{0}'.format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id="_bzW_jkH6l0", width=854, height=480, fs=1, rel=0)
  print('Video available at https://youtube.com/watch?v=' + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

display(out)_____no_output_____

---
# Setup
Please execute the cell(s) below to initialize the notebook environment._____no_output_____

# Imports
import numpy as np
import matplotlib.pyplot as plt
import os
from scipy import ndimage

import torch
from torch import nn, optim

from sklearn.datasets import fetch_openml_____no_output_____# @title Figure settings

%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/NMA2020/nma.mplstyle")_____no_output_____# @title Helper functions

def downloadMNIST():
  """
  Download MNIST dataset and transform it to torch.Tensor

  Args:
    None

  Returns:
    x_train : training images (torch.Tensor) (60000, 28, 28)
    x_test  : test images (torch.Tensor) (10000, 28, 28)
    y_train : training labels (torch.Tensor) (60000, )
    y_train : test labels (torch.Tensor) (10000, )
  """
  X, y = fetch_openml('mnist_784', version=1, return_X_y=True, as_frame=False)
  # Trunk the data
  n_train = 60000
  n_test = 10000

  train_idx = np.arange(0, n_train)
  test_idx = np.arange(n_train, n_train + n_test)

  x_train, y_train = X[train_idx], y[train_idx]
  x_test, y_test = X[test_idx], y[test_idx]

  # Transform np.ndarrays to torch.Tensor
  x_train = torch.from_numpy(np.reshape(x_train,
                                        (len(x_train),
                                         28, 28)).astype(np.float32))
  x_test = torch.from_numpy(np.reshape(x_test,
                                       (len(x_test),
                                        28, 28)).astype(np.float32))

  y_train = torch.from_numpy(y_train.astype(int))
  y_test = torch.from_numpy(y_test.astype(int))

  return (x_train, y_train, x_test, y_test)

def init_weights_kaiming_uniform(layer):
  """
  Initializes weights from linear PyTorch layer
  with kaiming uniform distribution.

  Args:
    layer (torch.Module)
        Pytorch layer

  Returns:
    Nothing.
  """
  # check for linear PyTorch layer
  if isinstance(layer, nn.Linear):
    # initialize weights with kaiming uniform distribution
    nn.init.kaiming_uniform_(layer.weight.data)

def init_weights_kaiming_normal(layer):
  """
  Initializes weights from linear PyTorch layer
  with kaiming normal distribution.

  Args:
    layer (torch.Module)
        Pytorch layer

  Returns:
    Nothing.
  """
  # check for linear PyTorch layer
  if isinstance(layer, nn.Linear):
    # initialize weights with kaiming normal distribution
    nn.init.kaiming_normal_(layer.weight.data)

def get_layer_weights(layer):
  """
  Retrieves learnable parameters from PyTorch layer.

  Args:
    layer (torch.Module)
        Pytorch layer

  Returns:
    list with learnable parameters
  """
  # initialize output list
  weights = []

  # check whether layer has learnable parameters
  if layer.parameters():
    # copy numpy array representation of each set of learnable parameters
    for item in layer.parameters():
      weights.append(item.detach().numpy())

  return weights

def eval_mse(y_pred, y_true):
  """
  Evaluates mean square error (MSE) between y_pred and y_true

  Args:
    y_pred (torch.Tensor)
        prediction samples

    v (numpy array of floats)
        ground truth samples

  Returns:
    MSE(y_pred, y_true)
  """

  with torch.no_grad():
    criterion = nn.MSELoss()
    loss = criterion(y_pred, y_true)

  return float(loss)

def eval_bce(y_pred, y_true):
  """
  Evaluates binary cross-entropy (BCE) between y_pred and y_true

  Args:
    y_pred (torch.Tensor)
        prediction samples

    v (numpy array of floats)
        ground truth samples

  Returns:
    BCE(y_pred, y_true)
  """

  with torch.no_grad():
    criterion = nn.BCELoss()
    loss = criterion(y_pred, y_true)

  return float(loss)

def plot_row(images, show_n=10, image_shape=None):
  """
  Plots rows of images from list of iterables (iterables: list, numpy array
  or torch.Tensor). Also accepts single iterable.
  Randomly selects images in each list element if item count > show_n.

  Args:
    images (iterable or list of iterables)
        single iterable with images, or list of iterables

    show_n (integer)
        maximum number of images per row

    image_shape (tuple or list)
        original shape of image if vectorized form

  Returns:
    Nothing.
  """

  if not isinstance(images, (list, tuple)):
    images = [images]

  for items_idx, items in enumerate(images):

    items = np.array(items)
    if items.ndim == 1:
      items = np.expand_dims(items, axis=0)

    if len(items) > show_n:
      selected = np.random.choice(len(items), show_n, replace=False)
      items = items[selected]

    if image_shape is not None:
      items = items.reshape([-1] + list(image_shape))

    plt.figure(figsize=(len(items) * 1.5, 2))
    for image_idx, image in enumerate(items):

      plt.subplot(1, len(items), image_idx + 1)
      plt.imshow(image, cmap='gray', vmin=image.min(), vmax=image.max())
      plt.axis('off')

    plt.tight_layout()

def to_s2(u):
  """
  Projects 3D coordinates to spherical coordinates (theta, phi) surface of
  unit sphere S2.
  theta: [0, pi]
  phi: [-pi, pi]

  Args:
    u (list, numpy array or torch.Tensor of floats)
        3D coordinates

  Returns:
    Sperical coordinates (theta, phi) on surface of unit sphere S2.
  """

  x, y, z = (u[:, 0], u[:, 1], u[:, 2])
  r = np.sqrt(x**2 + y**2 + z**2)
  theta = np.arccos(z / r)
  phi = np.arctan2(x, y)

  return np.array([theta, phi]).T

def to_u3(s):
  """
  Converts from 2D coordinates on surface of unit sphere S2 to 3D coordinates
  (on surface of S2), i.e. (theta, phi) ---> (1, theta, phi).

  Args:
    s (list, numpy array or torch.Tensor of floats)
        2D coordinates on unit sphere S_2

  Returns:
    3D coordinates on surface of unit sphere S_2
  """

  theta, phi = (s[:, 0], s[:, 1])
  x = np.sin(theta) * np.sin(phi)
  y = np.sin(theta) * np.cos(phi)
  z = np.cos(theta)

  return np.array([x, y, z]).T

def xy_lim(x):
  """
  Return arguments for plt.xlim and plt.ylim calculated from minimum
  and maximum of x.

  Args:
    x (list, numpy array or torch.Tensor of floats)
        data to be plotted

  Returns:
    Nothing.
  """
  x_min = np.min(x, axis=0)
  x_max = np.max(x, axis=0)

  x_min = x_min - np.abs(x_max - x_min) * 0.05 - np.finfo(float).eps
  x_max = x_max + np.abs(x_max - x_min) * 0.05 + np.finfo(float).eps

  return [x_min[0], x_max[0]], [x_min[1], x_max[1]]

def plot_generative(x, decoder_fn, image_shape, n_row=16, s2=False):
  """
  Plots images reconstructed by decoder_fn from a 2D grid in
  latent space that is determined by minimum and maximum values in x.

  Args:
    x (list, numpy array or torch.Tensor of floats)
        2D or 3D coordinates in latent space

    decoder_fn (integer)
        function returning vectorized images from 2D latent space coordinates

    image_shape (tuple or list)
        original shape of image

    n_row (integer)
        number of rows in grid

    s2 (boolean)
        convert 3D coordinates (x, y, z) to spherical coordinates (theta, phi)

  Returns:
    Nothing.
  """

  if s2:
    x = to_s2(np.array(x))

  xlim, ylim = xy_lim(np.array(x))

  dx = (xlim[1] - xlim[0]) / n_row
  grid = [np.linspace(ylim[0] + dx / 2, ylim[1] - dx / 2, n_row),
          np.linspace(xlim[0] + dx / 2, xlim[1] - dx / 2, n_row)]

  canvas = np.zeros((image_shape[0] * n_row, image_shape[1] * n_row))

  cmap = plt.get_cmap('gray')

  for j, latent_y in enumerate(grid[0][::-1]):
    for i, latent_x in enumerate(grid[1]):

      latent = np.array([[latent_x, latent_y]], dtype=np.float32)

      if s2:
        latent = to_u3(latent)

      with torch.no_grad():
        x_decoded = decoder_fn(torch.from_numpy(latent))

      x_decoded = x_decoded.reshape(image_shape)

      canvas[j * image_shape[0]: (j + 1) * image_shape[0],
             i * image_shape[1]: (i + 1) * image_shape[1]] = x_decoded

  plt.imshow(canvas, cmap=cmap, vmin=canvas.min(), vmax=canvas.max())
  plt.axis('off')

def plot_latent(x, y, show_n=500, s2=False, fontdict=None, xy_labels=None):
  """
  Plots digit class of each sample in 2D latent space coordinates.

  Args:
    x (list, numpy array or torch.Tensor of floats)
        2D coordinates in latent space

    y (list, numpy array or torch.Tensor of floats)
        digit class of each sample

    n_row (integer)
        number of samples

    s2 (boolean)
        convert 3D coordinates (x, y, z) to spherical coordinates (theta, phi)

    fontdict (dictionary)
        style option for plt.text

    xy_labels (list)
        optional list with [xlabel, ylabel]

  Returns:
    Nothing.
  """

  if fontdict is None:
    fontdict = {'weight': 'bold', 'size': 12}

  if s2:
    x = to_s2(np.array(x))

  cmap = plt.get_cmap('tab10')

  if len(x) > show_n:
    selected = np.random.choice(len(x), show_n, replace=False)
    x = x[selected]
    y = y[selected]

  for my_x, my_y in zip(x, y):
    plt.text(my_x[0], my_x[1], str(int(my_y)),
             color=cmap(int(my_y) / 10.),
             fontdict=fontdict,
             horizontalalignment='center',
             verticalalignment='center',
             alpha=0.8)

  xlim, ylim = xy_lim(np.array(x))
  plt.xlim(xlim)
  plt.ylim(ylim)

  if s2:
    if xy_labels is None:
      xy_labels = [r'$\varphi$', r'$\theta$']

    plt.xticks(np.arange(0, np.pi + np.pi / 6, np.pi / 6),
               ['0', '$\pi/6$', '$\pi/3$', '$\pi/2$',
                '$2\pi/3$', '$5\pi/6$', '$\pi$'])
    plt.yticks(np.arange(-np.pi, np.pi + np.pi / 3, np.pi / 3),
               ['$-\pi$', '$-2\pi/3$', '$-\pi/3$', '0',
                '$\pi/3$', '$2\pi/3$', '$\pi$'])

  if xy_labels is None:
    xy_labels = ['$Z_1$', '$Z_2$']

  plt.xlabel(xy_labels[0])
  plt.ylabel(xy_labels[1])

def plot_latent_generative(x, y, decoder_fn, image_shape, s2=False,
                           title=None, xy_labels=None):
  """
  Two horizontal subplots generated with encoder map and decoder grid.

  Args:
    x (list, numpy array or torch.Tensor of floats)
        2D coordinates in latent space

    y (list, numpy array or torch.Tensor of floats)
        digit class of each sample

    decoder_fn (integer)
        function returning vectorized images from 2D latent space coordinates

    image_shape (tuple or list)
        original shape of image

    s2 (boolean)
        convert 3D coordinates (x, y, z) to spherical coordinates (theta, phi)

    title (string)
        plot title

    xy_labels (list)
        optional list with [xlabel, ylabel]

  Returns:
    Nothing.
  """

  fig = plt.figure(figsize=(12, 6))

  if title is not None:
    fig.suptitle(title, y=1.05)

  ax = fig.add_subplot(121)
  ax.set_title('Encoder map', y=1.05)
  plot_latent(x, y, s2=s2, xy_labels=xy_labels)

  ax = fig.add_subplot(122)
  ax.set_title('Decoder grid', y=1.05)
  plot_generative(x, decoder_fn, image_shape, s2=s2)

  plt.tight_layout()
  plt.show()

def plot_latent_ab(x1, x2, y, selected_idx=None,
                   title_a='Before', title_b='After', show_n=500, s2=False):
  """
  Two horizontal subplots with encoder maps.

  Args:
    x1 (list, numpy array or torch.Tensor of floats)
        2D coordinates in latent space (left plot)

    x2 (list, numpy array or torch.Tensor of floats)
        digit class of each sample (right plot)

    y (list, numpy array or torch.Tensor of floats)
        digit class of each sample

    selected_idx (list of integers)
        indexes of elements to be plotted

    show_n (integer)
        maximum number of samples in each plot

    s2 (boolean)
        convert 3D coordinates (x, y, z) to spherical coordinates (theta, phi)

  Returns:
    Nothing.
  """

  fontdict = {'weight': 'bold', 'size': 12}

  if len(x1) > show_n:

    if selected_idx is None:
      selected_idx = np.random.choice(len(x1), show_n, replace=False)

    x1 = x1[selected_idx]
    x2 = x2[selected_idx]
    y = y[selected_idx]

  data = np.concatenate([x1, x2])

  if s2:
    xlim, ylim = xy_lim(to_s2(data))

  else:
    xlim, ylim = xy_lim(data)

  plt.figure(figsize=(12, 6))

  ax = plt.subplot(121)
  ax.set_title(title_a, y=1.05)
  plot_latent(x1, y, fontdict=fontdict, s2=s2)
  plt.xlim(xlim)
  plt.ylim(ylim)

  ax = plt.subplot(122)
  ax.set_title(title_b, y=1.05)
  plot_latent(x2, y, fontdict=fontdict, s2=s2)
  plt.xlim(xlim)
  plt.ylim(ylim)
  plt.tight_layout()

def runSGD(net, input_train, input_test, out_train=None, out_test=None,
           optimizer=None, criterion='bce', n_epochs=10, batch_size=32,
           verbose=False):
  """
  Trains autoencoder network with stochastic gradient descent with
  optimizer and loss criterion. Train samples are shuffled, and loss is
  displayed at the end of each opoch for both MSE and BCE. Plots training loss
  at each minibatch (maximum of 500 randomly selected values).

  Args:
    net (torch network)
        ANN network (nn.Module)

    input_train (torch.Tensor)
        vectorized input images from train set

    input_test (torch.Tensor)
        vectorized input images from test set

    criterion (string)
        train loss: 'bce' or 'mse'

    out_train (torch.Tensor)
        optional target images from train set

    out_test (torch.Tensor)
        optional target images from test set

    optimizer (torch optimizer)
        optional target images from train set

    criterion (string)
        train loss: 'bce' or 'mse'

    n_epochs (boolean)
        number of full iterations of training data

    batch_size (integer)
        number of element in mini-batches

    verbose (boolean)
        whether to print final loss

  Returns:
    Nothing.
  """

  if out_train is not None and out_test is not None:
    different_output = True
  else:
    different_output = False

  # Initialize loss function
  if criterion == 'mse':
    loss_fn = nn.MSELoss()
  elif criterion == 'bce':
    loss_fn = nn.BCELoss()
  else:
    print('Please specify either "mse" or "bce" for loss criterion')

  # Initialize SGD optimizer
  if optimizer is None:
    optimizer = optim.Adam(net.parameters())

  # Placeholder for loss
  track_loss = []

  print('Epoch', '\t', 'Loss train', '\t', 'Loss test')
  for i in range(n_epochs):

    shuffle_idx = np.random.permutation(len(input_train))
    batches = torch.split(input_train[shuffle_idx], batch_size)

    if different_output:
      batches_out = torch.split(out_train[shuffle_idx], batch_size)

    for batch_idx, batch in enumerate(batches):

      output_train = net(batch)

      if different_output:
        loss = loss_fn(output_train, batches_out[batch_idx])
      else:
        loss = loss_fn(output_train, batch)

      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

      # Keep track of loss at each epoch
      track_loss += [float(loss)]

    loss_epoch = f'{i+1}/{n_epochs}'
    with torch.no_grad():
      output_train = net(input_train)
      if different_output:
        loss_train = loss_fn(output_train, out_train)
      else:
        loss_train = loss_fn(output_train, input_train)

      loss_epoch += f'\t {loss_train:.4f}'

      output_test = net(input_test)
      if different_output:
        loss_test = loss_fn(output_test, out_test)
      else:
        loss_test = loss_fn(output_test, input_test)

      loss_epoch += f'\t\t {loss_test:.4f}'

    print(loss_epoch)

  if verbose:
    # Print loss
    if different_output:
      loss_mse = f'\nMSE\t {eval_mse(output_train, out_train):0.4f}'
      loss_mse += f'\t\t {eval_mse(output_test, out_test):0.4f}'
    else:
      loss_mse = f'\nMSE\t {eval_mse(output_train, input_train):0.4f}'
      loss_mse += f'\t\t {eval_mse(output_test, input_test):0.4f}'
    print(loss_mse)

    if different_output:
      loss_bce = f'BCE\t {eval_bce(output_train, out_train):0.4f}'
      loss_bce += f'\t\t {eval_bce(output_test, out_test):0.4f}'
    else:
      loss_bce = f'BCE\t {eval_bce(output_train, input_train):0.4f}'
      loss_bce += f'\t\t {eval_bce(output_test, input_test):0.4f}'
    print(loss_bce)

  # Plot loss
  step = int(np.ceil(len(track_loss)/500))
  x_range = np.arange(0, len(track_loss), step)
  plt.figure()
  plt.plot(x_range, track_loss[::step], 'C0')
  plt.xlabel('Iterations')
  plt.ylabel('Loss')
  plt.xlim([0, None])
  plt.ylim([0, None])
  plt.show()

def image_occlusion(x, image_shape):
  """
  Randomly selects on quadrant of images and sets to zeros.

  Args:
    x (torch.Tensor of floats)
        vectorized images

    image_shape (tuple or list)
        original shape of image

  Returns:
    torch.Tensor.
  """

  selection = np.random.choice(4, len(x))

  my_x = np.array(x).copy()
  my_x = my_x.reshape(-1, image_shape[0], image_shape[1])

  my_x[selection == 0, :int(image_shape[0] / 2), :int(image_shape[1] / 2)] = 0
  my_x[selection == 1, int(image_shape[0] / 2):, :int(image_shape[1] / 2)] = 0
  my_x[selection == 2, :int(image_shape[0] / 2), int(image_shape[1] / 2):] = 0
  my_x[selection == 3, int(image_shape[0] / 2):, int(image_shape[1] / 2):] = 0

  my_x = my_x.reshape(x.shape)

  return torch.from_numpy(my_x)

def image_rotation(x, deg, image_shape):
  """
  Randomly rotates images by +- deg degrees.

  Args:
    x (torch.Tensor of floats)
        vectorized images

    deg (integer)
        rotation range

    image_shape (tuple or list)
        original shape of image

  Returns:
    torch.Tensor.
  """

  my_x = np.array(x).copy()
  my_x = my_x.reshape(-1, image_shape[0], image_shape[1])

  for idx, item in enumerate(my_x):
    my_deg = deg * 2 * np.random.random() - deg
    my_x[idx] = ndimage.rotate(my_x[idx], my_deg,
                               reshape=False, prefilter=False)

  my_x = my_x.reshape(x.shape)

  return torch.from_numpy(my_x)

class AutoencoderClass(nn.Module):
  """
  Deep autoencoder network object (nn.Module) with optional L2 normalization
  of activations in bottleneck layer.

  Args:
    input_size (integer)
        size of input samples

    s2 (boolean)
        whether to L2 normalize activatinos in bottleneck layer

  Returns:
    Autoencoder object inherited from nn.Module class.
  """

  def __init__(self, input_size=784, s2=False):

    super().__init__()

    self.input_size = input_size
    self.s2 = s2

    if s2:
      self.encoding_size = 3

    else:
      self.encoding_size = 2

    self.enc1 = nn.Linear(self.input_size, int(self.input_size / 2))
    self.enc1_f = nn.PReLU()
    self.enc2 = nn.Linear(int(self.input_size / 2), self.encoding_size * 32)
    self.enc2_f = nn.PReLU()
    self.enc3 = nn.Linear(self.encoding_size * 32, self.encoding_size)
    self.enc3_f = nn.PReLU()
    self.dec1 = nn.Linear(self.encoding_size, self.encoding_size * 32)
    self.dec1_f = nn.PReLU()
    self.dec2 = nn.Linear(self.encoding_size * 32, int(self.input_size / 2))
    self.dec2_f = nn.PReLU()
    self.dec3 = nn.Linear(int(self.input_size / 2), self.input_size)
    self.dec3_f = nn.Sigmoid()

  def encoder(self, x):
    """
    Encoder component.
    """
    x = self.enc1_f(self.enc1(x))
    x = self.enc2_f(self.enc2(x))
    x = self.enc3_f(self.enc3(x))

    if self.s2:
        x = nn.functional.normalize(x, p=2, dim=1)

    return x

  def decoder(self, x):
    """
    Decoder component.
    """
    x = self.dec1_f(self.dec1(x))
    x = self.dec2_f(self.dec2(x))
    x = self.dec3_f(self.dec3(x))

    return x

  def forward(self, x):
    """
    Forward pass.
    """
    x = self.encoder(x)
    x = self.decoder(x)

    return x

def save_checkpoint(net, optimizer, filename):
  """
  Saves a PyTorch checkpoint.

  Args:
    net (torch network)
        ANN network (nn.Module)

    optimizer (torch optimizer)
        optimizer for SGD

    filename (string)
        filename (without extension)

  Returns:
    Nothing.
  """

  torch.save({'model_state_dict': net.state_dict(),
              'optimizer_state_dict': optimizer.state_dict()},
             filename+'.pt')

def load_checkpoint(url, filename):
  """
  Loads a PyTorch checkpoint from URL is local file not present.

  Args:
    url (string)
        URL location of PyTorch checkpoint

    filename (string)
        filename (without extension)

  Returns:
    PyTorch checkpoint of saved model.
  """

  if not os.path.isfile(filename+'.pt'):
    os.system(f"wget {url}.pt")

  return torch.load(filename+'.pt')

def reset_checkpoint(net, optimizer, checkpoint):
  """
  Resets PyTorch model to checkpoint.

  Args:
    net (torch network)
        ANN network (nn.Module)

    optimizer (torch optimizer)
        optimizer for SGD

    checkpoint (torch checkpoint)
        checkpoint of saved model

  Returns:
    Nothing.
  """

  net.load_state_dict(checkpoint['model_state_dict'])
  optimizer.load_state_dict(checkpoint['optimizer_state_dict'])_____no_output_____

---
# Section 1: Download  and prepare MNIST dataset
We use the helper function `downloadMNIST` to download the dataset and transform it into `torch.Tensor` and assign train and test sets to (`x_train`, `y_train`) and (`x_test`, `y_test`).

The variable `input_size` stores the length of *vectorized* versions of the images `input_train` and `input_test` for training and test images.

**Instructions:**
* Please execute the cell below_____no_output_____

# Download MNIST
x_train, y_train, x_test, y_test = downloadMNIST()

x_train = x_train / 255
x_test = x_test / 255

image_shape = x_train.shape[1:]

input_size = np.prod(image_shape)

input_train = x_train.reshape([-1, input_size])
input_test = x_test.reshape([-1, input_size])

test_selected_idx = np.random.choice(len(x_test), 10, replace=False)
train_selected_idx = np.random.choice(len(x_train), 10, replace=False)

test_subset_idx = np.random.choice(len(x_test), 500, replace=False)

print(f'shape image \t\t {image_shape}')
print(f'shape input_train \t {input_train.shape}')
print(f'shape input_test \t {input_test.shape}')_____no_output_____

---
# Section 2: Download a pre-trained model
The class `AutoencoderClass` implements the autoencoder architectures introduced in the previous tutorial. The design of this class follows the object-oriented programming (OOP) style from tutorial W3D4. Setting the boolean parameter `s2=True` specifies the model with projection onto the $S_2$ sphere.

We trained both models for `n_epochs=25` and saved the weights to avoid a lengthy initial training period - these will be our reference model states.

Experiments are run from the identical initial conditions by resetting the autoencoder to the reference state at the beginning of each exercise.

The mechanism for loading and storing models from PyTorch is the following:
```
model = nn.Sequential(...)
or
model = AutoencoderClass()

torch.save({'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict()},
           filename_path)

checkpoint = torch.load(filename_path)

model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
```
See additional [PyTorch instructions](https://pytorch.org/tutorials/recipes/recipes/saving_and_loading_a_general_checkpoint.html), and when to use `model.eval()` and `model.train()` for more complex models.

We provide the functions `save_checkpoint`, `load_checkpoint`, and `reset_checkpoint` to implement the steps above and download pre-trained weights from the GitHub repo.

If downloading from GitHub fails, please uncomment the 3rd cell bellow to train the model for `n_epochs=10` and save it locally.

**Instructions:**
* Please execute the cell(s) below_____no_output_____

root = 'https://github.com/mpbrigham/colaboratory-figures/raw/master/nma/autoencoders'
filename = 'ae_6h_prelu_bce_adam_25e_32b'
url = os.path.join(root, filename)
s2 = True

if s2:
  filename += '_s2'
  url += '_s2'_____no_output_____model = AutoencoderClass(s2=s2)
optimizer = optim.Adam(model.parameters())

encoder = model.encoder
decoder = model.decoder

checkpoint = load_checkpoint(url, filename)

model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])_____no_output_____# Please uncomment and execute this cell if download of
# pre-trained weights fail

# model = AutoencoderClass(s2=s2)
# encoder = model.encoder
# decoder = model.decoder
# n_epochs = 10
# batch_size = 128
# runSGD(model, input_train, input_test,
#        n_epochs=n_epochs, batch_size=batch_size)
# save_checkpoint(model, optimizer, filename)
# checkpoint = load_checkpoint(url, filename)_____no_output_____with torch.no_grad():
  output_test = model(input_test)
  latent_test = encoder(input_test)

plot_row([input_test[test_selected_idx], output_test[test_selected_idx]],
         image_shape=image_shape)

plot_latent_generative(latent_test, y_test, decoder,
                       image_shape=image_shape, s2=s2)_____no_output_____

---
# Section 3: Applications of autoencoders_____no_output_____## Application 1 - Image noise
Removing noise added to images is often showcased in dimensionality reduction techniques. The tutorial  *W1D5 Dimensionality reduction* illustrated this capability with PCA.

We first observe that autoencoders trained with noise-free images output noise-free images when receiving noisy images as input.  However, the reconstructed images will be different from the original images (without noise) since the added noise maps to different coordinates in latent space.

The ability to map noise-free and noisy versions to similar regions in latent space is known as *robustness* or *invariance* to noise. How can we build such functionality into the autoencoder?

The solution is to train the autoencoder with noise-free and noisy versions mapping to the noise-free version. A faster alternative is to re-train the autoencoder for few epochs with noisy images. These short training sessions fine-tune the weights to map noisy images to their noise-free versions from similar latent space coordinates.

Let's start by resetting to the reference state of the autoencoder.

**Instructions:**
* Please execute the cells below_____no_output_____

reset_checkpoint(model, optimizer, checkpoint)

with torch.no_grad():
  latent_test_ref = encoder(input_test)_____no_output_____

### Reconstructions before fine-tuning
Let's verify that an autoencoder trained on clean images will output clean images from noisy inputs. We visualize this by plotting three rows:
* Top row with noisy images inputs
* Middle row with reconstructions of noisy images
* Bottom row with reconstructions of the original images (noise-free)

![Noise task](https://github.com/mpbrigham/colaboratory-figures/raw/master/nma/autoencoders/applications_noise.png)

The bottom row helps identify samples with reconstruction issues before adding noise. This row shows the baseline reconstruction quality for these samples rather than the original images. (Why?)

**Instructions:**
* Please execute the cell(s) below_____no_output_____

noise_factor = 0.4

input_train_noisy = (input_train
                     + noise_factor * np.random.normal(size=input_train.shape))
input_train_noisy = np.clip(input_train_noisy, input_train.min(),
                            input_train.max(), dtype=np.float32)

input_test_noisy = (input_test
                    + noise_factor * np.random.normal(size=input_test.shape))
input_test_noisy = np.clip(input_test_noisy, input_test.min(),
                           input_test.max(), dtype=np.float32)_____no_output_____with torch.no_grad():
  output_test_noisy = model(input_test_noisy)
  latent_test_noisy = encoder(input_test_noisy)
  output_test = model(input_test)

plot_row([input_test_noisy[test_selected_idx],
          output_test_noisy[test_selected_idx],
          output_test[test_selected_idx]], image_shape=image_shape)_____no_output_____

### Latent space before fine-tuning
We investigate the origin of reconstruction errors by looking at how adding noise to input affects latent space coordinates. The decoder interprets significant coordinate changes as different digits.

The function `plot_latent_ab` compares latent space coordinates for the same set of samples between two conditions.  Here, we display coordinates for the ten samples from the previous cell before and after adding noise:
* The left plot shows the coordinates of the original samples (noise-free)
* The plot on the right shows the new coordinates after adding noise

**Instructions:**
* Please execute the cell below_____no_output_____

plot_latent_ab(latent_test, latent_test_noisy, y_test, test_selected_idx,
               title_a='Before noise', title_b='After noise', s2=s2)_____no_output_____

### Fine-tuning the autoencoder with noisy images
Let's re-train the autoencoder with noisy images on the input and original (noise-free) images on the output, and regenerate the previous plots.

We now see that both noisy and noise-free images match similar locations in latent space. The network denoises the input with a latent-space representation that is more robust to noise.

**Instructions:**
* Please execute the cell(s) below_____no_output_____

n_epochs = 3
batch_size = 32

model.train()

runSGD(model, input_train_noisy, input_test_noisy,
       out_train=input_train, out_test=input_test,
       n_epochs=n_epochs, batch_size=batch_size)_____no_output_____with torch.no_grad():
  output_test_noisy = model(input_test_noisy)
  latent_test_noisy = encoder(input_test_noisy)
  output_test = model(input_test)

plot_row([input_test_noisy[test_selected_idx],
          output_test_noisy[test_selected_idx],
          output_test[test_selected_idx]], image_shape=image_shape)

plot_latent_ab(latent_test, latent_test_noisy, y_test, test_selected_idx,
               title_a='Before fine-tuning',
               title_b='After fine-tuning', s2=s2)_____no_output_____

### Global latent space shift
The new latent space representation is more robust to noise and may result in a better internal representation of the dataset. We verify this by inspecting the latent space with clean images before and after fine-tuning with noisy images.

Fine-tuning the network with noisy images causes a *domain shift* in the dataset, i.e., a change in the distribution of images since the dataset was initially composed of noise-free images. Depending on the task and the extent of changes during re-train,  (number of epochs, optimizer characteristics, etc.), the new latent space representation may become less well adapted to the original data as a side-effect. How could we address *domain shift* and improve both noisy and noise-free images?

**Instructions:**
* Please execute the cell(s) below_____no_output_____

with torch.no_grad():
  latent_test = encoder(input_test)

plot_latent_ab(latent_test_ref, latent_test, y_test, test_subset_idx,
               title_a='Before fine-tuning',
               title_b='After fine-tuning', s2=s2)_____no_output_____

## Application 2 - Image occlusion
We now investigate the effects of image occlusion. Drawing from the previous exercise, we expect the autoencoder to reconstruct complete images since the train set does not contain occluded images (right?).

We visualize this by plotting three rows:
* Top row with occluded images
* Middle row with reconstructions of occluded images
* Bottom row with reconstructions of the original images

![Occlusion task](https://github.com/mpbrigham/colaboratory-figures/raw/master/nma/autoencoders/applications_occlusion.png)

Similarly, we investigate the source of this issue by looking at the representation of partial images in latent space and how it adjusts after fine-tuning.

**Instructions:**
* Please execute the cell(s) below_____no_output_____

reset_checkpoint(model, optimizer, checkpoint)

with torch.no_grad():
  latent_test_ref = encoder(input_test)_____no_output_____

### Before fine-tuning

**Instructions:**
* Please execute the cell(s) below_____no_output_____

input_train_mask = image_occlusion(input_train, image_shape=image_shape)
input_test_mask = image_occlusion(input_test, image_shape=image_shape)_____no_output_____with torch.no_grad():
  output_test_mask = model(input_test_mask)
  latent_test_mask = encoder(input_test_mask)
  output_test = model(input_test)

plot_row([input_test_mask[test_selected_idx],
          output_test_mask[test_selected_idx],
          output_test[test_selected_idx]], image_shape=image_shape)

plot_latent_ab(latent_test, latent_test_mask, y_test, test_selected_idx,
               title_a='Before occlusion', title_b='After occlusion', s2=s2)_____no_output_____

### After fine-tuning_____no_output_____

n_epochs = 3
batch_size = 32

model.train()

runSGD(model, input_train_mask, input_test_mask,
       out_train=input_train, out_test=input_test,
       n_epochs=n_epochs, batch_size=batch_size)_____no_output_____with torch.no_grad():
  output_test_mask = model(input_test_mask)
  latent_test_mask = encoder(input_test_mask)
  output_test = model(input_test)

plot_row([input_test_mask[test_selected_idx],
          output_test_mask[test_selected_idx],
          output_test[test_selected_idx]], image_shape=image_shape)

plot_latent_ab(latent_test, latent_test_mask, y_test, test_selected_idx,
               title_a='Before fine-tuning',
               title_b='After fine-tuning', s2=s2)_____no_output_____with torch.no_grad():
  latent_test = encoder(input_test)

plot_latent_ab(latent_test_ref, latent_test, y_test, test_subset_idx,
               title_a='Before fine-tuning',
               title_b='After fine-tuning', s2=s2)_____no_output_____

## Application 3 - Image rotation
Finally, we look at the effect of image rotation in latent space coordinates. This task is arguably more challenging since it may require a complete re-write of image reconstruction.

We visualize this by plotting three rows:
* Top row with rotated images
* Middle row with reconstructions of rotated images
* Bottom row with reconstructions of the original images

![Rotation task](https://github.com/mpbrigham/colaboratory-figures/raw/master/nma/autoencoders/applications_rotation.png)

We investigate the source of this issue by looking at the representation of rotated images in latent space and how it adjusts after fine-tuning.

**Instructions:**
* Please execute the cell(s) below_____no_output_____

reset_checkpoint(model, optimizer, checkpoint)

with torch.no_grad():
  latent_test_ref = encoder(input_test)_____no_output_____

### Before fine-tuning

**Instructions:**
* Please execute the cell(s) below_____no_output_____

input_train_rotation = image_rotation(input_train, 90, image_shape=image_shape)
input_test_rotation = image_rotation(input_test, 90, image_shape=image_shape)_____no_output_____with torch.no_grad():
  output_test_rotation = model(input_test_rotation)
  latent_test_rotation = encoder(input_test_rotation)
  output_test = model(input_test)

plot_row([input_test_rotation[test_selected_idx],
          output_test_rotation[test_selected_idx],
          output_test[test_selected_idx]], image_shape=image_shape)

plot_latent_ab(latent_test, latent_test_rotation, y_test, test_selected_idx,
               title_a='Before rotation', title_b='After rotation', s2=s2)_____no_output_____

### After fine-tuning

**Instructions:**
* Please execute the cell(s) below_____no_output_____

n_epochs = 5
batch_size = 32

model.train()

runSGD(model, input_train_rotation, input_test_rotation,
       out_train=input_train, out_test=input_test,
       n_epochs=n_epochs, batch_size=batch_size)_____no_output_____with torch.no_grad():
  output_test_rotation = model(input_test_rotation)
  latent_test_rotation = encoder(input_test_rotation)
  output_test = model(input_test)

plot_row([input_test_rotation[test_selected_idx],
          output_test_rotation[test_selected_idx],
          output_test[test_selected_idx]], image_shape=image_shape)

plot_latent_ab(latent_test, latent_test_rotation, y_test, test_selected_idx,
               title_a='Before fine-tuning',
               title_b='After fine-tuning', s2=s2)_____no_output_____with torch.no_grad():
  latent_test = encoder(input_test)

plot_latent_ab(latent_test_ref, latent_test, y_test, test_subset_idx,
               title_a='Before fine-tuning',
               title_b='After fine-tuning', s2=s2)_____no_output_____

## Application 4 - What would digit "6" look like if we had never seen it before?
Before we start melting our brains with such an impossible task, let's just ask the autoencoder to do it!

We train the autoencoder from scratch without digit class `6` and visualize reconstructions from digit `6`.

**Instructions:**
* Please execute the cell(s) below_____no_output_____

model = AutoencoderClass(s2=s2)
optimizer = optim.Adam(model.parameters())

encoder = model.encoder
decoder = model.decoder_____no_output_____missing = 6

my_input_train = input_train[y_train != missing]
my_input_test = input_test[y_test != missing]
my_y_test = y_test[y_test != missing]_____no_output_____n_epochs = 3
batch_size = 32

runSGD(model, my_input_train, my_input_test,
       n_epochs=n_epochs, batch_size=batch_size)

with torch.no_grad():
  output_test = model(input_test)
  my_latent_test = encoder(my_input_test)_____no_output_____plot_row([input_test[y_test == 6], output_test[y_test == 6]],
         image_shape=image_shape)

plot_latent_generative(my_latent_test, my_y_test, decoder,
                       image_shape=image_shape, s2=s2)_____no_output_____

### Exercise 1: Removing the most dominant digit classes
Digit classes `0` and `1` are dominant in the sense that these occupy large areas of the decoder grid, compared to other digit classes that occupy very little generative space.

How will latent space change when removing the two most dominant digit classes? Will latent space re-distribute evenly among remaining classes or choose another two dominant classes?

**Instructions:**
* Please execute the cell(s) below
* The intersection of two boolean arrays by condition is specified as `x[(cond_a)&(cond_b)]`_____no_output_____

model = AutoencoderClass(s2=s2)
optimizer = optim.Adam(model.parameters())

encoder = model.encoder
decoder = model.decoder_____no_output_____missing_a = 1
missing_b = 0
#################################################
## TODO for students:
#################################################
# input train data
# my_input_train = ...
# input test data
# my_input_test = ...
# model
# my_y_test = ...

# Uncomment to test your code
# print(my_input_train.shape)
# print(my_input_test.shape)
# print(my_y_test.shape)_____no_output_____

**SAMPLE OUTPUT**

```
torch.Size([47335, 784])
torch.Size([7885, 784])
torch.Size([7885])
```_____no_output_____[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/Bonus_Autoencoders/solutions/Bonus_Tutorial3_Solution_22e2c431.py)

_____no_output_____

n_epochs = 3
batch_size = 32

runSGD(model, my_input_train, my_input_test,
       n_epochs=n_epochs, batch_size=batch_size)

with torch.no_grad():
  output_test = model(input_test)
  my_latent_test = encoder(my_input_test)_____no_output_____plot_row([input_test[y_test == missing_a], output_test[y_test == missing_a]],
         image_shape=image_shape)

plot_row([input_test[y_test == missing_b], output_test[y_test == missing_b]],
         image_shape=image_shape)

plot_latent_generative(my_latent_test, my_y_test, decoder,
                       image_shape=image_shape, s2=s2)_____no_output_____

---
# Section 4: ANNs? Same but different!
"Same same but different" is an expression used in some parts of Asia to express differences between supposedly similar subjects. In this exercise, we investigate a fundamental difference in how fully-connected ANNs process visual information compared to human vision.

The previous exercises showed ANN autoencoder performing cognitive tasks with relative ease. However, there is a crucial aspect of ANN processing already encoded in the vectorization of images. This network architecture completely ignores the relative position of pixels. To illustrate this, we show that learning proceeds just as well with shuffled pixel locations.

First, we obtain a reversible shuffle map stored in `shuffle_image_idx` used to shuffle image pixels randomly.

![mnist_pixel_shuffle](https://github.com/mpbrigham/colaboratory-figures/raw/master/nma/autoencoders/mnist_pixel_shuffle.png)

The unshuffled image set `input_shuffle` is recovered as follows:
```
input_shuffle[:, shuffle_rev_image_idx]]
```

First, we set up the reversible shuffle map and visualize a few images with shuffled and unshuffled pixels, followed by their noisy versions.

**Instructions:**
* Please execute the cell(s) below_____no_output_____

# create forward and reverse indexes for pixel shuffling
shuffle_image_idx = np.arange(input_size)
shuffle_rev_image_idx = np.empty_like(shuffle_image_idx)

# shuffle pixel location
np.random.shuffle(shuffle_image_idx)

# store reverse locations
for pos_idx, pos in enumerate(shuffle_image_idx):
  shuffle_rev_image_idx[pos] = pos_idx

# shuffle train and test sets
input_train_shuffle = input_train[:, shuffle_image_idx]
input_test_shuffle = input_test[:, shuffle_image_idx]

input_train_shuffle_noisy = input_train_noisy[:, shuffle_image_idx]
input_test_shuffle_noisy = input_test_noisy[:, shuffle_image_idx]

# show samples with shuffled pixels
plot_row([input_test_shuffle,
          input_test_shuffle[:, shuffle_rev_image_idx]],
         image_shape=image_shape)_____no_output_____# show noisy samples with shuffled pixels
plot_row([input_train_shuffle_noisy[test_selected_idx],
          input_train_shuffle_noisy[:, shuffle_rev_image_idx][test_selected_idx]],
         image_shape=image_shape)_____no_output_____

We initialize and train the network in the denoising task with shuffled pixels.

**Instructions:**
* Please execute the cell below_____no_output_____

model = AutoencoderClass(s2=s2)

encoder = model.encoder
decoder = model.decoder

n_epochs = 3
batch_size = 32

# train the model to denoise shuffled images
runSGD(model, input_train_shuffle_noisy, input_test_shuffle_noisy,
       out_train=input_train_shuffle, out_test=input_test_shuffle,
       n_epochs=n_epochs, batch_size=batch_size)_____no_output_____

Finally, visualize reconstructions and latent space representation with the trained model.

We visualize reconstructions by plotting three rows:
* Top row with shuffled noisy images
* Middle row with reconstructions of shuffled denoised images
* Bottom row with unshuffled reconstructions of denoised images

![mnist_pixel_shuffle denoised](https://github.com/mpbrigham/colaboratory-figures/raw/master/nma/autoencoders/applications_ann_denoise.png)

We obtain the same organization in the encoder map as before. Sharing similar internal representations confirms the network to ignore the relative position of pixels. The decoder grid is different than before since it generates shuffled images.

**Instructions:**
* Please execute the cell below_____no_output_____

with torch.no_grad():
  latent_test_shuffle_noisy = encoder(input_test_shuffle_noisy)
  output_test_shuffle_noisy = model(input_test_shuffle_noisy)

plot_row([input_test_shuffle_noisy[test_selected_idx],
          output_test_shuffle_noisy[test_selected_idx],
          output_test_shuffle_noisy[:, shuffle_rev_image_idx][test_selected_idx]],
         image_shape=image_shape)

plot_latent_generative(latent_test_shuffle_noisy, y_test, decoder,
                       image_shape=image_shape, s2=s2)_____no_output_____

---
# Summary
Hooray! You have finished the last Tutorial of NMA 2020!

We hope you've enjoyed these tutorials and learned about the usefulness of autoencoders to model rich and non-linear representations of data. We hope you may find them useful in your research, perhaps to model certain aspects of cognition or even extend them to biologically plausible architectures - autoencoders of spiking neurons, anyone?

These are the key take away messages from these tutorials:

**Autoencoders trained in *learning by doing* tasks such as compression/decompression, removing noise, etc. can uncover rich lower-dimensional structure embedded in structured images and other cognitively relevant data.**

**The data domain seen during training imprints a "cognitive bias" - you only see what you expect to see, which can only be similar to what you saw before.**

Such bias is related to the concept [*What you see is all there is*](https://en.wikipedia.org/wiki/Thinking,_Fast_and_Slow) coined by Daniel Kahneman in psychology.

For additional applications of autoencoders to neuroscience, check the spike sorting application in the outro video, and also see [here](https://www.nature.com/articles/s41592-018-0109-9) how to replicate the input-output relationship of real networks of neurons with autoencoders._____no_output_____

# @title Video 2: Wrap-up
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  from IPython.display import IFrame
  class BiliVideo(IFrame):
      def __init__(self, id, page=1, width=400, height=300, **kwargs):
          self.id=id
          src = 'https://player.bilibili.com/player.html?bvid={0}&page={1}'.format(id, page)
          super(BiliVideo, self).__init__(src, width, height, **kwargs)

  video = BiliVideo(id="BV1ph411Z7uh", width=854, height=480, fs=1)
  print('Video available at https://www.bilibili.com/video/{0}'.format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id="ziiZK9P6AXQ", width=854, height=480, fs=1, rel=0)
  print('Video available at https://youtube.com/watch?v=' + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

display(out)_____no_output_____