%% ORR - lab 8 - Trajectory optimization
% In this seminar, you will get familiar with yet another way of finding an (sub) optimal trajectory for a given system. Specifically, you will implement a direct multiple shooting method, a simple method of transcription of an Optimal Control Problem (OCP) to a Nonlinear Programming Problem (NLP).
%% Toy problem
% Once again, we will use a pendulum as a sample problem for the trajectory optimization. The goal remains the same as we had two weeks ago; the goal is to get the pendulum from the stable downward position to the unstable upright position.
%
% We will consider the same mathematical model of the pendulum. For your convenience, the model is restated here. Considering the parameters and coordinates shown in the figure above, the equation of motion 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 input is u(t):=\tau(t). For ease of notation, let \mathbf{f}(x(t), u(t)) denote the right-hand side of the equation above and \mathbf{x}(t) denote the whole state vector, i.e. \mathbf{x}(t) = [x_1(t), x_2(t)]. Then the two equations above can be equivalently written as \dot{\mathbf{x}}=\mathbf{f}(\mathbf{x}(t), u(t)). Let us assume that the parameters are:

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);

%
%% Direct Transcription Method
% Direct transcription converts a continuous OCP to a NLP by discretization of the time interval at several time instants t_0, t_1, \dots, t_{N-1}. Let us denote \mathbf{x}(t_k) by \mathbf{x}_k and similarly u(t_k) by u_k. Then by this transcription, the NLP corresponding to a trajectory optimization problem is given by
%
% \begin{array}{rl}
%        \min_{\mathbf{x}_{k},\,\, u_k}& \sum_{k=0}^{N-1} g\left( \mathbf{x}_k, u_{k}\right) \\
%        \mathrm{s.t.}:  &\mathbf{x}_{k+1} = \mathbf{x}_{k} + T_{\mathrm{s}_k} \mathbf{f}\left(\mathbf{x}_k, u_k\right), \qquad k=0,\dots,N-1\\
%               & \mathbf{x}_N\ =  \mathbf{r}_N,\\
%               & \mathbf{x}_0\: \mathrm{given},
% \end{array}
% where \mathbf{r}_N is a vector of desired state values at the final time t_N and T_s is the length of the control period, i.e. T_{s_k } :=t_{k+1} -t_k.
% Of course, the constraints may differ from an optimization problem to an optimization problem. For instance, you can easily add inequality constraints on controls or states along the trajectory. Notice, that this optimization problem is similar to the one we encountered several weeks ago when simultaneous Model Predictive Control (MPC) was derived. The major difference is that here we have a general nonlinear continuous-time dynamics as opposed to linear discrete-time dynamics in MPC. In MPC, we knew that the corresponding optimization problem was an instance of a Quadratic Programming problem hence it was convex and easily solvable. Here, however, the problem is in general non-convex and thus rather difficult to solve to optimality as the solvers may easily get stuck in local minima.
%% Task
% Implement the direct transcription method using simple forward Euler integration for finding a trajectory getting the pendulum from the downward position to the upward position with minimal control effort. Consider the controls u_k to be limited by |u_k| \leq 2.75. Furthermore, consider the finite horizon to be [0,5] and discretize it at N=50 equidistantly spaced points. Altogether, by this transcription, the NLP for our problem is given by
%
% \begin{array}{rl}
%        \min_{\mathbf{x}_{k}, u_k}& \sum_{k=0}^{N-1} u^2_{k} \\
%        \mathrm{s.t.}:  &\mathbf{x}_{k+1} = \mathbf{x}_{k} + T_\mathrm{s} \mathbf{f}\left(\mathbf{x}_k, u_k\right), \qquad k=0,\dots,N-1\\
%               & |\mathbf{u}_k| \leq  2.75,\\
%               & \mathbf{x}_0\ = [0,0]^T,\\
%               & \mathbf{x}_N\ =  [\pi, 0]^T,
% \end{array}
% where T_\mathrm{s} = \frac{T_\mathrm{f}}{N} and T_\mathrm{f} = 5.
% Solve this NLP by functions from CasADi. CasADi is a framework for algorithmic differentiation and numerical optimization. The algorithmic differentatiation means, from user's point of view, that you do not have to provide the solver with derivatives of your objective function and/or constraints; CasADi computes them for you and hands them over to the solver. This makes solving an optimization problem at hand way more convenient.
% You can download CasADi from here and take some inspiration from the examples in the documentation.
% Fill in the missing parts of the code below.
% Hint: You will need only one function from CasADi and that its opti.subject_to().
% Questions
% Is the found trajectory any good? Can you improve it?

% Parameters
N = 50;
Tf = 5;
x0 = [0;0];

umin = -2.75;
umax = 2.75;

% Initialize the optimization problem
opti = casadi.Opti();

% define decision variables
X = opti.variable(2,N+1); % states
th = X(1, :);
Dth = X(2, :);
U = opti.variable(1,N); % controls

% Objective function (minimize the control effort)
opti.minimize(U*U');

% Dynamic constraints
f = @(x,u) pend_ode(x, u, a1, a2, a3); % dx/dt = f(x,u)
Ts = Tf/N; % control interval length

% Implement the dynamic constraints
%  - Your code goes here -

% Implement the initial and final-time constraints
%  - Your code goes here -

% Initialize decision variables
opti.set_initial(U, 0);
opti.set_initial(th, 0);
opti.set_initial(Dth, 0);

% Solve NLP
opti.solver('ipopt'); % use IPOPT solver
sol = opti.solve();
% Extract the states and control from the decision variables
x_nlp = sol.value(X)';
u_nlp = sol.value(U)';
t_nlp = (0:N)*Ts;

% Let us compare the obtained trajectory with the trajectory we obtained by solving the Two-Point Boundary Value Problem (BVP) two weeks ago.

load('optTraj.mat')
figure(1)
subplot(211)
plot(t_nlp, x_nlp);
hold on
plot(t_star, x_star, '--');
hold off
grid on
xlabel('Time [s]')
ylabel('States [-]')

subplot(212)
plot(t_nlp(1:end-1), u_nlp);
hold on
plot(t_star, u_star, 'r--');
hold off
grid on
xlabel('Time [s]')
ylabel('Controls [Nm]')
legend('NLP', 'TPBVP')

%
% Furthermore, let us try to apply the controls from the computed trajectory to the continous model of the pendulum and observe whether the pendulum gets to the upward position or not.

u_TS.time = t_nlp(1:end-1);
u_TS.signals.values = u_nlp;

sim('pendulumFF', [t_star(1) t_star(end)]);

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

% Run the following line to see an animation of the pendulum
% visu(simout.Time, simout.Data(:,1), l)