```python
%matplotlib inline
```

#Sequence Models and Long-Short Term Memory Networks

At this point, we have seen various feed-forward networks. That is,
there is no state maintained by the network at all. This might not be
the behavior we want. Sequence models are central to NLP: they are
models where there is some sort of dependence through time between your
inputs. The classical example of a sequence model is the Hidden Markov
Model for part-of-speech tagging. Another example is the conditional
random field.

A recurrent neural network is a network that maintains some kind of
state. For example, its output could be used as part of the next input,
so that information can propogate along as the network passes over the
sequence. In the case of an LSTM, for each element in the sequence,
there is a corresponding *hidden state* $h_t$, which in principle
can contain information from arbitrary points earlier in the sequence.
We can use the hidden state to predict words in a language model,
part-of-speech tags, and a myriad of other things.

##LSTM's in Pytorch

Before getting to the example, note a few things. Pytorch's LSTM expects
all of its inputs to be 3D tensors. The semantics of the axes of these
tensors is important. The first axis is the sequence itself, the second
indexes instances in the mini-batch, and the third indexes elements of
the input. We haven't discussed mini-batching, so lets just ignore that
and assume we will always have just 1 dimension on the second axis. If
we want to run the sequence model over the sentence "The cow jumped",
our input should look like

\begin{align}
\begin{bmatrix}
\overbrace{q_\text{The}}^\text{row vector} \\
q_\text{cow} \\
q_\text{jumped}
\end{bmatrix}\end{align}

Except remember there is an additional 2nd dimension with size 1.

In addition, you could go through the sequence one at a time, in which
case the 1st axis will have size 1 also.

Let's see a quick example.

```python
# Author: Robert Guthrie

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

torch.manual_seed(1)
```

    <torch._C.Generator at 0x7fa7dbe53198>

```python
lstm = nn.LSTM(3, 3)  # Input dim is 3, output dim is 3
inputs = [torch.randn(1, 3) for _ in range(5)]  # make a sequence of length 5

# initialize the hidden state.
hidden = (torch.randn(1, 1, 3),
          torch.randn(1, 1, 3))
for i in inputs:
    # Step through the sequence one element at a time.
    # after each step, hidden contains the hidden state.
    out, hidden = lstm(i.view(1, 1, -1), hidden)

# alternatively, we can do the entire sequence all at once.
# the first value returned by LSTM is all of the hidden states throughout
# the sequence. the second is just the most recent hidden state
# (compare the last slice of "out" with "hidden" below, they are the same)
# The reason for this is that:
# "out" will give you access to all hidden states in the sequence
# "hidden" will allow you to continue the sequence and backpropagate,
# by passing it as an argument  to the lstm at a later time
# Add the extra 2nd dimension
inputs = torch.cat(inputs).view(len(inputs), 1, -1)
hidden = (torch.randn(1, 1, 3), torch.randn(1, 1, 3))  # clean out hidden state
out, hidden = lstm(inputs, hidden)
print(out)
print(hidden)
```

    tensor([[[-0.0187,  0.1713, -0.2944]],

            [[-0.3521,  0.1026, -0.2971]],

            [[-0.3191,  0.0781, -0.1957]],

            [[-0.1634,  0.0941, -0.1637]],

            [[-0.3368,  0.0959, -0.0538]]], grad_fn=<StackBackward>)
    (tensor([[[-0.3368,  0.0959, -0.0538]]], grad_fn=<StackBackward>), tensor([[[-0.9825,  0.4715, -0.0633]]], grad_fn=<StackBackward>))

##Example: An LSTM for Part-of-Speech Tagging

In this section, we will use an LSTM to get part of speech tags. We will
not use Viterbi or Forward-Backward or anything like that, but as a
(challenging) exercise to the reader, think about how Viterbi could be
used after you have seen what is going on.

The model is as follows: let our input sentence be
$w_1, \dots, w_M$, where $w_i \in V$, our vocab. Also, let
$T$ be our tag set, and $y_i$ the tag of word $w_i$.
Denote our prediction of the tag of word $w_i$ by
$\hat{y}_i$.

This is a structure prediction, model, where our output is a sequence
$\hat{y}_1, \dots, \hat{y}_M$, where $\hat{y}_i \in T$.

To do the prediction, pass an LSTM over the sentence. Denote the hidden
state at timestep $i$ as $h_i$. Also, assign each tag a
unique index (like how we had word\_to\_ix in the word embeddings
section). Then our prediction rule for $\hat{y}_i$ is

\begin{align}\hat{y}_i = \text{argmax}_j \  (\log \text{Softmax}(Ah_i + b))_j\end{align}

That is, take the log softmax of the affine map of the hidden state,
and the predicted tag is the tag that has the maximum value in this
vector. Note this implies immediately that the dimensionality of the
target space of $A$ is $|T|$.

Prepare data:

```python
def prepare_sequence(seq, to_ix):
    idxs = [to_ix[w] for w in seq]
    return torch.tensor(idxs, dtype=torch.long)

training_data = [
    ("The dog ate the apple".split(), ["DET", "NN", "V", "DET", "NN"]),
    ("Everybody read that book".split(), ["NN", "V", "DET", "NN"])
]
word_to_ix = {}
for sent, tags in training_data:
    for word in sent:
        if word not in word_to_ix:
            word_to_ix[word] = len(word_to_ix)
print(word_to_ix)
tag_to_ix = {"DET": 0, "NN": 1, "V": 2}

# These will usually be more like 32 or 64 dimensional.
# We will keep them small, so we can see how the weights change as we train.
EMBEDDING_DIM = 6
HIDDEN_DIM = 6
```

    {'The': 0, 'dog': 1, 'ate': 2, 'the': 3, 'apple': 4, 'Everybody': 5, 'read': 6, 'that': 7, 'book': 8}

Create the model:

```python
class LSTMTagger(nn.Module):

    def __init__(self, embedding_dim, hidden_dim, vocab_size, tagset_size):
        super(LSTMTagger, self).__init__()
        self.hidden_dim = hidden_dim

        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)

        # The LSTM takes word embeddings as inputs, and outputs hidden states
        # with dimensionality hidden_dim.
        self.lstm = nn.LSTM(embedding_dim, hidden_dim)

        # The linear layer that maps from hidden state space to tag space
        self.hidden2tag = nn.Linear(hidden_dim, tagset_size)
        self.hidden = self.init_hidden()

    def init_hidden(self):
        # Before we've done anything, we dont have any hidden state.
        # Refer to the Pytorch documentation to see exactly
        # why they have this dimensionality.
        # The axes semantics are (num_layers, minibatch_size, hidden_dim)
        return (torch.zeros(1, 1, self.hidden_dim),
                torch.zeros(1, 1, self.hidden_dim))

    def forward(self, sentence):
        embeds = self.word_embeddings(sentence)
        lstm_out, self.hidden = self.lstm(
            embeds.view(len(sentence), 1, -1), self.hidden)
        tag_space = self.hidden2tag(lstm_out.view(len(sentence), -1))
        tag_scores = F.log_softmax(tag_space, dim=1)
        return tag_scores
```

Train the model:

