$\newcommand{\xv}{\mathbf{x}}
 \newcommand{\wv}{\mathbf{w}}
 \newcommand{\yv}{\mathbf{y}}
 \newcommand{\zv}{\mathbf{z}}
 \newcommand{\Chi}{\mathcal{X}}
 \newcommand{\R}{\rm I\!R}
 \newcommand{\sign}{\text{sign}}
 \newcommand{\Tm}{\mathbf{T}}
 \newcommand{\Xm}{\mathbf{X}}
 \newcommand{\Zm}{\mathbf{Z}}
 \newcommand{\I}{\mathbf{I}}
 \newcommand{\muv}{\boldsymbol\mu}
 \newcommand{\Sigmav}{\boldsymbol\Sigma}
$
### ITCS6155

# Reinforcement Learning

Along with supervised and unsupervised learning, reinforcement learning is one of the interesting fields of machine learining.
As we shortly discussed in the first week of class, reinforcement learning is different from other machine learning paradigms as summarized below:

- only reward signal as feedback,
- the feedback can be delayed,
- sequential data,
- interaction based on the actions taken.

The reinforcement learning resembles huuman learning or animal training that treats reward good behavior.
When the series of actions end up with good results, we can **reinforce** those actions by giving some rewards.

## Applictions

There are many possible applications:

- walking robot,
- playing boardgames,
- controlling joysticks for video games,
- smart thermostat,
- stock trading,
- music personalization,
- recommendation system,
- marketing,
- product delivery,
- and so on...

## Terminology

- Reward: $R_t$ is a scalar feedback signal at time $t$. It indicates how well agent is doing at the time.

- State: $S_t$ represents what happens currently and next.

- Action: $A_t$ is how an agent affects to the environment that can change the state.

- Observation: $O_t$ is what an agent recognizes the world for the state $S_t$.

- Policy: A function that determines an agent's behavior or a function that selects its action, $\pi(S_t)$.

 Q: Think about the rewards for the example applications above, and answer the possible rewards for them.
- walking robot: Positive reward for continuous successful walking; negative reward if it falls down.
- playing boardgames: Positive reward for wining; negative reward for loosing the game.
- controlling joysticks for video games: Positive reward for continuing the game; negative reward for loosing to oponent.
- smart thermostat: Not sure!
- stock trading: Positive reward for good prediction of profit; negative reward for loosing money.
- music personalization: Positive reward when user doesn't change the recommended music or play a lot from the recommended music; negative reward for not playing from recommendation or giving low ratting to the recommended music.
- recommendation system: Same as the `music personalization`.
- marketing: Positive reward for sales growth; negative reward for the oposit.
- product delivery: Positive reward for confirmation of product delivery; negative reward for not getting notification on-time.

## Modeling

How can we simplify the listed problems to solve? We have the sets of decision making problems or control problems.
With the assumption of **Markov property**, we can model the problems as Makrov Decision Process (MDP).

Here, the definition of Markov property is as follows,

**Definition**: A state $S_t$ is Markov if and only if

$$P [ S_{t+1} | S_t ] = P[ S_{t+1} | S_1, ..., S_t ].$$

That is, we can say the state transition model has Markov property when the future is dependent only on the current state not the past.

### Markov Decision Processes
So, the MDP can be defined as follows and the example state transition diagram has the transition probabilities and rewards along with discounting factor. The `discounting factor` ensures the mathematical simplicity, and it allows the modeling the uncertainty of the future as well.

 $ ( S, A, P, R, \gamma ) $
* $S$ : a finite set of states
* $A$ : a finite set of actions
* $P$ : a state transition probability
* $P^a_{ss^\prime} = P [ S_{t+1} = s^\prime | S_t = s, A_t = a ]$
* $R$ : a reward function
* $\gamma$ : a discount factor

 MDP with 3 states and 2 actions (wikipedia)

### Policy

The policy $\pi(S_t)$ is a function that maps from a state to an action. Thus, it can be easily represented by a probability:

$$ \pi(a | s) = P[A_t = a | S_t = s]. $$

Or, it can be a machine learning model to generate such probability for a certain behavior of an agent.

## Goal

Now with the MDP model, what we want to do is find a policy $\pi$ that `maximizes the long-term rewards`. The long-term rewards can be modeled as the sum of rewards. Here, considering the different importance on the actions (the current reward is not equivalent to future rewards), we can define the **Return**, $G_t$,

$$ G_t = R_{t+1} + \gamma R_{t+2} + \cdots = \sum_{k=0}^{\infty} \gamma^k R_{t+k+1}. $$

From this, we can model the long-term value of a certain state $s$ with the return:

$$ V(s) = E [ G_t | S_t = s ]. $$

We, we set our objective to maximize the return. Developing the optimal policy can be made by defining the evaluation function of state and action pair, the `Q funciton`:

$$
\begin{align}
    Q^\pi(s, a) &= E_\pi [ G_t | S_t = s, A_t = a ] \\
                &= E_\pi [ R_{t+1} + \gamma Q^\pi (S_{t+1}, A_{t+1}) | S_t = s, A_t = a ]
\end{align}
$$

## How to reach the goal? / How to find the optimal policy?

There are several ways to reach the goal such as `value iteration`, `policy iteration`, `linear programming`, and `temporal difference learning`.

First, let us examine the bootstraping representation of the value function to solve the problem.
This is called **Bellman equation**.

$$
\begin{aligned}
  V(s) &= E [ G_t | S_t = s ] \\
       &= E [ R_{t+1} + \gamma R_{t+2} + \cdots | S_t = s ] \\
       &= E [ R_{t+1} + \gamma ( R_{t+2} + \gamma R_{t+3} + \cdots ) | S_t = s ] \\
       &= E [ R_{t+1} + \gamma G_{t+1} | S_t = s ] \\
       &= E [ R_{t+1} + \gamma V(s_{t+1}) | S_t= s ]\\
      \\
  V(s) &= R(s) + \gamma \sum_{s^\prime \in S} P_{ss^\prime} V(s^\prime)
\end{aligned}
$$

With matrix representation, we can rewrite this as follows:

$$
 \begin{align}
  \begin{pmatrix}
  V_1 \\
  V_2 \\
  \vdots \\
  V_n
 \end{pmatrix} &=
  \begin{pmatrix}
  R_1 \\
  R_2 \\
  \vdots \\
  R_n
 \end{pmatrix} +
 \gamma
 \begin{pmatrix}
  P_{11} & P_{12} & \cdots & P_{1n} \\
  P_{21} & P_{22} & \cdots & P_{2n} \\
  \vdots  & \vdots  & \ddots & \vdots  \\
  P_{n1} & P_{n2} & \cdots & P_{nn}
 \end{pmatrix}
  \begin{pmatrix}
  V_1 \\
  V_2 \\
  \vdots \\
  V_n
 \end{pmatrix}
 \\
 \\
 V &= R + \gamma PV
 \end{align}
$$

This can be directly solved with

$$ V = (I - \gamma P)^{-1} R. $$

Considering the computational complexity with inverse matrix, however, we can consider iterative solutions such as `dynamic programming`, `Monte-Carlo evaluation`, or `temporal difference learning`.

# Value Iteration

* For state transition from $i$ to $j$ with action $k$
$$ V^{n+1} (s_i) = \max_k \Big[ R_i + \gamma \sum_{j=1}^N P^k_{ij} V^n (s_j) \Big] $$

* Value Iteration:

$$
\begin{align}
n=0 \quad \forall i, \quad &V^1 (s_i) = \max_k \Big[ R_i + \gamma \sum_{j=1}^N P^k_{ij} V^0 (s_j) \Big] \\
n=1 \quad \forall i, \quad &V^2 (s_i) = \max_k \Big[ R_i + \gamma \sum_{j=1}^N P^k_{ij} V^1 (s_j) \Big] \\
n=2 \quad \forall i, \quad &V^3 (s_i) = \max_k \Big[ R_i + \gamma \sum_{j=1}^N P^k_{ij} V^2 (s_j) \Big] \\
\vdots
\end{align}
$$

* Convergence can be tested as follows:

$$ \max_i \Big| V^{n+1}(s_i) - V^n(s_i) \Big| \lt \epsilon. $$

* Computing values for all $i$'s with Dynamic Programming.

Here follows an example from Wikipedia.

Let $\gamma = 1$. The values for each state can be iteratively updated as the following table.

| t | $V^\pi_{S_0}$ | $V^\pi_{S_1}$ | $V^\pi_{S_2}$ |
|:--|:-------------|:-------------|:-------------|
| 0 |      0       |       0.0      |      0       |
| 1 |      0       |       3.5      |    0.75      |
| 2 |      0       |       3.85     |    0.75      |
| 3 |      0.75       |      4.035      |    1.155      |
| 4 |      1.155       |      4.6595      |    1.5975      |
| 5 |      1.5975     |      5.09395   |    2.08335      |

In Python, we can write the codes.

```python
import collections
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
```

```python
gamma = 1.0
n_states = 3
n_actions = 2

# value vector
V = np.zeros(n_states).reshape((-1, 1))
# Reward dictionary
R = {(1,2,0): -1, (0,1,0): 5}
# transition probability: key (action, from, to)
P = {(0,0,0): 0.5, (0,0,2): 0.5, (1,0,2): 1.0,
     (0,1,0): 0.7, (0,1,1): 0.1, (0,1,2): 0.2,
     (1,1,1): 0.95, (1,1,2): 0.05,
     (0,2,0): 0.4, (0,2,2): 0.6,
     (1,2,0): 0.3, (1,2,1): 0.3, (1,2,2): 0.4 }

def Rf(s, sn, a):
    if (a, s, sn) in R:
        return R[(a,s,sn)]
    return 0

def Pf(s, sn, a):
    if (a, s, sn) in P:
        return P[(a,s,sn)]
    return 0.

""" with no matrix"""
def Rs(s, a):
    r = 0
    for j in range(3):
        if (a, s, j) in P:
            r += P[(a, s, j)] * Rf(a,s, j)
    return r

def Pfv(s, a):
    ret = np.zeros(3).reshape((-1, 1))
    for j in range(3):
        if (a, s, j) in P:
            ret[j] = P[(a, s, j)]
    return ret

```

```python
Rmat = np.zeros((n_states, n_actions))
Pmat = np.zeros((n_actions, n_states, n_states))

for i in range(n_states):
    for k in range(n_actions):
        for j in range(n_states):
            if (k, i, j) in P:
                Rmat[i, k] += P[(k, i, j)] * Rf(i, j, k)

for k, i, j in P.keys():
    Pmat[k,i,j] = P[(k, i, j)]

for t in range(5):
    print("t=", t, "\tV = ", V.T)
    # Bellman
    newV = Rmat + gamma * np.sum(Pmat @ V, axis=2).T
    V[:] = np.max(newV, axis=1, keepdims=True)

print("t=", t+1, "\tV = ", V.T)
```

    t= 0 	V =  [[0. 0. 0.]]
    t= 1 	V =  [[0.  3.5 0. ]]
    t= 2 	V =  [[0.   3.85 0.75]]
    t= 3 	V =  [[0.75  4.035 1.155]]
    t= 4 	V =  [[1.155  4.6595 1.5975]]
    t= 5 	V =  [[1.5975  5.09395 2.08335]]

```python
gamma=1
for t in range(50):
    print("t=", t, "\tV = ", V.T)
    for i in range(3):
        # Bellman
        V[i] = max(Rs(i, 0) + gamma * np.sum(Pfv(i, 0) * V),
                   Rs(i, 1) + gamma * np.sum(Pfv(i, 1) * V))
print("t=", t+1, "\tV = ", V.T)
```

    t= 0 	V =  [[1.5975  5.09395 2.08335]]
    t= 1 	V =  [[2.08335  4.94342  2.941371]]
    t= 2 	V =  [[2.941371   4.84331755 3.51195497]]
    t= 3 	V =  [[3.51195497 4.77674942 3.8913933 ]]
    t= 4 	V =  [[3.8913933  4.73248161 4.1437198 ]]
    t= 5 	V =  [[4.1437198  4.70304352 4.31151691]]
    t= 6 	V =  [[4.31151691 4.68346719 4.423102  ]]
    t= 7 	V =  [[4.423102   4.67044893 4.49730608]]
    t= 8 	V =  [[4.49730608 4.66179179 4.54665179]]
    t= 9 	V =  [[4.54665179 4.65603479 4.57946669]]
    t= 10 	V =  [[4.57946669 4.65220639 4.6012886 ]]
    t= 11 	V =  [[4.6012886  4.6496605  4.61580017]]
    t= 12 	V =  [[4.61580017 4.64796748 4.62545036]]
    t= 13 	V =  [[4.62545036 4.64684162 4.63186774]]
    t= 14 	V =  [[4.63186774 4.64609293 4.6361353 ]]
    t= 15 	V =  [[4.6361353  4.64559505 4.63897322]]
    t= 16 	V =  [[4.63897322 4.64526396 4.64086044]]
    t= 17 	V =  [[4.64086044 4.64504378 4.64211544]]
    t= 18 	V =  [[4.64211544 4.64489736 4.64295002]]
    t= 19 	V =  [[4.64295002 4.6448     4.64350501]]
    t= 20 	V =  [[4.64350501 4.64473525 4.64387408]]
    t= 21 	V =  [[4.64387408 4.64469219 4.64411952]]
    t= 22 	V =  [[4.64411952 4.64466356 4.64428273]]
    t= 23 	V =  [[4.64428273 4.64464452 4.64439126]]
    t= 24 	V =  [[4.64439126 4.64463185 4.64446344]]
    t= 25 	V =  [[4.64446344 4.64462343 4.64451144]]
    t= 26 	V =  [[4.64451144 4.64461783 4.64454336]]
    t= 27 	V =  [[4.64454336 4.64461411 4.64456458]]
    t= 28 	V =  [[4.64456458 4.64461163 4.6445787 ]]
    t= 29 	V =  [[4.6445787  4.64460999 4.64458808]]
    t= 30 	V =  [[4.64458808 4.64460889 4.64459433]]
    t= 31 	V =  [[4.64459433 4.64460816 4.64459848]]
    t= 32 	V =  [[4.64459848 4.64460768 4.64460124]]
    t= 33 	V =  [[4.64460124 4.64460736 4.64460307]]
    t= 34 	V =  [[4.64460307 4.64460714 4.64460429]]
    t= 35 	V =  [[4.64460429 4.644607   4.6446051 ]]
    t= 36 	V =  [[4.6446051  4.6446069  4.64460564]]
    t= 37 	V =  [[4.64460564 4.64460684 4.644606  ]]
    t= 38 	V =  [[4.644606   4.6446068  4.64460624]]
    t= 39 	V =  [[4.64460624 4.64460677 4.6446064 ]]
    t= 40 	V =  [[4.6446064  4.64460675 4.64460651]]
    t= 41 	V =  [[4.64460651 4.64460674 4.64460658]]
    t= 42 	V =  [[4.64460658 4.64460673 4.64460662]]
    t= 43 	V =  [[4.64460662 4.64460673 4.64460665]]
    t= 44 	V =  [[4.64460665 4.64460672 4.64460668]]
    t= 45 	V =  [[4.64460668 4.64460672 4.64460669]]
    t= 46 	V =  [[4.64460669 4.64460672 4.6446067 ]]
    t= 47 	V =  [[4.6446067  4.64460672 4.6446067 ]]
    t= 48 	V =  [[4.6446067  4.64460672 4.64460671]]
    t= 49 	V =  [[4.64460671 4.64460672 4.64460671]]
    t= 50 	V =  [[4.64460671 4.64460672 4.64460671]]

