```python
import sys
sys.path.append('../../')
import os
import dill
import numpy as np
import scipy as sc
import random as rand
from sklearn import preprocessing, linear_model
import matplotlib.pyplot as plt

from core.controllers import PDController
from core.dynamics import LinearSystemDynamics, ConfigurationDynamics

from koopman_core.controllers import OpenLoopController, MPCController,BilinearFBLinController, PerturbedController, LinearLiftedController
from koopman_core.dynamics import LinearLiftedDynamics, BilinearLiftedDynamics
from koopman_core.learning import Edmd, BilinearEdmd, KoopDnn
from koopman_core.basis_functions import PlanarQuadBasis
from koopman_core.learning.utils import differentiate_vec
from koopman_core.systems import PlanarQuadrotorForceInput

class QuadrotorPdOutput(ConfigurationDynamics):
    def __init__(self, dynamics, xd, t_d, n, m):
        ConfigurationDynamics.__init__(self, dynamics, 1)
        self.xd = xd
        self.t_d = t_d
        self.xd_dot = differentiate_vec(self.xd, self.t_d)
        self.n = n
        self.m = m

    def proportional(self, x, t):
        q, q_dot = x[:int(n/2)], x[int(n/2):]
        return self.y(q) - self.y_d(t)

    def derivative(self, x, t):
        q, q_dot = x[:int(n/2)], x[int(n/2):]
        return self.dydq(q)@q_dot - self.y_d_dot(t)

    def y(self, q):
        return q

    def dydq(self, q):
        return np.eye(int(self.n/2))

    def d2ydq2(self, q):
        return np.zeros((int(self.n/2), int(self.n/2), int(self.n/2)))

    def y_d(self, t):
        return self.desired_state_(t)[:int(self.n/2)]

    def y_d_dot(self, t):
        return self.desired_state_(t)[int(self.n/2):]

    def y_d_ddot(self, t):
        return self.desired_state_dot_(t)[int(self.n/2):]

    def desired_state_(self, t):
        return [np.interp(t, self.t_d.flatten(),self.xd[:,ii].flatten()) for ii in range(self.xd.shape[1])]

    def desired_state_dot_(self, t):
        return [np.interp(t, self.t_d.flatten(),self.xd_dot[:,ii].flatten()) for ii in range(self.xd_dot.shape[1])]

class PlanarQuadrotorForceInputDiscrete(PlanarQuadrotorForceInput):
    def __init__(self, mass, inertia, prop_arm, g=9.81, dt=1e-2):
        PlanarQuadrotorForceInput.__init__(self, mass, inertia, prop_arm, g=g)
        self.dt=dt

    def eval_dot(self, x, u, t):
        return x + self.dt*self.drift(x, t) + self.dt*np.dot(self.act(x, t),u)

    def get_linearization(self, x0, x1, u0, t):
        m, J, b, g = self.params
        A_lin = np.eye(self.n) + self.dt*np.array([[0, 0, 0, 1, 0, 0],
                                                   [0, 0, 0, 0, 1, 0],
                                                   [0, 0, 0, 0, 0, 1],
                                                   [0, 0, -(1/m)*np.cos(x0[2])*u0[0] -(1/m)*np.cos(x0[2])*u0[1], 0, 0, 0],
                                                   [0, 0, -(1/m)*np.sin(x0[2])*u0[0] -(1/m)*np.sin(x0[2])*u0[1], 0, 0, 0],
                                                   [0, 0, 0, 0, 0, 0],])

        B_lin = self.dt*np.array([[0, 0],
                                  [0, 0],
                                  [0, 0],
                                  [-(1/m)*np.sin(x0[2]), -(1/m)*np.sin(x0[2])],
                                  [(1/m)*np.cos(x0[2]), (1/m)*np.cos(x0[2])],
                                  [-b/J, b/J]])

        if x1 is None:
            x1 = A_lin@x0 + B_lin@u0

        f_d = self.eval_dot(x0,u0,t)
        r_lin = f_d - x1

        return A_lin, B_lin, r_lin

```

## Planar Quadrotor Example

Consider a planar quadrotor with states $\mathbf{x} = [y \, z \, \theta \, \dot{y} \, \dot{z} \, \dot{\theta}]^T$ and continuous-time dynamics

\begin{equation}
    \begin{bmatrix} \ddot{y} \\ \ddot{z} \\ \ddot{\theta} \end{bmatrix}
    = \begin{bmatrix}
    0\\-g\\0
    \end{bmatrix} +
    \begin{bmatrix}
    -\frac{1}{m}\text{sin}\theta & -\frac{1}{m}\text{sin}\theta\\
    \frac{1}{m}\text{cos}\theta & \frac{1}{m}\text{cos}\theta\\
    -\frac{l_{arm}}{I_{xx}} & \frac{l_{arm}}{I_{xx}}
    \end{bmatrix}
    \begin{bmatrix}
    T_1 \\ T_2
    \end{bmatrix}
\end{equation}

where $y,z$ describe the position of the vehicle in a fixed reference frame, $\theta$ is the orientation of the vehicle,
$T_1, T_2$ are the thrust from each of the propellers, $g$ is the gravitational acceleration, $m$ is the vehicle mass,
$l_{arm}$ is the distance from the vehicle's center of mass to the center of the propeller, and $I_{xx}$ is the inertia
around the x-axis.

```python
mass = 2.
inertia = 1.
prop_arm = 0.2
gravity = 9.81
sys_name = 'planar_quad'
system = PlanarQuadrotorForceInput(mass, inertia, prop_arm, g=gravity)

# Linearized system specification:
n, m = 6, 2                                                         # Number of states, number of control inputs
A_nom = np.array([[0., 0., 0., 1., 0., 0.],                         # Linearization of the true system around the origin
                  [0., 0., 0., 0., 1., 0.],
                  [0., 0., 0., 0., 0., 1.],
                  [0., 0., -gravity, 0., 0., 0.],
                  [0., 0., 0., 0., 0., 0.],
                  [0., 0., 0., 0., 0., 0.]])
B_nom = np.array([[0., 0.],                                         # Linearization of the true system around the origin
                  [0., 0.],
                  [0., 0.],
                  [0., 0.],
                  [1./mass, 1./mass],
                  [-prop_arm/inertia, prop_arm/inertia]])

hover_thrust = mass*gravity/m
```

### Collect data for learning

To collect data, a nominal controller is designed with LQR on the dynamics's linearization around hover. However, any
controller can be used and the method does not require the knowledge of model's linearization. In addition, a
exploratory white noise is added to the controller to ensure that the data is sufficiently excited. Note that the system
is underactuated and that trajectory optimization is necessary to control the position of the vehicle. We use a
simplified trajectory generator based on a model predictive controller for the linearized dynamics. More careful design
of the desired trajectory may be necessary for more demanding applications and this is readily compatible with our method.

```python
q_dc, r_dc = 5e2, 1                                                 # State and actuation penalty values, data collection
Q_dc = q_dc * np.identity(n)                                        # State penalty matrix, data collection
R_dc = r_dc*np.identity(m)                                          # Actuation penalty matrix, data collection
P_dc = sc.linalg.solve_continuous_are(A_nom, B_nom, Q_dc, R_dc)     # Algebraic Ricatti equation solution, data collection
K_dc = np.linalg.inv(R_dc)@B_nom.T@P_dc                             # LQR feedback gain matrix, data collection
K_dc_p = K_dc[:,:int(n/2)]                                          # Proportional control gains, data collection
K_dc_d = K_dc[:,int(n/2):]                                          # Derivative control gains, data collection
nominal_sys = LinearLiftedDynamics(A_nom, B_nom, np.eye(n), lambda x: x)

# Data collection parameters:
collect_data = True
dt = 1.0e-2                                                         # Time step length
traj_length_dc = 2.                                                 # Trajectory length, data collection
n_pred_dc = int(traj_length_dc/dt)                                  # Number of time steps, data collection
t_train = dt * np.arange(n_pred_dc + 1)                              # Simulation time points
n_traj_train = 100                                                     # Number of trajectories to execute, data collection
n_traj_val = int(0.2*n_traj_train)
noise_var = 5.                                                      # Exploration noise to perturb controller, data collection

xmax = np.array([2, 2, np.pi/3, 2.,2.,2.])                          # State constraints, trajectory generation
xmin = -xmax
umax = np.array([2*hover_thrust, 2*hover_thrust]) - hover_thrust    # Actuation constraint, trajectory generation
umin = np.array([0., 0.]) - hover_thrust
x0_max = np.array([xmax[0], xmax[1], xmax[2], 1., 1., 1.])          # Initial value limits
Q_trajgen = sc.sparse.diags([0,0,0,0,0,0])                          # State penalty matrix, trajectory generation
QN_trajgen = sc.sparse.diags([5e1,5e1,5e1,1e1,1e1,1e1])             # Final state penalty matrix, trajectory generation
R_trajgen = sc.sparse.eye(m)                                        # Actuation penalty matrix, trajectory generation
sub_sample_rate = 1                                                 # Rate to subsample data for training
model_fname = 'examples/planar_quad_models'                         # Path to save learned models
n_cols = 10                                                         # Number of columns in training data plot
save_figures = True
directory = os.path.abspath("")                                     # Path to save learned models
```

```python
if collect_data:
    xd = np.empty((n_traj_train + n_traj_val, n_pred_dc + 1, n))
    xs = np.empty((n_traj_train + n_traj_val, n_pred_dc + 1, n))
    us = np.empty((n_traj_train + n_traj_val, n_pred_dc, m))

    plt.figure(figsize=(12, 12 * (n_traj_train + n_traj_val) / (n_cols ** 2)))
    for ii in range(n_traj_train+n_traj_val):
        x0 = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
        set_pt_dc = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
        mpc_trajgen = MPCController(nominal_sys, n_pred_dc, dt, umin, umax, xmin, xmax, QN_trajgen, R_trajgen,
                                    QN_trajgen, set_pt_dc)
        mpc_trajgen.eval(x0, 0)
        xd[ii, :, :] = mpc_trajgen.parse_result().T
        while abs(x0[0]) + abs(x0[1]) < 1 or np.any(np.isnan(xd[ii, :, :])):
            x0 = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
            set_pt_dc = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
            mpc_trajgen = MPCController(nominal_sys, n_pred_dc, dt, umin, umax, xmin, xmax, QN_trajgen, R_trajgen,
                                        QN_trajgen, set_pt_dc)
            mpc_trajgen.eval(x0, 0)
            xd[ii, :, :] = mpc_trajgen.parse_result().T

        output = QuadrotorPdOutput(system, xd[ii, :, :], t_train, n, m)
        pd_controller = PDController(output, K_dc_p, K_dc_d)
        perturbed_pd_controller = PerturbedController(system, pd_controller, noise_var, const_offset=hover_thrust)
        xs[ii, :, :], us[ii, :, :] = system.simulate(x0, perturbed_pd_controller, t_train)

        plt.subplot(int(np.ceil((n_traj_train + n_traj_val) / n_cols)), n_cols, ii + 1)
        plt.plot(t_train, xs[ii, :, 0], 'b', label='$y$')
        plt.plot(t_train, xs[ii, :, 1], 'g', label='$z$')
        plt.plot(t_train, xs[ii, :, 2], 'r', label='$\\theta$')
        plt.plot(t_train, xd[ii, :, 0], '--b', label='$y_d$')
        plt.plot(t_train, xd[ii, :, 1], '--g', label='$z_d$')
        plt.plot(t_train, xd[ii, :, 2], '--r', label='$\\theta_d$')

    plt.suptitle(
        'Training data \nx-axis: time (sec), y-axis: state value, $x$ - blue, $xd$ - dotted blue, $\\theta$ - red, $\\theta_d$ - dotted red',
        y=0.94)
    plt.show()

    xs_train, us_train = xs[:n_traj_train,:,:], us[:n_traj_train, :, :]
    xs_val, us_val = xs[n_traj_train:,:,:], us[n_traj_train:, :, :]

    data_list = [xs_train, us_train, t_train, n_traj_train, xs_val, us_val, t_train, n_traj_val]
    outfile = open(directory + '/data/' + sys_name + '_data.pickle', 'wb')
    dill.dump(data_list, outfile)
    outfile.close()
else:
    infile = open(directory + '/data/' + sys_name + '_data.pickle', 'rb')
    xs_train, us_train, t_train, n_traj_train, xs_val, us_val, t_val, n_traj_val = dill.load(infile)
    infile.close()
```

### Learn a linear model with dynamic mode decomposition (DMD)

To compare our method with existing techniques, we first learn a linear state space model from data. This is dubbed
dynamic mode decomposition. I.e. we use linear regression with LASSO regularization to learn an approximate linear model
with model structure

\begin{equation}
    \mathbf{\dot{x}} = A_{dmd}\mathbf{x} + B_{dmd}\mathbf{u}
\end{equation}

```python
#DMD parameters:
alpha_dmd = 9.8e-5                                                  # Regularization strength (LASSO) DMD
tune_mdl_dmd = False
```

```python
basis = lambda x: x
C_dmd = np.eye(n)

optimizer_dmd = linear_model.MultiTaskLasso(alpha=alpha_dmd, fit_intercept=False, selection='random')
cv_dmd = linear_model.MultiTaskLassoCV(fit_intercept=False, n_jobs=-1, cv=3, selection='random')
standardizer_dmd = preprocessing.StandardScaler(with_mean=False)

model_dmd = Edmd(n, m, basis, n, n_traj_train, optimizer_dmd, cv=cv_dmd, standardizer=standardizer_dmd, C=C_dmd,
                 first_obs_const=False, continuous_mdl=False, dt=dt)
xdmd, y_dmd = model_dmd.process(xs_train, us_train-hover_thrust, np.tile(t_train,(n_traj_train,1)),
                                downsample_rate=sub_sample_rate)
model_dmd.fit(xdmd, y_dmd, cv=tune_mdl_dmd, override_kinematics=True)
sys_dmd = LinearLiftedDynamics(model_dmd.A, model_dmd.B, model_dmd.C, model_dmd.basis, continuous_mdl=False, dt=dt)
if tune_mdl_dmd:
    print('$\\alpha$ DMD: ',model_dmd.cv.alpha_)
```

