%% ORR - lab 6 - Trajectory tracking
% In this worksheet, you will implement a trajectory tracking regulator for a pendulum. The term trajectory tracking (or sometimes also trajectory stabilization) refers to a situation where you are given a feasible trajectory--time-sequence of control inputs and states consistent with the dynamics of a given model--and you want to implement a feedback controller ensuring that the system follows this given trajectory. For the pendulum, the trajectory can be, for instance, a sequence of controls getting the pendulum from the downward position to the upright position. If you find pendulum too boring, you may easily imagine using the same techniques for a rocket returning from space to Earth. Again, you precalculate a trajectory consistent with a model of the rocket and then you just follow this precalculated trajectory in real time. Nevertheless, we stick to the pendulum.
%
% Let us postpone for a moment the discussion of the trajectory stabilization and derive the state-space model of the pendulum. Considering the parameters and coordinates shown in the figure above, the equation of motion of  a pendulum is
%
%   ml^2\ddot{\theta}(t) + b\dot{\theta}(t) + mgl\sin{\theta}(t) = \tau(t).
% Grouping the physical parameters to a_1:=\frac{g}{l}, a_2:=\frac{b}{ml^2} and a_3:=\frac{1}{ml^2} this equation of motion can be rewritten to the following nonlinear state-space model
%
% \begin{array}{rl}
%        \dot{x}_1(t) &= x_2(t), \\
%        \dot{x}_2(t) &= -a_1 \sin (x_1(t)) - a_2 x_2(t) + a_3 u(t) ,
% \end{array}
% where the states are [x_1(t), x_2(t)] := [\theta(t), \dot{\theta}(t)] and the control input is u(t):=\tau(t). We will use the following physical parameters:

m = 1; % mass of the pendulum
l = 1; % length of the pendulum
g = 9.81; % gravitational acceleration
b = 0.1; % friction coefficient

a1 = g/l;
a2 = b/(m*l*l);
a3 = 1/(m*l*l);

% Since from now on, we will be working solely in discrete-time, let us discretize the continuous dynamics of the pendulum. We will do so in the simplest possible way: by Forward Euler method. This way, we obtain the following discrete-time nonlinear model of a pendulum
%
% \begin{array}{rl}
%        x_1[k+1] &= x_1[k] + T_\mathrm{s}\,x_2[k], \\
%        x_2[k+1] &= x_2[k] + T_\mathrm{s}\left( -a_1\,\sin(x_1[k]) - a_2\,x_2[k] + a_3\,u[k] \right),
% \end{array}
% where T_\mathrm{s} is the sampling period.
% For ease of notation, let f(\mathbf{x}[k], u[k]) denote the right-hand side of the equation above and \mathbf{x}[k] the whole state vector, i.e. \mathbf{x}[k] = [x_1[k], x_2[k]]. Then the two equations above can be equivalently written as \mathbf{x}[k+1]=f(\mathbf{x}[k], u[k]).
%
% Now, let us load a trajecotory which, when the control inputs are applied, gets the pendulum from the downward position to the upright positon. The trajectory consists of sampled controls u^\ast_k and resulting states \mathbf{x}^\ast[k]. The samples are separated by sampling period T_\mathrm{s}=0.1\,\mathrm{s}. By the way, how to find such a trajectory will be your task in the homework assignment for this week.

% time vector - t_star, controls - u_star, state - x_star
load('optTraj.mat')

figure(1)
subplot(211)
stairs(t_star, x_star)
grid on
xlabel('Time [s]')
h = legend('$\theta[k]$ - precalculated', '$\dot{\theta}[k]$ - precalculated', 'Location', 'eastoutside');
set(h,'Interpreter','latex');
xlim([t_star(1) t_star(end)])

subplot(212)
stairs(t_star, u_star)
grid on
xlabel('Time [s]')
h = legend('$\tau[k]$ - precalculated', 'Location', 'eastoutside');
set(h,'Interpreter','latex');
xlim([t_star(1) t_star(end)])