# Policy Iteration

* Start with an initial policy $\pi^0$.
* Iteratively,
    * Evaluate policy ($V^n(x) = V^{\pi^n}(x)$):
    $$ V^n(s) = R(s, a = \pi^n(s)) + \gamma \sum_{s^\prime} P(s^\prime | s, a = \pi^n(s)) V^{n}(s^\prime) $$
    * Improve policy:
    $$ \pi_{n+1}(s) = \arg \max_a \Big[ R(s, a) + \gamma \sum_{s^\prime} P(s^\prime | s, a) V^{n}(s^\prime) \Big] $$

* Stop condition:
    * Policy does not change anymore (for about 10 iterations)
    * The changes in evaluation of values are minor

```python
n_iter = 10
gamma = 1.
V = np.zeros(3)

pi = np.random.randint(2, size=n_states)

for n in range(n_iter):
    print("n = ", n, "\t", pi, end="")

    # evaluate
    for s in range(n_states):
        V[s] = Rs(s, pi[s]) + gamma * np.sum([Pf(s, sn, pi[s]) * V[sn] for sn in range(n_states)])
    print("\t", V)

    # improve
    for s in range(n_states):
        pi[s] = np.argmax([Rs(s, a) + gamma * np.sum([Pf(s, sn, a) * V[sn] for sn in range(n_states)])  for a in range(n_actions)] )
```

    n =  0 	 [1 0 0]	 [0. 0. 0.]
    n =  1 	 [0 0 0]	 [0. 0. 0.]
    n =  2 	 [0 0 0]	 [0. 0. 0.]
    n =  3 	 [0 0 0]	 [0. 0. 0.]
    n =  4 	 [0 0 0]	 [0. 0. 0.]
    n =  5 	 [0 0 0]	 [0. 0. 0.]
    n =  6 	 [0 0 0]	 [0. 0. 0.]
    n =  7 	 [0 0 0]	 [0. 0. 0.]
    n =  8 	 [0 0 0]	 [0. 0. 0.]
    n =  9 	 [0 0 0]	 [0. 0. 0.]

# Linear Programming

                                                                                [Manne '60]

\begin{equation*}
\begin{aligned}
& \underset{x}{\text{minimize}}
& & \sum_x V(x) \\
& \text{subject to}
& & V(x) \geq R(x, u) + \gamma \sum_{x^\prime} P(x^\prime | x, u) V^{n}(x^\prime) &\forall x, a.
\end{aligned}
\end{equation*}

# Temporal Difference Learning

When there are large number of states, the memory requirement grows exponentially.
*Temporal difference (TD) learning* considers that the agent knows only the partial information of the MDP.
With only current and next state transition and without any model transition probability, TD lets the agent explore the environment to examine the random policy.
With an estiate of the value function $V(s)$, $\hat{V}(s)$,

   $$
        \begin{align}
        V(s_t) &= R_{t+1} + \gamma V(s_{t+1}) \\
        V(s_t) &\sim R_{t+1} + \gamma \hat{V}(s_{t+1}) \\
        \Rightarrow \quad \delta_t &= R_{t+1} + \gamma \hat{V}(s_{t+1}) - V(s_t).
        \end{align}
    $$

Here, $\delta$ represents the *temporal diffrence error*.
We can use this error as a gradient to update the value estimation.

$$
    \begin{align}
        V(s_t) &\leftarrow R_{t+1} + \alpha \delta_t \\
        V(s_t) &\leftarrow R_{t+1} + \alpha (R_{t+1} + \gamma \hat{V}(s_{t+1}) - V(s_t))
    \end{align}
$$

## Example Model

```python
# model (problem to solve)
# a: 0 - loop  1 - to next state
# r:       +1                 +0      +100 when reaches 4
#
n_states = 5
n_actions = 2

def transition(s, a):
    s1 = s + a
    if s == s1:
        r = 1
    elif s1 == n_states-1:
        r = 100
    else:
        r = 0
    return s1, r

def pause_print():
    sys.stdout.write('\r')
    sys.stdout.write("\033[F")
    sys.stdout.flush()
    time.sleep(0.5)
```

```python
# Example TD(0) Update
alpha = 1.0
gamma = 0.9
V = np.zeros(n_states)  # 5 states
pi = np.random.randint(n_actions, size=n_states)

np.set_printoptions(precision=2)
for k in range(2):
    if k == 1:
        pi = np.ones(n_states, dtype=np.int)
    print("Policy:", pi)
    # Evaluation of the random policy
    for e in range(20):
        s = 0

        print("\r\tTraj: ", s, end=" ")
        for step in range(10):
            a = pi[s]
            s1, r1 = transition(s, a)
            V[s] += alpha * (r1 + gamma * V[s1] - V[s])
            s = s1

            print(a, s, end=" ")
            if s == n_states-1:
                break

        print("\t", V, end="\n")
        #pause_print()
    print()
```

    Policy: [1 0 0 0 1]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [0.   6.13 0.   0.   0.  ]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [5.51 8.5  0.   0.   0.  ]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [7.65 9.42 0.   0.   0.  ]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [8.48 9.77 0.   0.   0.  ]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [8.8  9.91 0.   0.   0.  ]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [8.92 9.97 0.   0.   0.  ]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [8.97 9.99 0.   0.   0.  ]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [8.99 9.99 0.   0.   0.  ]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]
    	Traj:  0 1 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 	 [ 9. 10.  0.  0.  0.]

    Policy: [1 1 1 1 1]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [  9.   0.   0. 100.   0.]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [  0.   0.  90. 100.   0.]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [  0.  81.  90. 100.   0.]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]
    	Traj:  0 1 1 1 2 1 3 1 4 	 [ 72.9  81.   90.  100.    0. ]

# Control Problems and Q function

For control problems, we have defined Q function above to evaluate the state and action altogether.
Updating the Q values with TD learning is similar to previous update with two different considerations.
First, we update the Q with assumption that we follow a certain behavior policy. Thus, we call this as *on-policy control*, or **SARSA**.

$$
    Q(s_t, a_t) \leftarrow Q(s_t, a_t) + \alpha ( R_{t+1} + \gamma Q(s_{t+1}, a_{t+1}) - Q(s_t, a_t))
$$

Next, without making assumption of behavior policy, we can explore other possible policies to update the Q. We call this as *off-policy control*, or **Q-learning**.

$$
    Q(s_t, a_t) \leftarrow Q(s_t, a_t) + \alpha ( R_{t+1} + \gamma \max_a Q(s_{t+1}, a) - Q(s_t, a_t))
$$

Here describes the psedocode for each algorithm.

**[Algorithm: TD Learning]**

**[Algorithm: SARSA]**

**[Algorithm: Q-learning]**

# Choosing an Action

Picking an action can be simple by selecting one with maximum Q value, so *Greedy!*.

$$
a^* = \arg \max_a Q(S_t, a)
$$

However, this can cause limited experience to develop good Q estimation, and eventually a good policy.
Without new data, greedy action selection will repeat the same actions, or repeatedly *exploit* your current knowledge. Thus, you need to *explore* other non-greedy actions to increase the experience to improve the Q estimation.

This is called "exploration-exploitation dilemma."

One of the way for this dilemma is $\epsilon$-greedy action selection. With a parameter $\epsilon \in [0, 1]$, we can control the exploration and exploitation level. When $\epsilon = 0$, the actions are selected in greedy manner, but when $\epsilon = 1$, the actions are selected randomly.

```python
# action selection
def greedy(Q, s):
    return np.argmax(Q[s])  # greedy action selection

def e_greedy(Q, s, e):
    if np.random.rand() < e:
        return np.random.randint(n_actions)
    else:
        return greedy(Q,s)
```

```python
# SARSA example
alpha = 1.  # learning rate
gamma = 0.9 # discount factor
epsilon = 0.1

# tabular approximation
Q = np.random.rand(n_states, n_actions)  # 5 states and 2 actions

for e in range(20):
    s = 0
    a = e_greedy(Q, s, epsilon)  # greedy action selection

    for step in range(100):
        s1, r1 = transition(s, a)
        a1 = e_greedy(Q, s, epsilon)

        # TODO: update Q table here!
        Q[s, a] += alpha * (r1 + gamma * Q[s1, a1] - Q[s,a])

        #print("s: ", s, "a: ", a, "s1:", s1, "a1:", a1, "Q: ", Q[s,a])
        s, a = s1, a1
        if s == n_states-1:
            break

print("Final Q: ", Q)
print("Policy:", np.argmax(Q, axis=1))
```

    Final Q:  [[  2.52  73.17]
     [ 74.32  81.47]
     [ 82.47  90.52]
     [ 91.52 100.57]
     [  0.42   0.64]]
    Policy: [1 1 1 1 1]

```python
# Q-learning example
alpha = 1.  # learning rate
gamma = 0.9 # discount factor
epsilon = 0.1

# tabular approximation
Q = np.random.rand(n_states, n_actions)  # 5 states and 2 actions

for e in range(20):
    s = 0

    for step in range(100):
        a = e_greedy(Q, s, epsilon)  # greedy action selection
        s1, r1 = transition(s, a)

        # TODO: update Q table here!
        Q[s, a] += alpha * (r1 + gamma * np.max(Q[s1, :]) - Q[s,a])

        print("s: ", s, "a: ", "s1:", s1, "a1:", a1, "Q: ", Q[s,a])
        s = s1
        if s == n_states-1:
            break

print("Final Q: ", Q)
print("Policy:", np.argmax(Q, axis=1))
```

    s:  0 a:  s1: 1 a1: 1 Q:  0.7561608580140299
    s:  1 a:  s1: 1 a1: 1 Q:  1.75616085801403
    s:  1 a:  s1: 1 a1: 1 Q:  2.580544772212627
    s:  1 a:  s1: 1 a1: 1 Q:  3.3224902949913644
    s:  1 a:  s1: 1 a1: 1 Q:  3.990241265492228
    s:  1 a:  s1: 1 a1: 1 Q:  4.591217138943005
    s:  1 a:  s1: 1 a1: 1 Q:  5.132095425048704
    s:  1 a:  s1: 1 a1: 1 Q:  5.618885882543834
    s:  1 a:  s1: 2 a1: 1 Q:  0.7505308201058045
    s:  2 a:  s1: 3 a1: 1 Q:  0.42173196249345
    s:  3 a:  s1: 3 a1: 1 Q:  1.4217319624934501
    s:  3 a:  s1: 3 a1: 1 Q:  2.2795587662441053
    s:  3 a:  s1: 3 a1: 1 Q:  3.0516028896196947
    s:  3 a:  s1: 3 a1: 1 Q:  3.7464426006577254
    s:  3 a:  s1: 3 a1: 1 Q:  4.371798340591953
    s:  3 a:  s1: 3 a1: 1 Q:  4.9346185065327575
    s:  3 a:  s1: 3 a1: 1 Q:  5.441156655879482
    s:  3 a:  s1: 3 a1: 1 Q:  5.897040990291534
    s:  3 a:  s1: 3 a1: 1 Q:  6.307336891262381
    s:  3 a:  s1: 3 a1: 1 Q:  6.676603202136143
    s:  3 a:  s1: 3 a1: 1 Q:  7.008942881922528
    s:  3 a:  s1: 3 a1: 1 Q:  7.308048593730275
    s:  3 a:  s1: 3 a1: 1 Q:  7.577243734357248
    s:  3 a:  s1: 3 a1: 1 Q:  7.819519360921523
    s:  3 a:  s1: 3 a1: 1 Q:  8.037567424829371
    s:  3 a:  s1: 3 a1: 1 Q:  8.233810682346434
    s:  3 a:  s1: 3 a1: 1 Q:  8.410429614111791
    s:  3 a:  s1: 3 a1: 1 Q:  8.569386652700612
    s:  3 a:  s1: 3 a1: 1 Q:  8.71244798743055
    s:  3 a:  s1: 3 a1: 1 Q:  8.841203188687496
    s:  3 a:  s1: 3 a1: 1 Q:  8.957082869818747
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  5.0569972942894506
    s:  1 a:  s1: 1 a1: 1 Q:  6.0569972942894506
    s:  1 a:  s1: 1 a1: 1 Q:  6.451297564860505
    s:  1 a:  s1: 1 a1: 1 Q:  6.806167808374455
    s:  1 a:  s1: 1 a1: 1 Q:  7.12555102753701
    s:  1 a:  s1: 1 a1: 1 Q:  7.412995924783309
    s:  1 a:  s1: 1 a1: 1 Q:  7.671696332304978
    s:  1 a:  s1: 1 a1: 1 Q:  7.904526699074481
    s:  1 a:  s1: 1 a1: 1 Q:  8.114074029167032
    s:  1 a:  s1: 1 a1: 1 Q:  8.302666626250328
    s:  1 a:  s1: 1 a1: 1 Q:  8.472399963625296
    s:  1 a:  s1: 1 a1: 1 Q:  8.625159967262768
    s:  1 a:  s1: 1 a1: 1 Q:  8.762643970536491
    s:  1 a:  s1: 1 a1: 1 Q:  8.886379573482841
    s:  1 a:  s1: 1 a1: 1 Q:  8.997741616134558
    s:  1 a:  s1: 1 a1: 1 Q:  9.097967454521102
    s:  1 a:  s1: 1 a1: 1 Q:  9.188170709068991
    s:  1 a:  s1: 1 a1: 1 Q:  9.269353638162093
    s:  1 a:  s1: 1 a1: 1 Q:  9.342418274345883
    s:  1 a:  s1: 2 a1: 1 Q:  0.379558766244105
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  8.408176446911295
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 0 a1: 1 Q:  66.98755171443602
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 3 a1: 1 Q:  91.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 2 a1: 1 Q:  82.4661132276988
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 0 a1: 1 Q:  66.98755171443602
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    s:  0 a:  s1: 1 a1: 1 Q:  73.31950190492891
    s:  1 a:  s1: 2 a1: 1 Q:  81.46611322769878
    s:  2 a:  s1: 3 a1: 1 Q:  90.51790358633198
    s:  3 a:  s1: 4 a1: 1 Q:  100.57544842925776
    Final Q:  [[ 66.99  73.32]
     [  9.34  81.47]
     [ 82.47  90.52]
     [ 91.52 100.58]
     [  0.64   0.38]]
    Policy: [1 1 1 1 0]