```python
model = LSTMTagger(EMBEDDING_DIM, HIDDEN_DIM, len(word_to_ix), len(tag_to_ix))
loss_function = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# See what the scores are before training
# Note that element i,j of the output is the score for tag j for word i.
# Here we don't need to train, so the code is wrapped in torch.no_grad()
with torch.no_grad():
    inputs = prepare_sequence(training_data[0][0], word_to_ix)
    tag_scores = model(inputs)
    print(tag_scores)

for epoch in range(300):  # again, normally you would NOT do 300 epochs, it is toy data
    print(epoch)
    for sentence, tags in training_data:
        # Step 1. Remember that Pytorch accumulates gradients.
        # We need to clear them out before each instance
        model.zero_grad()

        # Also, we need to clear out the hidden state of the LSTM,
        # detaching it from its history on the last instance.
        model.hidden = model.init_hidden()

        # Step 2. Get our inputs ready for the network, that is, turn them into
        # Tensors of word indices.
        sentence_in = prepare_sequence(sentence, word_to_ix)
        targets = prepare_sequence(tags, tag_to_ix)

        # Step 3. Run our forward pass.
        tag_scores = model(sentence_in)

        # Step 4. Compute the loss, gradients, and update the parameters by
        #  calling optimizer.step()
        loss = loss_function(tag_scores, targets)
        loss.backward()
        optimizer.step()

# See what the scores are after training
with torch.no_grad():
    inputs = prepare_sequence(training_data[0][0], word_to_ix)
    tag_scores = model(inputs)

    # The sentence is "the dog ate the apple".  i,j corresponds to score for tag j
    # for word i. The predicted tag is the maximum scoring tag.
    # Here, we can see the predicted sequence below is 0 1 2 0 1
    # since 0 is index of the maximum value of row 1,
    # 1 is the index of maximum value of row 2, etc.
    # Which is DET NOUN VERB DET NOUN, the correct sequence!
    print(tag_scores)
```

    tensor([[-1.1673, -1.2724, -0.8949],
            [-1.0838, -1.3116, -0.9357],
            [-0.9590, -1.3362, -1.0388],
            [-1.0451, -1.3285, -0.9585],
            [-1.0733, -1.3155, -0.9422]])
    0
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    88
    89
    90
    91
    92
    93
    94
    95
    96
    97
    98
    99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    206
    207
    208
    209
    210
    211
    212
    213
    214
    215
    216
    217
    218
    219
    220
    221
    222
    223
    224
    225
    226
    227
    228
    229
    230
    231
    232
    233
    234
    235
    236
    237
    238
    239
    240
    241
    242
    243
    244
    245
    246
    247
    248
    249
    250
    251
    252
    253
    254
    255
    256
    257
    258
    259
    260
    261
    262
    263
    264
    265
    266
    267
    268
    269
    270
    271
    272
    273
    274
    275
    276
    277
    278
    279
    280
    281
    282
    283
    284
    285
    286
    287
    288
    289
    290
    291
    292
    293
    294
    295
    296
    297
    298
    299
    tensor([[-0.0447, -3.3470, -4.7651],
            [-4.8805, -0.0112, -5.6515],
            [-3.8234, -4.7204, -0.0312],
            [-0.0561, -3.2038, -4.2751],
            [-4.3101, -0.0176, -5.5201]])

##Exercise: Augmenting the LSTM part-of-speech tagger with character-level features

In the example above, each word had an embedding, which served as the
inputs to our sequence model. Let's augment the word embeddings with a
representation derived from the characters of the word. We expect that
this should help significantly, since character-level information like
affixes have a large bearing on part-of-speech. For example, words with
the affix *-ly* are almost always tagged as adverbs in English.

To do this, let $c_w$ be the character-level representation of
word $w$. Let $x_w$ be the word embedding as before. Then
the input to our sequence model is the concatenation of $x_w$ and
$c_w$. So if $x_w$ has dimension 5, and $c_w$
dimension 3, then our LSTM should accept an input of dimension 8.

To get the character level representation, do an LSTM over the
characters of a word, and let $c_w$ be the final hidden state of
this LSTM. Hints:

* There are going to be two LSTM's in your new model.
  The original one that outputs POS tag scores, and the new one that
  outputs a character-level representation of each word.
* To do a sequence model over characters, you will have to embed characters.
  The character embeddings will be the input to the character LSTM.


```python
from __future__ import print_function
import sisl
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
```

Analyzing output from Siesta makes analysis much easier since many things are intrinsically enabled through `sisl`, i.e. orbital numbering etc. may easily be handled with `sisl`.

In this example we will use Siesta to calculate the *heavy* part, i.e. the projected density of states for a fine $k$-point sampling. The formula used to calculate the projected density of states is:
\begin{align}
  \mathrm{PDOS}_\nu(E) &= \int \mathrm d k \sum_i \psi^*_{i,k,\nu} [\mathbf S(k) | \psi_{i,k}\rangle]_\nu D(E-\epsilon_i),
  \\
  \mathrm{DOS}(E) &= \int \mathrm d k \sum_i D(E-\epsilon_i) = \sum_\nu \mathrm{PDOS}_\nu(E),
\end{align}
with $D(\Delta E)$ being the distribution function (in Siesta this is the Gaussian distribution function).
Finally we will also compare with the same calculation using `sisl` to ensure that correctness of both methods.
Note that `sisl` is (for now) not MPI parallelized and thus for large systems it may be *way* more efficient to use Siesta to calculate PDOS, but use `sisl` to post-process the PDOS data.

The system will again be graphene.

```python
graphene = sisl.geom.graphene(1.44)
graphene.write('STRUCT.fdf')
graphene.write('STRUCT.xyz')
```

To tell Siesta to calculate the projected DOS one has to add these flags to the input (if you want you can play with the numbers):

```python
open('PDOS.fdf', 'w').write("""
# K-point sampling
%block PDOS.kgrid.MonkhorstPack
  63  0 0
   0 63 0
   0  0 1
%endblock
# Energy grid to calculate the PDOS on
# E_min = -20 eV
# E_max = 15 eV
# Gaussian width: \sigma 2 ^{1/2} = 0.2 eV
# Number of points: 3500 (spacing of 0.01 eV)
%block Projected.DensityOfStates
-20.00 15.00 0.200 3500 eV
%endblock
""");
```

Please open `PDOS.fdf` and `RUN.fdf`.
- Why do you see two different $k$-point samplings?

Now run Siesta to calculate the electronic structure and the PDOS:

     siesta RUN.fdf | tee RUN.out

Subsequently we will read in the PDOS information and post-process it.
First we will data found in the `siesta.PDOS.xml` file, 3 quantities will be returned:
1. The geometry (as found in the XML file), this also contains all atomic species etc.
2. The energy grid used
3. Orbital and energy resolved DOS

Before analyzing, please examine the data returned by viewing the shapes and printing the geometry.

```python
geometry, E, PDOS = sisl.get_sile('siesta.PDOS.xml').read_data()
```

```python
# Examine output by printing shapes and geometry

```

## Exercises

1. Plot the total DOS (invoke Eq. 2 in the top paragraph)
2. Plot the total PDOS on the first atom
   *HINT*: Sum all orbitals on the first atom
3. Figure out which orbital contributes to the Dirac cone by plotting the PDOS for all orbitals
4. Read in the Hamiltonian (as done in [S 1](../S_01/run.ipynb)) and create a Monkhorst-Pack grid as the one in Siesta input. Use this snippet to calculate the PDOS using `sisl`

       mp = sisl.MonkhorstPack(H, [nx, ny, 1])
       pdos = mp.asaverage().PDOS(E, eta=True)

    Search the API documentation for the `MonkhorstPack.asaverage` method and figure out what it does. Note, there are other `MonkhorstPack.as*` methods, these are all extremely usefull when calculating a large quantity of data in a Brillouin-zone object.
5. Compare the Siesta PDOS with sisl PDOS, why are they different?
   *HINT*: $\sigma$. Check the API for the `PDOS` method.

```python
plt.plot(E, PDOS.sum(0), label='DOS')
plt.xlabel(r'$E - E_F$ [eV]')
plt.ylabel(r'DOS [1/eV]')

# Add more plots for individual PDOS

# Add legend(s)
plt.legend();
```


# Neuromatch Academy: Week 2, Day 5, Tutorial 3
# Learning to Act: Temporal Difference Learning

# Tutorial Objectives

In this tutorial you learn how to act in the more realistic setting of sequential decisions, formalized by Markov Decision Processes (MDPs). In a sequential decision problem, the actions executed in one state not only may lead to immediate rewards (as in a bandit problem), but may also affect the states experienced next (unlike a bandit problem). In other words, each individual action can ultimately affect affect all future rewards. Thus, making decisions in this setting requires considering each action in terms of their expected **cumulative** future reward.

We will consider here the example of spatial navigation, where actions (movements) in one state (location) affect the states experienced next, and an agent might need to execute a whole sequence of actions before a reward is obtained.

* You will learn the basics of the Q-learning algorithm for estimating action values
* You will understand the concept of on-policy and off-policy value estimation
* You will learn the differences between Q-learning (an off-policy algorithm) and SARSA (an on-policy algorithm)
* You will learn how the concept of exploration and exploitation, reviewed in the bandit case, also applies to the sequential decision setting.