### Learn a lifted linear model with extended dynamic mode decomposition (EDMD)

In addition, we compare our method with the current state of the art of Koopman based learning, the extended dynamic mode
decomposition. We use a dictionary of nonlinear functions $\boldsymbol{\phi(x)}$ to lift the state variables and learn a lifted state space model
of the dynamics. I.e. we first lift and then use linear regression with LASSO regularization to learn an approximate
lifted linear model with model structure

\begin{equation}
    \mathbf{\dot{z}} = A_{edmd}\mathbf{z} + B_{edmd}\mathbf{u}, \qquad \mathbf{z} = \boldsymbol{\phi(x)}
\end{equation}

```python
#EDMD parameters:
alpha_edmd = 2.22e-4                                                 # Regularization strength (LASSO) EDMD
tune_mdl_edmd = False
```

```python
basis = PlanarQuadBasis(n, poly_deg=3)
basis.construct_basis()
planar_quad_features = preprocessing.FunctionTransformer(basis.basis)
planar_quad_features.fit(np.zeros((1,n)))
n_lift_edmd = planar_quad_features.transform((np.zeros((1,n)))).shape[1]
C_edmd = np.zeros((n,n_lift_edmd))
C_edmd[:,1:n+1] = np.eye(n)

optimizer_edmd = linear_model.MultiTaskLasso(alpha=alpha_edmd, fit_intercept=False, selection='random',
                                             max_iter=2000)
cv_edmd = linear_model.MultiTaskLassoCV(fit_intercept=False, n_jobs=-1, cv=3, selection='random', max_iter=2000)
standardizer_edmd = preprocessing.StandardScaler(with_mean=False)

model_edmd = Edmd(n, m, basis.basis, n_lift_edmd, n_traj_train, optimizer_edmd, cv=cv_edmd,
                  standardizer=standardizer_edmd, C=C_edmd, continuous_mdl=False, dt=dt)
X_edmd, y_edmd = model_edmd.process(xs_train, us_train, np.tile(t_train,(n_traj_train,1)),
                                    downsample_rate=sub_sample_rate)
model_edmd.fit(X_edmd, y_edmd, cv=tune_mdl_edmd, override_kinematics=True)
model_edmd.reduce_mdl()
sys_edmd = LinearLiftedDynamics(model_edmd.A, model_edmd.B, model_edmd.C, model_edmd.basis_reduced,
                                continuous_mdl=False, dt=dt)
if tune_mdl_edmd:
    print('$\\alpha$ EDMD: ',model_edmd.cv.alpha_)

```

### Learn a lifted bilinear model with bilinear extended mode decomposition (bEDMD)

Finally, we use the method developed in the paper to learn a lifted bilinear model of the dynamics, dubbed bilinear
extended mode decomposition (bEDMD). I.e. we first lift and then use linear regression with LASSO regularization to learn an approximate
lifted linear model with model structure

\begin{equation}
    \mathbf{\dot{z}}=F\mathbf{z}+\sum_{i=1}^m G_i\mathbf{z}\mathbf{u}_i, \qquad \mathbf{z} = \boldsymbol{\phi(x)}
\end{equation}

```python
#Bilinear EDMD parameters:
alpha_bedmd = 6.9e-5  # Regularization strength (LASSO) bEDMD
tune_mdl_bedmd = False
```

```python
n_lift_bedmd = n_lift_edmd
C_bedmd = np.zeros((n,n_lift_bedmd))
C_bedmd[:,1:n+1] = np.eye(n)

basis_bedmd = lambda x: planar_quad_features.transform(x)
optimizer_bedmd = linear_model.MultiTaskLasso(alpha=alpha_bedmd, fit_intercept=False, selection='random')
cv_bedmd = linear_model.MultiTaskLassoCV(fit_intercept=False, n_jobs=-1, cv=3, selection='random')
standardizer_bedmd = preprocessing.StandardScaler(with_mean=False)

model_bedmd = BilinearEdmd(n, m, basis_bedmd, n_lift_bedmd, n_traj_train, optimizer_bedmd, cv=cv_bedmd,
                           standardizer=standardizer_bedmd, C=C_bedmd, continuous_mdl=False, dt=dt)
X_bedmd, y_bedmd = model_bedmd.process(xs_train, us_train, np.tile(t_train,(n_traj_train,1)),
                                       downsample_rate=sub_sample_rate)
model_bedmd.fit(X_bedmd, y_bedmd, cv=tune_mdl_bedmd, override_kinematics=True)
sys_bedmd = BilinearLiftedDynamics(model_bedmd.n_lift, m, model_bedmd.A, model_bedmd.B, model_bedmd.C,
                                   model_bedmd.basis, continuous_mdl=False, dt=dt)
if tune_mdl_bedmd:
    print('$\\alpha$ bilinear EDMD: ', model_bedmd.cv.alpha_)

```

    /Users/carlaxelfolkestad/.conda/envs/keedmd/lib/python3.7/site-packages/sklearn/linear_model/_coordinate_descent.py:1790: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Duality gap: 1.3172523546308825, tolerance: 0.01812327679309411
      check_random_state(self.random_state), random)

### Learn a lifted bilinear model with Koopman NN

```python
import dill, os, torch

load_tuned_params = False

if load_tuned_params:
    infile = open(os.path.abspath('') + '/data/planar_quad_best_params.pickle', 'rb')
    net_params, val_loss, test_loss, open_loop_mse, open_loop_std = dill.load(infile)
    infile.close()

else:
    net_params = {}
    net_params['state_dim'] = n
    net_params['ctrl_dim'] = m
    net_params['encoder_hidden_width'] = 100
    net_params['encoder_hidden_depth'] = 2
    net_params['encoder_output_dim'] = 10
    net_params['optimizer'] = 'adam'
    net_params['activation_type'] = 'relu'
    net_params['lr'] = 1e-3
    net_params['epochs'] = 500
    net_params['batch_size'] = 256
    net_params['lin_loss_penalty'] = 0.5
    net_params['l2_reg'] = 5e-6
    net_params['l1_reg'] = 5e-6
    net_params['first_obs_const'] = True
    net_params['override_kinematics'] = True
    net_params['dt'] = dt

print(net_params)
```

    {'state_dim': 6, 'ctrl_dim': 2, 'encoder_hidden_width': 100, 'encoder_hidden_depth': 2, 'encoder_output_dim': 10, 'optimizer': 'adam', 'activation_type': 'relu', 'lr': 0.001, 'epochs': 500, 'batch_size': 256, 'lin_loss_penalty': 0.5, 'l2_reg': 5e-06, 'l1_reg': 5e-06, 'first_obs_const': True, 'override_kinematics': True, 'dt': 0.01}