# Maze Example

Now, let us solve some maze problems. To deal with various maze configuration, we define our class that can handle a file input of maze. Here the file is expected to have 3 letters: 'O' for open space, 'H' for obstacles (walls), and 'G' for a goal position. The example grid.txt file looks as follows.

Note: this bash command runs only on Linux or Mac. The shell is to show the content of file only. You can skip running the cell and creat a file with the same content.

```bash
%%bash
cat grid.txt
```

    OOOHOOOOO
    OOOHOOHOO
    OOOOOOHOO
    OOOOHHHOO
    OOHOOOOOH
    OOHOOOGOO
    OOOOOOOOO

Now, let us define the environment class. The environment constructor gets only one input, the filename. Along with some utility functions, the main simulation functions are init() and next() for initiating and running one step simulation. Please take a look at the example use and try them in a separate cell.

```python
# maze example

class GridWorld:
    """ Grid World environment
            there are four actions (left, right, up, and down) to move an agent
            In a grid, if it reaches a goal, it get 30 points of reward.
            If it falls in a hole or moves out of the grid world, it gets -5.
            Each step costs -1 point.

        to test GridWorld, run the following sample codes:

            env = GridWorld('grid.txt')

            env.print_map()
            print [2,3], env.check_state([2,3])
            print [0,0], env.check_state([0,0])
            print [3,4], env.check_state([3,4])
            print [10,3], env.check_state([10,3])

            env.init([0,0])
            print env.next(1)  # right
            print env.next(3)  # down
            print env.next(0)  # left
            print env.next(2)  # up
            print env.next(2)  # up

        Parameters
        ==========
        _map        ndarray
                    string array read from a file input
        _size       1d array
                    the size of _map in ndarray
        goal_pos    tuple
                    the index for the goal location
        _actions    list
                    list of actions for 4 actions
        _s          1d array
                    current state
    """
    def __init__(self, fn):
        # read a map from a file
        self._map = self.read_map(fn)
        self._size = np.asarray(self._map.shape)
        self.goal_pos = np.where(self._map == 'G')

        # definition of actions (left, right, up, and down repectively)
        self._actions = [[0, -1], [0, 1], [-1, 0], [1, 0]]
        self._s = None

    def get_cur_state(self):
        return self._s

    def get_size(self):
        return self._size

    def read_map(self, fn):
        grid = []
        with open(fn) as f:
            for line in f:
               grid.append(list(line.strip()))
        return np.asarray(grid)

    def print_map(self):
        print( self._map )

    def check_state(self, s):
        if isinstance(s, collections.Iterable) and len(s) == 2:
            if s[0] < 0 or s[1] < 0 or\
               s[0] >= self._size[0] or s[1] >= self._size[1]:
               return 'N'
            return self._map[tuple(s)].upper()
        else:
            return 'F'  # wrong input

    def init(self, state=None):
        if state is None:
            s = [0, 0]
        else:
            s = state

        if self.check_state(s) == 'O':
            self._s = np.asarray(state)
        else:
            raise ValueError("Invalid state for init")

    def next(self, a):
        s1 = self._s + self._actions[a]
        # state transition
        curr = self.check_state(s1)

        if curr == 'H' or curr == 'N':
            return -5
        elif curr == 'F':
            warnings.warn("invalid state " + str(s1))
            return -5
        elif curr == 'G':
            self._s = s1
            return 30
        else:
            self._s = s1
            return -1

    def is_goal(self):
        return self.check_state(self._s) == 'G'

    def get_actions(self):
        return self._actions
```

```python
env = GridWorld("grid.txt")
env.print_map()
```

    [['O' 'O' 'O' 'H' 'O' 'O' 'O' 'O' 'O']
     ['O' 'O' 'O' 'H' 'O' 'O' 'H' 'O' 'O']
     ['O' 'O' 'O' 'O' 'O' 'O' 'H' 'O' 'O']
     ['O' 'O' 'O' 'O' 'H' 'H' 'H' 'O' 'O']
     ['O' 'O' 'H' 'O' 'O' 'O' 'O' 'O' 'H']
     ['O' 'O' 'H' 'O' 'O' 'O' 'G' 'O' 'O']
     ['O' 'O' 'O' 'O' 'O' 'O' 'O' 'O' 'O']]

## Agent

Let us define an RL agent class. Two main methods are train() and test() as other ML classes. Unlike supervised learining, train function actively collect data while interacting with an environment and update its Q estimation with SARSA (or Q-learning). The next cell defines a utility to adjust coordianates to start from left-top.

```python
def coord_convert(s, sz):
    return [s[1], sz[0]-s[0]-1]
```

```python
class RLAgent:
    """
        Reinforcement Learning Agent Model for training/testing
        with Tabular function approximation

    """

    def __init__(self, evn):
        self.env = env
        self.size = env.get_size()
        self.n_a = len(env.get_actions())
        # self.Q table including the surrounding border
        self.Q = np.zeros((self.size[0], self.size[1], self.n_a))

    def epsilon_greed(self, epsilon, s):
        if np.random.uniform() < epsilon:
            a = np.random.randint(self.n_a)
        else:
            i_max = np.where(self.Q[s[0], s[1], :] == np.max(self.Q[s[0], s[1], :]))[0]
            a = int(np.random.choice(i_max))
        return a

    def train(self, start, **params):

        # parameters
        gamma = params.pop('gamma', 0.99)
        alpha = params.pop('alpha', 0.1)
        epsilon= params.pop('epsilon', 0.1)
        maxiter= params.pop('maxiter', 1000)
        maxstep= params.pop('maxstep', 1000)

        # init self.Q matrix
        self.Q[...] = 0
        self.Q[self.env._map == 'H'] = -np.inf

        # online train
        # rewards and step trace
        rtrace = []
        steps = []
        for j in range(maxiter):

            env.init(start)
            s = env.get_cur_state()
            # selection an action
            a = self.epsilon_greed(epsilon, s)

            rewards = []
            trace = np.array(coord_convert(s, self.size))
            # run simulation for max number of steps
            for step in range(maxstep):
                # move
                r = env.next(a)
                s1 = env.get_cur_state()
                a1 = self.epsilon_greed(epsilon, s1)

                rewards.append(r)
                trace = np.vstack((trace, coord_convert(s1, self.size)))

                # ToDo: update self.Q table (SARSA)
                self.Q[s[0], s[1], a] += alpha * (r + gamma * self.Q[s1[0], s1[1], a1] -\
                                             self.Q[s[0], s[1], a])

                if env.is_goal(): # reached the goal
                    # ToDo: update Q table
                    self.Q[s1[0], s1[1], a1] = 0
                    break

                s = s1
                a = a1

            rtrace.append(np.sum(rewards))
            steps.append(step+1)
        return rtrace, steps, trace # last trace of trajectory

    def test(self, start, maxstep=1000):

        #### ToDo: Here is one episode of simulation for testing.
        ####       Add one line to make this code test the trained agent
        epsilon = 0.0
        env.init(start)
        s = env.get_cur_state()
        # selection an action
        a = self.epsilon_greed(epsilon, s)

        trace = np.array(coord_convert(s, self.size))
        # run simulation for max number of steps
        for step in range(maxstep):
            # move
            r = env.next(a)
            s1 = env.get_cur_state()
            a1 = self.epsilon_greed(epsilon, s1)

            trace = np.vstack((trace, coord_convert(s1, self.size)))

            if env.is_goal(): # reached the goal
                break

            s = s1
            a = a1

        return trace
```

```python
### Plotting tools

def plot_trace(agent, start, trace, title="test trajectory"):
    plt.plot(trace[:, 0], trace[:, 1], "ko-")
    plt.text(env.goal_pos[1], agent.size[0]-env.goal_pos[0]-1, 'G')
    plt.text(start[1], agent.size[0]-start[0]-1, 'S')
    plt.xlim([0, agent.size[1]])
    plt.ylim([0, agent.size[0]])

def plot_train(agent, rtrace, steps, trace, start):

    fig = plt.figure(figsize=(10,8))
    ax = fig.add_subplot(221)
    plt.plot(rtrace, "b-")
    plt.ylabel("sum of rewards")

    ax1 = fig.add_subplot(222)
    plt.plot(steps)
    plt.ylabel("# steps")

    # contour plot for agent.Q
    ax2 = fig.add_subplot(223)
    xs = range(agent.size[1])
    ys = range(agent.size[0])
    maxQ = np.max(agent.Q, axis=2)
    h_b = (maxQ==-np.inf)
    maxQ[h_b] = 0
    maxQ[h_b] = np.min(maxQ) - 100
    cs = plt.contourf(xs, ys[::-1], maxQ)
    plt.colorbar(cs)
    plt.text(env.goal_pos[1], agent.size[0]-env.goal_pos[0]-1, 'G')
    plt.text(start[1], agent.size[0]-start[0]-1, 'S')
    plt.ylabel("max agent.Q")

    # plot traces
    ax3 = fig.add_subplot(224)
    plot_trace(agent, start, trace, "trace of the last episode")

    plt.plot()
```

```python
agent = RLAgent(env)
start = [0,0]
rtrace, steps, trace = agent.train(start,
                                   amma=0.99,
                                   alpha=0.1,
                                   epsilon=0.1,
                                   maxiter=100,
                                   maxstep=1000)
```

```python
plot_train(agent, rtrace, steps, trace, start)
```

```python
test_start = [0,2]
test_trace = agent.test(test_start)
plot_trace(agent, test_start, test_trace)
```

```python
test_start = [1,8]
test_trace = agent.test(test_start)
plot_trace(agent, test_start, test_trace)
```

Q: What do you think about the test results above? Does the agent run fine?


# Лабораторная работа по методам Монте-Карло

В рамках данной работы предлагается:
- Реализовать динамику Ланжевена для сэмплирования из апостериорного распределения
- Реализовать стохастическую динамику Ланжевена для сэмплирования из апостериорного распределения
- Применить оба метода на модельной задаче оценки параметров смеси нормальных распределений
- *(бонусная часть)* Применить динамику Ланжевена для генерации изображений

Оценка ставится на основе дополненного кода в ячейках (вычисление градиентов в модельной задаче, динамика Ланжевена и стохастическая динамика Ланжевена), а также ответов на вопросы по экспериментам. В бонусной части задания в качестве ненормированной плотности используется выход сверточной нейронной сети, поэтому для её выполнения вам может потребоваться GPU, облачный сервер или немного терпения.

Полезные источники:

- Welling M., Teh Y. W. Bayesian learning via stochastic gradient Langevin dynamics // Proceedings of the 28th International Conference on Machine Learning (ICML-11). – 2011. – С. 681-688.
- Neal R. M. et al. MCMC using Hamiltonian dynamics // Handbook of Markov Chain Monte Carlo. – 2011. – Т. 2. – №. 11
- Grathwohl W. et al. Your Classifier is Secretly an Energy Based Model and You Should Treat it Like One // ICLR 2020