# Setup

```python
import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import convolve as conv
```

```python
#@title Figure Settings
%matplotlib inline
fig_w, fig_h = (8, 6)
plt.rcParams.update({'figure.figsize': (fig_w, fig_h)})
%config InlineBackend.figure_format = 'retina'
```

```python
#@title Helper Functions
def epsilon_greedy(q, epsilon):
  """Epsilon-greedy policy: selects the maximum value action with probabilty
  (1-epsilon) and selects randomly with epsilon probability.

  Args:
    q (ndarray): an array of action values
    epsilon (float): probability of selecting an action randomly

  Returns:
    int: the chosen action
  """
  if np.random.random() > epsilon:
    action = np.argmax(q)
  else:
    action = np.random.choice(len(q))

  return action

class CliffWorld:
  """
  World: Cliff world.
  40 states (4-by-10 grid world).
  The mapping from state to the grids are as follows:
  30 31 32 ... 39
  20 21 22 ... 29
  10 11 12 ... 19
  0  1  2  ...  9
  0 is the starting state (S) and 9 is the goal state (G).
  Actions 0, 1, 2, 3 correspond to right, up, left, down.
  Moving anywhere from state 9 (goal state) will end the session.
  Taking action down at state 11-18 will go back to state 0 and incur a
      reward of -100.
  Landing in any states other than the goal state will incur a reward of -1.
  Going towards the border when already at the border will stay in the same
      place.
  """
  def __init__(self):
    self.name = "cliff_world"
    self.n_states = 40
    self.n_actions = 4
    self.dim_x = 10
    self.dim_y = 4
    self.init_state = 0

  def get_outcome(self, state, action):
    if state == 9:  # goal state
      reward = 0
      next_state = None
      return next_state, reward
    reward = -1  # default reward value
    if action == 0:  # move right
      next_state = state + 1
      if state % 10 == 9:  # right border
        next_state = state
      elif state == 0:  # start state (next state is cliff)
        next_state = None
        reward = -100
    elif action == 1:  # move up
      next_state = state + 10
      if state >= 30:  # top border
        next_state = state
    elif action == 2:  # move left
      next_state = state - 1
      if state % 10 == 0:  # left border
        next_state = state
    elif action == 3:  # move down
      next_state = state - 10
      if state >= 11 and state <= 18:  # next is cliff
        next_state = None
        reward = -100
      elif state <= 9:  # bottom border
        next_state = state
    else:
      print("Action must be between 0 and 3.")
      next_state = None
      reward = None
    return int(next_state) if next_state is not None else None, reward

  def get_all_outcomes(self):
    outcomes = {}
    for state in xrange(self.n_states):
      for action in xrange(self.n_actions):
        next_state, reward = self.get_outcome(state, action)
        outcomes[state, action] = [(1, next_state, reward)]
    return outcomes

def learn_environment(env, learning_rule, params, max_steps, n_episodes):
  # Start with a uniform value function
  value = np.ones((env.n_states, env.n_actions))

  # Run learning
  reward_sums = np.zeros(n_episodes)

  # Loop over episodes
  for episode in range(n_episodes):
    state = env.init_state  # initialize state
    reward_sum = 0

    for t in range(max_steps):
      # choose next action
      action = epsilon_greedy(value[state], params['epsilon'])

      # observe outcome of action on environment
      next_state, reward = env.get_outcome(state, action)

      # update value function
      value = learning_rule(state, action, reward, next_state, value, params)

      # sum rewards obtained
      reward_sum += reward

      if next_state is None:
          break  # episode ends
      state = next_state

      reward_sums[episode] = reward_sum

  return value, reward_sums

def plot_state_action_values(env, value, ax=None):
  """
  Generate plot showing value of each action at each state.
  """
  if ax is None:
    fig, ax = plt.subplots()

  for a in range(env.n_actions):
    ax.plot(range(env.n_states), value[:, a], marker='o', linestyle='--')
  ax.set(xlabel='States', ylabel='Values')
  ax.legend(['R','U','L','D'], loc='lower right')

def plot_quiver_max_action(env, value, ax=None):
  """
  Generate plot showing action of maximum value or maximum probability at
    each state (not for n-armed bandit or cheese_world).
  """
  if ax is None:
    fig, ax = plt.subplots()

  X = np.tile(np.arange(env.dim_x), [env.dim_y,1]) + 0.5
  Y = np.tile(np.arange(env.dim_y)[::-1][:,np.newaxis], [1,env.dim_x]) + 0.5
  which_max = np.reshape(value.argmax(axis=1), (env.dim_y,env.dim_x))
  which_max = which_max[::-1,:]
  U = np.zeros(X.shape)
  V = np.zeros(X.shape)
  U[which_max == 0] = 1
  V[which_max == 1] = 1
  U[which_max == 2] = -1
  V[which_max == 3] = -1

  ax.quiver(X, Y, U, V)
  ax.set(
      title='Maximum value/probability actions',
      xlim=[-0.5, env.dim_x+0.5],
      ylim=[-0.5, env.dim_y+0.5],
  )
  ax.set_xticks(np.linspace(0.5, env.dim_x-0.5, num=env.dim_x))
  ax.set_xticklabels(["%d" % x for x in np.arange(env.dim_x)])
  ax.set_xticks(np.arange(env.dim_x+1), minor=True)
  ax.set_yticks(np.linspace(0.5, env.dim_y-0.5, num=env.dim_y))
  ax.set_yticklabels(["%d" % y for y in np.arange(0, env.dim_y*env.dim_x, env.dim_x)])
  ax.set_yticks(np.arange(env.dim_y+1), minor=True)
  ax.grid(which='minor',linestyle='-')

def plot_heatmap_max_val(env, value, ax=None):
  """
  Generate heatmap showing maximum value at each state
  """
  if ax is None:
    fig, ax = plt.subplots()

  if value.ndim == 1:
      value_max = np.reshape(value, (env.dim_y,env.dim_x))
  else:
      value_max = np.reshape(value.max(axis=1), (env.dim_y,env.dim_x))
  value_max = value_max[::-1,:]

  im = ax.imshow(value_max, aspect='auto', interpolation='none', cmap='afmhot')
  ax.set(title='Maximum value per state')
  ax.set_xticks(np.linspace(0, env.dim_x-1, num=env.dim_x))
  ax.set_xticklabels(["%d" % x for x in np.arange(env.dim_x)])
  ax.set_yticks(np.linspace(0, env.dim_y-1, num=env.dim_y))
  # ax.set_yticklabels(["%d" % y for y in np.arange(0, env.dim_y*env.dim_x, env.dim_x)])
  if env.name != 'windy_cliff_grid':
      ax.set_yticklabels(
          ["%d" % y for y in np.arange(
              0, env.dim_y*env.dim_x, env.dim_x)][::-1])
  return im
  # ax.colorbar(im)
  # return fig

def plot_rewards(n_episodes, rewards, average_range=10, ax=None):
  """
  Generate plot showing total reward accumulated in each episode.
  """
  if ax is None:
    fig, ax = plt.subplots()

  smoothed_rewards = (conv(rewards, np.ones(average_range), mode='same')
                      / average_range)

  ax.plot(range(0, n_episodes, average_range),
          smoothed_rewards[0:n_episodes:average_range],
          marker='o', linestyle='--')
  ax.set(xlabel='Episodes', ylabel='Total reward')

# # Plot results
# fig = plot_state_action_values(env, value_qlearning)
# fig = plot_quiver_max_action(env, value_qlearning)
# fig = plot_heatmap_max_val(env, value_qlearning)
# fig = plot_rewards(n_episodes, reward_sums_qlearning, average_range=10)

def plot_performance(env, value, reward_sums):
  fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(16, 12))
  plot_state_action_values(env, value, ax=axes[0,0])
  plot_quiver_max_action(env, value, ax=axes[0,1])
  plot_rewards(n_episodes, reward_sums, ax=axes[1,0])
  im = plot_heatmap_max_val(env, value, ax=axes[1,1])
  fig.colorbar(im)
```

# Grid Worlds

As pointed out, bandits only have a single state and immediate rewards for our actions. Many problems we are interested in have multiple states and delayed rewards, i.e. we won't know if the choices we made will pay off over time, or which actions we took contributed to the outcomes we observed.