```python
from koopman_core.learning import KoopDnn, KoopmanNetCtrl

#standardizer_x_kdnn = preprocessing.StandardScaler(with_mean=False)
standardizer_x_kdnn = preprocessing.StandardScaler(with_mean=True)
#standardizer_x_kdnn = None
standardizer_u_kdnn = preprocessing.StandardScaler(with_mean=True)
#standardizer_u_kdnn = None
n_tot = n + net_params['encoder_output_dim'] + int(net_params['first_obs_const'])

net = KoopmanNetCtrl(net_params, standardizer_x=standardizer_x_kdnn, standardizer_u=standardizer_u_kdnn)
model_koop_dnn = KoopDnn(net)
model_koop_dnn.set_datasets(xs_train, t_train, u_train=us_train, x_val=xs_val, u_val=us_val, t_val=t_train)
model_koop_dnn.model_pipeline(net_params)
model_koop_dnn.construct_koopman_model()
sys_koop_dnn = BilinearLiftedDynamics(n_tot, m, model_koop_dnn.A, model_koop_dnn.B, model_koop_dnn.C, model_koop_dnn.basis_encode,
                                    continuous_mdl=False, dt=dt, standardizer_x=standardizer_x_kdnn,
                                      standardizer_u=standardizer_u_kdnn)
```

    Epoch   1: train loss: 5.79225588, validation loss: 1.24986596
    Epoch   2: train loss: 0.83037537, validation loss: 0.42855491
    Epoch   3: train loss: 0.33319420, validation loss: 0.22175703
    Epoch   4: train loss: 0.17941733, validation loss: 0.13503779
    Epoch   5: train loss: 0.11287631, validation loss: 0.09434711
    Epoch   6: train loss: 0.07859979, validation loss: 0.07143931
    Epoch   7: train loss: 0.06068058, validation loss: 0.06006999
    Epoch   8: train loss: 0.05046217, validation loss: 0.05041111
    Epoch   9: train loss: 0.04282654, validation loss: 0.04312954
    Epoch  10: train loss: 0.03758250, validation loss: 0.03760370
    Epoch  11: train loss: 0.03320860, validation loss: 0.03351406
    Epoch  12: train loss: 0.02992171, validation loss: 0.03069416
    Epoch  13: train loss: 0.02762811, validation loss: 0.02898422
    Epoch  14: train loss: 0.02494172, validation loss: 0.02538829
    Epoch  15: train loss: 0.02307066, validation loss: 0.02423847
    Epoch  16: train loss: 0.02080001, validation loss: 0.02142645
    Epoch  17: train loss: 0.01942540, validation loss: 0.01949529
    Epoch  18: train loss: 0.01802165, validation loss: 0.01930718
    Epoch  19: train loss: 0.01709201, validation loss: 0.01737142
    Epoch  20: train loss: 0.01577595, validation loss: 0.01687637
    Epoch  21: train loss: 0.01500758, validation loss: 0.01545202
    Epoch  22: train loss: 0.01415862, validation loss: 0.01484914
    Epoch  23: train loss: 0.01325533, validation loss: 0.01420603
    Epoch  24: train loss: 0.01251044, validation loss: 0.01378286
    Epoch  25: train loss: 0.01194354, validation loss: 0.01248725
    Epoch  26: train loss: 0.01133224, validation loss: 0.01197852
    Epoch  27: train loss: 0.01088677, validation loss: 0.01174055
    Epoch  28: train loss: 0.01062767, validation loss: 0.01093675
    Epoch  29: train loss: 0.01004568, validation loss: 0.01284418
    Epoch  30: train loss: 0.00963226, validation loss: 0.00994158
    Epoch  31: train loss: 0.00913601, validation loss: 0.01057452
    Epoch  32: train loss: 0.00882153, validation loss: 0.00948092
    Epoch  33: train loss: 0.00862754, validation loss: 0.00979661
    Epoch  34: train loss: 0.00832926, validation loss: 0.00872535
    Epoch  35: train loss: 0.00781409, validation loss: 0.00877404
    Epoch  36: train loss: 0.00759831, validation loss: 0.00772222
    Epoch  37: train loss: 0.00753533, validation loss: 0.00999367
    Epoch  38: train loss: 0.00745700, validation loss: 0.00742849
    Epoch  39: train loss: 0.00681914, validation loss: 0.00712146
    Epoch  40: train loss: 0.00692503, validation loss: 0.00735920
    Epoch  41: train loss: 0.00652081, validation loss: 0.00686995
    Epoch  42: train loss: 0.00644565, validation loss: 0.00637342
    Epoch  43: train loss: 0.00623332, validation loss: 0.00646325
    Epoch  44: train loss: 0.00629828, validation loss: 0.00893176
    Epoch  45: train loss: 0.00606772, validation loss: 0.00609800
    Epoch  46: train loss: 0.00572035, validation loss: 0.00597551
    Epoch  47: train loss: 0.00569322, validation loss: 0.00573133
    Epoch  48: train loss: 0.00542325, validation loss: 0.00770620
    Epoch  49: train loss: 0.00534481, validation loss: 0.00521132
    Epoch  50: train loss: 0.00534199, validation loss: 0.00574045
    Epoch  51: train loss: 0.00528706, validation loss: 0.00579366
    Epoch  52: train loss: 0.00517862, validation loss: 0.00473424
    Epoch  53: train loss: 0.00480173, validation loss: 0.00472433
    Epoch  54: train loss: 0.00482118, validation loss: 0.00568945
    Epoch  55: train loss: 0.00480158, validation loss: 0.00439994
    Epoch  56: train loss: 0.00464682, validation loss: 0.00471954
    Epoch  57: train loss: 0.00455467, validation loss: 0.00434539
    Epoch  58: train loss: 0.00482891, validation loss: 0.00447735
    Epoch  59: train loss: 0.00428220, validation loss: 0.00427874
    Epoch  60: train loss: 0.00430363, validation loss: 0.00446903
    Epoch  61: train loss: 0.00421611, validation loss: 0.00429145
    Epoch  62: train loss: 0.00431762, validation loss: 0.00405835
    Epoch  63: train loss: 0.00404505, validation loss: 0.00370175
    Epoch  64: train loss: 0.00386445, validation loss: 0.00462665
    Epoch  65: train loss: 0.00401839, validation loss: 0.00415844
    Epoch  66: train loss: 0.00378527, validation loss: 0.00348224
    Epoch  67: train loss: 0.00377309, validation loss: 0.00451359
    Epoch  68: train loss: 0.00396709, validation loss: 0.00389344
    Epoch  69: train loss: 0.00372084, validation loss: 0.00455662
    Epoch  70: train loss: 0.00375983, validation loss: 0.00659405
    Epoch  71: train loss: 0.00411369, validation loss: 0.00319588
    Epoch  72: train loss: 0.00343619, validation loss: 0.00310895
    Epoch  73: train loss: 0.00337245, validation loss: 0.00313441
    Epoch  74: train loss: 0.00344687, validation loss: 0.00337706
    Epoch  75: train loss: 0.00364262, validation loss: 0.00305629
    Epoch  76: train loss: 0.00325462, validation loss: 0.00292753
    Epoch  77: train loss: 0.00310593, validation loss: 0.00289130
    Epoch  78: train loss: 0.00333657, validation loss: 0.00315406
    Epoch  79: train loss: 0.00327399, validation loss: 0.00302356
    Epoch  80: train loss: 0.00326501, validation loss: 0.00458416
    Epoch  81: train loss: 0.00319644, validation loss: 0.00321232
    Epoch  82: train loss: 0.00297700, validation loss: 0.00314786
    Epoch  83: train loss: 0.00315560, validation loss: 0.00364954
    Epoch  84: train loss: 0.00303592, validation loss: 0.00494946
    Epoch  85: train loss: 0.00295322, validation loss: 0.00404852
    Epoch  86: train loss: 0.00314734, validation loss: 0.00264067
    Epoch  87: train loss: 0.00288175, validation loss: 0.00269493
    Epoch  88: train loss: 0.00299053, validation loss: 0.00250429
    Epoch  89: train loss: 0.00265284, validation loss: 0.00277654
    Epoch  90: train loss: 0.00287037, validation loss: 0.00652205
    Epoch  91: train loss: 0.00320621, validation loss: 0.00252735
    Epoch  92: train loss: 0.00270760, validation loss: 0.00252588
    Epoch  93: train loss: 0.00265602, validation loss: 0.00243585
    Epoch  94: train loss: 0.00267965, validation loss: 0.00249292
    Epoch  95: train loss: 0.00251295, validation loss: 0.00258033
    Epoch  96: train loss: 0.00266445, validation loss: 0.00291114
    Epoch  97: train loss: 0.00286263, validation loss: 0.00256747
    Epoch  98: train loss: 0.00245153, validation loss: 0.00314988
    Epoch  99: train loss: 0.00262023, validation loss: 0.00292697
    Epoch 100: train loss: 0.00249519, validation loss: 0.00226936
    Epoch 101: train loss: 0.00248802, validation loss: 0.00224856
    Epoch 102: train loss: 0.00248330, validation loss: 0.00253684
    Epoch 103: train loss: 0.00240903, validation loss: 0.00251325
    Epoch 104: train loss: 0.00237967, validation loss: 0.00243032
    Epoch 105: train loss: 0.00227549, validation loss: 0.00206167
    Epoch 106: train loss: 0.00218241, validation loss: 0.00246376
    Epoch 107: train loss: 0.00224916, validation loss: 0.00209173
    Epoch 108: train loss: 0.00240168, validation loss: 0.00218054
    Epoch 109: train loss: 0.00217923, validation loss: 0.00198222
    Epoch 110: train loss: 0.00230895, validation loss: 0.00228715
    Epoch 111: train loss: 0.00215953, validation loss: 0.00234274
    Epoch 112: train loss: 0.00213709, validation loss: 0.00211844
    Epoch 113: train loss: 0.00210973, validation loss: 0.00203843
    Epoch 114: train loss: 0.00209313, validation loss: 0.00264051
    Epoch 115: train loss: 0.00244121, validation loss: 0.00218282
    Epoch 116: train loss: 0.00200982, validation loss: 0.00212742
    Epoch 117: train loss: 0.00201141, validation loss: 0.00201828
    Epoch 118: train loss: 0.00203203, validation loss: 0.00190810
    Epoch 119: train loss: 0.00205742, validation loss: 0.00182529
    Epoch 120: train loss: 0.00201869, validation loss: 0.00202707
    Epoch 121: train loss: 0.00199597, validation loss: 0.00180952
    Epoch 122: train loss: 0.00189482, validation loss: 0.00209971
    Epoch 123: train loss: 0.00245030, validation loss: 0.00235107
    Epoch 124: train loss: 0.00211455, validation loss: 0.00180742
    Epoch 125: train loss: 0.00176874, validation loss: 0.00178815
    Epoch 126: train loss: 0.00175496, validation loss: 0.00217022
    Epoch 127: train loss: 0.00189208, validation loss: 0.00206719
    Epoch 128: train loss: 0.00184735, validation loss: 0.00236503
    Epoch 129: train loss: 0.00183261, validation loss: 0.00191410
    Epoch 130: train loss: 0.00185281, validation loss: 0.00189244
    Epoch 131: train loss: 0.00164966, validation loss: 0.00160544
    Epoch 132: train loss: 0.00174577, validation loss: 0.00319794
    Epoch 133: train loss: 0.00202358, validation loss: 0.00205673
    Epoch 134: train loss: 0.00178098, validation loss: 0.00174782
    Epoch 135: train loss: 0.00162571, validation loss: 0.00157392
    Epoch 136: train loss: 0.00224547, validation loss: 0.00197109
    Epoch 137: train loss: 0.00181257, validation loss: 0.00155943
    Epoch 138: train loss: 0.00156519, validation loss: 0.00157839
    Epoch 139: train loss: 0.00155262, validation loss: 0.00177314
    Epoch 140: train loss: 0.00158969, validation loss: 0.00176624
    Epoch 141: train loss: 0.00161264, validation loss: 0.00201728
    Epoch 142: train loss: 0.00179312, validation loss: 0.00190871
    Epoch 143: train loss: 0.00182507, validation loss: 0.00155508
    Epoch 144: train loss: 0.00163039, validation loss: 0.00441373
    Epoch 145: train loss: 0.00170034, validation loss: 0.00183307
    Epoch 146: train loss: 0.00144177, validation loss: 0.00172955
    Epoch 147: train loss: 0.00172906, validation loss: 0.00338051
    Epoch 148: train loss: 0.00186909, validation loss: 0.00201458
    Epoch 149: train loss: 0.00160922, validation loss: 0.00152451
    Epoch 150: train loss: 0.00145538, validation loss: 0.00150882
    Epoch 151: train loss: 0.00164302, validation loss: 0.00200427
    Epoch 152: train loss: 0.00152015, validation loss: 0.00177693
    Epoch 153: train loss: 0.00148017, validation loss: 0.00172292
    Epoch 154: train loss: 0.00149707, validation loss: 0.00173788
    Epoch 155: train loss: 0.00155261, validation loss: 0.00159063
    Epoch 156: train loss: 0.00153494, validation loss: 0.00218051
    Epoch 157: train loss: 0.00170063, validation loss: 0.00138974
    Epoch 158: train loss: 0.00133955, validation loss: 0.00164791
    Epoch 159: train loss: 0.00167261, validation loss: 0.00141801
    Epoch 160: train loss: 0.00135793, validation loss: 0.00163650
    Epoch 161: train loss: 0.00133069, validation loss: 0.00157772
    Epoch 162: train loss: 0.00141800, validation loss: 0.00220752
    Epoch 163: train loss: 0.00142252, validation loss: 0.00160189
    Epoch 164: train loss: 0.00139955, validation loss: 0.00141898
    Epoch 165: train loss: 0.00129019, validation loss: 0.00132710
    Epoch 166: train loss: 0.00140015, validation loss: 0.00150085
    Epoch 167: train loss: 0.00153812, validation loss: 0.00175513
    Epoch 168: train loss: 0.00162252, validation loss: 0.00144463
    Epoch 169: train loss: 0.00140141, validation loss: 0.00131737
    Epoch 170: train loss: 0.00125315, validation loss: 0.00133116
    Epoch 171: train loss: 0.00128587, validation loss: 0.00160716
    Epoch 172: train loss: 0.00148668, validation loss: 0.00148051
    Epoch 173: train loss: 0.00133957, validation loss: 0.00211417
    Epoch 174: train loss: 0.00153662, validation loss: 0.00176566
    Epoch 175: train loss: 0.00132164, validation loss: 0.00136360
    Epoch 176: train loss: 0.00125364, validation loss: 0.00142442
    Epoch 177: train loss: 0.00121827, validation loss: 0.00132256
    Epoch 178: train loss: 0.00123788, validation loss: 0.00157878
    Epoch 179: train loss: 0.00138961, validation loss: 0.00139521
    Epoch 180: train loss: 0.00127230, validation loss: 0.00147543
    Epoch 181: train loss: 0.00132980, validation loss: 0.00146676
    Epoch 182: train loss: 0.00131916, validation loss: 0.00130245
    Epoch 183: train loss: 0.00132318, validation loss: 0.00157621
    Epoch 184: train loss: 0.00122836, validation loss: 0.00237077
    Epoch 185: train loss: 0.00176852, validation loss: 0.00211424
    Epoch 186: train loss: 0.00120928, validation loss: 0.00208305
    Epoch 187: train loss: 0.00206180, validation loss: 0.00211075
    Epoch 188: train loss: 0.00123358, validation loss: 0.00118391
    Epoch 189: train loss: 0.00112289, validation loss: 0.00114369
    Epoch 190: train loss: 0.00109211, validation loss: 0.00115871
    Epoch 191: train loss: 0.00111325, validation loss: 0.00135410
    Epoch 192: train loss: 0.00117998, validation loss: 0.00125950
    Epoch 193: train loss: 0.00132033, validation loss: 0.00116435
    Epoch 194: train loss: 0.00123164, validation loss: 0.00119297
    Epoch 195: train loss: 0.00123628, validation loss: 0.00114117
    Epoch 196: train loss: 0.00119157, validation loss: 0.00200442
    Epoch 197: train loss: 0.00119164, validation loss: 0.00126214
    Epoch 198: train loss: 0.00135684, validation loss: 0.00128723
    Epoch 199: train loss: 0.00110434, validation loss: 0.00145214
    Epoch 200: train loss: 0.00114354, validation loss: 0.00121961
    Epoch 201: train loss: 0.00123845, validation loss: 0.00122548
    Epoch 202: train loss: 0.00119689, validation loss: 0.00129707
    Epoch 203: train loss: 0.00129549, validation loss: 0.00172619
    Epoch 204: train loss: 0.00130671, validation loss: 0.00119561
    Epoch 205: train loss: 0.00113093, validation loss: 0.00114040
    Epoch 206: train loss: 0.00107067, validation loss: 0.00134028
    Epoch 207: train loss: 0.00106126, validation loss: 0.00130506
    Epoch 208: train loss: 0.00115197, validation loss: 0.00128711
    Epoch 209: train loss: 0.00121351, validation loss: 0.00119133
    Epoch 210: train loss: 0.00108956, validation loss: 0.00110090
    Epoch 211: train loss: 0.00103745, validation loss: 0.00119024
    Epoch 212: train loss: 0.00115014, validation loss: 0.00359916
    Epoch 213: train loss: 0.00180269, validation loss: 0.00110165
    Epoch 214: train loss: 0.00100940, validation loss: 0.00115686
    Epoch 215: train loss: 0.00122132, validation loss: 0.00118690
    Epoch 216: train loss: 0.00102829, validation loss: 0.00105288
    Epoch 217: train loss: 0.00113035, validation loss: 0.00161083
    Epoch 218: train loss: 0.00120304, validation loss: 0.00109303
    Epoch 219: train loss: 0.00107262, validation loss: 0.00101189
    Epoch 220: train loss: 0.00103222, validation loss: 0.00108345
    Epoch 221: train loss: 0.00107369, validation loss: 0.00126005
    Epoch 222: train loss: 0.00115262, validation loss: 0.00120875
    Epoch 223: train loss: 0.00112930, validation loss: 0.00112006
    Epoch 224: train loss: 0.00102221, validation loss: 0.00096797
    Epoch 225: train loss: 0.00108443, validation loss: 0.00144499
    Epoch 226: train loss: 0.00120224, validation loss: 0.00100510
    Epoch 227: train loss: 0.00099819, validation loss: 0.00113718
    Epoch 228: train loss: 0.00100571, validation loss: 0.00213398
    Epoch 229: train loss: 0.00112004, validation loss: 0.00135513
    Epoch 230: train loss: 0.00109633, validation loss: 0.00092874
    Epoch 231: train loss: 0.00099205, validation loss: 0.00119561
    Epoch 232: train loss: 0.00134267, validation loss: 0.00103823
    Epoch 233: train loss: 0.00172646, validation loss: 0.00162787
    Epoch 234: train loss: 0.00129124, validation loss: 0.00100946
    Epoch 235: train loss: 0.00091647, validation loss: 0.00095498
    Epoch 236: train loss: 0.00095084, validation loss: 0.00123949
    Epoch 237: train loss: 0.00100373, validation loss: 0.00098741
    Epoch 238: train loss: 0.00092436, validation loss: 0.00102499
    Epoch 239: train loss: 0.00098492, validation loss: 0.00101096
    Epoch 240: train loss: 0.00105150, validation loss: 0.00170253
    Epoch 241: train loss: 0.00111537, validation loss: 0.00104615
    Epoch 242: train loss: 0.00094953, validation loss: 0.00093894
    Epoch 243: train loss: 0.00105159, validation loss: 0.00119488
    Epoch 244: train loss: 0.00091832, validation loss: 0.00096785
    Epoch 245: train loss: 0.00100732, validation loss: 0.00123458
    Epoch 246: train loss: 0.00111832, validation loss: 0.00116431
    Epoch 247: train loss: 0.00096419, validation loss: 0.00126087
    Epoch 248: train loss: 0.00098255, validation loss: 0.00120453
    Epoch 249: train loss: 0.00103609, validation loss: 0.00297266
    Epoch 250: train loss: 0.00124524, validation loss: 0.00114636
    Epoch 251: train loss: 0.00102004, validation loss: 0.00121987
    Epoch 252: train loss: 0.00099135, validation loss: 0.00096586
    Epoch 253: train loss: 0.00097994, validation loss: 0.00094773
    Epoch 254: train loss: 0.00103071, validation loss: 0.00094039
    Epoch 255: train loss: 0.00087902, validation loss: 0.00108757
    Epoch 256: train loss: 0.00116164, validation loss: 0.00104841
    Epoch 257: train loss: 0.00093574, validation loss: 0.00096843
    Epoch 258: train loss: 0.00089855, validation loss: 0.00092135
    Epoch 259: train loss: 0.00103182, validation loss: 0.00096735
    Epoch 260: train loss: 0.00093665, validation loss: 0.00107030
    Epoch 261: train loss: 0.00121914, validation loss: 0.00109531
    Epoch 262: train loss: 0.00093579, validation loss: 0.00097617
    Epoch 263: train loss: 0.00087939, validation loss: 0.00103489
    Epoch 264: train loss: 0.00094577, validation loss: 0.00184936
    Epoch 265: train loss: 0.00116374, validation loss: 0.00090025
    Epoch 266: train loss: 0.00086100, validation loss: 0.00092744
    Epoch 267: train loss: 0.00092596, validation loss: 0.00110370
    Epoch 268: train loss: 0.00094005, validation loss: 0.00100758
    Epoch 269: train loss: 0.00091714, validation loss: 0.00135915
    Epoch 270: train loss: 0.00092583, validation loss: 0.00110098
    Epoch 271: train loss: 0.00096976, validation loss: 0.00123970
    Epoch 272: train loss: 0.00127745, validation loss: 0.00093362
    Epoch 273: train loss: 0.00099395, validation loss: 0.00099756
    Epoch 274: train loss: 0.00085863, validation loss: 0.00104818
    Epoch 275: train loss: 0.00092693, validation loss: 0.00121846
    Epoch 276: train loss: 0.00099986, validation loss: 0.00097108
    Epoch 277: train loss: 0.00097060, validation loss: 0.00091235
    Epoch 278: train loss: 0.00088086, validation loss: 0.00092233
    Epoch 279: train loss: 0.00085567, validation loss: 0.00094297
    Epoch 280: train loss: 0.00099869, validation loss: 0.00241286
    Epoch 281: train loss: 0.00096627, validation loss: 0.00088798
    Epoch 282: train loss: 0.00090142, validation loss: 0.00108350
    Epoch 283: train loss: 0.00098239, validation loss: 0.00091818
    Epoch 284: train loss: 0.00083357, validation loss: 0.00096326
    Epoch 285: train loss: 0.00090031, validation loss: 0.00089485
    Epoch 286: train loss: 0.00089467, validation loss: 0.00499492
    Epoch 287: train loss: 0.00152132, validation loss: 0.00112584
    Epoch 288: train loss: 0.00094044, validation loss: 0.00087090
    Epoch 289: train loss: 0.00081436, validation loss: 0.00096446
    Epoch 290: train loss: 0.00082687, validation loss: 0.00099858
    Epoch 291: train loss: 0.00084631, validation loss: 0.00096369
    Epoch 292: train loss: 0.00086067, validation loss: 0.00106924
    Epoch 293: train loss: 0.00089214, validation loss: 0.00087873
    Epoch 294: train loss: 0.00089185, validation loss: 0.00088478
    Epoch 295: train loss: 0.00104592, validation loss: 0.00109929
    Epoch 296: train loss: 0.00085723, validation loss: 0.00078156
    Epoch 297: train loss: 0.00083699, validation loss: 0.00087235
    Epoch 298: train loss: 0.00086460, validation loss: 0.00085079
    Epoch 299: train loss: 0.00099904, validation loss: 0.00106085
    Epoch 300: train loss: 0.00098141, validation loss: 0.00092075
    Epoch 301: train loss: 0.00081920, validation loss: 0.00099346
    Epoch 302: train loss: 0.00094096, validation loss: 0.00083690
    Epoch 303: train loss: 0.00106907, validation loss: 0.00193305
    Epoch 304: train loss: 0.00096057, validation loss: 0.00085076
    Epoch 305: train loss: 0.00081898, validation loss: 0.00096874
    Epoch 306: train loss: 0.00077259, validation loss: 0.00099447
    Epoch 307: train loss: 0.00080688, validation loss: 0.00080888
    Epoch 308: train loss: 0.00082848, validation loss: 0.00082018
    Epoch 309: train loss: 0.00085099, validation loss: 0.00095222
    Epoch 310: train loss: 0.00079871, validation loss: 0.00075801
    Epoch 311: train loss: 0.00086551, validation loss: 0.00093802
    Epoch 312: train loss: 0.00099750, validation loss: 0.00082952
    Epoch 313: train loss: 0.00082630, validation loss: 0.00089408
    Epoch 314: train loss: 0.00118076, validation loss: 0.00081055
    Epoch 315: train loss: 0.00075265, validation loss: 0.00082435
    Epoch 316: train loss: 0.00076296, validation loss: 0.00092707
    Epoch 317: train loss: 0.00084758, validation loss: 0.00085345
    Epoch 318: train loss: 0.00075547, validation loss: 0.00078662
    Epoch 319: train loss: 0.00088533, validation loss: 0.00216657
    Epoch 320: train loss: 0.00163788, validation loss: 0.00083488
    Epoch 321: train loss: 0.00075495, validation loss: 0.00085455
    Epoch 322: train loss: 0.00076250, validation loss: 0.00079223
    Epoch 323: train loss: 0.00076549, validation loss: 0.00083350
    Epoch 324: train loss: 0.00077615, validation loss: 0.00080033
    Epoch 325: train loss: 0.00076336, validation loss: 0.00076928
    Epoch 326: train loss: 0.00075855, validation loss: 0.00079906
    Epoch 327: train loss: 0.00075086, validation loss: 0.00107829
    Epoch 328: train loss: 0.00100992, validation loss: 0.00095608
    Epoch 329: train loss: 0.00079915, validation loss: 0.00085811
    Epoch 330: train loss: 0.00085024, validation loss: 0.00081684
    Epoch 331: train loss: 0.00076079, validation loss: 0.00078721
    Epoch 332: train loss: 0.00077217, validation loss: 0.00104725
    Epoch 333: train loss: 0.00156179, validation loss: 0.00074761
    Epoch 334: train loss: 0.00073946, validation loss: 0.00147021
    Epoch 335: train loss: 0.00089103, validation loss: 0.00075132
    Epoch 336: train loss: 0.00071530, validation loss: 0.00086720
    Epoch 337: train loss: 0.00074078, validation loss: 0.00078406
    Epoch 338: train loss: 0.00074182, validation loss: 0.00092171
    Epoch 339: train loss: 0.00078194, validation loss: 0.00086428
    Epoch 340: train loss: 0.00081480, validation loss: 0.00078926
    Epoch 341: train loss: 0.00076006, validation loss: 0.00080109
    Epoch 342: train loss: 0.00078100, validation loss: 0.00074002
    Epoch 343: train loss: 0.00073423, validation loss: 0.00083269
    Epoch 344: train loss: 0.00072414, validation loss: 0.00103983
    Epoch 345: train loss: 0.00101394, validation loss: 0.00106113
    Epoch 346: train loss: 0.00082145, validation loss: 0.00085152
    Epoch 347: train loss: 0.00078141, validation loss: 0.00080331
    Epoch 348: train loss: 0.00069564, validation loss: 0.00082213
    Epoch 349: train loss: 0.00077503, validation loss: 0.00131843
    Epoch 350: train loss: 0.00085349, validation loss: 0.00075777
    Epoch 351: train loss: 0.00067773, validation loss: 0.00078871
    Epoch 352: train loss: 0.00066166, validation loss: 0.00070740
    Epoch 353: train loss: 0.00103420, validation loss: 0.00078153
    Epoch 354: train loss: 0.00070738, validation loss: 0.00071634
    Epoch 355: train loss: 0.00067165, validation loss: 0.00075299
    Epoch 356: train loss: 0.00092475, validation loss: 0.00294023
    Epoch 357: train loss: 0.00082633, validation loss: 0.00073762
    Epoch 358: train loss: 0.00068358, validation loss: 0.00068444
    Epoch 359: train loss: 0.00102637, validation loss: 0.00188669
    Epoch 360: train loss: 0.00082894, validation loss: 0.00077172
    Epoch 361: train loss: 0.00076420, validation loss: 0.00087094
    Epoch 362: train loss: 0.00065216, validation loss: 0.00071045
    Epoch 363: train loss: 0.00070877, validation loss: 0.00067883
    Epoch 364: train loss: 0.00063403, validation loss: 0.00069902
    Epoch 365: train loss: 0.00068384, validation loss: 0.00071971
    Epoch 366: train loss: 0.00063582, validation loss: 0.00078469
    Epoch 367: train loss: 0.00080967, validation loss: 0.00067853
    Epoch 368: train loss: 0.00063606, validation loss: 0.00072633
    Epoch 369: train loss: 0.00060736, validation loss: 0.00113455
    Epoch 370: train loss: 0.00076479, validation loss: 0.00062143
    Epoch 371: train loss: 0.00065260, validation loss: 0.00079042
    Epoch 372: train loss: 0.00073676, validation loss: 0.00071641
    Epoch 373: train loss: 0.00075117, validation loss: 0.00063622
    Epoch 374: train loss: 0.00061058, validation loss: 0.00068547
    Epoch 375: train loss: 0.00073223, validation loss: 0.00078174
    Epoch 376: train loss: 0.00079426, validation loss: 0.00099509
    Epoch 377: train loss: 0.00072176, validation loss: 0.00084782
    Epoch 378: train loss: 0.00061672, validation loss: 0.00064491
    Epoch 379: train loss: 0.00061419, validation loss: 0.00075005
    Epoch 380: train loss: 0.00112049, validation loss: 0.00063819
    Epoch 381: train loss: 0.00057194, validation loss: 0.00059208
    Epoch 382: train loss: 0.00055642, validation loss: 0.00068847
    Epoch 383: train loss: 0.00057515, validation loss: 0.00074646
    Epoch 384: train loss: 0.00065456, validation loss: 0.00061999
    Epoch 385: train loss: 0.00061405, validation loss: 0.00092507
    Epoch 386: train loss: 0.00067584, validation loss: 0.00061326
    Epoch 387: train loss: 0.00070584, validation loss: 0.00099933
    Epoch 388: train loss: 0.00071693, validation loss: 0.00090925
    Epoch 389: train loss: 0.00069401, validation loss: 0.00074272
    Epoch 390: train loss: 0.00057821, validation loss: 0.00060813
    Epoch 391: train loss: 0.00058593, validation loss: 0.00055908
    Epoch 392: train loss: 0.00062935, validation loss: 0.00057303
    Epoch 393: train loss: 0.00055196, validation loss: 0.00064129
    Epoch 394: train loss: 0.00126625, validation loss: 0.00062421
    Epoch 395: train loss: 0.00056335, validation loss: 0.00074653
    Epoch 396: train loss: 0.00056355, validation loss: 0.00062698
    Epoch 397: train loss: 0.00053850, validation loss: 0.00056042
    Epoch 398: train loss: 0.00052870, validation loss: 0.00056581
    Epoch 399: train loss: 0.00053052, validation loss: 0.00059617
    Epoch 400: train loss: 0.00053607, validation loss: 0.00088505
    Epoch 401: train loss: 0.00066666, validation loss: 0.00186690
    Epoch 402: train loss: 0.00119052, validation loss: 0.00067151
    Epoch 403: train loss: 0.00054852, validation loss: 0.00063310
    Epoch 404: train loss: 0.00052614, validation loss: 0.00054202
    Epoch 405: train loss: 0.00053641, validation loss: 0.00072213
    Epoch 406: train loss: 0.00055394, validation loss: 0.00055267
    Epoch 407: train loss: 0.00055834, validation loss: 0.00054785
    Epoch 408: train loss: 0.00055462, validation loss: 0.00082774
    Epoch 409: train loss: 0.00065174, validation loss: 0.00053766
    Epoch 410: train loss: 0.00060951, validation loss: 0.00060815
    Epoch 411: train loss: 0.00052364, validation loss: 0.00057097
    Epoch 412: train loss: 0.00053981, validation loss: 0.00055151
    Epoch 413: train loss: 0.00052171, validation loss: 0.00052500
    Epoch 414: train loss: 0.00059586, validation loss: 0.00059979
    Epoch 415: train loss: 0.00173218, validation loss: 0.00109936
    Epoch 416: train loss: 0.00075883, validation loss: 0.00052755
    Epoch 417: train loss: 0.00051721, validation loss: 0.00050778
    Epoch 418: train loss: 0.00049349, validation loss: 0.00051720
    Epoch 419: train loss: 0.00050352, validation loss: 0.00058975
    Epoch 420: train loss: 0.00050543, validation loss: 0.00054280
    Epoch 421: train loss: 0.00049380, validation loss: 0.00053791
    Epoch 422: train loss: 0.00050584, validation loss: 0.00055699
    Epoch 423: train loss: 0.00048801, validation loss: 0.00051431
    Epoch 424: train loss: 0.00053522, validation loss: 0.00058704
    Epoch 425: train loss: 0.00058705, validation loss: 0.00056832
    Epoch 426: train loss: 0.00054814, validation loss: 0.00050946
    Epoch 427: train loss: 0.00048262, validation loss: 0.00051916
    Epoch 428: train loss: 0.00053806, validation loss: 0.00060748
    Epoch 429: train loss: 0.00056472, validation loss: 0.00072669
    Epoch 430: train loss: 0.00052885, validation loss: 0.00053735
    Epoch 431: train loss: 0.00052596, validation loss: 0.00060280
    Epoch 432: train loss: 0.00068272, validation loss: 0.00081625
    Epoch 433: train loss: 0.00088503, validation loss: 0.00119153
    Epoch 434: train loss: 0.00080632, validation loss: 0.00054437
    Epoch 435: train loss: 0.00048190, validation loss: 0.00054693
    Epoch 436: train loss: 0.00049144, validation loss: 0.00053110
    Epoch 437: train loss: 0.00050435, validation loss: 0.00053237
    Epoch 438: train loss: 0.00051228, validation loss: 0.00078325
    Epoch 439: train loss: 0.00079294, validation loss: 0.00050444
    Epoch 440: train loss: 0.00050953, validation loss: 0.00050212
    Epoch 441: train loss: 0.00052415, validation loss: 0.00059089
    Epoch 442: train loss: 0.00055400, validation loss: 0.00052182
    Epoch 443: train loss: 0.00052482, validation loss: 0.00052627
    Epoch 444: train loss: 0.00050417, validation loss: 0.00051749
    Epoch 445: train loss: 0.00053329, validation loss: 0.00103713
    Epoch 446: train loss: 0.00061278, validation loss: 0.00052381
    Epoch 447: train loss: 0.00052905, validation loss: 0.00097861
    Epoch 448: train loss: 0.00060001, validation loss: 0.00061146
    Epoch 449: train loss: 0.00053926, validation loss: 0.00050589
    Epoch 450: train loss: 0.00049979, validation loss: 0.00048631
    Epoch 451: train loss: 0.00048773, validation loss: 0.00050589
    Epoch 452: train loss: 0.00057042, validation loss: 0.00057504
    Epoch 453: train loss: 0.00054516, validation loss: 0.00055242
    Epoch 454: train loss: 0.00052804, validation loss: 0.00050957
    Epoch 455: train loss: 0.00059130, validation loss: 0.00053183
    Epoch 456: train loss: 0.00076837, validation loss: 0.00058264
    Epoch 457: train loss: 0.00050944, validation loss: 0.00048359
    Epoch 458: train loss: 0.00064279, validation loss: 0.00056913
    Epoch 459: train loss: 0.00053545, validation loss: 0.00051711
    Epoch 460: train loss: 0.00047284, validation loss: 0.00045717
    Epoch 461: train loss: 0.00047856, validation loss: 0.00052342
    Epoch 462: train loss: 0.00054410, validation loss: 0.00051787
    Epoch 463: train loss: 0.00049290, validation loss: 0.00056659
    Epoch 464: train loss: 0.00061312, validation loss: 0.00047896
    Epoch 465: train loss: 0.00055769, validation loss: 0.00047728
    Epoch 466: train loss: 0.00046093, validation loss: 0.00049182
    Epoch 467: train loss: 0.00055168, validation loss: 0.00064541
    Epoch 468: train loss: 0.00100547, validation loss: 0.00050028
    Epoch 469: train loss: 0.00045833, validation loss: 0.00046892
    Epoch 470: train loss: 0.00047499, validation loss: 0.00055394
    Epoch 471: train loss: 0.00058058, validation loss: 0.00048047
    Epoch 472: train loss: 0.00058883, validation loss: 0.00057088
    Epoch 473: train loss: 0.00050628, validation loss: 0.00046559
    Epoch 474: train loss: 0.00046228, validation loss: 0.00044857
    Epoch 475: train loss: 0.00047975, validation loss: 0.00062482
    Epoch 476: train loss: 0.00052015, validation loss: 0.00049028
    Epoch 477: train loss: 0.00069017, validation loss: 0.00084841
    Epoch 478: train loss: 0.00054894, validation loss: 0.00052269
    Epoch 479: train loss: 0.00050262, validation loss: 0.00056461
    Epoch 480: train loss: 0.00055599, validation loss: 0.00063470
    Epoch 481: train loss: 0.00053296, validation loss: 0.00060813
    Epoch 482: train loss: 0.00048785, validation loss: 0.00048495
    Epoch 483: train loss: 0.00048136, validation loss: 0.00050461
    Epoch 484: train loss: 0.00050100, validation loss: 0.00062679
    Epoch 485: train loss: 0.00051928, validation loss: 0.00052477
    Epoch 486: train loss: 0.00054245, validation loss: 0.00051511
    Epoch 487: train loss: 0.00047895, validation loss: 0.00049170
    Epoch 488: train loss: 0.00065683, validation loss: 0.00063497
    Epoch 489: train loss: 0.00055582, validation loss: 0.00046024
    Epoch 490: train loss: 0.00045986, validation loss: 0.00049115
    Epoch 491: train loss: 0.00045358, validation loss: 0.00049144
    Epoch 492: train loss: 0.00052923, validation loss: 0.00052943
    Epoch 493: train loss: 0.00046403, validation loss: 0.00046569
    Epoch 494: train loss: 0.00048530, validation loss: 0.00045099
    Epoch 495: train loss: 0.00051436, validation loss: 0.00044140
    Epoch 496: train loss: 0.00055119, validation loss: 0.00056624
    Epoch 497: train loss: 0.00062318, validation loss: 0.00111081
    Epoch 498: train loss: 0.00077277, validation loss: 0.00044776
    Epoch 499: train loss: 0.00046284, validation loss: 0.00044364
    Epoch 500: train loss: 0.00043643, validation loss: 0.00054867
    Finished Training

