function optimal_control
%
% NSCI 850 - Computational Approaches to Neuroscience
%
% optimal feedback control example: point-mass movement in 2D space
% using iterative LQR (linear quadratic regulator)
% based on Li & Todorov (2006)
%
% written by Tim Lillicrap (2006)
% adapted by G. Blohm (2010)
%--------------------------------------------------------------------------

%% simulations
u = zeros(50,2); % initial guess of command u: no command
u = ilqr(u, [0 0 0 0], [3 4 0 0], diag([1 1 1 1]),0*ones(4),0.001*eye(2), 0.05, 100, 0); % find optimal command using ILQR
x = ilqr_openloop([0 0 0 0], u, 0.05); % simulate system with optimal command
% state x is composed of X/Y position and X/Y velocity

%% plot results
figure
subplot(221)
plot(1:length(x(:,1)), x(:,1:2));
legend('X','Y')
xlabel('time')
ylabel('position')

subplot(223)
plot(1:length(x(:,1)), x(:,3:4));
xlabel('time')
ylabel('velocity')

subplot(222)
plot(1:length(u(:,1)), u(:,1:2));
xlabel('time')
ylabel('command u')

subplot(224)
plot(x(:,1), x(:,2),'k.')
xlabel('X')
ylabel('Y')
axis equal;

%% ========================================================================
%% USED FUNCTIONS
%% ========================================================================

%% ILQR algorithm
function newu = ilqr(u, x0, xf, Qf, Q, R, dt, iterations, ethreshold)

% du = ilqr(u, x0, xf, Qf, Q, R, dt, iterations, ethreshold)
%
% run the ilqr algorithm using some initial state, nominal comman signal,
% and cost function.  iterate until the maximum iterations are reached
% or until the change in the command signal goes below some threshold.
%
% given -> u
%	-> x0 the state of the system at some initial time
%       -> xf the requested goal state of the system
%       -> Qf the quadratic final cost matrix
%	-> Q the quadratic integrated cost matrix on the state
%	-> R the quadratic integrated cost matrix on the command
%	-> dt the current time step used
%	-> iterations the maximum number of iterations to run for
%	-> ethreshold stop when the change in u goes below this scalar
%
% returns -> newu the new nominal trajectory given by the algorithm
%
% created by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005
% last edited by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005

currentu = u;

oldcost = ilqr_cost(u, x0, xf, Qf, Q, R, dt); % computes the cost of the current control

for i = 1:iterations % iterative adjustment of the controls
	
	du = ilqr_iteration(currentu, x0, xf, Qf, Q, R, dt); % compute control update
	currentu = currentu + .25*du; % update of u; uses damped version of du for numerical stability
	
	newcost = ilqr_cost(currentu, x0, xf, Qf, Q, R, dt); % compute cost with updated control
	
	if(abs(newcost - oldcost) < ethreshold) 
		break;
	end;
		
	oldcost = newcost;	

end;

newu = currentu;

return;


%% ILQR interation
function du = ilqr_iteration(u, x0, xf, Qf, Q, R, dt)

% du = ilqr_iteration(u, x0, xf, Qf, Q, R, dt)
%
% run through one iteration of the ilqr algorithm.  the du returned
% can be used to update the nominal trajectory u given to the function.
%
% given -> x0 the state of the system at some initial time
%       -> xf the desired final state of the system
%       -> Qf the quadratic final cost matrix
%	-> Q the quadratic integrated cost matrix on the state
%	-> R the quadratic integrated cost matrix on the command
%	-> dt the current time step used
%
% returns -> du the update to be applied to the nominal trajectory in ilqr.
%
% created by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005
% last edited by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005

du = zeros(size(u)); % memory pre-allocation

nx = length(x0);
nu = length(u(1,:));

x = ilqr_openloop(x0, u, dt); % simulate current system
N = length(x);
dx = 0.00000000000001*ones(size(x));

A = zeros(nx, nx, N-1);
B = zeros(nx, nu, N-1);

for i = 1:N-1
	[A(:,:,i) B(:,:,i)] = ilqr_linearization(x(i,:), u(i,:), dt);
end;

S = zeros(nx, nx, N);
S(:,:,N) = Qf;