In order to explore these ideas, we turn the a common problem setting: the gird world. Grid worlds are simple environments where each state corresponds to a tile on a 2D grid, and the only actions the agent can take are to move up, down, left, or right across the grid tiles. The agent's job is almost always to find a way to a goal tile in the most direct way possible while overcoming some maze or other obstacles, either static or dynamic.

For our discussion we will be looking at the classic Cliff World, or Cliff Walker, environment. This is a 4x10 grid with a starting position in the lower-left and the goal position in the lower-right. Every tile between these two is the "cliff", and should the agent enter the cliff, they will receive a -100 reward and be sent back to the starting position. Every tile other than the cliff produces a -1 reward when entered. The goal tile ends the episode after taking any action from it.

Given these conditions, the maximum achievable reward is -11 (1 up, 9 right, 1 down). Using negative rewards is a common technique to encourage the agent to move and seek out the goal state as fast as possible.

# Q-learning

Now that we have our environment, how can we solve it?

One of the most famous algorithms for estimating action values (aka Q-values) is the Temporal Differences (TD) **control** algorithm known as *Q-learning* (Watkins, 1989).

\begin{align}
Q(s_t,a_t) \leftarrow Q(s_t,a_t) + \alpha \big(r_t + \gamma\max_\limits{a} Q(s_{t+1},a_{t+1}) - Q(s_t,a_t)\big)
\end{align}

where $Q(s,a)$ is the value function for action $a$ at state $s$, $\alpha$ is the learning rate, $r$ is the reward, and $\gamma$ is the temporal discount rate.

The expression $r_t + \gamma\max_\limits{a} Q(s_{t+1},a_{t+1})$ is referred to as the TD target while the full expression
\begin{align}
r_t + \gamma\max_\limits{a} Q(s_{t+1},a_{t+1}) - Q(s_t,a_t),
\end{align}
i.e. the difference between the TD target and the current Q-value, is referred to as the TD error, or reward prediction error.

Because of the max operator used to select the optimal Q-value in the TD target, Q-learning directly estimates the optimal action value, i.e. the cumulative future reward that would be obtained if the agent behaved optimally, regardless of the policy currently followed by the agent. For this reason, Q-learning is referred to as an **off-policy** method.

## Exercise: Implement the Q-learning algorithm

In this exercise you will implement the Q-learning update rule described above. It takes in as arguments the previous state $s_t$, the action $a_t$ taken, the reward received $r_t$, the current state $s_{t+1}$, the Q-value table, and a dictionary of parameters that contain the learning rate $\alpha$ and discount factor $\gamma$. The method returns the updated Q-value table. For the parameter dictionary, $\alpha$: `params['alpha']` and $\gamma$: `params['gamma']`.

```python
def q_learning(prev_state, action, reward, state, value, params):
  """Q-learning: updates the value function and returns it.

  Args:
    prev_state (int): the previous state identifier
    action (int): the action taken
    reward (float): the reward received
    state (int): the current state identifier
    value (ndarray): current value function of shape (n_states, n_actions)
    params (dict): a dictionary containing the default parameters

  Returns:
    ndarray: the updated value function of shape (n_states, n_actions)
  """
  # value of previous state-action pair
  prev_value = value[prev_state, action]

  ##########################################################
  ## TODO for students: implement the Q-learning update rule
  ##########################################################
  # comment this out when you've filled
  raise NotImplementedError("Student excercise: implement the Q-learning update rule")

  # write an expression for finding the maximum Q-value at the current state
  if state is None:
    max_value = 0
  else:
    max_value = ...

  # write the expression for the reward prediction error
  delta = ...
  # write the expression of that updates the Q-value of previous state-action pair
  value[prev_state, action] = ...

  return value
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D5_ReinforcementLearning/solutions/W2D5_Tutorial3_Solution_a0a54043.py)

Now that we have our Q-learning algorithm, let's see how it handles learning to solve the Cliff World environment.

```python
# set for reproducibility, comment out / change seed value for different results
np.random.seed(1)

# parameters needed by our policy and learning rule
params = {
  'epsilon': 0.1,  # epsilon-greedy policy
  'alpha': 0.1,  # learning rate
  'gamma': 1.0,  # discount factor
}

# episodes/trials
n_episodes = 500
max_steps = 1000

# environment initialization
env = CliffWorld()

# solve Cliff World using Q-learning
results = learn_environment(env, q_learning, params, max_steps, n_episodes)
value_qlearning, reward_sums_qlearning = results

# Plot results
plot_performance(env, value_qlearning, reward_sums_qlearning)
```

If all went well, we should see some interesting plots. The top left is a representation of the Q-table itself, showing the values for different actions in different states. Notably, going right from the starting state or down when above the cliff is clearly very bad. The top right figure shows the greedy policy based on the Q-table, while the bottom right shows the best reward we ever received in a given state.

Finally, and most importantly, the bottom left is the actual proof of learning, as we see the total reward steadily increasing until asymptoting at the maximum possible reward of -11.

Feel free to try changing the parameters or random seed and see how the agent's behavior changes.

# SARSA

An alternative to Q-learning, the SARSA algorithm also estimates action values. However, rather than estimating the optimal (off-policy) values, SARSA estimates the **on-policy** action value, i.e. the cumulative future reward that would be obtained if the agent behaved according to its current beliefs.

\begin{align}
Q(s_t,a_t) \leftarrow Q(s_t,a_t) + \alpha \big(r_t + \gamma Q(s_{t+1},a_{t+1}) - Q(s_t,a_t)\big)
\end{align}

where, once again, $Q(s,a)$ is the value function for action $a$ at state $s$, $\alpha$ is the learning rate, $r$ is the reward, and $\gamma$ is the temporal discount rate.

In fact, you will notices that the *only* difference between Q-learning and SARSA is the TD target calculation uses the policy to select the next action (in our case epsilon-greedy) rather than using the action that maximizes the Q-value.

## Exercise: Implement the SARSA algorithm

In this exercise you will implement the SARSA update rule described above. Just like Q-learning, it takes in as arguments the previous state $s_t$, the action $a_t$ taken, the reward received $r_t$, the current state $s_{t+1}$, the Q-value table, and a dictionary of parameters that contain the learning rate $\alpha$ and discount factor $\gamma$. The method returns the updated Q-value table. You may use the `epsilon_greedy` function to acquire the next action. For the parameter dictionary, $\alpha$: `params['alpha']`, $\gamma$: `params['gamma']`, and $\epsilon$: `params['epsilon']`.

```python
def sarsa(prev_state, action, reward, state, value, params):
  """SARSA: updates the value function and returns it.

  Args:
    prev_state (int): the previous state identifier
    action (int): the action taken
    reward (float): the reward received
    state (int): the current state identifier
    value (ndarray): current value function of shape (n_states, n_actions)
    params (dict): a dictionary containing the default parameters

  Returns:
    ndarray: the updated value function of shape (n_states, n_actions)
  """
  # value of previous state-action pair
  prev_value = value[prev_state, action]

  ##########################################################
  ## TODO for students: implement the SARSA update rule
  ##########################################################
  # comment this out when you've filled
  raise NotImplementedError("Student excercise: implement the SARSA update rule")

  # select the expected value at current state based on our policy by sampling
  # from it
  if state == None:
    policy_value = 0
  else:
    # write an expression for selecting an action using epsilon-greedy
    policy_action = ...
    # write an expression for obtaining the value of the policy action at the
    # current state
    policy_value = ...

  # write the expression for the reward prediction error
  delta = ...
  # write the expression of that updates the value of previous state-action pair
  value[prev_state, action] = ...

  return value
```

```python
def sarsa(prev_state, action, reward, state, value, params):
  """SARSA: updates the value function and returns it.

  Args:
    prev_state (int): the previous state identifier
    action (int): the action taken
    reward (float): the reward received
    state (int): the current state identifier
    value (ndarray): current value function of shape (n_states, n_actions)
    params (dict): a dictionary containing the default parameters

  Returns:
    ndarray: the updated value function of shape (n_states, n_actions)
  """
  # value of previous state-action pair
  prev_value = value[prev_state, action]

  # select the expected value at current state based on our policy by sampling
  # from it
  if state == None:
    policy_value = 0
  else:
    # sample action from the policy for the next state
    policy_action = epsilon_greedy(value[state], params['epsilon'])
    # get the value based on the action sampled from the policy
    policy_value = value[state, policy_action]

  # reward prediction error
  delta = reward + params['gamma'] * policy_value - prev_value

  # update value of previous state-action pair
  value[prev_state, action] = prev_value + params['alpha'] * delta

  return value
