% OFC_NSCI850.m
%
% optimal control of a car in 1D
% x_t+1 = Ax_t + Bu_t
%
% written by G. Blohm for NSCI 850 (2017)

clear

% parameters
dt = .1; % time increment
T = 100; % # time steps

% define state matrices
A = [1 dt; 0 1]; % state propagation
B = [dt^2/2; dt]; % augmented control matrix
Qf = [200 0; 0 0.2]; % final cost weight on state
Qt = 0; % state cost  weight during movement
R = 1e-5; % cost of motor command weight

% augment system description for initial conditions
A = [A zeros(2);
    zeros(2) eye(2)];
B = [B;zeros(2,1)];
Qf = [Qf -Qf; -Qf Qf];

% backward recursion for cost to go
S{T} = Qf;
for t = (T-1):-1:1,
    S{t} = Qt + A'*S{t+1}*(eye(4) - B*inv(R)*B'*S{t+1})^(-1)*A;
end

%% Optimal control (no feedback)
% note how this open-loop control does very well even with perturbations!
% but it does not account for noise, i.e. optimally combine forward and
% feedback info

% forward recursion for x and u
x(:,1) = [0; 0; 10; 0];
for t = 1:(T-1),
    x(:,t+1) = (eye(4) - B*inv(R)*B'*S{t+1})^(-1)*A*x(:,t);
    u(:,t) = inv(R)*B'*S{t+1}*x(:,t+1);
    if t == 50, x(1,t+1) = x(1,t+1)+2; end % add perturbation to x
end

% plot result
figure
subplot(3,1,1)
plot(x')
ylabel('State')
subplot(3,1,2)
plot(diff(x'))
ylabel('State derivative')
subplot(3,1,3)
plot(u')
ylabel('motor command')

%% Optimal feedback control
% incorporate optimal state estimation
% define additional matrices
H = eye(4); % observation matrix (perfect observation)
Ow = 1e-5*H; % measurment noise
Oe = 1e-10*eye(4); % process noise
SIG{1} = 1e-10*eye(4); % covariance for Kalman filter

% backward recursion for cost to go
S{T} = Qf;
for t = (T-1):-1:1,
    L{t} = inv(R + B'*S{t+1}*B)*B'*S{t+1}*A;
    S{t} = Qt + A'*S{t+1}*(A - B*L{t});
end

% forward recursion for x and u
x(:,1) = [0; 0; 10; 0];
xe = x;
for t = 1:(T-1),
    y(:,t) = H*x(:,t) + sqrt(diag(Ow)).*[randn(2,1);zeros(2,1)];
    u(:,t) = -L{t}*xe(:,t);
    x(:,t+1) = A*x(:,t) + B*u(:,t) + sqrt(diag(Oe)).*[randn(2,1);zeros(2,1)];
    K{t} = A*SIG{t}*H'*inv(H*SIG{t}*H' + Ow);
    SIG{t+1} = Oe + (A - K{t}*H)*SIG{t}*A';
    xe(:,t+1) = (A - B*L{t})*xe(:,t) + K{t}*(y(:,t) - H*xe(:,t));
    if t == 50, x(1,t+1) = x(1,t+1)+2; end % add perturbation to x
end

% plot result
figure
subplot(3,1,1)
plot(xe'); hold on; plot(x')
ylabel('State')
subplot(3,1,2)
plot(diff(xe')); hold on; plot(diff(x'))
ylabel('State derivative')
subplot(3,1,3)
plot(u')
ylabel('motor command')