```python
sys_koop_dnn.A
```

    array([[ 1.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00],
           [ 0.0000000e+00,  1.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             9.5007028e-03,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00],
           [ 0.0000000e+00,  0.0000000e+00,  1.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  9.1959769e-03,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00],
           [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  1.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  2.6239881e-02,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00],
           [ 1.1247686e-03, -5.5586908e-04,  0.0000000e+00, -7.9869367e-03,
             1.0000000e+00,  0.0000000e+00, -2.0529462e-04, -8.7566916e-03,
            -3.1576012e-03,  3.0395496e-03, -3.7833580e-03, -3.7106494e-03,
            -1.7403871e-03, -4.8968536e-03, -5.9845792e-03, -1.1257441e-02,
            -1.3937379e-03],
           [ 8.7711617e-04,  0.0000000e+00,  0.0000000e+00, -2.9414403e-04,
             0.0000000e+00,  1.0000000e+00, -3.0430441e-05,  4.2124046e-03,
             3.4934634e-04, -2.0425599e-03, -3.5547593e-03,  1.0319303e-02,
            -2.8737064e-03, -2.4449585e-03,  2.3388036e-02, -1.9125523e-02,
            -2.1245836e-03],
           [ 0.0000000e+00, -2.5778736e-05,  0.0000000e+00,  1.4070771e-04,
            -1.3597561e-05,  0.0000000e+00,  9.9972099e-01, -4.4351106e-04,
             1.5781533e-03, -9.1366016e-04, -2.2260526e-03,  1.3512012e-04,
            -2.3370169e-04,  7.7043049e-04,  0.0000000e+00,  9.0118170e-05,
            -1.1500117e-03],
           [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
            -8.4711217e-05,  0.0000000e+00,  1.1088921e-02,  9.9909866e-01,
             0.0000000e+00,  0.0000000e+00, -1.4925176e-04,  5.0574187e-03,
             0.0000000e+00,  0.0000000e+00, -2.5576926e-03,  0.0000000e+00,
             0.0000000e+00],
           [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             1.4601325e-04,  0.0000000e+00,  1.4200504e-03, -4.7820009e-05,
             1.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00],
           [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00, -5.2168230e-03,  2.3085156e-03,
             0.0000000e+00,  9.9836934e-01,  0.0000000e+00, -9.0420031e-04,
             0.0000000e+00,  0.0000000e+00, -9.1821078e-04, -2.0684330e-03,
             0.0000000e+00],
           [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  4.2709764e-03,  0.0000000e+00,
             0.0000000e+00, -9.9564779e-05,  1.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00],
           [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  1.2819145e-02,
             0.0000000e+00, -9.5829647e-03,  5.1245140e-04,  1.0027224e+00,
             1.8274348e-04,  0.0000000e+00, -1.7213659e-02, -1.2209680e-02,
             0.0000000e+00],
           [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  1.2023482e-04,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             1.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00],
           [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
            -1.1200932e-04,  0.0000000e+00,  6.1595715e-03, -2.1398674e-04,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  1.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00],
           [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  2.6545203e-03, -7.2301994e-03,
             0.0000000e+00,  1.2142119e-02,  0.0000000e+00,  3.2542937e-02,
             0.0000000e+00, -2.5940083e-03,  9.9354249e-01,  0.0000000e+00,
            -4.8921015e-03],
           [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  7.3437497e-04,
             0.0000000e+00,  0.0000000e+00,  1.2775570e-02, -2.2624282e-03,
             0.0000000e+00, -8.0450280e-03, -9.2068862e-04, -2.0239200e-02,
             0.0000000e+00,  4.1580144e-03,  8.5503170e-03,  1.0023123e+00,
             5.4848753e-03],
           [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  1.5258008e-03,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
             0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  5.7066313e-06,
             1.0000000e+00]], dtype=float32)