Критерии оценки:
1. Функции для вычисления плотностей в модельной задаче - 2 балла
2. Стохастическая динамика Ланжевена - 2 балла
3. Динамика Ланжевена - 2 балла
4. Комментарии к экспериментам 1 - 1.5 балл
5. Комментарии к экспериментам 2 - 1.5 балл
6. Комментарии к экспериментам 3 - 1 балл
7. Бонусная часть - 2 балла максимум

```python
# Библиотеки, использованные при разработке задания:
# numpy==1.19.2
# jax==0.2.3
# matplotlib==3.3.2
# tqdm==4.50.2
# torch==1.6.0

import matplotlib.pylab as plt
import numpy as onp
import jax.numpy as np
import jax

import copy
import gzip
import os

from jax.scipy.stats import norm
from tqdm.auto import tqdm

# убираем warning если у вас нет gpu или tpu
jax.config.update("jax_platform_name", "cpu")
%matplotlib inline
```

### Пара слов о библиотеках для лабораторной работы

Поскольку для динамики Ланжевена нужны градиенты плотности, в этой работе мы воспользуемся библиотекой JAX для их подсчета. JAX переобределяет большую часть библиотеки numpy в модуле jax.numpy и добавляет возможности для
- Автоматическго построения градиентов функций;

```python
def foo(x):
    return np.sin(x)

grad_foo = jax.grad(foo)
x = 0.0
print("f(0) = {}, f'(0) = {}".format(foo(x), grad_foo(x)))
```

    f(0) = 0.0, f'(0) = 1.0

- Векторизации функций;

```python
# np.cos и так векторизован, но jax.vmap добавлять векторную размерность и для более сложных питоновских функций
vector_grad_foo = jax.vmap(grad_foo)
# vector_grad_foo(0.) теперь не заработает
vector_grad_foo(0.5 * np.pi * np.arange(8))
```

    DeviceArray([ 1.0000000e+00, -4.3711388e-08, -1.0000000e+00,
                  1.1924881e-08,  1.0000000e+00, -3.3776624e-07,
                 -1.0000000e+00,  6.6360758e-07], dtype=float32)

- Выполнения вычислений на ускорителях и компиляции функций "just-in-time"