nv = nx;
v = zeros(N, nx);
v(N,:) = (Qf*(x(N,:) - xf)')';

K = zeros(nu, nx, N-1);
Kv = zeros(nu, nv, N-1);
Ku = zeros(nu, nu, N-1);

% following algorithm comes straight out of Li and Todorov (2006).
for k = N-1:-1:1
	
 	K(:,:,k) = inv(B(:,:,k)' * S(:,:,k+1) * B(:,:,k) + R) * B(:,:,k)' * S(:,:,k+1) * A(:,:,k);
	Kv(:,:,k) = inv(B(:,:,k)' * S(:,:,k+1) * B(:,:,k) + R) * B(:,:,k)';
	Ku(:,:,k) = inv(B(:,:,k)' * S(:,:,k+1) * B(:,:,k) + R) * R;

	S(:,:,k) = A(:,:,k)' * S(:,:,k+1) * (A(:,:,k) - B(:,:,k) * K(:,:,k)) + Q;
	v(k,:) = ((A(:,:,k) - B(:,:,k) * K(:,:,k))' * v(k+1,:)' - K(:,:,k)' * R * u(k,:)' + Q * x(k,:)')';

	du(k,:) = (-1) * (K(:,:,k)* dx(k,:)' + Kv(:,:,k) * v(k+1,:)' + Ku(:,:,k) * u(k,:)')';
end;

return;


%% cost function
function cost = ilqr_cost(u, x0, xf, Qf, Q, R, dt)

% newx = ilqr_cost(u, x0, xf, Qf, Q, R, dt)
%
% run the system under consideration forward using the command signal u
% and calculate the associated cost defined by xf, Qf, Q, R.
%
% given -> u the command signal with which to control the system
%	-> x0 the state of the system at some initial time
%	-> xf the desired final state of the system
%	-> Qf the quadratic final cost matrix
%	-> Q the quadratic integrated cost matrix on the state
%	-> R the quadratic integrated cost matrix on the command
%       -> dt the current time step used
%
% returns -> the associated cost.
%
% created by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005
% last edited by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005

x = ilqr_openloop(x0, u, dt); % simulate system
N = length(x);

cost = (x(N,:) - xf) * Qf * (x(N,:) - xf)'; % compute final cost

for i = 1:N-1
	cost = cost + x(i,:) * Q * x(i,:)' + u(i,:) * R * u(i,:)'; % add instantaneous cost
end;

cost = (1/2)*cost; % this is for conventions but doesn't make a difference

return;


%% locally linear model
function [A, B] = ilqr_linearization(x, u, dt)

% [A, B] = ilqr_linearization(x, u, dt)
%
% linearizes the system in preparation to apply the ilqr techniques.
%
% given	-> x the state of the system at some time
%	-> u the command into the system at some time
%	-> dt the current time step used
%
% returns -> A the current jacobian Dxf(x,u) as defined by todorov
%	  -> B the current jacobian Duf(x,u) as defined by todorov
%
% created by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005
% last edited by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005

A = zeros(length(x), length(x));
B = zeros(length(x), length(u));

mass = [1 1];

A = diag([1 1 1 1]);
A(1,3) = dt;
A(2,4) = dt;

B(1,1) = (1/(2*mass(1)))*dt^2;
B(2,2) = (1/(2*mass(2)))*dt^2;
B(3,1) = (1/mass(1))*dt;
B(4,2) = (1/mass(2))*dt;
% these matrix entries come from the linearization of the 2D movving mass
% problem (which actually is linear) and can be determined straight-away from
% the equations of motions of a point mass

return;


%% open-loop computations
function trajectory = ilqr_openloop(x0, u, dt)

% trajectory = ilqr_openloop(x, u, dt)
%
% run the system forward in open loop from an initial state x0 using
% the control signal u.  note that length(trajectory) == length(u) + 1.
%
% given -> x0 the state of the system at some initial time
%       -> u the command signal into the system 
%       -> dt the current time step used
%
% returns -> the trajectory through state space taken by the system.
%
% created by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005
% last edited by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005

trajectory = zeros(length(u)+1, length(x0));

trajectory(1,:) = x0; % initial state

for i = 2:length(trajectory)
	trajectory(i,:) = ilqr_system(trajectory(i-1,:), u(i-1,:), dt); % compute states over time
end;

return;


%% the 2-D mass system
function newx = ilqr_system(x, u, dt)

% newx = ilqr_system(x, u, dt)
%
% run the system under consideration forward and return the new state.
%
% given -> x the state of the system at some time
%       -> u the command into the system at some time
%       -> dt the current time step used
%
% returns -> newx the new state of the system
%
% created by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005
% last edited by timothy lillicrap (tim@biomed.queensu.ca); december 1, 2005

newx = zeros(size(x));
A = zeros(length(x), length(x));
B = zeros(length(x), length(u));

[A B] = ilqr_linearization(x, u, dt); % compute linear system

newx = (A * x' + B * u')'; % simulate

return;