### Evaluate open loop prediction performance

We first evaluate the open loop prediction performance of the proposed method.
This is done by generating a new data set in the same way as the training set, predicting the evolution of the system
with the control sequence of each trajectory executed in the data set with each of the models, and finally comparing
the mean and standard deviation of the error between the true and predicted evolution over the trajectories. The
experimental results support what is to be expected from the theory as the error in the $y$ and $z$ terms are
significantly lower for the bEDMD method than both DMD and EDMD. The reason for this
improvement is that the bEDMD method can capture the nonlinearities present in the actuation matrix of the
$(y,z)$-dynamics.

```python
# Prediction performance evaluation parameters:
folder_plots = 'figures/'                                  # Path to save plots
n_traj_ol = 50                                                     # Number of trajectories to execute, open loop
t_eval = t_train
```

```python
from koopman_core.util import evaluate_ol_pred
from tabulate import tabulate

xs_test = np.empty((n_traj_ol, t_eval.shape[0], n))
us_test = np.empty((n_traj_ol, t_eval.shape[0]-1, m))

for ii in range(n_traj_ol):
    x0 = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
    set_pt_dc = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
    mpc_trajgen = MPCController(nominal_sys, n_pred_dc, dt, umin, umax, xmin, xmax, QN_trajgen, R_trajgen,
                                QN_trajgen, set_pt_dc)
    mpc_trajgen.eval(x0, 0)
    xd = mpc_trajgen.parse_result().T

    while xd[0,0] is None:
        x0 = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
        set_pt_dc = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
        mpc_trajgen = MPCController(nominal_sys, n_pred_dc, dt, umin, umax, xmin, xmax, QN_trajgen, R_trajgen,
                                    QN_trajgen, set_pt_dc)
        mpc_trajgen.eval(x0, 0)
        xd = mpc_trajgen.parse_result().T

    output = QuadrotorPdOutput(system, xd, t_eval, n, m)
    pd_controller = PDController(output, K_dc_p, K_dc_d)
    perturbed_pd_controller = PerturbedController(system, pd_controller, noise_var, const_offset=hover_thrust)

    xs_test[ii,:,:], us_test[ii,:,:] = system.simulate(x0, perturbed_pd_controller, t_eval)

#    mdl_lst = [sys_koop_dnn, sys_bedmd, sys_edmd, sys_dmd]
    mdl_lst = [sys_koop_dnn, sys_bedmd]
    mdl_names = ['Koop DNN', 'bEDMD', 'EDMD', 'DMD']
    error, mse, std = [], [], []

    for sys in mdl_lst:
        err_tmp, mse_tmp, std_tmp = evaluate_ol_pred(sys, xs_test, t_eval, us=us_test)
        error.append(err_tmp)
        mse.append(mse_tmp)
        std.append(std_tmp)

print('\nOpen loop performance statistics:')
table_data = []
for name, mse_mdl, std_mdl in zip(mdl_names, mse, std):
    table_data.append([name, "{:.5f}".format(mse_mdl), "{:.5f}".format(std_mdl)])

print(tabulate(table_data, headers=['Mean squared error', 'Standard deviation']))
```

    Open loop performance statistics:
                Mean squared error    Standard deviation
    --------  --------------------  --------------------
    Koop DNN               0.00916               0.09368
    bEDMD                  0.00242               0.04922

```python
import matplotlib.pyplot as plt
import matplotlib

figwidth = 12
lw = 2
fs = 14
y_lim_gain = 1.2
row = 2
col = n/row

#Plot open loop results:
plt.figure(figsize=(figwidth,4))
axs = [plt.subplot(row,col,jj+1) for jj in range(n)]

for ii, err in enumerate(error):
    err_mean = np.mean(err, axis=0)
    err_std = np.std(err, axis=0)

    for jj in range(n):
        axs[jj].plot(t_eval[1:], err_mean[:,jj], label=mdl_names[ii])
        axs[jj].fill_between(t_eval[1:], err_mean[:,jj]-err_std[:,jj], err_mean[:,jj]+err_std[:,jj], alpha=0.1)

for jj in range(n):
    axs[jj].set_xlabel('Time (sec)', fontsize=fs)
    axs[jj].set_ylabel('$x_'+ str(jj+1) + '$', fontsize=fs)

plt.legend(frameon=False, fontsize=fs)
stitle=plt.suptitle('Open loop prediction performance of learned models', fontsize=fs+2)
matplotlib.rcParams['pdf.fonttype'] = 42
matplotlib.rcParams['ps.fonttype'] = 42
plt.savefig(folder_plots + 'koop_sys_prediction.pdf', format='pdf', dpi=2400, bbox_extra_artists=(stitle,), bbox_inches="tight")
plt.show()
```

# Design trajectories based on learned models

We now study the closed loop performance of the control design.

```python
#Closed loop performance evaluation parameters:
traj_length=250
t_eval = dt * np.arange(traj_length+1)                       # Simulation time points, closed loop
Q_mpc = sc.sparse.diags([0,0,0,0,0,0])                       # State penalty matrix, trajectory generation
#QN_mpc = sc.sparse.diags([1e5,1e5,1e5,1e5,1e5,1e5])         # Final state penalty matrix, trajectory generation
QN_mpc = Q_mpc
R_mpc = sc.sparse.eye(m)                                     # Actuation penalty matrix, trajectory generation

ctrl_offset = np.array([hover_thrust, hover_thrust])

# Design trajectory:
x0_cl = np.array([0., 0., 0., 0., 0., 0.])                   # Initial value, closed loop trajectory
set_pt_cl = np.array([1.5, 0.5, 0., 0., 0., 0.])              # Desired final value, closed loop trajectory
xmax = np.array([2, 2, np.pi/3, 1.,1.,1.])                          # State constraints, trajectory generation
xmin = -xmax
term_constraint=True

# Define initial solution for SQP algorithm:
x_init = np.linspace(x0_cl, set_pt_cl, int(traj_length)+1)
u_init = np.zeros((m,traj_length)).T
```