```python
x = 0.5 * np.pi * np.arange(100)

%timeit vector_grad_foo(x)
vector_grad_foo = jax.jit(vector_grad_foo)  # компилируем функцию
%timeit vector_grad_foo(x)
%timeit np.cos(x) # так работает библиотечная функция
```

    2.45 ms ± 507 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
    5.53 µs ± 502 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
    4.71 µs ± 905 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Модуль jax.numpy можно использовать аналогично библиотеке numpy. Главным исключением является генерация случайных чисел: для этого в jax есть модуль jax.random, о котором можно прочитать [тут](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#%F0%9F%94%AA-Random-Numbers). **Tl;dr: для генерации каждый раз нужно передавать уникальный ключ.**

```python
rng = jax.random.PRNGKey(seed=42)
rng, key = jax.random.split(rng)
n = jax.random.normal(key, (3,))
print("Random sample:\n{}".format(n))
n = jax.random.normal(key, (3,))
print("Same key, same sample:\n{}".format(n))
rng, key = jax.random.split(rng)
n = jax.random.normal(key, (3,))
print("New key, new sample:\n{}".format(n))
```

    Random sample:
    [-0.5675503   0.2843918  -0.93206096]
    Same key, same sample:
    [-0.5675503   0.2843918  -0.93206096]
    New key, new sample:
    [ 0.6790333  -1.2206062   0.94670594]

Рекомендуем использовать jax.random при сэмплировании в динамике Ланжевена: он быстрее и "случайней". Тем не менее, для простоты мы будем инициализировать веса с помощью *numpy.random*.

# Задача оценки параметров смеси распределений

> Рассмотрим вероятностную модель, в которой данные приходят из смеси нормальных распределений:

>\begin{align}
& \theta_1 \sim N(0,\sigma_1^2); \quad \theta_2 \sim N(0, \sigma_2^2) \\
& \\
& x_i \overset{i.i.d.}{\sim} \frac{1}{2} N(\theta_1, \sigma_x^2) + \frac{1}{2} N(\theta_1 + \theta_2, \sigma_x^2) \quad i=1,\dots,N \\
\end{align}

>На $\theta_1$ и $\theta_2$ введены априорные нормальные распределения, а все остальные параметры полагаются равными
\begin{align}
& \sigma_x^2 = 2 \\
& \sigma_1^2 = 10 \\
& \sigma_2^2 = 1 \\
\end{align}

>Требуется по сгеренированной с параметрами $\theta_1 = 0$, $\theta_2 = 1$ выборке $X$ из $N=100$ элементов построить сэмплы из апостериорного распределения $p({\bar \theta} | X, \sigma_x, \sigma_1, \sigma_2)$.

Построим выборку из данного распределения:

```python
sigma_x = onp.sqrt(2.0)
sigma_1 = onp.sqrt(10.0)
sigma_2 = 1.0

def generate_mixture_data(N, theta_1=None, theta_2=None, seed=6):
    """
    Функция геренирует выобрку данного размера из описанной выше вероятностной модели.
    В случае, когда для theta_1 или theta_2 передано скалярное значение, вместо случайного
    значения параметра из априорного распределения используется переданное значение.
    """
    onp.random.seed(seed)
    if theta_1 is None:
        theta_1 = sigma_1 * onp.random.randn()
    if theta_2 is None:
        theta_2 = sigma_2 * onp.random.randn()
    mixture_component = onp.random.randint(0, 2, N).astype(bool)
    first_component = theta_1 + sigma_x * onp.random.randn(N)
    second_component = (theta_1 + theta_2) + sigma_x * onp.random.randn(N)
    samples = onp.where(mixture_component, first_component, second_component)
    return samples

# Пример выборки из описанной выше модели
fig, ax = plt.subplots(figsize=(7, 4))
ax.hist(generate_mixture_data(2048), bins=50, density=True)
ax.set_xlabel("X")
ax.set_ylabel("Empirical probability");
```

Наша главная цель - по полученной при фиксированных значениях $\theta_1$ и $\theta_2$ выборке построить апостериорное распределение на их значения.

\begin{equation}
p({\bar \theta} | X, \sigma_x, \sigma_1, \sigma_2) \propto \left[ \prod_{i=1}^{N} p(X| {\bar \theta}, \sigma_x) \right] p(\theta | \sigma_1, \sigma_2)
\end{equation}

Хотя задача и допускает аналитический подсчет апостериорного распределения, после раскрытия скобок оно будет представлять из себя смесь $2^N$ нормальных распределений. Поэтому даже для такой простой задачи аналитический вывод может оказаться неэффективным на практике.

Методы Монте-Карло позволяют получить сэмплы из апостериорного распределения.

> Для наших целей нам понядобятся функции для подсчета плотности априорного распределения и логарифма правдоподобия. Реализуйте их:

```python
def log_p_prior(weights):
    """
    Логарифм плотности априорного распределения.
    Вход:
        weights - вектор из двух вещественных чисел \theta_1, \theta_2
    Выход: число
    """
    theta_1 = weights[0]
    theta_2 = weights[1]
    return norm.logpdf(theta_1, 0, sigma_1) + norm.logpdf(theta_2, 0, sigma_2)

def log_p_likelihood(weights, x):
    """
    Возвращает логарифм правдоподобия данных x относительно параметров weights
    Вход:
        weights - вектор из двух вещественных чисел \theta_1, \theta_2,
        x - данные, представленные вектором фиксированной длины
    Выход: число
    """
    theta_1 = weights[0]
    theta_2 = weights[1]
    gauss_1 = norm.pdf(x, theta_1, sigma_x)
    gauss_2 = norm.pdf(x, theta_1 + theta_2, sigma_x)
    return np.log(0.5 * gauss_1 + 0.5 * gauss_2).sum(axis=0)

def log_p_joint(weights, x):
    return log_p_likelihood(weights, x) + log_p_prior(weights)

# Определим функцию для градиента и векторизуем обе функции для подсчета по батчу параметров:
grad_log_p_joint = jax.grad(log_p_joint)
# векторизуем только по первому аргументу
log_p_joint = jax.vmap(log_p_joint, (0, None))
grad_log_p_joint = jax.vmap(grad_log_p_joint, (0, None))
# компиляция для ускорения, можете попробовать опустить этот шаг
log_p_joint = jax.jit(log_p_joint)
grad_log_p_joint = jax.jit(grad_log_p_joint)
```

Согласно задаче, к нам поступает полученная со значениями $\theta_1=0$ и $\theta_2=1$ выборка из ста наблюдений $x_i$.

Изобразим приближенно плотность аспостериорного распределения в зависимости от числа поступивших наблюдений:

```python
rng = jax.random.PRNGKey(seed=42)
rng, key = jax.random.split(rng)
data = generate_mixture_data(N=100, theta_1=0, theta_2=1)

def plot_posterior(axes, data, resolution=100):
    theta_1 = np.linspace(-1.0, 2.0, resolution)
    theta_2 = np.linspace(-2.0, 2.0, resolution)
    X, Y = np.meshgrid(theta_1, theta_2)
    weights = np.stack((X.reshape(-1), Y.reshape(-1)), 1)
    log_probs = log_p_joint(weights, data)
    log_probs = log_probs - log_probs.max()
    Z = np.exp(log_probs - log_probs.max())
    norm_const = np.sum(Z) * 3 * 4 / resolution ** 2
    Z /= norm_const
    Z = Z.reshape((resolution, resolution))

    CS = axes.contour(X, Y, Z)
    axes.set_xlabel(r"$\theta_1$")
    axes.set_ylabel(r"$\theta_2$")

Ns = (
    10,
    20,
    50,
    100,
)
fig, axes = plt.subplots(ncols=len(Ns), figsize=(15, 3))
for n, ax in zip(Ns, axes):
    plot_posterior(ax, data[:n])
    ax.set_title("N = %d" % n)
```

> У апостериорного распределения на последнем графике должны быть две ярко выраженные моды. Чем это можно объяснить?

Так как $x$ из смеси гауссиан, то 2 моды всегда есть. Однако только для последнего эксперимента было выбрано достаточное количество точек, чтобы явно получить разделение.

> Ниже необходимо реализовать функции для построения сэмплов из апостериорного распределения:

> 1. Стохастическая динамика Ланжевена. Для реализации достаточно слегка модифицировать предложенную реализацию градиентного подъема
> 2. Динамика Ланжевена с приятием точек по схеме Метрополиса-Гастингса
    - Вычисление новых весов
    - Подсчет вероятности их принятия
    - Случайное принятие части весов

```python
def gradient_ascent_update(epsilon, gradient, weights, rng):
    return weights + 0.5 * epsilon * gradient

def stochastic_langevin_update(epsilon, gradient, weights, rng):
    """
    Стохастическая динамика Ланжевена
    Вход:
        epsilon - размер шага градиента
        gradient - вещественная матрица с градиентами для пар $\theta$
        weights -  вещественная матрица с параметрами $\theta$
        rng - сид для генерации случайных чисел
    Выход:
        вещественная матрица, соответствующая обновленным параметрам $\theta$
    """
    noise = jax.random.normal(rng, shape=gradient.shape) * np.sqrt(epsilon)
    return gradient_ascent_update(epsilon, gradient, weights, rng) + noise

def mh_acceptance_ratio(epsilon, weights, new_weights):
    """
    Динамика Ланжевена
    Вход:
        epsilon - размер шага градиента
        weights -  вещественная матрица с параметрами $\theta$
        new_weigts - параметры, полученные после шага схемы Ланжевена
    Выход:
        вещественный вектор, задающий логарифм вероятности принятия каждой пары из батча
    """
    return log_p_joint(new_weights, data) - log_p_joint(weights, data)

def langevin_update(epsilon, gradient, weights, rng):
    """
    Динамика Ланжевена
    Вход:
        epsilon - размер шага градиента
        gradient - вещественная матрица с градиентами для пар $\theta$
        weights -  вещественная матрица с параметрами $\theta$
        rng - сид для генерации случайных чисел
    Выход:
        вещественная матрица, соответствующая обновленным параметрам $\theta$
        вещественное число, доля принятых точек из батча
    """
    new_weights = stochastic_langevin_update(epsilon, gradient, weights, rng)

    accept_proba = np.minimum(1, np.exp(mh_acceptance_ratio(epsilon, weights, new_weights)))
    accept_mask = onp.random.binomial(1, p=accept_proba).astype(bool)

    new_weights.at[~accept_mask].set(weights[~accept_mask])

    return new_weights, np.mean(accept_mask)
```

## Динамика Ланжевена на модельной задаче

Для проведения эксперимента определим несколько вспомогательных функций

```python
def append_ema(array, value, alpha=0.05):
    """
    Добавление элементов в массив с экспоненциальным сглаживанием
    """
    if not array:
        array.append(value)
    else:
        array.append((1 - alpha) * array[-1] + alpha * value)
    return array

def train_mixture(data, weights, epsilon, n_epochs=2000, n_trajectories=5):
    """
    Построение точек из апостериорного распределения с помощью динамики Ланжевена
    Параллельно с построением точек сохраняются траектории n_trajectories точек и динамика
    acceptance_rate
    """
    rng = jax.random.PRNGKey(27)
    acceptance_rate = []
    theta_dynamics = []
    for epoch in tqdm(range(n_epochs)):
        theta_dynamics.append(copy.copy(weights[:n_trajectories]))
        gradient = grad_log_p_joint(weights, data)
        rng, key = jax.random.split(rng)
        weights, accepted = langevin_update(epsilon, gradient, weights, key)
        append_ema(acceptance_rate, accepted)
    theta_dynamics = np.asarray(theta_dynamics)
    return weights, acceptance_rate, theta_dynamics
```

Запустим M схем Ланжевена с шагом $\varepsilon = 10^{-3}$ и с шагом $\varepsilon = 10^{-4}$

```python
M = 128

weights_a = onp.random.randn(M, 2)
weights_b = onp.random.randn(M, 2)

weights_a, acceptance_rates_a, theta_dynamics_a = train_mixture(data, weights_a, 1e-3)
weights_b, acceptance_rates_b, theta_dynamics_b = train_mixture(data, weights_b, 1e-4)
```

      0%|          | 0/2000 [00:00<?, ?it/s]

      0%|          | 0/2000 [00:00<?, ?it/s]

В следующей ячейке строится четыре графика:
1. Два множества M точек из параллельно запущенных динамик Ланжевена после 2000 итераций
2. Динамика пяти точек из схемы с шагом $\varepsilon=10^{-3}$
3. Динамика пяти точек из схемы с шагом $\varepsilon=10^{-4}$
4. Доля принятых точек (из M параллельных динамик) после каждой итерации для разных шагов

```python
fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(14, 14))

plot_posterior(axes[0][0], data)
axes[0][0].scatter(weights_a[..., 0], weights_a[..., 1], marker="o", c="goldenrod", label=r"$\varepsilon=10^{-2}$")
axes[0][0].scatter(weights_b[..., 0], weights_b[..., 1], marker="v", c="deepskyblue", label=r"$\varepsilon=10^{-4}$")
axes[0][0].set_title("GLD samples after 2000 steps")
axes[0][0].legend(loc="upper right")

plot_posterior(axes[0][1], data)
plot_posterior(axes[1][0], data)
axes[0][1].set_title(r"GLD point dynamics, $\varepsilon=10^{-3}$")
axes[1][0].set_title(r"GLD point dynamics, $\varepsilon=10^{-4}$")

for i in range(5):
    axes[0][1].plot(theta_dynamics_a[:, i, 0], theta_dynamics_a[:, i, 1])
    axes[1][0].plot(theta_dynamics_b[:, i, 0], theta_dynamics_b[:, i, 1])

axes[1][1].plot(acceptance_rates_a, label=r"$\varepsilon=10^{-3}$", ls=":", c="goldenrod")
axes[1][1].plot(acceptance_rates_b, label=r"$\varepsilon=10^{-4}$", ls="-", c="deepskyblue")
axes[1][1].set_title("Metropolis-Hastings acceptance rate")
axes[1][1].set_xlabel("Epoch")
axes[1][1].set_ylabel(r"EMA of acceptance rate, $\alpha = 0.05$")
axes[1][1].legend(loc="lower right");
```

> На основе этих графиков ответьте на следующие вопросы:
- Хорошо ли точки покрывают апостериорное распределение?
Визуально покрытие кажется хорошим, однако с меньшим $\varepsilon$ покрытие хуже.
- Удается ли точкам "перепрыгнуть" с одной моды на другую?
Некоторым удаётся, например, красная и синия траектория.
- Как меняется доля принятых точек в зависимости от величины шага?
С меньшим шагом доля принятых выше.
- Исходя из каких соображений стоит выбирать длину шага в динамике Ланжевена?
Для более разнообразных точек следует выбирать шаг меньшего размера. Однако больший шаг может быть полезен для прогрева генерации.

## Стохастическая динамика Ланжевена на модельной задаче

Для проведения следующего эксперимента определим еще несколько вспомогательных функций

```python
# стохастическая оценка совместного правдоподобия
def log_p_joint_estimate(weights, x, N_over_n):
    return N_over_n * log_p_likelihood(weights, x) + log_p_prior(weights)

# стохастический градиент
grad_log_p_joint_estimate = jax.grad(log_p_joint_estimate)

log_p_joint_estimate = jax.jit(jax.vmap(log_p_joint_estimate, (0, None, None)))
grad_log_p_joint_estimate = jax.jit(jax.vmap(grad_log_p_joint_estimate, (0, None, None)))

# определим вектор градиентов $\nablda log p(x_i), i=1, \dots, N$ для подсчета дисперсии относительно выбора данных
batch_gradients = jax.vmap(jax.grad(log_p_likelihood), (None, 0))
```

```python
def iterate_data(rng, length, batchsize=100, shuffle=True):
    """
    Функция для прохода по выборке
    """
    indices = onp.arange(length)
    if shuffle:
        onp.random.shuffle(indices)
    for start_idx in range(0, length - batchsize + 1, batchsize):
        yield indices[start_idx : start_idx + batchsize]

def stochastic_train_mixture(rng, data, weights, method, n_epochs=2000, n=10, n_trajectories=5):
    """
    Построение точек из апостериорного распределения с помощью стохастической динамики Ланжевена
    Параллельно с построением точек сохраняются траектории n_trajectories точек после каждой эпохи,
    оценки дисперий стохастического градиента и дисперсия добавочного шума динамики Ланжевена
    """
    noise_vars = []
    theta_vars = []
    theta_dynamics = []
    N_over_n = len(data) / n
    steps_counter = 0
    for epochs in tqdm(range(n_epochs)):
        epsilon = 1e-1 * (steps_counter + 1e1) ** (-0.55)
        rng, key = jax.random.split(rng)
        for indices in iterate_data(key, len(data), n):
            rng, key = jax.random.split(rng)
            gradient = grad_log_p_joint_estimate(weights, data[indices], N_over_n)
            weights = method(epsilon, gradient, weights, key)
            steps_counter += 1

        theta_dynamics.append(copy.copy(weights[:n_trajectories, :]))
        # также как и в статье, оцениваем дисперсии по подвыборке
        rng, key = jax.random.split(rng)
        theta_grad = batch_gradients(weights, data)
        theta_vars.append(theta_grad.var(0).max(0) * (0.5 * epsilon * N_over_n) ** 2 * n)
        noise_vars.append(epsilon)

    theta_dynamics = np.asarray(theta_dynamics)
    theta_vars = np.asarray(theta_vars)
    noise_vars = np.asarray(noise_vars)

    return weights, noise_vars, theta_vars, theta_dynamics
```

Запустим стохастическую динамику Ланжевена и стохастический градиентный подъем

```python
M = 128
n = 10
n_trajectories = 5

weights_a = onp.random.randn(M, 2)
weights_b = onp.random.randn(M, 2)

weights_a, noise_vars_a, theta_vars_a, theta_dynamics_a = stochastic_train_mixture(
    rng, data, weights_a, stochastic_langevin_update, n_epochs=200
)

weights_b, _, theta_vars_b, theta_dynamics_b = stochastic_train_mixture(
    rng, data, weights_b, gradient_ascent_update, n_epochs=200
)
```

      0%|          | 0/200 [00:00<?, ?it/s]

      0%|          | 0/200 [00:00<?, ?it/s]

В следующей ячейке строятся четыре графика:

1. M точек из параллельно запущенных динамик Ланжевена после 2000 эпох и M точек из параллельно замущенных стохастических градиентных спусков после 2000 эпох
2. Траектории пяти точек при движении согласно стохастической динамике Ланжевена
3. Траектории пять точек при движении в направлении стохастического градиента
4. Дисперсии стохастических градиентов и добавочного шума стохастической динамики Ланжевена в зависимости от эпохи

```python
fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(14, 14))

plot_posterior(axes[0][0], data, 200)
axes[0][0].scatter(weights_a[:, 0], weights_a[:, 1], c="goldenrod", label="Stochastic Gradient Langevin Dynamics")
axes[0][0].scatter(weights_b[:, 0], weights_b[:, 1], marker="v", c="deepskyblue", label="Stochastic Gradient Descent")
axes[0][0].set_title("SGD vs. SGLD")
axes[0][0].legend(loc="upper right")

plot_posterior(axes[0][1], data)
plot_posterior(axes[1][0], data)
axes[0][1].set_title(r"SGLD point dynamics")
axes[1][0].set_title(r"SGD point dynamics")

for i in range(5):
    axes[0][1].plot(theta_dynamics_a[:, i, 0], theta_dynamics_a[:, i, 1])
    axes[1][0].plot(theta_dynamics_b[:, i, 0], theta_dynamics_b[:, i, 1])

axes[1][1].set_title("Stochstic gradient variance")
axes[1][1].set_yscale("log", nonpositive="clip")
axes[1][1].set_xscale("log", nonpositive="clip")
axes[1][1].plot(theta_vars_a[:, 0], label=r"SGLD $\nabla \theta_1$", c="goldenrod", ls=":")
axes[1][1].plot(theta_vars_a[:, 1], label=r"SGLD $\nabla \theta_2$", c="deepskyblue", ls=":")
axes[1][1].plot(theta_vars_b[:, 0], label=r"SGD $\nabla \theta_1$", c="goldenrod", ls="-")
axes[1][1].plot(theta_vars_b[:, 1], label=r"SGD $\nabla \theta_2$", c="deepskyblue", ls="-")
axes[1][1].plot(noise_vars_a, label=r"$\varepsilon$", c="black")
axes[1][1].set_xlabel("Epoch")
axes[1][1].set_ylabel("log variance")
axes[1][1].legend(loc="lower left");
```

> На основе графиков ответьте на следующие вопросы:
1. Хорошо ли представленные точки покрывают апостериорное распределение? Где собираются точки после градиентного подъема?
Точки хорошо покрывают апостериорное распределение, после градиентного подъема они собрались в модах.
2. Удается ли стохастической динамике Ланжевена "перепрыгнуть" с одной моды на другую?
По синей и оранжевой траектории видно, что удалось перепрыгнуть.
3. Как соотносятся дисперсии стохастических градиентов с дисперсией добавленного шума динамики Ланжевена?
Изначально дисперсии градиентов были больше дисперсии шума, однако они значительно быстрее уменьшались, поэтому довольно быстро стали меньше.

# Бонусная часть

В работе "Your Classifier is Secretly an Energy Based Model and You Should Treat it Like One" авторы предлагают приближать ненормированную плотность $\hat{p}(x, y)$ данных CIFAR-10 с помощью нейронной сети, разработанной для дискриминативной классификации. Обученная ими модель оказывается одновременно применима и для классификации, и для генерации данных. Для последнего как раз нужна схема Ланжевена.

В этой части задания мы предлагаем реализовать схему для сэмплирования из предобученной модели, выложенной [в репозитории авторов](https://github.com/wgrathwohl/JEM/). Там также можно найти приближенную схему сэмплирования: без поправки Метрополиса-Гастингса и с несогласованными множителями перед градиентом и добавочным шумом.

Попробуйте реализовать максимально обоснованную схему сэмплирования и получить сэмплы из CIFAR-10. Подумайте как выбор начальных точек влияет на скорость сэмплирования? Насколько хорошо определена плотность для аргумнетов, не похожих на изображения из CIFAR-10? Как реализовать условную и безусловную генерацию?

Для удобства к заданию мы прикладываем скрипт, позволяющий загрузить предобученную модель, вычислить значение логарифма ненормированной плотности и градиентов плотности по входному изображению. Пример их использования дан ниже.

```python
from wideresnet import load_pretrained_model, get_grads_wrt_input

path_to_pretrained_model = "./CIFAR10_MODEL.pt"
model = load_pretrained_model(path_to_pretrained_model)

batch_size = 4
x = onp.random.randn(batch_size, 3, 32, 32)
y = onp.random.randint(10, size=(batch_size,))

log_p, input_grads = get_grads_wrt_input(model, x)
log_p, input_grads = get_grads_wrt_input(model, x, y)
```


# Setup

```python
import numpy as np
import matplotlib.pyplot as plt
from sympy import *
import sympy.physics.mechanics as me
from pydy.system import System
from pydy.viz import Cylinder, Plane, VisualizationFrame, Scene
from sympy import sin, cos, symbols, solve
from pydy.codegen.ode_function_generators import generate_ode_function
from scipy.integrate import odeint
from IPython.display import SVG
SVG(filename='openrov2.svg')
```

```python
%matplotlib nbagg
```

```python
me.init_vprinting(use_latex='mathjax')
```

# Define Variables

constants:

- $m_A$: Mass of the Rov.
- $T_1$: Distance from COM to back thruster axis.
- $T_2$: Distance from COM to upper thruster axis.
- $W_x$: distance from center axis to back thruster axis.
- $W_h$: distance from center axis to upper thruster.
- $B_h$: the z messure of COB (center of bouyency) distance from COG (center of gravity)
- $B_h$: the x messure of COB distance from COG.

```python
# Inertial Reference Frame
N = me.ReferenceFrame('N')

# Define a world coordinate origin
O = me.Point('O')
O.set_vel(N, 0)
```

```python
rot = list(me.dynamicsymbols('r0:3'))
#rot = list(symbols('r0:3'))
drot = list(me.dynamicsymbols('dr0:3'))
x = list(me.dynamicsymbols('v0:3')) # Coordinates of robot in World Frame
dx = list(me.dynamicsymbols('dv0:3'))
kin_diff=Matrix(x+rot).diff()-Matrix(dx+drot)
kin_diff
```

$$\left[\begin{matrix}- dv_{0} + \dot{v}_{0}\\- dv_{1} + \dot{v}_{1}\\- dv_{2} + \dot{v}_{2}\\- dr_{0} + \dot{r}_{0}\\- dr_{1} + \dot{r}_{1}\\- dr_{2} + \dot{r}_{2}\end{matrix}\right]$$

```python
#xxx = me.dynamicsymbols('xxx')
#from sympy import *
#eval(sm.srepr(kin_diff))
#import pickle
#pickle.dump(xxx,open('/tmp/forcing_vector.pkl','wb'))
```

```python
# Constants for the Robot Body
Wx = symbols('W_x')                                     # 2*w is the width between thrusters
Wh = symbols('W_h')
T1 = symbols('T_1')                                     # Distance between thruster base and center of mass
T2 = symbols('T_2')
Bh = symbols('B_h')
Bw = symbols('B_w')

m_b = symbols('m_b')                                 # Mass of the body
v_b = symbols('v_b') # Volume of the body
mu = symbols('\mu') #drag
mu_r = symbols('\mu_r') #rotational drag
g = symbols('g')
I = list(symbols('Ixx, Iyy, Izz'))             # Moments of inertia of body

```

```python
# Robot Reference Frame
Rz=N.orientnew('R_z', 'Axis', (rot[2], N.z))
Rz.set_ang_vel(N,drot[2]*N.z)

Rx=Rz.orientnew('R_x', 'Axis', (rot[0], Rz.x))
Rx.set_ang_vel(Rz,drot[0]*Rz.x)

R=Rx.orientnew('R', 'Axis', (rot[1], Rx.y))
R.set_ang_vel(Rx,drot[1]*Rx.y)

#### adding dumping Torqe for each rotation
T_z=(R,-drot[2]*N.z*mu_r) #rotaional dumping Torqe
T_x=(R,-drot[0]*Rz.x*mu_r) #rotaional dumping Torqe
T_y=(R,-drot[1]*Rx.y*mu_r) #rotaional dumping Torqe

# Center of mass of body
COM = O.locatenew('COM', x[0]*N.x + x[1]*N.y + x[2]*N.z)

# Set the velocity of COM
COM.set_vel(N, dx[0]*N.x + dx[1]*N.y + dx[2]*N.z)

#center of bouyency
COB = COM.locatenew('COB', R.x*Bw+R.z*Bh)
COB.v2pt_theory(COM, N, R);
R.ang_vel_in(N)
```

$$dr_{1}\mathbf{\hat{r_x}_y} + dr_{0}\mathbf{\hat{r_z}_x} + dr_{2}\mathbf{\hat{n}_z}$$

```python
# Calculate inertia of body
Ib = me.inertia(R, *I)
# Create a rigid body object for body
Body = me.RigidBody('Body', COM, R, m_b, (Ib, COM))

```

```python
# Points of thrusters
L1 = COM.locatenew('L_1', -R.x*T1-Wx*R.y)
L2 = COM.locatenew('L_2', -R.x*T1+Wx*R.y)
L3 = COM.locatenew('L_3', -R.x*T2-Wh*R.z)

# Set the velocity of points
L1.v2pt_theory(COM, N, R)
L2.v2pt_theory(COM, N, R)
L3.v2pt_theory(COM, N, R);

```

## Calculating hydrodynamic drag

under sphire assumption and ignoring inertia forces

$F_{D}\,=\,{\tfrac {1}{2}}\,\rho \,v^{2}\,C_{D}\,A$

https://en.wikipedia.org/wiki/Drag_(physics)

we define $\mu$ as:

$\mu=\,{\tfrac {1}{2}}\,\rho \,C_{D}\,A$

then:

$F_{D}\,=\mu \,v^{2}$

```python
#dCw=Cw.diff()
v=N.x*dx[0]+N.y*dx[1]+N.z*dx[2]
Fd=-v.normalize()*v.magnitude()**2*mu
Fd
```

$$-  \mu \sqrt{dv^{2}_{0} + dv^{2}_{1} + dv^{2}_{2}} dv_{0}\mathbf{\hat{n}_x} -  \mu \sqrt{dv^{2}_{0} + dv^{2}_{1} + dv^{2}_{2}} dv_{1}\mathbf{\hat{n}_y} -  \mu \sqrt{dv^{2}_{0} + dv^{2}_{1} + dv^{2}_{2}} dv_{2}\mathbf{\hat{n}_z}$$

```python
#rotational drags
#Fr=-R.ang_vel_in(R_static)*mu
#Fr=(N,-drot[2]**2*N.z*mu),-drot[0]**2*Rz.x-drot[1]**2*Rx.y
#Fr=Fr*mu
#Fr
```

```python
#thrust forces symbols
F1, F2, F3 = symbols('f_1, f_2, f_3')
Fg = -N.z *m_b * g
Fb = N.z * v_b * 1e3 *g #whight of 1m^3 water in kg (MKS units)
```

```python
kane = me.KanesMethod(N, q_ind=x+rot, u_ind=dx+drot, kd_eqs=kin_diff)
```

```python
bodies = (Body,)
loads = (
    (L1, F1 * R.x),
    (L2, F2 * R.x),
    (L3, F3 * R.z),
    (COM, Fg ),
    (COB, Fb ),
    (COM, Fd ),
    T_x,
    T_y,
    T_z
    )

fr, frstar = kane.kanes_equations(loads=loads, bodies=bodies)
```

```python
mu_r
```

$$\mu_r$$

```python
mass_matrix = trigsimp(kane.mass_matrix_full)
mass_matrix
```

$$\left[\begin{array}{cccccccccccc}1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\0 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0\\0 & 0 & 0 & 0 & 0 & 0 & m_{b} & 0 & 0 & 0 & 0 & 0\\0 & 0 & 0 & 0 & 0 & 0 & 0 & m_{b} & 0 & 0 & 0 & 0\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & m_{b} & 0 & 0 & 0\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & Ixx \operatorname{cos}^{2}\left(r_{1}\right) + Izz \operatorname{sin}^{2}\left(r_{1}\right) & 0 & \left(- Ixx + Izz\right) \operatorname{sin}\left(r_{1}\right) \operatorname{cos}\left(r_{0}\right) \operatorname{cos}\left(r_{1}\right)\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & Iyy & Iyy \operatorname{sin}\left(r_{0}\right)\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \left(- Ixx + Izz\right) \operatorname{sin}\left(r_{1}\right) \operatorname{cos}\left(r_{0}\right) \operatorname{cos}\left(r_{1}\right) & Iyy \operatorname{sin}\left(r_{0}\right) & Ixx \operatorname{sin}^{2}\left(r_{1}\right) \operatorname{cos}^{2}\left(r_{0}\right) + Iyy \operatorname{sin}^{2}\left(r_{0}\right) + Izz \operatorname{cos}^{2}\left(r_{0}\right) \operatorname{cos}^{2}\left(r_{1}\right)\end{array}\right]$$

```python
forcing_vector = trigsimp(kane.forcing_full)
```

```python
#open('/tmp/mass_matrix.srepr','wb').write(mass_matrix)

```

```python

```

```python
coordinates = tuple(x+rot)
coordinates
```

$$\left ( v_{0}, \quad v_{1}, \quad v_{2}, \quad r_{0}, \quad r_{1}, \quad r_{2}\right )$$

```python
speeds = tuple(dx+drot)
speeds
```

$$\left ( dv_{0}, \quad dv_{1}, \quad dv_{2}, \quad dr_{0}, \quad dr_{1}, \quad dr_{2}\right )$$

```python
specified = (F1, F2, F3)
constants = [Wx,Wh,T1,T2,Bh,Bw,m_b,v_b,mu,mu_r,g]+I
```

```python
if 1:
    open('./forcing_vector.srepr','wb').write(srepr((\
                                                    forcing_vector,
                                                   coordinates,
                                                    mass_matrix,
                                                   speeds,
                                                    constants,
                                                   specified,
                                                   )).encode())

```

```python
if 1:
    (forcing_vector,\
    coordinates,
    mass_matrix,
    speeds,
    constants,
    specified)=eval(open('./forcing_vector.srepr','rb').read())

```

```python
right_hand_side = generate_ode_function(forcing_vector, coordinates,
                                        speeds, constants,
                                        mass_matrix=mass_matrix,
                                        specifieds=specified)
```

```python
help(right_hand_side)
```

    Help on function rhs in module pydy.codegen.ode_function_generators:

    rhs(*args)
        Returns the derivatives of the states, i.e. numerically evaluates the right
        hand side of the first order differential equation.

        x' = f(x, t, r, p)

        Parameters
        ==========
        x : ndarray, shape(12,)
            The state vector is ordered as such:
                - v0(t)
                - v1(t)
                - v2(t)
                - r0(t)
                - r1(t)
                - r2(t)
                - dv0(t)
                - dv1(t)
                - dv2(t)
                - dr0(t)
                - dr1(t)
                - dr2(t)
        t : float
            The current time.
        r : dictionary; ndarray, shape(3,); function

            There are three options for this argument. (1) is more flexible but
            (2) and (3) are much more efficient.

            (1) A dictionary that maps the specified functions of time to floats,
            ndarrays, or functions that produce ndarrays. The keys can be a single
            specified symbolic function of time or a tuple of symbols. The total
            number of symbols must be equal to 3. If the value is a
            function it must be of the form g(x, t), where x is the current state
            vector ndarray and t is the current time float and it must return an
            ndarray of the correct shape. For example::

              r = {a: 1.0,
                   (d, b) : np.array([1.0, 2.0]),
                   (e, f) : lambda x, t: np.array(x[0], x[1]),
                   c: lambda x, t: np.array(x[2])}

            (2) A ndarray with the specified values in the correct order and of the
            correct shape.

            (3) A function that must be of the form g(x, t), where x is the current
            state vector and t is the current time and it must return an ndarray of
            the correct shape.

            The specified inputs are, in order:
                - f_1
                - f_2
                - f_3
        p : dictionary len(14) or ndarray shape(14,)
            Either a dictionary that maps the constants symbols to their numerical
            values or an array with the constants in the following order:
                - W_x
                - W_h
                - T_1
                - T_2
                - B_h
                - B_w
                - m_b
                - v_b
                - \mu
                - \mu_r
                - g
                - Ixx
                - Iyy
                - Izz

        Returns
        =======
        dx : ndarray, shape(12,)
            The derivative of the state vector.

```python
x0 = np.zeros(12)
#MKS units
#constants = [Wx,Wh,T1,T2,Bh,Bw,m_b,v_b,mu,g]+I
numerical_constants = np.array([
                            0.1,  # Wx [m]
                            0.15,  # Wh [m]
                            0.1,  # T1 [m]
                            0.05,  # T2 [m]
                            0.08,  # Bh [m]
                            0.01,  # Bw [m]
                            1.0 *3,  # m_b [kg]
                            0.001 *3 ,  # v_b [M^3]
                            5.9,  # mu
                            0.2,  # mu_r
                            9.8,  # g MKS
                            0.5,  # Ixx [kg*m^2]
                            0.5,  # Iyy [kg*m^2]
                            0.5,  # Izz [kg*m^2]
                               ]
                            )
```

```python
#args = {'constants': numerical_constants,
numerical_specified=[0.8,0.5,0]
frames_per_sec = 60.0
final_time = 60.0

t = np.linspace(0.0, final_time, int(final_time * frames_per_sec))
```

```python
right_hand_side(x0, 0.0, numerical_specified, numerical_constants)
```

    array([ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ,
            0.        ,  0.43333333,  0.        ,  0.        ,  0.        ,
           -0.588     ,  0.06      ])

```python
def controller(x, t):
    if t<20 or t>35:
        #return [0.8,0.5,0]
        return [0.0,0.0,0]
    else:
        return [-0.55*30,-0.5*30,0]
#def controller(x, t):
#    return [0.0,0.0,0]
y = odeint(right_hand_side, x0, t, args=(controller, numerical_constants))

#y = odeint(right_hand_side, x0, t, args=(numerical_specified, numerical_constants))
y.shape
```

$$\left ( 3600, \quad 12\right )$$

```python
def plot():
    plt.figure()
    #plt.plot(sys.times, np.rad2deg(x[:, :3]))
    plt.subplot(2,3,1)
    #plt.plot(t, np.rad2deg(y[:, 0]))
    plt.plot(t, y[:, :3])
    plt.legend([latex(s, mode='inline') for s in coordinates[:3]])
    plt.subplot(2,3,2)
    plt.plot(t, np.rad2deg(y[:, 3:6]))
    plt.legend([latex(s, mode='inline') for s in coordinates[3:6]])
    plt.subplot(2,3,3)
    plt.title('XY')
    plt.plot(y[:,0],y[:,1])
    plt.axis('equal')
    plt.subplot(2,3,6)
    plt.title('Z')
    plt.plot(y[:,2])
    plt.axis('equal')
    plt.subplot(2,3,4)
    plt.plot(t, np.rad2deg(y[:, 9:12]))
    plt.legend([latex(s, mode='inline') for s in coordinates[9:12]])
plot()
```

    <IPython.core.display.Javascript object>

# another method

```python
x0 = np.zeros(12)
xx=x0
y=[]
for ct in t:
    x_dot=right_hand_side(xx, ct, controller, numerical_constants)
    y.append(xx)
    xx=xx+x_dot*1/frames_per_sec
y=np.array(y)
y.shape
```

$$\left ( 2400, \quad 12\right )$$

```python
plot()
```

    <IPython.core.display.Javascript object>

```python

```

```python

```


# some configs for pyplot
# https://stackoverflow.com/questions/332289/how-do-you-change-the-size-of-figures-drawn-with-matplotlib

```python
experimento = 'experimento_02'
sp_str      = '100'
```

```python
# %ls ../etc
```

```python
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
```

```python
import matplotlib
import scipy
print('Matplotlib Version:{}'.format(matplotlib.__version__))
print('Numpy Version:{}'.format(np.__version__))
print('Pandas Version:{}'.format(pd.__version__))
print('Scipy Version:{}'.format(scipy.__version__))
```

    Matplotlib Version:3.3.2
    Numpy Version:1.19.2
    Pandas Version:1.0.5
    Scipy Version:1.5.0

```python
motor_d = pd.read_csv('../etc/motor_direito_'+sp_str+ '_'+experimento+ '.csv')
motor_e = pd.read_csv('../etc/motor_esquerdo_'+sp_str+'_'+experimento+'.csv')

motor_d_c = pd.read_csv('../etc/control_motor_direito_'+sp_str+ '_'+experimento+ '.csv')
motor_e_c = pd.read_csv('../etc/control_motor_esquerdo_'+sp_str+'_'+experimento+'.csv')

sp = motor_e_c['SET_POINT'][0]

t  = motor_e['TIME']
we_raw = motor_e['OMEGA_RAW']
wd_raw = motor_d['OMEGA_RAW']

we = motor_e['OMEGA_FILTERED']
wd = motor_d['OMEGA_FILTERED']

t_c = motor_e_c['TIME']
we_c= motor_e_c['OMEGA_FILTERED']
wd_c= motor_d_c['OMEGA_FILTERED']

wref = motor_e_c['SET_POINT']*motor_e_c['OMEGA_MAX']

R_e  = motor_e['OMEGA_RAW'][motor_e['TIME'] >= motor_e['TAU']*10].var()
R_d  = motor_d['OMEGA_RAW'][motor_d['TIME'] >= motor_d['TAU']*10].var()
```

```python
def func(t,K,Tm):
    return K*(1.0 - np.exp(-t/Tm))
```

```python
we_opt, _   = curve_fit(func, t, we, bounds=([-10000., 0.01], [10000., 1000]))
we_c_opt,_ = curve_fit(func, t, we_c, bounds=([-10000., 0.01], [10000., 1000]))

wd_opt,_   = curve_fit(func, t, wd, bounds=([-10000., 0.01], [10000., 1000]))
wd_c_opt,_ = curve_fit(func, t, wd_c, bounds=([-10000., 0.01], [10000., 1000]))

we_opt_raw    = [motor_e['K'][0], motor_e['TAU'][0]]
wd_opt_raw    = [motor_d['K'][0], motor_d['TAU'][0]]

we_c_opt_raw,_  = curve_fit(func, t, motor_d_c['OMEGA_RAW'], bounds=([-10000., 0.01], [10000., 1000]))
wd_c_opt_raw,_  = curve_fit(func, t, motor_d_c['OMEGA_RAW'], bounds=([-10000., 0.01], [10000., 1000]))
```

```python
R_e
```

    1189.8554546739656

```python
R_d
```

    3349.449539155261

```python
# data = np.array([
# [1905.72803 , 1.00000],
# [1821.21313 , 0.95750],
# [1729.47572 , 0.91500],
# [1653.03481 , 0.87250],
# [1616.04560 , 0.83000],
# [1580.67555 , 0.78750],
# [1480.13788 , 0.74500],
# [1376.98560 , 0.70250],
# [1230.54941 , 0.66000],
# [1111.08493 , 0.61750],
# [989.78975 , 0.57500],
# [877.41730 , 0.53250],
# [785.88934 , 0.49000],
# [678.45646 , 0.44750],
# [604.61752 , 0.40500],
# [507.48609 , 0.36250],
# [371.34665 , 0.32000],
# [210.34399 , 0.27750],
# [90.20955 , 0.23500],
# [0.00000 , 0.19250]])
# x = data[:,0]
# y = data[:,1]
```

```python
# def _calcTau(t,w, wss):
#     Sxx = 0.0
#     Sxy = 0.0
#     y = 0.0
#     for i in range(len(t)):
#         if (abs(w[i]) > 0.2*abs(wss)) and (abs(w[i]) < 0.8*abs(wss)):
#             y = -np.log(1.0 - w[i]/wss);
#             Sxy+= t[i]*y;
#             Sxx+= t[i]*t[i];
#     return Sxx/Sxy

# def line(x, a,b):
#     return x*a + b

# pop, pcov = curve_fit(line, x, y, bounds=([0, 0], [np.inf, 1.0]))
# print(pop,pcov)

# plt.plot(x,y, label='data')
# plt.plot(x,line(x,pop[0], pop[1]), '--k')
# plt.ylim([0,1])
# plt.show()
```

```python
# def _calcTau(t,w, wss):
#     Sxx = 0.0
#     Sxy = 0.0
#     y = 0.0
#     for i in range(len(t)):
#         if (abs(w[i]) > 0.2*abs(wss)) and (abs(w[i]) < 0.8*abs(wss)):
#             y = -np.log(1.0 - w[i]/wss);
#             Sxy+= t[i]*y;
#             Sxx+= t[i]*t[i];
#     return Sxx/Sxy

# data = np.array([[0.00001, 0.00000],
# [0.00419, 0.00000],
# [0.00919, 1.39470],
# [0.01418, 1.40516],
# [0.01919, 325.16614],
# [0.02419, 325.16614],
# [0.02918, 414.23954],
# [0.03419, 709.00308],
# [0.03918, 841.12253],
# [0.04419, 953.29773]])

# x = data[:,0]
# y = data[:,1]
```

## TESTAGEM DA CALIBRAÇÃO

```python
motor_e.head()
```

      MOTOR
      CONTROLLER
      SET_POINT
      OMEGA_MAX
      K
      TAU
      FORWARD_KP
      BACK_KP
      FORWARD_ANG_COEF
      FORWARD_LIN_COEF
      BACK_ANG_COEF
      BACK_LIN_COEF
      TIME
      OMEGA_RAW
      OMEGA_FILTERED
      OMEGA_PREDICTED
      K_GAIN
      PREDIC_ERR
      MEASURE_ERR

      0
      0
      0
      1
      2625.29
      3236.8
      0.076151
      0.000137
      0.000185
      0.000287
      0.09487
      0.000274
      -0.087521
      0.000002
      0.0
      0.0
      0.0
      0.0
      0.0
      0

      1
      0
      0
      1
      2625.29
      3236.8
      0.076151
      0.000137
      0.000185
      0.000287
      0.09487
      0.000274
      -0.087521
      0.001037
      0.0
      0.0
      0.0
      0.0
      0.0
      0

      2
      0
      0
      1
      2625.29
      3236.8
      0.076151
      0.000137
      0.000185
      0.000287
      0.09487
      0.000274
      -0.087521
      0.003037
      0.0
      0.0
      0.0
      0.0
      0.0
      0

      3
      0
      0
      1
      2625.29
      3236.8
      0.076151
      0.000137
      0.000185
      0.000287
      0.09487
      0.000274
      -0.087521
      0.005037
      0.0
      0.0
      0.0
      0.0
      0.0
      0

      4
      0
      0
      1
      2625.29
      3236.8
      0.076151
      0.000137
      0.000185
      0.000287
      0.09487
      0.000274
      -0.087521
      0.007038
      0.0
      0.0
      0.0
      0.0
      0.0
      0

```python
motor_d.head()
```

      MOTOR
      CONTROLLER
      SET_POINT
      OMEGA_MAX
      K
      TAU
      FORWARD_KP
      BACK_KP
      FORWARD_ANG_COEF
      FORWARD_LIN_COEF
      BACK_ANG_COEF
      BACK_LIN_COEF
      TIME
      OMEGA_RAW
      OMEGA_FILTERED
      OMEGA_PREDICTED
      K_GAIN
      PREDIC_ERR
      MEASURE_ERR

      0
      1
      0
      1
      2625.29
      3047.72
      0.065716
      0.000124
      0.000082
      0.000293
      0.1044
      0.000303
      -0.079345
      0.000003
      0.0
      0.0
      0.0
      0.055118
      66.1417
      1200

      1
      1
      0
      1
      2625.29
      3047.72
      0.065716
      0.000124
      0.000082
      0.000293
      0.1044
      0.000303
      -0.079345
      0.001404
      0.0
      0.0
      0.0
      0.055118
      66.1417
      1200

      2
      1
      0
      1
      2625.29
      3047.72
      0.065716
      0.000124
      0.000082
      0.000293
      0.1044
      0.000303
      -0.079345
      0.003404
      0.0
      0.0
      0.0
      0.055118
      66.1417
      1200

      3
      1
      0
      1
      2625.29
      3047.72
      0.065716
      0.000124
      0.000082
      0.000293
      0.1044
      0.000303
      -0.079345
      0.005404
      0.0
      0.0
      0.0
      0.055118
      66.1417
      1200

      4
      1
      0
      1
      2625.29
      3047.72
      0.065716
      0.000124
      0.000082
      0.000293
      0.1044
      0.000303
      -0.079345
      0.007404
      0.0
      0.0
      0.0
      0.055118
      66.1417
      1200

```python
print('Motor Esquerdo')
print('| Km | Tm | R | Kp |')
print('| {:.2f} | {:.2e} | {:.2f} | {:.2e}|'.format(motor_e['K'][0],  motor_e['TAU'][0], R_e, motor_e['FORWARD_KP'][0]))
```

    Motor Esquerdo
    | Km | Tm | R | Kp |
    | 3236.80 | 7.62e-02 | 1189.86 | 1.37e-04|

```python
print('Motor Direito')
print('| Km | Tm | R | Kp |')
print('| {:.2f} | {:.2e} | {:.2f} | {:.2e}|'.format(motor_d['K'][0],  motor_d['TAU'][0], R_d, motor_d['FORWARD_KP'][0]))
```

    Motor Direito
    | Km | Tm | R | Kp |
    | 3047.72 | 6.57e-02 | 3349.45 | 1.24e-04|

```python
figsize=(10,5)
# plt.title('Comportamento Esperado vs Velocidade Medida')
plt.plot(t, we_raw, '-b', label=r'$\omega_{E_{medido}}$')
plt.plot(t, func(t,sp*we_opt_raw[0],we_opt_raw[1]), '--k', label=r'$\omega_E(t) = %.2f\left( 1 - e^{-t/%.2f}\right)$'%(sp*we_opt_raw[0],we_opt_raw[1]))

plt.xlim([0, motor_e_c['TIME'].max()])
plt.ylabel(r'$rad.s^{-1}$')
plt.xlabel(r'$t$ [s]')
plt.grid(True)
plt.legend()

plt.savefig('regressao_vs_medido_direito.eps', format='eps')
plt.show()

figsize=(10,5)
# plt.title('Comportamento Esperado vs Velocidade Medida')
plt.plot(t, wd_raw, '-b', label=r'$\omega_{D_{medido}}$')
plt.plot(t, func(t,sp*wd_opt_raw[0],wd_opt_raw[1]), '--k', label=r'$\omega_D(t) = %.2f\left( 1 - e^{-t/%.2f}\right)$'%(sp*wd_opt_raw[0],wd_opt_raw[1]))

plt.xlim([0, motor_e_c['TIME'].max()])
plt.ylabel(r'$rad.s^{-1}$')
plt.xlabel(r'$t$ [s]')
plt.grid(True)
plt.legend()

plt.savefig('regressao_vs_medido_direito' + sp_str + '.eps', format='eps')
plt.show()
```

```python
def reta(x, a, b):
    return a*x + b
```

```python
# Curva PWM x Omega
# Motor esquerdo
wmax = motor_e['OMEGA_MAX'][0]
w    = np.linspace(0,wmax, 10)
ae_f = motor_e['FORWARD_ANG_COEF'][0]
be_f = motor_e['FORWARD_LIN_COEF'][0]
ae_b = motor_e['BACK_ANG_COEF'][0]
be_b = motor_e['BACK_LIN_COEF'][0]

ue_f = reta(w,ae_f, be_f)
ue_b = reta(-w,ae_b, be_b)

plt.figure(figsize=(8,4))
plt.title(r'Curva $u(\omega_{ss})$ do Motor Esquerdo')

plt.plot(w, ue_f, 'b', label= r'$u_{f}(\omega) = %.2e*\omega + %.2f$'%(ae_f, be_f))
plt.plot(-w, ue_b, 'r',label= r'$u_{b}(\omega) = %.2e*\omega  %.2f$'%(ae_b, be_b))

plt.ylabel('PWM')
plt.xlabel(r'$\omega_{ss}$ [rad/s]')
plt.grid(True)
plt.ylim([-1.1,1.1])
plt.xlim([-wmax-100,wmax+100])
plt.legend()

# plt.grid(linestyle="--", linewidth=0.5, color='.25', zorder=-10)

plt.savefig('curva_feedforward_esquerdo' + sp_str + '.eps', format='eps')
plt.show()
```

```python
# Curva PWM x Omega
# Motor direito
wmax = motor_d['OMEGA_MAX'][0]
w    = np.linspace(0,wmax, 10)
ad_f = motor_d['FORWARD_ANG_COEF'][0]
bd_f = motor_d['FORWARD_LIN_COEF'][0]
ad_b = motor_d['BACK_ANG_COEF'][0]
bd_b = motor_d['BACK_LIN_COEF'][0]

ud_f = reta(w,ad_f, bd_f)
ud_b = reta(-w,ad_b, bd_b)

plt.figure(figsize=(8,4))
plt.title(r'Curva $u(\omega_{ss})$ do Motor Direito')

plt.plot(w, ud_f, 'b', label= r'$u_{f}(\omega) = %.2e*\omega + %.2f$'%(ad_f, bd_f))
plt.plot(-w, ud_b, 'r',label= r'$u_{b}(\omega) = %.2e*\omega  %.2f$'%(ad_b, bd_b))

plt.ylabel('PWM')
plt.xlabel(r'$\omega_{ss}$ [rad/s]')
plt.grid(True)
plt.ylim([-1.1,1.1])
plt.xlim([-wmax-100,wmax+100])
plt.legend()

# plt.grid(linestyle="--", linewidth=0.5, color='.25', zorder=-10)

plt.savefig('curva_feedforward_direito' + sp_str + '.eps', format='eps')
plt.show()
```

## TESTAGEM DO FILTRO

## Teste offline: Filtro de Kalman

```python
# plt.figure(num=1,figsize=(15,5))
# plt.title('Função de aproximação para o motor Esquerdo e Ganho do sistema filtrado')
# plt.plot(t, motor_e['OMEGA_RAW'], '-y', label=r'$\omega_{e_{raw}}$')
# plt.plot(t, we, '-b', label=r'$\omega_e$')
# plt.plot(t, func(t,Ke,Tme), '--k', label=r'$\omega(t) = %.2f\left( 1 - e^{-t/%.2f}\right)$'%(Ke,Tme))
# # plt.plot(t, func(t,x_k,motor_e['TAU'][0]), '--g', label=r'$\omega_{cal}(t) = %.2f\left( 1 - e^{-t/%.2f}\right)$'%(x_k,motor_e['TAU'][0]))
# plt.plot(t, func(t,motor_e['SET_POINT']*motor_e['K'][0],motor_e['TAU'][0]), '--r', label=r'$\omega_{cal_{orig}}(t) = %.2f\left( 1 - e^{-t/%.2f}\right)$'%(motor_e['K'][0]*motor_e['SET_POINT'][0],motor_e['TAU'][0]))

# plt.xlim([0, motor_e_c['TIME'].max()])
# plt.ylabel(r'$rad.s^{-1}$')
# plt.xlabel(r'$t$ [s]')
# plt.grid(True)
# plt.legend()
# plt.show()
```

```python
p0 = 100
r  = 1200
q  = 10
```

```python
# # inicialização
# # r = motor_e['OMEGA_RAW'][t >= Tme*5].var() #incerteza da medição
# #############################################################
# w_mean  = np.zeros_like(t)    #omega predito
# w_check = np.zeros_like(t)    #omega predito
# w_hat   = np.zeros_like(t)    #melhor estimativa de oemga (omega filtrado)
# p_check = np.zeros_like(t)    #incerteza de omega
# p_hat   = np.zeros_like(t)
# K = np.zeros_like(t)          #ganho do filtro

# p_check[0] = p_hat[0] = p0
# ##############################################################
# #input
# Tm= motor_e['TAU'][0]
# Kgain= motor_e['K'][0]
# u = motor_e['SET_POINT'][0]*Kgain

# for i in range(1, len(t)):
#     # medição
#     wz = motor_e['OMEGA_RAW'][i]
#     w_mean[i] = (w_mean[i-1] + wz)/2.0

#     predição
#     w_check[i] = w_hat[i-1] + (u - w_hat[i-1])*(1.0 - np.exp(-(t[i]- t[i-1])/Tm))
#     p_check[i] = p_hat[i-1] + q

#     # atualização
#     K[i] = p_check[i]/(p_check[i]+r)
#     w_hat[i] = w_check[i] + K[i]*(wz - w_check[i])
#     p_hat[i] = (1 - K[i])*p_check[i]
```

```python
# plt.figure(figsize=(10, 5), dpi=100)

# # plt.plot(t, w_check, '-b', label=r'$\check{\omega}(t)$')

# plt.plot(t, w_hat, '-g', label=r'$\hat{\omega}(t)$')

# plt.plot(t, motor_e['OMEGA_RAW'], '-k', label=r'$\omega_{raw}(t)$')

# # plt.plot(t, w_mean, '-y', label=r'$\omega_{mean}(t)$')

# # plt.plot(t, motor_e['OMEGA_FILTERED'], '-r', label=r'$\omega_{filtered}(t)$')

# plt.title('Teste filtro de Kalman para o Motor Esquerdo')
# plt.xlim([0, motor_e_c['TIME'].max()])
# plt.ylabel(r'$rad/s$');
# plt.xlabel(r'$t(s)$');
# plt.legend();
# plt.grid(True)
# plt.show();
```

```python
# # inicialização
# # p0 = 60.0
# # r = motor_d['OMEGA_RAW'][t >= Tme*5].var() #incerteza da medição
# # q = 10  #bias da incerteza (procurar uma definição mais adequada)
# #############################################################
# w_mean  = np.zeros_like(t)    #omega predito
# w_check = np.zeros_like(t)    #omega predito
# w_hat   = np.zeros_like(t)    #melhor estimativa de oemga (omega filtrado)
# p_check = np.zeros_like(t)    #incerteza de omega
# p_hat   = np.zeros_like(t)
# K = np.zeros_like(t)          #ganho do filtro

# p_check[0] = p_hat[0] = p0
# ##############################################################
# #input
# Tm= motor_d['TAU'][0]
# Kgain= motor_d['K'][0]
# u = motor_d['SET_POINT'][0]*Kgain

# for i in range(1, len(t)):
#     # medição
#     wz = motor_d['OMEGA_RAW'][i]
#     w_mean[i] = (w_mean[i-1] + wz)/2.0

#     # predição
#     w_check[i] = w_hat[i-1] + (u - w_hat[i-1])*(1.0 - np.exp(-(t[i]- t[i-1])/Tm))
#     p_check[i] = p_hat[i-1] + q

#     # atualização
#     K[i] = p_check[i]/(p_check[i]+r)
#     w_hat[i] = w_check[i] + K[i]*(wz - w_check[i])
#     p_hat[i] = (1 - K[i])*p_check[i]
```

```python
# plt.figure(figsize=(10, 5), dpi=100)

# # plt.plot(t, w_check, '-b', label=r'$\check{\omega}(t)$')

# plt.plot(t, w_hat, '-g', label=r'$\hat{\omega}(t)$')

# plt.plot(t, wd_raw, '-k', label=r'$\omega_{raw}(t)$')

# plt.plot(t, w_mean, '-y', label=r'$\omega_{mean}(t)$')

# plt.plot(t, wd, '-r', label=r'$\omega_{filtered}(t)$')

# plt.title('Teste filtro de Kalman para o Motor Direito')
# plt.xlim([0, motor_e_c['TIME'].max()])
# plt.ylabel(r'$rad/s$');
# plt.xlabel(r'$t(s)$');
# plt.legend();
# plt.grid(True)
# plt.show();
```

```python
# Teste Com vs Sem Filtro
# Sem controlador
plt.figure(figsize=(10,5))
# plt.title(r'$\hat{\omega}$ vs $\omega_{medido}$')
plt.title('Com Filtro vs Sem Filtro')
plt.plot(t, we_raw, '-b', label=r'$\omega_{E_{medido}}$')
plt.plot(t, wd_raw, '-g', label=r'$\omega_{D_{medido}}$')
plt.plot(t, we, '--b', label=r'$\hat{\omega}_{E}$')
plt.plot(t, wd, '--g', label=r'$\hat{\omega}_{D}$')

plt.xlim([0, motor_e['TIME'].max()])
plt.ylabel(r'$rad.s^{-1}$')
plt.xlabel(r'$t$ [s]')
plt.grid(True)
plt.legend()

plt.savefig('filtro_vs_sem_filtro' + sp_str + '.eps', format='eps')
plt.show()
```

## TESTAGEM DO Controlador

### CONTROLE COM FILTRO VS CONTROLE SEM FILTRO

```python
# Esse grafico seria informativo se fosse o controle usando o filtro vs o controle sem usar o filtro

# plt.figure(figsize=(15,5))
# plt.title('Ambos os motores | controlador ligado | Com Filtro x Sem Filtro')
# plt.plot(t_c, we_c, '--b', label=r'$\omega_e$')
# plt.plot(t_c, wd_c, '--g', label=r'$\omega_d$')
# plt.plot(t_c, motor_e_c['OMEGA_RAW'], '-b', label=r'$\omega_{e_{raw}}$')
# plt.plot(t_c, motor_d_c['OMEGA_RAW'], '-g', label=r'$\omega_{d_{raw}}$')

# plt.xlim([0, motor_e_c['TIME'].max()])
# plt.ylabel(r'$rad.s^{-1}$')
# plt.xlabel(r'$t$ [s]')
# plt.grid(True)
# plt.legend()
# plt.show()
```

### CONTROLADOR VS SEM CONTROLADOR

```python
plt.figure(figsize=(8,4))
plt.title('Com Controlador vs Sem Controlador')

plt.plot(t_c, we_c, '--b', label=r'$\omega_{E_{control}}(t) = %.2f(1 - e^{-t/%.2f})$'%(sp*we_c_opt[0],we_c_opt[1]))
plt.plot(t_c, wd_c, '--g', label=r'$\omega_{D_{control}}(t) = %.2f(1 - e^{-t/%.2f})$'%(sp*wd_c_opt[0],wd_c_opt[1]))

plt.plot(t_c, we, '-b', label=r'$\omega_{E}(t) = %.2f(1 - e^{-t/%.2f})$'%(sp*we_opt[0],we_opt[1]))
plt.plot(t_c, wd, '-g', label=r'$\omega_{D}(t) = %.2f(1 - e^{-t/%.2f})$'%(sp*wd_opt[0],wd_opt[1]))

plt.plot(t_c, wref, '-k', label=r'$\omega_{ref} = %.2f rad.s^{-1}$'%(wref[0]))

plt.xlim([0, motor_e_c['TIME'].max()])
plt.ylabel(r'$rad.s^{-1}$')
plt.xlabel(r'$t$ [s]')
plt.grid(True)
plt.legend()

plt.savefig('controlador_vs_sem_controlador' + sp_str + '.eps', format='eps')
plt.show()
```

```python
plt.figure(figsize=(8,4))
plt.title('Antes vs Depois')

plt.plot(t_c, we_c, '--b', label=r'$\omega_{E}(t) = %.2f(1 - e^{-t/%.2f})$'%(sp*we_opt[0],we_opt[1]))
plt.plot(t_c, wd_c, '--g', label=r'$\omega_{D}(t) = %.2f(1 - e^{-t/%.2f})$'%(sp*wd_opt[0],wd_opt[1]))

plt.plot(t_c, we_raw, '-b', label=r'$\omega_{{E}}(t) = %.2f(1 - e^{-t/%.2f})$'%(sp*we_opt_raw[0],we_opt_raw[1]))
plt.plot(t_c, wd_raw, '-g', label=r'$\omega_{{D}}(t) = %.2f(1 - e^{-t/%.2f})$'%(sp*wd_opt_raw[0],wd_opt_raw[1]))

plt.plot(t_c, wref, '-k', label=r'$\omega_{ref} = %.2f rad.s^{-1}$'%(wref[0]))

plt.xlim([0, motor_e_c['TIME'].max()])
plt.ylabel(r'$rad.s^{-1}$')
plt.xlabel(r'$t$ [s]')
plt.grid(True)
plt.legend()

plt.savefig('antes_vs_depois' + sp_str + '.eps', format='eps')
plt.show()
```

# Ref. Teórico
## Sistema(Planta):
\begin{equation}
G(s) = \frac{Y(s)}{U(s)} = \frac{K}{T_ms + 1}
\end{equation}

\begin{equation}
y(t) = K(1 - e^{-t/T_m})
\end{equation}

\begin{equation}
\omega(t) = u(t)K\left( 1 - e^{-t/T_m} \right); u(t) \in \left[-1, 1\right]
\end{equation}

## Filtro de Kalman
### Sistema
\begin{equation}
\textbf{x}_k = F_x x_{k-1} + B_k u_k + w_k; w_k \sim N(0, Q_k)
\end{equation}
### Medição
\begin{equation}
z_k = H_k x_k + v_k; v_k \sim N(0, R_k)
\end{equation}

### Predição
$\check{x}_k = F_k \hat{x}_{k-1} + B_k u_k$

$\check{P}_k = F_k \hat{P}_{k-1} F^T_k + Q_k$
### Atualização
$K_k = \check{P}_k H^T \left( H_k \check{P}_k H^T_k + R_k\right)^{-1}$

$\hat{x}_k = \check{x}_k + K_k\left( z_k - H_k \check{x}_k \right)$

$\hat{P}_k = \left(I - KH_k \right)\check{P}_k$

## Filtro de Kalman para a nossa planta
\begin{cases}
\textbf{x}_k = \left[ \omega_k \right]\\
z_k = x_k = \omega_k\\
F_k = 1\\
H_k = 1
\end{cases}
### Medição
\begin{align*}
z_k = \omega_z = \frac{N}{\Delta t}.k
\end{align*}

$N$: Números de pulsos;

$k$: Constante para converter de quantidade de pulsos para $rad/s$.
### Predição
\begin{align*}
\check{x}_k &= \hat{x}_{k-1} + (u_k - \hat{x}_{k-1}) \left( 1 - e^{-\frac{\Delta t}{T_m}} \right)\\
\check{x}_k &= \hat{x}_{k-1} +  u^{'}_k \left( 1 - e^{-\frac{\Delta t}{T_m}} \right)\\
\check{P}_k &= \hat{P}_{k-1} + Q_k
\end{align*}

**Obs.:** $\Delta t = t_f - t_0$. $\Delta t$ é relativo ao sinal de entrada $u_k$, sendo $t_0$ o instante que o sinal é aplicado e $t_f$ o instante atual $k$.
### Atualização
\begin{align*}
K_k &= \check{P}_k \left( \check{P}_k + R_k \right)^{-1} = \frac{\check{P}_k}{\check{P}_k + R_k}\\
\hat{x}_k &= \check{x}_k + K_k \left( z_k - \check{x}_k \right)\\
\hat{P}_k &= \left( 1 - K_k \right) \check{P}_k
\end{align*}