%
% In order to verify that the controls u^\ast[k] truly get the pendulum from the downward position (\theta[k]=0) to the upward position (\theta[k]=\pi), let us simulate the response of the pendulum model to the controls u^\ast[k]. The pendulum is modelled as a continuous-time model in Simulink model pendulumFF.slx. The FF suffix suggests that here we are using feedforward control; there is no feedback from the states or other measured quantities, we only apply the precalculted controls u^\ast[k] and hope that the system follows the precalculated state trajectory \mathbf{x}^\ast[k].

u_TS.time = t_star;
u_TS.signals.values = u_star;
sim('pendulumFF', [t_star(1) t_star(end)]);

figure(1)
subplot(211)
stairs(t_star, x_star)
hold on
stairs(simout.Time, simout.Data, '--')
hold off
h = legend('$\theta[k]$ - precalculated', '$\dot{\theta}[k]$ - precalculated', ...
           '$\theta[k]$ - simulated FF', '$\dot{\theta}[k]$ - simulated FF', ...
           'Location', 'eastoutside');
set(h,'Interpreter','latex');
grid on
xlabel('Time [s]')
xlim([t_star(1) t_star(end)])
% Call the following line in the Maltab command line for an animated visualization.
% visu(simout.Time, simout.Data(:,1), l)

% Apparently, controls u^\ast[k] successfully get the pendulum from the downward position to the upright position. Nevertheless, we can see some small discrepancies between the simulated response and the precalculated trajectory. Even though the trajectory was calculated on the exactly same model as we are using here for the simulation, the trajectory contains only samples of the controls and states at cetain time instants whereas the model evolves continuously in time. Furthermore, as there is no feedback even a small deviation of the model with respect to the real system makes the real trajectory different. For instance, try to re-run the script with friction coefficient b:=0.2 instead of the original b:=0.1.
% The discrepancies can be rectified by using trajectory stabilization that is by using a feedback controller ensuring that the system follows a given trajectory.
%
%% Trajectory stabilization
% Finally, let us design a trajectory stabilizing controller. At first, we show that the problem of trajectory stabilization can be solved by a finite-horizon time-varying LQR.
% We linearize the model of a system around the trajectory and thus obtain a linear time-varying deviation model
%
%   \delta\mathbf{x}[k+1] = \mathbf{A}[k]\delta\mathbf{x}[k] +
%   \mathbf{B}[k]\delta u[k],
% where \delta\mathbf{x}[k] = \mathbf{x}[k] - \mathbf{x}^\ast[k] and \delta u[k] = u[k] - u^\ast[k] and the time-varying matrices are
%
% \begin{array}{rl}
%   \mathbf{A}[k] &= \frac{\partial f \left(\mathbf{x}[k], u[k]\right)}{\partial \mathbf{x}[k]}\Big|_{\mathbf{x}[k] = \mathbf{x}^\ast[k], u[k]=u^\ast[k]}, \\
%   \mathbf{B}[k] &= \frac{\partial f \left(\mathbf{x}[k], u[k]\right)}{\partial u[k]}\Big|_{\mathbf{x}[k] = \mathbf{x}^\ast[k], u[k]=u^\ast[k]}.
% \end{array}
% Now, it remains to design a LQR stabilizing the linear time-varying devation model. In other words, you need to design a LQR for the model
%
%   \delta\mathbf{x}[k+1] = \mathbf{A}[k]\delta\mathbf{x}[k] +
%   \mathbf{B}[k]\delta u[k],
% on finite-time control horizon minimizing the quadratic cost function over the duration of the trajectory to be followed
%
%   J = \frac{1}{2}\delta\mathbf{x}^T[N] S_N \delta\mathbf{x}^T[N] +
%   \frac{1}{2}\sum_{k=0}^{N-1}\left( \delta\mathbf{x}^T[k] \mathbf{Q}
%   \delta\mathbf{x}[k] + \delta u^T[k] \mathbf{R} \delta u[k] \right).
% This is an instance of discrete-time time-varying LQR. The design procedure here is exactly the same as with the discrete-time time-invariant LQR, the only difference is that matrices \mathbf{A} and \mathbf{B} are not constant in time in the Difference Riccati Equation. Therefore, by iterating backward in time the Difference Riccati Equation, you calculate the matrices \mathbf{S}[k] and also the optimal state-feeback gain \mathbf{K}[k]. You can find the details in the lecture notes for week 4. Having the state-feedback gain the trajectory stabilizing controller is given by
% u[k] = u^\ast[k] - \mathbf{K}[k]\delta\mathbf{x}[k],
% that the control input is given by the nominal control input u^\ast[k] from the precalculated trajectory plus the term -\mathbf{K}[k]\delta\mathbf{x}[k] stabilizing the devition from the precalculated state trajectory.
% Fill in the gaps in the following code. Pretty much everything has been already implemented. The only things missing are the linearization of the discretized model along the given trajectory (calculation of the matrices \mathbf{A}[k] and \mathbf{B}[k]) and solution of the Difference Riccati Equation.
% If you succeed in implementing these missing parts, you can also implement an infinite-horizon LQR stabilizing the pendulum in the upward position. The trajectory controller only ensures that the pendulum ends up in the upward position but then it stops regulating and the pendulum falls down. That is why it is desirable to design yet another LQR. Having the second LQR, the trajectory regulator can hand over the regulation to the second LQR just after the pendulum reach the upward position.