```

Now that we have an implementation for SARSA, let's see how it tackles Cliff World. We will again use the same setup we tried with Q-learning.

```python
np.random.seed(1)

# parameters needed by our policy and learning rule
params = {
  'epsilon': 0.1,  # epsilon-greedy policy
  'alpha': 0.1,  # learning rate
  'gamma': 1.0,  # discount factor
}

# episodes/trials
n_episodes = 500
max_steps = 1000

# environment initialization
env = CliffWorld()

# learn Cliff World using Sarsa
results = learn_environment(env, sarsa, params, max_steps, n_episodes)
value_sarsa, reward_sums_sarsa = results

# Plot results
plot_performance(env, value_sarsa, reward_sums_sarsa)
```

We should see that SARSA also solves the task with similar looking outcomes to Q-learning. One notable difference is that SARSA seems to be skittsh around the cliff edge and often goes further away before coming back down to the goal.

Again, feel free to try changing the parameters or random seed and see how the agent's behavior changes.

# On-Policy vs Off-Policy

We have now seen an example of both on- and off-policy learning algorithms. Let's compare both Q-learning and SARSA reward results again, side-by-side, to see how they stack up.

```python
# parameters needed by our policy and learning rule
params = {
  'epsilon': 0.1,  # epsilon-greedy policy
  'alpha': 0.1,  # learning rate
  'gamma': 1.0,  # discount factor
}

# episodes/trials
n_episodes = 500
max_steps = 1000

# environment initialization
env = CliffWorld()

# learn Cliff World using Sarsa
np.random.seed(1)
results = learn_environment(env, q_learning, params, max_steps, n_episodes)
value_qlearning, reward_sums_qlearning = results
np.random.seed(1)
results = learn_environment(env, sarsa, params, max_steps, n_episodes)
value_sarsa, reward_sums_sarsa = results
```

```python
fig, ax = plt.subplots()
ax.plot(reward_sums_qlearning, label='Q-learning')
ax.plot(reward_sums_sarsa, label='SARSA')
ax.set(xlabel='Episodes', ylabel='Total reward')
plt.legend(loc='lower right');
```

On this simple Cliff World task, Q-learning and SARSA are almost indistinguisable from a performance standpoint, but we can see that Q-learning has a slight-edge within the 500 episode time horizon. Let's look at the illustrated "greedy policy" plots again.

```python
fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(16, 6))
plot_quiver_max_action(env, value_qlearning, ax=ax1)
ax1.set(title='Q-learning maximum value/probability actions')
plot_quiver_max_action(env, value_sarsa, ax=ax2)
ax2.set(title='SARSA maximum value/probability actions');
```

What immidiately jumps out is that Q-learning learned to go up, then immediately go to the right, skirting the cliff edge, until it hits the wall and goes down to the goal. The policy further away from the cliff is less certain.

SARSA, on the other hand, appears to avoid the cliff edge, going up one more tile before starting over to the goal side. This also clearly solves the challenge of getting to the goal, but does so at an additional -2 cost over the truly optimal route.

Why do you think these behaviors emerged the way they did?

```python

```


# Import modules

```python
from common import colors, savefig, sbarlw, insetlw
```

```python
import os
os.chdir(os.path.join('..', '..', '..'))
print(os.getcwd())
```

    /home/eharkin/projects/5HTGIF-dev

```python
from __future__ import division

import pickle

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gs
from mpl_toolkits.axes_grid1.inset_locator import inset_axes, mark_inset
import seaborn as sns
from scipy.signal import find_peaks
import scipy.optimize as optimize
import scipy.stats as stats

from grr.cell_class import Cell
from ezephys import pltools
from ezephys.stimtools import BiexponentialSynapticKernel
```

# Load data

```python
MINI_SAVE_PATH = './data/processed/5HT/GABA_synapses/detected_minis/'

mini_detectors = []

for fname in os.listdir(MINI_SAVE_PATH):
    with open(MINI_SAVE_PATH + fname, 'rb') as f:
        mini_detectors.append(pickle.load(f))

SAMPLE_TR_PATH = './data/processed/5HT/GABA_synapses/sample_traces/'

with open(SAMPLE_TR_PATH + 'averaged_traces.dat', 'rb') as f:
    averaged_traces = pickle.load(f)

with open(SAMPLE_TR_PATH + 'sample_decay_fit.dat', 'rb') as f:
    sample_tr = pickle.load(f)
```

# Perform linear fit on I/V data

Extract reversal potential of Cl based on linear fit of IV curves.
Use extracted reversal to calculate mIPSC peak conductance from recorded
currents.

```python
IV_data = {
    'amplitudes': [],
    'voltages': []
}

for i in range(len(averaged_traces['t_mats'])):
    mask = 0 < np.gradient(averaged_traces['voltages'][i])
    IV_data['voltages'].append(averaged_traces['voltages'][i][mask])
    IV_data['amplitudes'].append(averaged_traces['traces'][i][250, mask])

X = np.concatenate(IV_data['voltages'])
X = np.array([X, np.ones_like(X)]).T
Y = np.concatenate(IV_data['amplitudes'])

def linear_fit(X, Y):
    coeffs = np.dot(np.linalg.inv(np.dot(X.T, X)), np.dot(X.T, Y))
    return coeffs

m, b = linear_fit(X, Y)
IV_coeffs = {
    'slope': m,
    'intercept': b
}
del m, b

reversal = -IV_coeffs['intercept']/IV_coeffs['slope']

def I_to_g(I, V, reversal):
    """Convert current (amperes) to conductance (siemens).

    Note: 1 pA / 1 mV = 1nS

    Inputs:
        I -- float or array of floats
            Current value(s) to convert.
        V -- float
            Voltage at which current was measured.
        reversal -- float
            Reversal potential of the current.

    Returns:
        I / (V - reversal)
    """
    return I/(V - reversal)
```

# Fit gamma distribution to mIPSC amplitudes

```python
for MiDe in mini_detectors:
    print('no. minis: {}'.format(MiDe.no_minis))

# Since all cells have a similar number of minis, we'll pool all amplitudes.

amplis = np.concatenate([MiDe.amplitudes(5) for MiDe in mini_detectors])
conds = I_to_g(amplis, 20, reversal)
a, loc, scale = stats.gamma.fit(conds, floc = 0)

gamma_params = {
    'a': a,
    'loc': loc,
    'scale': scale
}
del a, loc, scale

plt.figure()

ax = plt.subplot(111)
plt.hist(conds)

ax2 = ax.twinx()
x = np.arange(0, 1.4, 0.01)
ax2.plot(
    x,
    stats.gamma.pdf(x, a = gamma_params['a'], loc = gamma_params['loc'], scale = gamma_params['scale']),
    'k-'
)
ax2.set_ylim(0, ax2.get_ylim()[1])

plt.show()
```

```python
print('mIPSC amplitude: {:.5f}nS +/- {:.5f}nS'.format(np.mean(conds), stats.sem(conds)))
```

    mIPSC amplitude: 0.30849nS +/- 0.00366nS

# Fit exponential distribution to mIPSC IEIs

```python
# Again, we'll pool all IEIs here since the count per cell is similar.

IEIs = np.concatenate([MiDe.inter_mini_intervals() * 1e-3 for MiDe in mini_detectors])
loc, scale = stats.expon.fit(IEIs)

expon_params = {
    'loc': loc,
    'scale': scale
}
del loc, scale

ax = plt.subplot(111)
plt.hist(IEIs)

ax2 = ax.twinx()
x = np.arange(0, 2.5, 0.01)
ax2.plot(
    x,
    stats.expon.pdf(x, loc = expon_params['loc'], scale = expon_params['scale']),
    'k-'
)
ax2.set_ylim(0, ax2.get_ylim()[1])