#### Design controllers for learned DMD, EDMD, and bEDMD models

```python
from koopman_core.controllers import MPCController, NonlinearMPCController, BilinearMPCController

# Define DMD-based controller:
controller_dmd = MPCController(sys_dmd, traj_length, dt, umin, umax, xmin, xmax, Q_mpc, R_mpc, QN_mpc, set_pt_cl, terminal_constraint=term_constraint, const_offset=ctrl_offset)

# Define EDMD-based controller:
controller_edmd = MPCController(sys_edmd, traj_length, dt, umin, umax, xmin, xmax, Q_mpc, R_mpc, QN_mpc, set_pt_cl, terminal_constraint=term_constraint, const_offset=ctrl_offset)

# Define bEDMD-based controller:
controller_bedmd = BilinearMPCController(sys_bedmd, traj_length, dt, umin, umax, xmin, xmax, Q_mpc, R_mpc, QN_mpc, set_pt_cl, terminal_constraint=term_constraint, const_offset=ctrl_offset)
z0_cl = sys_bedmd.basis(x0_cl.reshape((1,-1))).squeeze()
z_init = sys_bedmd.basis(x_init)
controller_bedmd.construct_controller(z_init, u_init)
```

#### Design controller using full knowledge of nonlinear controller

```python
quadrotor_d = PlanarQuadrotorForceInputDiscrete(mass, inertia, prop_arm, g=gravity, dt=dt)
controller_nmpc = NonlinearMPCController(quadrotor_d, traj_length, dt, umin+hover_thrust, umax+hover_thrust, xmin, xmax, Q_mpc, R_mpc, QN_mpc, set_pt_cl, terminal_constraint=term_constraint)
controller_nmpc.construct_controller(x_init, u_init+hover_thrust)
```

#### Design trajectories with the contructed MPCs

```python
max_iter = 50

controller_dmd.eval(x0_cl, 0)
xr_dmd = controller_dmd.parse_result()
ur_dmd = controller_dmd.get_control_prediction() + hover_thrust

controller_edmd.eval(x0_cl, 0)
xr_edmd = sys_edmd.C@controller_edmd.parse_result()
ur_edmd = controller_edmd.get_control_prediction() + hover_thrust

controller_bedmd.solve_to_convergence(z0_cl, 0., z_init, u_init, max_iter=max_iter)
xr_bedmd = sys_bedmd.C@controller_bedmd.get_state_prediction().T
ur_bedmd = controller_bedmd.get_control_prediction().T + hover_thrust
```

```python

controller_nmpc.solve_to_convergence(x0_cl, 0., x_init, u_init + ctrl_offset.reshape(1,-1), max_iter=max_iter)
xr_nmpc = controller_nmpc.get_state_prediction().T
ur_nmpc = controller_nmpc.get_control_prediction().T
```

#### Simulate designed trajectories open loop

```python
ol_controller_dmd = OpenLoopController(quadrotor, ur_dmd.T, t_eval[:-1])
xs_dmd, us_dmd = quadrotor.simulate(x0_cl, ol_controller_dmd, t_eval)
xs_dmd, us_dmd = xs_dmd.T, us_dmd.T

ol_controller_edmd = OpenLoopController(quadrotor, ur_edmd.T, t_eval[:-1])
xs_edmd, us_edmd = quadrotor.simulate(x0_cl, ol_controller_edmd, t_eval)
xs_edmd, us_edmd = xs_edmd.T, us_edmd.T

ol_controller_bedmd = OpenLoopController(quadrotor, ur_bedmd.T, t_eval[:-1])
xs_bedmd, us_bedmd = quadrotor.simulate(x0_cl, ol_controller_bedmd, t_eval)
xs_bedmd, us_bedmd = xs_bedmd.T, us_bedmd.T

ol_controller_nmpc = OpenLoopController(quadrotor, ur_nmpc.T, t_eval[:-1])
xs_nmpc, us_nmpc = quadrotor.simulate(x0_cl, ol_controller_nmpc, t_eval)
xs_nmpc, us_nmpc = xs_nmpc.T, us_nmpc.T
```

##### Compare performance

```python
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator

plot_inds = [0, 1, 2, 3, 4, 5, 0, 1]
subplot_inds = [1, 2, 3, 5, 6, 7, 4, 8]
labels = ['$y$ (m)', '$z$ (m)', '$\\theta$ (rad)', '$\\dot{y}$ (m/s)','$\\dot{z}$ (m/s)', '$\\dot{\\theta}$', '$T_1$ (N)','$T_2$ (N)']
titles = ['y-coordinates', 'z-coordinates', '$\\theta$-coordinates', 'Control inputs']
colors = ['tab:blue', 'tab:orange', 'tab:brown', 'tab:cyan']

plt.figure(figsize=(12,4))
#plt.suptitle('Trajectory designed with model predictive controllers\nsolid lines - designed trajectory | dashed lines - open loop simulated trajectory | black dotted lines - state/actuation bounds')
for ii in range(8):
    ind = plot_inds[ii]
    if ii < 6:
        ax = plt.subplot(2,4,subplot_inds[ii])
        plt.plot(t_eval, xr_dmd[ind,:], colors[0], label='DMD MPC')
        plt.plot(t_eval, xr_edmd[ind, :], colors[1], label='EDMD MPC')
        plt.plot(t_eval, xr_bedmd[ind, :], colors[2], label='K-MPC')
        plt.plot(t_eval, xr_nmpc[ind,:], colors[3], label='NMPC')

        plt.plot(t_eval, xs_dmd[ind,:], '--', color=colors[0], linewidth=1)
        plt.plot(t_eval, xs_edmd[ind, :], '--', color=colors[1], linewidth=1)
        plt.plot(t_eval, xs_bedmd[ind, :], '--', color=colors[2], linewidth=1)
        plt.plot(t_eval, xs_nmpc[ind,:], '--', color=colors[3], linewidth=1)

        plt.scatter(t_eval[0], x0_cl[ind], color='g')
        plt.scatter(t_eval[-1], set_pt_cl[ind], color='r')
        plt.ylabel(labels[ind])
        ax.yaxis.set_major_locator(MaxNLocator(integer=True))
        if ii >= 3:
            plt.plot([0, t_eval[-1]], [xmax[ind], xmax[ind]], ':k')
            plt.plot([0, t_eval[-1]], [xmin[ind], xmin[ind]], ':k')
            #plt.ylim(xmin[ind]-0.1,xmax[ind]+0.1)
        if subplot_inds[ii]==1:
            plt.legend(loc='upper left', frameon=False)
    elif ii < 8:
        ax = plt.subplot(2,4,subplot_inds[ii])
        plt.plot(t_eval[:-1],ur_dmd[ind,:], color=colors[0], label='DMD MPC')
        plt.plot(t_eval[:-1], ur_edmd[ind, :], color=colors[1], label='EDMD MPC')
        plt.plot(t_eval[:-1], ur_bedmd[ind, :], color=colors[2], label='K-NMPC')
        plt.plot(t_eval[:-1],ur_nmpc[ind,:], color=colors[3], label='NMPC')
        plt.plot([0, t_eval[-1]], [umax[ind]+hover_thrust, umax[ind]+hover_thrust], ':k')
        plt.plot([0, t_eval[-1]], [umin[ind]+hover_thrust, umin[ind]+hover_thrust], ':k')
        plt.ylabel(labels[ii])
        ax.yaxis.set_major_locator(MaxNLocator(integer=True))

    if subplot_inds[ii] > 4:
        plt.xlabel('Time (sec)')
    else:
        plt.title(titles[subplot_inds[ii]-1])

if save_figures:
    matplotlib.rcParams['pdf.fonttype'] = 42
    matplotlib.rcParams['ps.fonttype'] = 42
    plt.tight_layout()
    plt.savefig(dropbox_folder + 'planar_quad_trajectory.pdf', format='pdf', dpi=2400)
plt.show()

cost_ref_dmd = (xr_dmd[:,-1]-set_pt_cl).T@QN_mpc@(xr_dmd[:,-1]-set_pt_cl) + np.sum(np.diag(ur_dmd.T@R_mpc@ur_dmd))
cost_ref_edmd = (xr_edmd[:,-1]-set_pt_cl).T@QN_mpc@(xr_edmd[:,-1]-set_pt_cl) + np.sum(np.diag(ur_edmd.T@R_mpc@ur_edmd))
cost_ref_bedmd = (xr_bedmd[:,-1]-set_pt_cl).T@QN_mpc@(xr_bedmd[:,-1]-set_pt_cl) + np.sum(np.diag(ur_bedmd.T@R_mpc@ur_bedmd))
cost_ref_nmpc = (xr_nmpc[:,-1]-set_pt_cl).T@QN_mpc@(xr_nmpc[:,-1]-set_pt_cl) + np.sum(np.diag(ur_nmpc.T@R_mpc@ur_nmpc))

dist_ol_dmd = np.linalg.norm(xs_dmd[:,-1] - set_pt_cl)
dist_ol_edmd = np.linalg.norm(xs_edmd[:,-1] - set_pt_cl)
dist_ol_bedmd = np.linalg.norm(xs_bedmd[:,-1] - set_pt_cl)
dist_ol_nmpc = np.linalg.norm(xs_nmpc[:,-1] - set_pt_cl)

print('Solution statistics:\n')
print(tabulate([['DMD MPC', "{:.4f}".format(cost_ref_dmd/cost_ref_nmpc), "{:.4f}".format(dist_ol_dmd), '-','-',sum(controller_dmd.comp_time)],
                ['EDMD MPC', "{:.4f}".format(cost_ref_edmd/cost_ref_nmpc), "{:.4f}".format(dist_ol_edmd),'-','-',sum(controller_edmd.comp_time)],
                ['bEDMD MPC', "{:.4f}".format(cost_ref_bedmd/cost_ref_nmpc), "{:.4f}".format(dist_ol_bedmd), len(controller_bedmd.x_iter), "{:.4f}".format(np.mean(controller_bedmd.comp_time)), sum(controller_bedmd.comp_time)],
                ['NMPC (benchmark)', 1, "{:.4f}".format(dist_ol_nmpc), len(controller_nmpc.x_iter), "{:.4f}".format(np.mean(controller_nmpc.comp_time)), sum(controller_nmpc.comp_time)]],
               headers=['Normalized cost,\ndesigned trajectory', 'Realized terminal,\nerror', '# of SQP\niterations','Mean comp. time\nper iteration (secs)', 'Total comp.\ntime (secs)']))
```

#### Study evolution of the solution after each iteration of the SQP-algorithm

```python
n_iter = min(len(controller_nmpc.x_iter),len(controller_bedmd.x_iter))

# Calculate cost after each iteration:
iter_cost_bedmd, iter_cost_nmpc = [], []

ol_controller_init = OpenLoopController(quadrotor, u_init, t_eval[:-1])
xs_init, _ = quadrotor.simulate(x0_cl, ol_controller_init, t_eval)
xs_init, us_init = xs_init.T, u_init.T+hover_thrust
init_cost = (xs_init[:,-1]-set_pt_cl).T@QN_mpc@(xs_init[:,-1]-set_pt_cl) + np.sum(np.diag(us_init.T@R_mpc@us_init))
iter_cost_bedmd = [init_cost]
iter_cost_nmpc = [init_cost]
iter_norm_dist_bedmd = [np.linalg.norm(xs_init[:,-1]-set_pt_cl)]
iter_norm_dist_nmpc = [np.linalg.norm(xs_init[:,-1]-set_pt_cl)]

for ii in range(len(controller_bedmd.x_iter)):
    ur_bedmd_iter = controller_bedmd.u_iter[ii].T+hover_thrust
    ol_controller_bedmd_iter = OpenLoopController(quadrotor, ur_bedmd_iter, t_eval[:-1])
    xs_bedmd_iter, _ = quadrotor.simulate(x0_cl, ol_controller_bedmd_iter, t_eval)
    xs_bedmd_iter, us_bedmd_iter = xs_bedmd_iter.T, ur_bedmd_iter.T
    iter_cost_bedmd.append((xs_bedmd_iter[:,-1]-set_pt_cl).T@QN_mpc@(xs_bedmd_iter[:,-1]-set_pt_cl) + np.sum(np.diag(us_bedmd_iter.T@R_mpc@us_bedmd_iter)))
    iter_norm_dist_bedmd.append(np.linalg.norm(xs_bedmd_iter[:,-1]-set_pt_cl))

for ii in range(len(controller_nmpc.x_iter)):
    ur_nmpc_iter = controller_nmpc.u_iter[ii].T
    ol_controller_nmpc_iter = OpenLoopController(quadrotor, ur_nmpc_iter, t_eval[:-1])
    xs_nmpc_iter, _ = quadrotor.simulate(x0_cl, ol_controller_nmpc_iter, t_eval)
    xs_nmpc_iter, us_nmpc_iter = xs_nmpc_iter.T, ur_nmpc_iter.T
    iter_cost_nmpc.append((xs_nmpc_iter[:,-1]-set_pt_cl).T@QN_mpc@(xs_nmpc_iter[:,-1]-set_pt_cl) + np.sum(np.diag(us_nmpc_iter.T@R_mpc@us_nmpc_iter)))
    iter_norm_dist_nmpc.append(np.linalg.norm(xs_nmpc_iter[:,-1]-set_pt_cl))
```