% Weight matrices
Q = diag([100 100]); % state
R = 1; % control

% Linearization around the optimal trajectory
% x := [th; Dth]
A = zeros(2, 2, numel(t_star));
B = zeros(2, 1, numel(t_star));

for i=1:numel(t_star)
    % your code goes here
    A(:,:,i) = zeros(2,2);
    B(:,:,i) = zeros(2,1);
end

% Design a LQR stabilizing the trajectory
S = zeros(2, 2, numel(t_star));
K = zeros(1, 2, numel(t_star));

S(:,:, end) = Q;
K(:,:, end) = R\B(:,:, end)'*S(:,:, end);

% Iterate the Difference Riccati Equation. You may find equations (71) and
% (72) in the lecture notes
% (https://moodle.fel.cvut.cz/pluginfile.php/234279/mod_resource/content/6/L4_discrete_time_optimal_control_indirect_LQ_ARE.pdf)
% quite useful here.
for i = (numel(t_star) - 1):-1:1
    % your code goes here     
    K(:,:,i) = zeros(2,1);
    S(:,:,i) = zeros(2,2);
end

K = -squeeze(K)';

% Simulate the response of the pendulum model with the trajectory regulator
% implemented.
x0 = [0; 0];

K_TS = timeseries(K, t_star);
x_star_TS = timeseries(x_star, t_star);
u_star_TS = timeseries(u_star, t_star);

sim('pendulumFB', [t_star(1) t_star(end)+1]);

figure(1)
subplot(211)
stairs(t_star, x_star)
hold on
stairs(simout.Time, simout.Data(:,1:2), '--')
hold off
h = legend('$\theta[k]$ - precalculated', '$\dot{\theta}[k]$ - precalculated', ...
           '$\theta[k]$ - simulated FB', '$\dot{\theta}[k]$ - simulated FB', ...
           'Location', 'eastoutside');
set(h,'Interpreter','latex');
grid on
xlabel('Time [s]')
xlim([simout.Time(1) simout.Time(end)])

subplot(212)
stairs(t_star, u_star)
hold on
stairs(simout.Time, simout.Data(:,3), '--')
hold off
h = legend('$\tau[k]$ - precalculated', ...
           '$\tau[k]$ - simulated FB', ...
           'Location', 'eastoutside');
set(h,'Interpreter','latex');
grid on
xlabel('Time [s]')
xlim([simout.Time(1) simout.Time(end)])
       
% Call the following line in the Maltab command line for an animated visualization.
% visu(simout.Time, simout.Data(:,1), l)