plt.show()
```

# Make figure

Define some reference timepoints. `stim_time` is the beginning of the stimulus artefact, and `foot_time` is an approximate index of the IPSC start (i.e., `stim_time` + propagation delay).

```python
stim_time = 162
foot_time = 180
```

```python
first_peak_times = []

for voltages, tr in zip(averaged_traces['voltages'], averaged_traces['traces']):
    ind = np.argmin(np.abs(voltages - 60.))
    locs, details = find_peaks(tr[(stim_time + 10):, ind], (10, 500), width=2, prominence=2)
    first_peak_times.append((locs[0] + 10 + (stim_time - foot_time)) * 0.1)

    plt.plot(tr[:, ind], 'k-')
    plt.plot(locs + 162 + 10, details['peak_heights'], 'ro')
    plt.legend()
    plt.ylim(-10, 170)
    plt.xlim(150, 300)

plt.axvline(stim_time, color='g', ls='--', lw=4, label='Stimulus time')
plt.axvline(foot_time, color='b', ls='--', lw=4, label='IPSC foot time')
plt.legend()
```

```python
print('First peak times: {}'.format(first_peak_times))
print('First peak times: {:.1f} +/- {:.1f}'.format(np.mean(first_peak_times), stats.sem(first_peak_times)))
print(
    'Decay time constants: {:.1f} +/- {:.1f}'.format(
        np.mean(averaged_traces['decay']),
        stats.sem(averaged_traces['decay'])
    )
)
```

    First peak times: [2.1, 5.7, 5.300000000000001]
    First peak times: 4.4 +/- 1.1
    Decay time constants: 26.0 +/- 2.1

## IPSC model

The GIFnet models GABA synapses with a biexponential synaptic kernel.
The decay of the synaptic kernel in GIFnet simulations should obviously be the same as fitted here (26.0ms).
We can calculate the rising time constant of the synaptic kernel as follows.

The biexponential kernel used in GIFnet simulations is defined as
\begin{align}
    I &= e^{-t/\tau_d} - e^{-t/\tau_r} \\
    \frac{\partial I}{\partial t} &= \frac{1}{\tau_r} e^{-t/\tau_r} - \frac{1}{\tau_d} e^{-t/\tau_d}
\end{align}
where $\tau_d$ is the decay time constant and $\tau_r$ is the rising time constant. Setting the derivative to zero gives us a partial solution for $\tau_r$ in terms of $\tau_d$ and the peak time (which we have---ish):
$$ \frac{t}{\tau_d} + \log \tau_d = \frac{t}{\tau_r} + \log \tau_r. $$

**Substituting in $\tau_d = 26.0$ms and $t = 4.4$ms we get $\tau_r = 1.44$ms.**

P.S. The solution for the time to peak $t$ in terms of $\tau_r$ and $\tau_d$ is
$$ t = \frac{1}{\frac{1}{\tau_d} - \frac{1}{\tau_r}} \log \frac{\tau_r}{\tau_d} $$

```python
kernel = BiexponentialSynapticKernel(125, 1.4, 26.0)
_ = kernel.generate(100., 0.1)
```

```python
def plot_kernel_vs_data(ax=None):
    if ax is None:
        ax = plt.gca()

    for voltages, tr in zip(averaged_traces['voltages'], averaged_traces['traces']):
        ind = np.argmin(np.abs(voltages - 60.))

        peak_amplitude = tr[stim_time + 10:, ind].max()
        ax.plot(
            np.arange(0, (tr.shape[0] - stim_time) * 0.1, 0.1),
            tr[stim_time:, ind] / peak_amplitude,
            'k-'
        )

    ax.plot(
        kernel.time_supp + (foot_time - stim_time) * 0.1,
        kernel.kernel / kernel.kernel.max(),
        color=colors['gif'],
        label='Biexponential\nmodel',
        lw=4,
        alpha=0.8
    )
    ax.set_ylim(0, 1.05)
    ax.set_xlabel('Time (ms)')
    ax.set_ylabel('Normalized amplitude (AU)')
    ax.legend()

    pltools.hide_border('tr', trim=True, ax=ax)
```

```python
plot_kernel_vs_data()
```

The simplified model gives a good fit.

| Parameter | Value |
| :-------: | :---: |
| $\tau_r$  | 1.44ms |
| $\tau_d$  | 26.0ms |
| $g_{syn}$ | 0.3nS  |

```python
SAVE_PATH = './figs/ims/writeup/gaba_synapses/'

sample_minis_tr = 1

def subtract_baseline(arr, samples = 10):
    return np.copy(arr) - arr[:samples, :].mean(axis = 0)

bl_sub_sample_minis = subtract_baseline(mini_detectors[sample_minis_tr].minis)

spec_outer = gs.GridSpec(2, 4, wspace = 0.5, hspace = 0.5, top = 0.9, right = 0.95, left = 0.05)
spec_IV_extracted = gs.GridSpecFromSubplotSpec(1, 2, spec_outer[0, 2:4], width_ratios = [1, 0.4], wspace = 0.6)
spec_mini_tr = gs.GridSpecFromSubplotSpec(3, 1, spec_outer[1, 1], hspace = 0)

plt.style.use('./figs/scripts/writeup/publication_figure_style.dms')

plt.figure(figsize = (7, 5))

plt.subplot(spec_outer[1, 0])
plot_kernel_vs_data()

for i in range(3):
    plt.subplot(spec_mini_tr[i, :])
    if i == 0:
        plt.text(0.05, 0.9, 'TTX', ha = 'left', va = 'top', transform = plt.gca().transAxes)

    plt.plot(mini_detectors[i].t_vec, mini_detectors[i].I[:, 16], 'k-', lw = 0.5)
    plt.ylim(-12, 50)
    plt.xlim(500, 2500)

    if i == 0:
        pltools.add_scalebar(
            x_units = 'ms', y_units = 'pA', x_size = 200, y_size = 20,
            anchor = (0.7, 0.6), bar_space = 0, x_label_space = -0.05)
    else:
        plt.xticks([])
        plt.yticks([])
        pltools.hide_border()

plt.subplot(spec_outer[1, 2])
plt.text(
    0.95, 0, 'N = {}'.format(mini_detectors[sample_minis_tr].no_minis),
    ha = 'right', va = 'center', transform = plt.gca().transAxes
)
t_vec = np.arange(
    0,
    mini_detectors[sample_minis_tr].minis.shape[0] * mini_detectors[sample_minis_tr].dt,
    mini_detectors[sample_minis_tr].dt
)
plt.plot(
    np.tile(
        t_vec[:, np.newaxis],
        (1, mini_detectors[sample_minis_tr].no_minis)
    ),
    bl_sub_sample_minis,
    'k-', lw = 0.5, alpha = 0.05
)
plt.plot(
    t_vec, np.nanmean(bl_sub_sample_minis, axis = 1),
    '-', color=colors['gif'], lw=1, label = 'Mean'
)
pltools.add_scalebar(
    x_units = 'ms', y_units = 'pA', anchor = (0.35, 0.), bar_space = 0,
    x_size = 10, x_label_space = -0.03
)
plt.legend(loc = 'upper right')

cond_ax = plt.subplot(spec_outer[1, 3])
for sorted_maxes in [MiDe.amplitudes(5, sort = True) for MiDe in mini_detectors]:
    sorted_gs = I_to_g(sorted_maxes, 20, reversal)
    pctiles = np.linspace(0, 100, len(sorted_maxes))
    plt.plot(sorted_gs, pctiles, 'k-')
cond_ax.set_xticks([0, 0.5, 1])
plt.xlabel('Conductance (nS)')
plt.ylabel('Percentile')
pltools.hide_border('tr')

cond_ins = inset_axes(cond_ax, '35%', '40%', loc = 'center right')
for sorted_maxes in [MiDe.amplitudes(5, sort = True) for MiDe in mini_detectors]:
    sorted_gs = I_to_g(sorted_maxes, 20, reversal)
    pctiles = np.linspace(0, 100, len(sorted_maxes))
    cond_ins.plot(sorted_gs, pctiles, 'k-')
cond_ins.set_xlim(0.15, 0.4)
cond_ins.set_ylim(0, 100)