```python
plt.figure(figsize=(6,4))
#plt.suptitle('Control solution after each iteration of the SQP-algorithm for NMPC and K-NMPC')
plt.subplot(2,1,1)
plt.plot(np.arange(n_iter), iter_cost_bedmd[:n_iter]/iter_cost_nmpc[-1], color=colors[2], label='K-NMPC')
plt.plot(np.arange(n_iter), iter_cost_nmpc[:n_iter]/iter_cost_nmpc[-1], color=colors[3], label='NMPC')
plt.title('Control effort')
plt.ylabel('$||u||$')
plt.legend(loc='upper right', frameon=False)
plt.xlabel('SQP iteration')

plt.subplot(2,1,2)
plt.plot(np.arange(n_iter), iter_norm_dist_bedmd[:n_iter], color=colors[2], label=labels[2])
plt.plot(np.arange(n_iter), iter_norm_dist_nmpc[:n_iter], color=colors[3], label=labels[3])
plt.ylim(0,5)
plt.title('Realized terminal distance from setpoint')
plt.ylabel('$||x_N - x_d||$')
plt.xlabel('SQP iteration')

if save_figures:
    matplotlib.rcParams['pdf.fonttype'] = 42
    matplotlib.rcParams['ps.fonttype'] = 42
    plt.tight_layout()
    plt.savefig(dropbox_folder + 'planar_quad_sqp_iterations.pdf', format='pdf', dpi=2400)

plt.show()

print('Solution statistics\n')
print(tabulate([['Nonlinear MPC', len(controller_nmpc.x_iter), np.mean(controller_nmpc.comp_time), np.std(controller_nmpc.comp_time), sum(controller_nmpc.comp_time)],
                ['Koopman bilinear MPC', len(controller_bedmd.x_iter), np.mean(controller_bedmd.comp_time), np.std(controller_bedmd.comp_time), sum(controller_bedmd.comp_time)]],
               headers=['Number of SQP\niterations','Mean comp. time per\niteration (secs)', 'Std comp. time per\niteration (secs)', 'Total comp.\ntime (secs)']))
```

# Evaluate performance of controllers for closed-loop control

#### Design finite horizon controllers

```python
from koopman_core.controllers import PerturbedController

Q_mpc_cl = 5e2*np.diag([1, 1, 1, 2e-1, 2e-1, 2e-1])
QN_mpc_cl = Q_mpc_cl
R_mpc_cl = np.eye(m)
traj_duration = 1.
N_cl = int(traj_duration/dt)
t_eval_cl=np.arange(300)*dt

# Solver settings closed loop:
polish_osqp = False
max_iter_osqp = 10
linsys_solver_osqp = 'qdldl'
warm_start_osqp = True

controller_dmd_cl = MPCController(sys_dmd, N_cl, dt, umin, umax, xmin, xmax, Q_mpc_cl, R_mpc_cl, QN_mpc_cl, set_pt_cl, add_slack=True)
controller_dmd_cl = PerturbedController(sys_dmd,controller_dmd_cl,0.,const_offset=hover_thrust, umin=umin, umax=umax)

controller_edmd_cl = MPCController(sys_edmd, N_cl, dt, umin, umax, xmin, xmax, Q_mpc_cl, R_mpc_cl, QN_mpc_cl, set_pt_cl, add_slack=True)
controller_edmd_cl = PerturbedController(sys_edmd,controller_edmd_cl,0.,const_offset=hover_thrust, umin=umin, umax=umax)

controller_bedmd_cl = BilinearMPCController(sys_bedmd, N_cl, dt, umin, umax, xmin, xmax, Q_mpc_cl, R_mpc_cl, QN_mpc_cl, set_pt_cl, add_slack=True)
controller_bedmd_cl.construct_controller(controller_bedmd.cur_z[:N_cl+1,:], controller_bedmd.cur_u[:N_cl,:])
controller_bedmd_cl.solve_to_convergence(z0_cl, 0., controller_bedmd.cur_z[:N_cl+1,:], controller_bedmd.cur_u[:N_cl,:], max_iter=max_iter)
controller_bedmd_cl = PerturbedController(sys_bedmd,controller_bedmd_cl,0.,const_offset=hover_thrust, umin=umin, umax=umax)
controller_bedmd_cl.nom_controller.update_solver_settings(polish=polish_osqp, max_iter=max_iter_osqp, linsys_solver=linsys_solver_osqp, warm_start=warm_start_osqp)

controller_nmpc_cl = NonlinearMPCController(quadrotor_d, N_cl, dt, umin+hover_thrust, umax+hover_thrust, xmin, xmax, Q_mpc_cl, R_mpc_cl, QN_mpc_cl, set_pt_cl, add_slack=True)
controller_nmpc_cl.construct_controller(controller_nmpc.cur_z[:N_cl+1,:], controller_nmpc.cur_u[:N_cl,:])
controller_nmpc_cl.solve_to_convergence(x0_cl, 0., controller_nmpc.cur_z[:N_cl+1,:], controller_nmpc.cur_u[:N_cl,:], max_iter=max_iter)
controller_nmpc_cl.update_solver_settings(polish=polish_osqp, max_iter=max_iter_osqp, linsys_solver=linsys_solver_osqp, warm_start=warm_start_osqp)
```

#### Simulate designed trajectories closed-loop

```python
xs_dmd_cl, us_dmd_cl = quadrotor.simulate(x0_cl, controller_dmd_cl, t_eval_cl)
xs_dmd_cl, us_dmd_cl = xs_dmd_cl.T, us_dmd_cl.T

xs_edmd_cl, us_edmd_cl = quadrotor.simulate(x0_cl, controller_edmd_cl, t_eval_cl)
xs_edmd_cl, us_edmd_cl = xs_edmd_cl.T, us_edmd_cl.T

controller_bedmd_cl.comp_time = []
xs_bedmd_cl, us_bedmd_cl = quadrotor.simulate(x0_cl, controller_bedmd_cl, t_eval_cl)
xs_bedmd_cl, us_bedmd_cl = xs_bedmd_cl.T, us_bedmd_cl.T

controller_nmpc_cl.comp_time = []
xs_nmpc_cl, us_nmpc_cl = quadrotor.simulate(x0_cl, controller_nmpc_cl, t_eval_cl)
xs_nmpc_cl, us_nmpc_cl = xs_nmpc_cl.T, us_nmpc_cl.T
```

#### Plot/analyze the results

```python
plot_inds = [0, 1, 2, 0, 1]
subplot_inds = [1, 2, 3, 4, 8]

plt.figure(figsize=(12,2.5))
for ii in range(5):
    ind = plot_inds[ii]
    if ii < 3:
        ax = plt.subplot(1,4,subplot_inds[ii])
        plt.plot(t_eval_cl, xs_dmd_cl[ind,:], colors[0], label='DMD MPC')
        plt.plot(t_eval_cl, xs_edmd_cl[ind, :], colors[1], label='EDMD MPC')
        plt.plot(t_eval_cl, xs_bedmd_cl[ind, :], colors[2], label='K-NMPC')
        plt.plot(t_eval_cl, xs_nmpc_cl[ind,:], colors[3], label='NMPC')

        plt.scatter(t_eval_cl[0], x0_cl[ind], color='g')
        plt.scatter(t_eval_cl[-1], set_pt_cl[ind], color='r')
        plt.ylabel(labels[ind])
        ax.yaxis.set_major_locator(MaxNLocator(integer=True))
        plt.title(titles[subplot_inds[ii]-1])
        plt.xlabel('Time (sec)')
        if subplot_inds[ii]==1:
            plt.legend(loc='upper left', frameon=False)
            plt.ylim(-0.15,2)
    else:
        bx = plt.subplot(2,4,subplot_inds[ii])
        plt.plot(t_eval_cl[:-1],us_dmd_cl[ind,:], color=colors[0], label='DMD MPC')
        plt.plot(t_eval_cl[:-1], us_edmd_cl[ind, :], color=colors[1], label='EDMD MPC')
        plt.plot(t_eval_cl[:-1], us_bedmd_cl[ind, :], color=colors[2], label='K-NMPC')
        plt.plot(t_eval_cl[:-1],us_nmpc_cl[ind,:], color=colors[3], label='NMPC')
        plt.plot([0, t_eval_cl[-1]], [umax[ind]+hover_thrust, umax[ind]+hover_thrust], ':k')
        plt.plot([0, t_eval_cl[-1]], [umin[ind]+hover_thrust, umin[ind]+hover_thrust], ':k')
        plt.ylabel(labels[ii+3])
        ax.yaxis.set_major_locator(MaxNLocator(integer=True))
        if subplot_inds[ii] == 4:
            plt.title('Control inputs')
        else:
            plt.xlabel('Time (sec)')
if save_figures:
    matplotlib.rcParams['pdf.fonttype'] = 42
    matplotlib.rcParams['ps.fonttype'] = 42
    plt.tight_layout()
    plt.savefig(dropbox_folder + 'planar_quad_closed_loop.pdf', format='pdf', dpi=2400)
plt.show()

from matplotlib._png import read_png
from matplotlib.offsetbox import AnnotationBbox, OffsetImage
from scipy import ndimage

draw_inds = np.arange(0,t_eval_cl.size)[::50]

plt.figure(figsize=(12,2))
ax = plt.subplot(1,1,1, frameon=False)
plt.plot(xs_bedmd_cl[0,:], xs_bedmd_cl[1,:], color=colors[2], label='Koopman NMPC closed loop trajectory with quadrotor orientation sampled at 2 hz')
plt.xlabel('y (m)')
plt.ylabel('z (m)')
plt.ylim(-0.1, xs_bedmd[1,-1]+0.45)
plt.legend(loc='upper left',frameon=False)
for ii in draw_inds:
    im_quad = read_png('working_files/figures/quad_figure.png')
    im_quad = ndimage.rotate(im_quad, xs_bedmd_cl[2,ii]*180)
    imagebox_quad = OffsetImage(im_quad, zoom=.11)
    ab = AnnotationBbox(imagebox_quad, [xs_bedmd_cl[0,ii], xs_bedmd_cl[1,ii]], frameon=False)
    ax.add_artist(ab)

if save_figures:
    matplotlib.rcParams['pdf.fonttype'] = 42
    matplotlib.rcParams['ps.fonttype'] = 42
    plt.tight_layout()
    plt.savefig(dropbox_folder + 'planar_quad_closed_loop_2.pdf', format='pdf', dpi=2400)
plt.show()

cost_cl_dmd = np.sum(np.diag((xs_dmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)).T@Q_mpc_cl@(xs_dmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)))) + (xs_dmd_cl[:,-1]-set_pt_cl).T@QN_mpc_cl@(xs_dmd_cl[:,-1]-set_pt_cl) + np.sum(np.diag(us_dmd_cl.T@R_mpc_cl@us_dmd_cl))
cost_cl_edmd = np.sum(np.diag((xs_edmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)).T@Q_mpc_cl@(xs_edmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)))) + (xs_edmd_cl[:,-1]-set_pt_cl).T@QN_mpc_cl@(xs_edmd_cl[:,-1]-set_pt_cl) + np.sum(np.diag(us_edmd_cl.T@R_mpc_cl@us_edmd_cl))
cost_cl_bedmd = np.sum(np.diag((xs_bedmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)).T@Q_mpc_cl@(xs_bedmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)))) + (xs_bedmd_cl[:,-1]-set_pt_cl).T@QN_mpc_cl@(xs_bedmd_cl[:,-1]-set_pt_cl) + np.sum(np.diag(us_bedmd_cl.T@R_mpc_cl@us_bedmd_cl))
cost_cl_nmpc = np.sum(np.diag((xs_nmpc_cl[:,:-1]-set_pt_cl.reshape(-1,1)).T@Q_mpc_cl@(xs_nmpc_cl[:,:-1]-set_pt_cl.reshape(-1,1)))) + (xs_nmpc_cl[:,-1]-set_pt_cl).T@QN_mpc_cl@(xs_nmpc_cl[:,-1]-set_pt_cl) + np.sum(np.diag(us_nmpc_cl.T@R_mpc_cl@us_nmpc_cl))

print('Solution statistics:\n')
print(tabulate([['DMD MPC', "{:.4f}".format(cost_cl_dmd/cost_cl_nmpc), np.mean(controller_dmd_cl.nom_controller.comp_time), np.std(controller_dmd_cl.nom_controller.comp_time)],
                ['EDMD MPC', "{:.4f}".format(cost_cl_edmd/cost_cl_nmpc),np.mean(controller_edmd_cl.nom_controller.comp_time), np.std(controller_edmd_cl.nom_controller.comp_time)],
                ['bEDMD MPC', "{:.4f}".format(cost_cl_bedmd/cost_cl_nmpc), np.mean(controller_bedmd_cl.nom_controller.comp_time), np.std(controller_bedmd_cl.nom_controller.comp_time)],
                ['NMPC (benchmark, known model)',1, np.mean(controller_nmpc_cl.comp_time), np.std(controller_nmpc_cl.comp_time)]],
               headers=['Normalized cost,\nrealized trajectory', 'Mean comp. time (secs)', 'std comp. time (secs)']))
```