mark_inset(cond_ax, cond_ins, loc1 = 1, loc2 = 4, ls = '--', color = 'gray', lw = 0.5)

plt.subplot(spec_outer[0, :2])
plt.plot(
    averaged_traces['t_mats'][sample_tr['ind']][:, :-2],
    averaged_traces['traces'][sample_tr['ind']][:, :-2],
    'k-', lw = 0.5
)
plt.plot(
    sample_tr['x_fit'],
    sample_tr['y_fit'],
    '--',
    lw=2,
    color=colors['gif'], label = 'Monoexponential fit'
)
#plt.plot(kernel.time_supp + foot_time * 0.1, kernel.kernel, label='Synthetic IPSC', lw=4, alpha=0.5)
plt.ylim(-50, 175)
pltools.add_scalebar(
    x_units = 'ms', y_units = 'pA', x_size = 20, y_size = 50, bar_space = 0,
    x_on_left = False, anchor = (0.7, 0.6)
)
plt.legend(loc = 'lower right')

plt.subplot(spec_IV_extracted[:, 0])
plt.axhline(0, color = 'k', lw = 0.5)
plt.axvline(0, color = 'k', lw = 0.5)
for i in range(len(IV_data['amplitudes'])):
    plt.plot(
        IV_data['voltages'][i],
        IV_data['amplitudes'][i],
        'k-'
    )
x_linear = np.array([[-80, 1], [60, 1]])
plt.plot(
    x_linear[:, 0],
    np.dot(x_linear, np.array([IV_coeffs['slope'], IV_coeffs['intercept']])),
    '--', color=colors['gif'], lw=1, label = 'Linear fit'
)
plt.annotate(
    'Reversal:\n${:.1f}$mV'.format(reversal),
    (reversal, 1),
    (0.05, 0.5),
    xycoords = 'data',
    textcoords = 'axes fraction',
    arrowprops = {'arrowstyle': '->'}
)
plt.xlabel('$V$ (mV)')
plt.ylabel('$I$ (pA)')
plt.legend()

plt.subplot(spec_IV_extracted[:, 1])
plt.ylim(0, 35)
sns.swarmplot(y = averaged_traces['decay'], color = colors['gif'])
plt.xticks([])
plt.ylabel('$\\tau$ (ms)')
pltools.hide_border('trb')

if SAVE_PATH is not None:
    savefig(SAVE_PATH + 'gaba_synapses.png')

plt.show()
```

```python

```


# Constrained Optimization

All methods to compute paths or configurations solve constrained optimization problems. To use them, you need to learn to define constrained optimization problems. Some definitions:

* An objective defines either a sum-of-square cost term, or an equality constraint, or an inequality constraint in the optimization problem. An objective is defined by its type and its feature. The type can be `sos` (sum-of-squares), `eq`, or `ineq`, referring to the three cases.
* A feature is a (differentiable mapping) from the decision variable (the full path, or robot configuration) to a feature space. If the feature space is, e.g., 3-dimensional, this defines 3 sum-of-squares terms, or 3 inequality, or 3 equality constraints, one for each dimension. For instance, the feature can be the 3-dim robot hand position in the 15th time slice of a path optimization problem. If you put an equality constraint on this feature, then this adds 3 equality constraints to the overall path optimization problem.
* A feature is defined by the keyword for the feature map (e.g., `pos` for position), typically by a set of frame names that tell which objects we refer to (e.g., `pr2L` for the left hand of the pr2), optionally some modifiers (e.g., a scale or target, which linearly transform the feature map), and the set of configuration IDs or time slices the feature is to be computed from (e.g., `confs=[15]` if the feat is to be computed from the 15th time slice in a path optimization problem).
* In path optimization problems, we often want to add objectives for a whole time interval rather than for a single time slice or specific configuration. E.g., avoid collisions from start to end. When adding objectives to the optimization problem we can specify whole intervals, with `times=[1., 2.]`, so that the objective is added to each configuration in this time interval.
* Some features, especially velocity and acceleration, refer to a tuple of (consecutive) configurations. E.g., when you impose an acceleration feature, you need to specify `confs=[13, 14, 15]`. Or if you use `times=[1.,1.]`, the acceleration features is computed from the configuration that corresponds to time=1 and the two configurations *before*.
* All kinematic feature maps (that depend on only one configuration) can be modified to become a velocity or acceleration features. E.g., the position feature map can be modified to become a velocity or acceleration feature.
* The `sos`, `eq`, and `ineq` always refer to the feature map to be *zero*, e.g., constraining all features to be equal to zero with `eq`. This is natural for many features, esp. when they refer to differences (e.g. `posDiff` or `posRel`, which compute the relative position between two frames). However, when one aims to constrain the feature to a non-zero constant value, one can modify the objective with a `target` specification.
* Finally, all features can be rescaled with a `scale` specification. Rescaling changes the costs that arise from `sos` objectives. Rescaling also has significant impact on the convergence behavior for `eq` and `ineq` constraints. As a guide: scale constraints so that if they *would* be treated as squared penalties (squaredPenalty optim mode, to be made accessible) convergence to reasonable approximate solutions is efficient. Then the AugLag will also converge efficiently to precise constraints.

## Concrete features

The currently implemented feature keywords are:
  FS_none,
  FS_position,
  FS_positionDiff,
  FS_positionRel,
  FS_quaternion,
  FS_quaternionDiff,
  FS_quaternionRel,
  FS_pose,
  FS_poseDiff,
  FS_poseRel,
  FS_vectorX,
  FS_vectorXDiff,
  FS_vectorXRel,
  FS_vectorY,
  FS_vectorYDiff,
  FS_vectorYRel,
  FS_vectorZ,
  FS_vectorZDiff,
  FS_vectorZRel,
  FS_scalarProductZ,
  FS_gazeAt,
  FS_accumulatedCollisions,
  FS_jointLimits,
  FS_distance,
  FS_qItself,
  FS_aboveBox,
  FS_insideBox,
  FS_standingAbove,

The table is a bit outdated.. The expressions in the left of the table declare features. All of these are functions of a \emph{single} configuration---differences and accelerations can tie multiple configurations together. Arguments typically indicate objects (actually \emph{coordinate frame identifies}), which we write as o1, o2, etc. Another typical argument is a vector (e.g., offset) interpreted in the object's coordinate frame; e.g.\ v1=$(0,0,1)$ to indicate the object's z-axis.

| feature | dim | description |
|:---:|:---:|:---:|
  | `pos(o1)` | 3 | 3D position of o1 in world coordinates |
  | `posRel(o1,o2)` | 3 | 3D position of o1 in o2 coordinates |
  | `posDiff(o1,o2)` | 3 | difference of 3D positions of o1 and o2 in world coordinates |
  | `vec(o1,v1)` | 3 | 3D vector v1 transformed from o1 to world |
  | `vecRel Diff(o1,v1,o2)` | 3 | ... |
  | `prod(o1,v1,o2,v2)` | 1 | scalar product of v1 (in o1) and v2 (in o2) |
  | `quat(o1)` | 4 | 4D quaternion of o1 in world coordinates\footnote{There is ways to handle the invariance w.r.t.\ quaternion sign properly.} |
  | `quatRel Diff(o1,o2)` | 4 | ... |
  | `pose(o1)` | 7 | 7D pose of o1 in world coordinates |
  | `poseRel Diff(o1,o2)` | 7 | ...|
| `dist(o1,o2)` | 1 | NEGATIVE distance between convex meshes o1 and o2, positive for penetration|
  | `above(o1,o2)` | ? | when all negative, o1 is above (inside support of) o2 |
  | `placed(o1,o2)` | ? | o1 is above and straightly placed on o2 |
| `coll` | ? | when negative, nothing is colliding |
  | `limits` | ? | when negative, all robot joint limits are ok |

## Demo of features in Inverse Kinematics

Let's setup a standard configuration. (Lock the window with "Always on Top".)

```python
import sys
sys.path.append('../src/ry')
```

```python
from libry import *
from numpy import *
```

```python
K = Configuration()
D = K.camera()
K.addFile('../rai-robotModels/pr2/pr2.g')
K.addFile('../test/kitchen.g')
```

For simplicity, let's add a frame that represents goals

```python
K.addFrame("goal", "", "shape:marker size:[.3] color:[.5 1 1]" )
K.setFrameState([1,.5,1,1,0,0,0], ["goal"])
x0 = K.getFrameState().flatten().tolist()
```

The `posDiff` (position difference in world coordinates) between `pr2L` (the yellow blob in the left hand) and `goal` is equal to zero:

```python
K.setFrameState(x0)
IK = K.komo_IK()
IK.addObjective(type='eq', feature='positionDiff', frames=['pr2L', 'goal'])
IK.optimize()
IK.getConfiguration(0)
```

```python
IK.getReport()
```

    [{'feature': 'Transition',
      'scale': 1.0,
      'sos_value': 3.6657745735746645,
      'type': 'sos'},
     {'description': 'QuaternionNorms',
      'scale': 3.0,
      'sos_value': 0.0,
      'type': 'sos'},
     {'eq_sumOfAbs': 0.0003536309427615292,
      'feature': 'posDiff',
      'o1': 'pr2L',
      'o2': 'goal',
      'scale': 10.0,
      'type': 'eq'}]

The `posDiff` between `pr2R` and `goal` is equal to zero:

```python
K.setFrameState(x0)
IK = K.komo_IK()
IK.addObjective(type='eq', feature='positionDiff', frames=['pr2R', 'goal'])
IK.optimize()
IK.getConfiguration(0)
```

The relative position of `goal` in `pr2R` coordinates equals [0,0,-.2] (which is 20cm straight in front of the yellow blob)

```python
K.setFrameState(x0)
IK = K.komo_IK()
IK.addObjective(type='eq', feature='positionRel', frames=['goal','pr2R'], target=[0,0,-.2])
IK.optimize()
IK.getConfiguration(0)
```

The distance between `pr2R` and `pr2L` is zero:

```python
K.setFrameState(x0)
IK = K.komo_IK()
IK.addObjective(type='eq', feature='distance', frames=['pr2L','pr2R'])
IK.optimize()
IK.getConfiguration(0)
```

```python
IK.getReport()
```

    [{'feature': 'Transition',
      'scale': 1.0,
      'sos_value': 0.2077689310624419,
      'type': 'sos'},
     {'description': 'QuaternionNorms',
      'scale': 3.0,
      'sos_value': 0.0,
      'type': 'sos'},
     {'eq_sumOfAbs': 0.01132489500129505,
      'feature': 'dist',
      'o1': 'pr2L',
      'o2': 'pr2R',
      'scale': 100.0,
      'type': 'eq'}]

The 3D difference between the z-vector of `pr2R` and the z-vector of `goal`:

```python
K.setFrameState(x0)
IK = K.komo_IK()
IK.addObjective(type='eq', feature='vectorZDiff', frames=['pr2R', 'goal'])
IK.optimize()
IK.getConfiguration(0)
```

The scalar product between the z-vector of `pr2R` and the z-vector of `goal` is zero:

```python
K.setFrameState(x0)
IK = K.komo_IK()
IK.addObjective(type='eq', feature='scalarProductZ', frames=['pr2R', 'goal'])
IK.optimize()
IK.getConfiguration(0)
```

etc etc

## Graph-structured NLP

We have a set $\{x_1,..,x_n\}$ of $n$ variables, each with domain $x_i \in \mathbb{R}^{d_i}$.

We have a set $\{\phi_1,..,\phi_K\}$ of $K$ features, each of which is a differentiable mapping $\phi_k: x_{\pi_k} \mapsto \mathbb{R}^{D_k}$. Here $\pi_k \subseteq \{1,..,n\}$ indicates a subset of the variables, and $x_{\pi_k}$ is the value of this subset of variables. In addition,  $\rho_k \in\{\texttt{ineq, eq, sos}\}$ indicates for each feature whether it is an inequality, an equality, or a sum-of-square cost objective.

This defines the mathematical program
\begin{align}
  \min_{x_1,..,x_n} \sum_{k : \rho_k=\texttt{sos}} \phi_k(x_{\pi_k})^T \phi_k(x_{\pi_k})
  ~\text{s.t.}~ \mathop\forall_{k : \rho_k=\texttt{ineq}} \phi_k(x_{\pi_k}) \le 0 ~,\quad
  \mathop\forall_{k : \rho_k=\texttt{eq}} \phi_k(x_{\pi_k}) = 0 ~,\quad
\end{align}

In summary, a graph-structured NLP is specified as a tuple $(n, \{d_i\}, K, \{\phi_k\}, \{\pi_k\}, \{\rho_k\})$.

Note: I assumed here that all cose features are sum-of-squares... That's not necessary, of course. But typically the most relevant case; and we only need the Jacobian $\partial_x \phi$ to approximate the Gauss-Newton Hessian.

Note: If all features are cost objectives, then
\begin{align}
  \prod_k \exp\{ -  \phi_k(x_{\pi_k}) \}
\end{align}
defines a factor graph. If all features are actually sum-of-square
costs and we use Gaussian belief approximations, solving the problem
(inference, which now also is MAP, which is also minimization) can be
tackled with Gaussian message passing, as in GTSAM. (But in principle,
this does not exploit more structure than a sparse matrix method to
compute Newton steps...) In this sense, a graph-structured NLP is a
generalization of factor graphs.

Note: If the structure is $k$-order Markovian (meaning that the variables $x_i$ can be sorted into a chain and cliques $\pi_k$ contain only $k+1$  consecutive variables), then the Hessian of the Augmented Lagrangian becomes banded diagonal, and Newton steps can be computed in $O(n)$. This is what I call KOMO ($k$-order Markov Optimization). BUT FOR SMALL $n$ ONE CAN TACKLE THE PROBLEM ALSO WITHOUT THE $k$-ORDER ASSUMPTION.

## Generic modifiers for all features: linear rescaling, difference/velocity, acceleration

We assume each variable $x_i$ is a world configuration, that is, the vector of all DOFs of the world in a specific time slice.

%\paragraph{\texttt{order} modifier: velocities and acceleration}

Most features are defined as functions of a single configuration $x_i$. E.g. $\phi(x_i)$ could be the 3D position of endeffector in configuration $x_i$. However, for such features we automatically also define the difference and ``acceleration'' as
\begin{align}
  \phi(x_1,x_2) &\equiv \frac{1}{\tau}(\phi(x_2) - \phi(x_1)) \\
  \phi(x_1,x_2,x_3) &\equiv \frac{1}{\tau^2}(\phi(x_1) + \phi(x_3) - 2 \phi(x_2))
\end{align}
When in the specification such a $\phi$ is indicated to depend on two configurations, it automatically means that we impose an objective on the velocity/difference; if on three configurations, then we impose an objective on the acceleration.

%\paragraph{\texttt{scale, target} modifiers: linear rescaling}

The (in)equalities always refer to zero. To constrain $\phi(x) = \textit{target}$ you can modify any defined feature by specifying a \texttt{target} and/or \texttt{scale}, which implies
\begin{align}
  \tilde \phi(x) \gets \texttt{scale} \cdot (\phi(x) - \texttt{target})
\end{align}

## Example Specification of a robot problem

We have two configurations $x_1, x_2$; here is a specification of features that define the NLP for grasping an object in $x_1$ and placing it on a tray in $x_2$:
```
(1,2): { sos, trans, scale:0.1 }
(1,2): { eq, pose(object, hand) }
  (1):   { eq, dist(object, hand) }
  (1):   { ineq, coll }
  (2):   { ineq, coll }
  (2):   { eq, dist(object, tray) }
  (2):   { ineq, above(object, tray) }
  (2):   { eq, vec(object, (0,0,1)), target:(0,0,1) }
```
Here, every tuple $(i,j)$ indicates an objective that depends on variables $x_i,x_j$. The @{..}@ desribes the objective. The objectives read:
* We penalize configuration difference between $x_1, x_2$ with a weak sum-of-square term
* The relative pose of object in hand has zero difference between configurations $x_1, x_2$
* The distance between object and hand is zero in configuration $x_1$
* There are no collisions in configuration $x_1$
* There are no collisions in configuration $x_2$
* The object touches tray (distance=zero) in config 2
* The object is above the tray in config 2
* The object's z-axis equals the world z-axis in config 2

```python

```