* This notebook implements the VAE model in [Kingma (2013)](https://arxiv.org/abs/1312.6114).

## Variational Auto-encoder

### Motivation

In addition to the values of observable and of its latent variable, we are to encode the uncertainties of them also.

### Definitions

* Let $X$ be the observed random variable and $Z$ the latent random variable. Let

  \begin{equation}
    Z \sim P_{Z}
  \end{equation}

  be some given prior.

* And let

  \begin{equation}
    X \mid Z \sim P_{X \mid Z; \phi}
  \end{equation}

  for some distribution in a family parameterized by $\phi$.

For instance, if $X \in \mathbb{R}$, it is general to assume a multivariate Gaussian distribution, diagonalized or semi-diagonalized, then we have

\begin{equation}
  X \mid Z \sim \mathcal{N} \left( \mu(Z; \phi_1), \sigma(Z; \phi_2) \right),
\end{equation}

for some universality functions $\mu(\cdot; \phi_2)$ and $\sigma(\cdot; \phi_2)$ and $\phi := (\phi_1, \phi_2)$.

### Variational Inference

Utilization of the dataset of $x$ forces us to employ some distribution $q(z \mid x)$ (to be determined) to fit $p(z \mid x)$ (and if the dataset is of $z$, then we are forced to $p(x \mid z)$ instead), then we have the bound by KL-divergence

\begin{equation}
  \text{KL} \left( q(z \mid x) \| p(z \mid x) \right)
  = \text{KL} \left( q(z \mid x) \| p(x, z) \right) + \ln p(x)
  \geq 0.
\end{equation}

Thus, a loss

\begin{align}
  L(x)
  & := \text{KL} \left( q(z \mid x) \| p(y, z) \right) \\
  & = \mathbb{E}_{z \sim q(z \mid x)} \left[
         \ln q(z \mid x) - \ln p(x \mid z) - \ln p(z) \right] \\
  & \geq - \ln p(x).
\end{align}

The equality can be reached if and only if $q(z \mid x) = p(z \mid x)$.

### Evaluation of Inference

The evaluation of this fitting is comparing the $\ln p(x)$ and the $-L(x)$. Given any $p(x \mid z)$, we have

\begin{align}
  p(x) & = \int dz p(x \mid z) p(z) \\
       & = \mathbb{E}_{z \sim p(z)} \left[ p(x \mid z) \right].
\end{align}

### Relation with Autoencoder

* The $\mathbb{E}_{z \sim q(z \mid x)} \left[ - \ln p(x \mid z) \right]$ part can be interpreted as a reconstruction loss.

* The $\mathbb{E}_{z \sim q(z \mid x)} \left[ - \ln p(z) \right]$ part serves as regularization terms.

* It seems that the $\mathbb{E}_{z \sim q(z \mid x)} \left[ \ln q(z \mid x) \right]$ part has no correspondence in auto-encoder.

### Example

In this very example on the MNIST dataset:

\begin{align}
  Z & \sim \mathcal{N} (0, 1); \\
  X \mid Z & \sim \text{Bernoulli}\left( p(z; \theta) \right).
\end{align}

And inference distributions:

\begin{align}
  Z_0 \mid X & \sim \mathcal{N}\left( \mu(x; \phi_1), \sigma(x; \phi_2) \right); \\
  Z \mid Z_0 & \sim T(z_0).
\end{align}

where $Z_0$ obeys the "base inference distribution", as a Gaussian; and $T$ is bijective with universality.

## Implementation

```python
%matplotlib inline

from IPython.display import display
import matplotlib.pyplot as plt
from PIL import Image
from tqdm import tqdm
import numpy as np
import tensorflow as tf
import tensorflow.contrib.distributions as tfd
from tensorflow.contrib.distributions.python.ops import bijectors as tfb
from tensorflow.examples.tutorials.mnist import input_data

# For reproducibility
seed = 42
np.random.seed(seed)
tf.set_random_seed(seed)
```

### Functions

```python
def get_p_X_z(z, X_dim, hidden_layers=None,
              name='p_X_z', reuse=None):
  """Returns the distribution of P(X|Z).

  X | Z ~ Bernoulli( p(Z) ).

  Args:
    z: Tensor of the shape `[batch_size, z_dim]`.
    X_dim: Positive integer.
    hidden_layers: List of positive integers. Defaults to
      `[128, 256, 512]`.

  Returns:
    An instance of `tfd.Distribution`, with batch-shape `batch_size`
    and event-shape `X_dim`.
  """
  if hidden_layers is None:
    hidden_layers = [128, 256, 512]

  with tf.variable_scope(name, reuse=reuse):
    hidden = z
    for hidden_layer in hidden_layers:
      hidden = tf.layers.dense(hidden, hidden_layer,
                               activation=tf.nn.relu)
    logits = tf.layers.dense(hidden, X_dim, activation=None)

    p_X_z = tfd.Independent(tfd.Bernoulli(logits=logits))
    return p_X_z
```

```python
def get_q_z_X(X, z_dim, hidden_layers=None, bijectors=None,
              dtype='float32', name='q_z_X', reuse=None):
  """Returns the distribution of Z | X.

  Z = bijector(Z_0), and
  Z_0 | X ~ Normal(mu(X;phi), sigma(X;phi)).

  Args:
    X: Tensor with shape `[batch_size, X_dim]`.
    z_dim: Positive integer.
    hidden_layers: List of positive integers. Defaults to
      `[512, 256, 128]`.
    bijectors: List of `tfb.Bijector`s. Defaults to an empty
      list.

  Returns:
    An instance of `tfd.Distribution`, with batch-shape `batch_size`
    and event-shape `z_dim`.
  """
  if bijectors is None:
    bijectors = []
  if hidden_layers is None:
    hidden_layers = [512, 256, 128]

  with tf.variable_scope(name, reuse=reuse):
    hidden = X
    for hidden_layer in hidden_layers:
      hidden = tf.layers.dense(hidden, hidden_layer,
                               activation=tf.nn.relu)
    # Outputs in the fiber-bundle space
    output = tf.layers.dense(hidden, z_dim * 2, activation=None)
    # shape: [batch_size, z_dim]
    mu, log_var = tf.split(output, [z_dim, z_dim], axis=1)

    q_z0_X = tfd.MultivariateNormalDiag(mu, tf.exp(log_var))
    chain = tfb.Chain(bijectors)
    q_z_X = tfd.TransformedDistribution(q_z0_X, chain)
    return q_z_X
```

```python
def get_bijectors(name='bjiectors', reuse=None):
  """Complexify the inference distribution by extra-bijectors like
  normalizing flows.

  Returns:
    List of `Bijector`s.
  """
  with tf.variable_scope(name, reuse=reuse):
    bijectors = []
    #return bijectors  # test!
    for i in range(10):
      # Get one bijector
      shift_and_log_scale_fn = \
        tfb.masked_autoregressive_default_template([128])
      # MAP is extremely slow in training. Use IAF instead.
      bijector = tfb.Invert(
          tfb.MaskedAutoregressiveFlow(shift_and_log_scale_fn))
      bijectors.append(bijector)
    return bijectors
```

```python
def get_loss_X(get_q_z_X, get_p_X_z, p_z, reuse=None):
  """L(X) := E_{z ~ q(z|X)} [ log_q(z|X) - log_p(z) - log_p(X|z) ].

  Args:
    get_q_z_X: Callable with the signature:
      Args:
        X: Tensor with shape `[batch_size, X_dim]`.
        reuse: Boolean.
      Returns:
        An instance of `tfd.Distribution`, with batch-shape `batch_size`
        and event-shape `z_dim`.
    get_p_X_z: Callable with the signature:
      Args:
        z: Tensor of the shape `[batch_size, z_dim]`.
        reuse: Boolean.
      Returns:
        An instance of `tfd.Distribution`, with batch-shape `batch_size`
        and event-shape `X_dim`.
    p_z: An instance of `tfd.Distribution`, with batch-shape `batch_size`
      and event-shape `z_dim`.
    reuse: If reuse the variables in `get_q_z_X` and `get_p_X_z`.

  Returns:
    Callable with the signature:
      Args:
        X: Tensor of the shape `[batch_size, X_dim]`.
      Returns:
        Tensor of the shape `[batch_size]`.
  """

  def loss_X(X, name='loss_X'):
    """
    Args:
      X: Tensor of the shape `[batch_size, X_dim]`.

    Returns:
      Tensor of the shape `[batch_size]`.
    """
    with tf.name_scope(name):
      # Get the distribution q(z|X)
      q_z_X = get_q_z_X(X, reuse=reuse)

      # Get the distribution p(X|z)
      z_samples = q_z_X.sample()
      p_X_z = get_p_X_z(z_samples, reuse=reuse)

      # Compute the tensor of L(X)
      loss_X_tensor = tf.zeros([batch_size])  # initialize.
      # E_{z ~ q(z|X)} [ log_q(z|X) ]
      loss_X_tensor += q_z_X.log_prob(z_samples)
      # E_{z ~ q(z|X)} [ - log_p(z) ]
      loss_X_tensor += -1 * p_z.log_prob(z_samples)
      # E_{z ~ q(z|X)} [ - log_p(X|z) ]
      loss_X_tensor += -1 * p_X_z.log_prob(X)
      return loss_X_tensor

  return loss_X
```

```python
def get_log_p_X(get_p_X_z, p_z, n_samples=100, reuse=None):
  """Returns the function ln p(X) by Monte-Carlo integral.

  p(X) = E_{z~p(z)} [ p(X|z) ].

  Args:
    get_p_X_z: Callable with the signature:
      Args:
        z: Tensor of the shape `[batch_size, z_dim]`.
        reuse: Boolean.
      Returns:
        An instance of `tfd.Distribution`, with batch-shape `batch_size`
        and event-shape `X_dim`.
    p_z: An instance of `tfd.Distribution`, with batch-shape `batch_size`
      and event-shape `z_dim`.
    n_samples: Positive integer.
    reuse: If reuse the variables in `get_p_X_z`.

  Returns:
    Callable with the signature:
      Args:
        X: Tensor of the shape `[batch_size, X_dim]`.
        name: String.
      Returns:
        Tensor of the shape `[batch_size]`.
  """

  def log_p_X(X, name='log_p_X'):
    """Returns the tensor of ln p(X).

    Args:
      X: Tensor of the shape `[batch_size, X_dim]`.
      name: String.

    Returns:
      Tensor of the shape `[batch_size]`.
    """
    with tf.name_scope(name):

      def log_p_X_z(z_sample):
        """Returns the tensor of ln p(X|z)."""
        p_X_z = get_p_X_z(z_sample, reuse=reuse)
        return p_X_z.log_prob(X)  # [batch_size]

      z_samples = p_z.sample(n_samples)  # [n_samples, batch_size, z_dim]
      expected = tf.map_fn(log_p_X_z, z_samples)  # [n_samples, batch_size]
      log_p_X_tensor = (tf.reduce_logsumexp(expected, axis=0) -
                        tf.log(float(n_samples)))
      return log_p_X_tensor

  return log_p_X
```

### Loss

```python
batch_size = 128
X_dim = 28 * 28
X = tf.placeholder(shape=[batch_size, X_dim], dtype='float32', name='X')
```

```python
z_dim = 64
p_z = tfd.MultivariateNormalDiag(tf.zeros([batch_size, z_dim]), name='p_z')
```

```python
def _get_q_z_X(X, reuse):
  bijectors = get_bijectors(reuse=reuse)
  #bijectors = []  # test!
  return get_q_z_X(X, z_dim, bijectors=bijectors, reuse=reuse)

def _get_p_X_z(z, reuse):
  return get_p_X_z(z, X_dim=X_dim, reuse=reuse)

loss_X = get_loss_X(_get_q_z_X, _get_p_X_z, p_z, reuse=tf.AUTO_REUSE)
loss_X_tensor = loss_X(X)
loss_X_scalar = tf.reduce_mean(loss_X_tensor)
```

### Generating

```python
z_samples = tf.placeholder(shape=[batch_size, z_dim],
                           dtype='float32',
                           name='z_samples')
X_samples = _get_p_X_z(z_samples, reuse=tf.AUTO_REUSE).sample()
```

```python
def get_image(array):
  """
  Args:
    array: Numpy array with shape `[28*28]`.

  Returns:
    An image.
  """
  array = 255 * array
  array = array.reshape([28, 28])
  array = array.astype(np.uint8)
  return Image.fromarray(array)
```

### Training

```python
optimizer = tf.train.AdamOptimizer(epsilon=1e-3)
train_op = optimizer.minimize(loss_X_scalar)
```

```python
sess = tf.Session()
sess.run(tf.global_variables_initializer())
```

```python
mnist = input_data.read_data_sets('dat/MNIST', one_hot=True)

def get_X_y_batch():
    X_batch, y_batch = mnist.train.next_batch(batch_size)
    # Since X | Y, Z ~ Bernoulli, the observed value of X shall
    # either be 0 or 1,
    X_batch = np.where(X_batch >= 0.5, np.ones_like(X_batch),
                       np.zeros_like(X_batch))
    return X_batch, y_batch
```

    Extracting dat/MNIST/train-images-idx3-ubyte.gz
    Extracting dat/MNIST/train-labels-idx1-ubyte.gz
    Extracting dat/MNIST/t10k-images-idx3-ubyte.gz
    Extracting dat/MNIST/t10k-labels-idx1-ubyte.gz

```python
loss_vals = []
for i in tqdm(range(100000)):
  X_batch, y_batch = mnist.train.next_batch(batch_size)
  _, loss_val = sess.run([train_op, loss_X_scalar], {X: X_batch})
  if np.isnan(loss_val):
    raise ValueError('Loss has been NaN.')
  loss_vals.append(loss_val)

# Visualization
plt.plot(loss_vals)
plt.xlabel('steps')
plt.ylabel('loss')
plt.show()

# Zoomed in
last_steps = 20000
plt.plot(loss_vals[-last_steps:])
plt.xlabel('last steps')
plt.ylabel('loss')
plt.show()

print('Final loss:', np.mean(loss_vals[-100:]))
```

```python
z_sample_vals = np.random.normal(size=[batch_size, z_dim])
X_sample_vals = sess.run(X_samples, {z_samples: z_sample_vals})

# Display the results
n_display = 5
for i in range(n_display):
  print('Gnerated:')
  display(get_image(X_sample_vals[i]))
  print()
```

### Effect of Normalizing-flow

* While not employing normalizing flow, 100000 iterations give the final loss about 97.8.
* With 10 normalizing flows, 100000 iterations give the final loss 93.5 (improved 4.3).

```python
np.exp(4)
```

    54.598150033144236

### Evaluation

```python
# Compute the lower bound of loss for any (trained) p(X|z).

log_p_X = get_log_p_X(_get_p_X_z, p_z,
                      n_samples=10**5,
                      reuse=tf.AUTO_REUSE)
log_p_X_tensor = log_p_X(X)
loss_lower_bound = -1 * tf.reduce_mean(log_p_X_tensor)
```

```python
print(sess.run(loss_lower_bound, {X: X_batch}))
```

    106.46511

It's strange that the lower bound of loss is even greater than the trained loss! And the Monte-Carlo integral is not that in-sensitive to the `n_samples`, as it should be.