%% Train Custom LQR Agent
% This example shows how to train a custom linear quadratic regulation (LQR) agent to control a discrete-time linear system modeled in MATLAB®.
%
%% Create Linear System Environment
% The reinforcement learning environment for this example is a discrete-time linear system. The dynamics for the system are given by
% x_{t+1} = Ax_t+Bu_t
% The feedback control law is
% u_t = -Kx_t
% The control objective is to minimize the quadratic cost: J=\sum_{t=0}^{\infty } \left({x_t }^{\prime } {\textrm{Qx}}_t +{u_t }^{\prime } {\textrm{Ru}}_t \right).
% In this example, the system matrices are
% \begin{array}{l}
% A=\left\lbrack \begin{array}{ccc}
% 1\ldotp 05 & 0\ldotp 05 & 0\ldotp 05\\
% 0\ldotp 05 & 1\ldotp 05 & 0\ldotp 05\\
% 0 & 0\ldotp 05 & 1\ldotp 05
% \end{array}\right\rbrack \\
% B=\left\lbrack \begin{array}{ccc}
% 0\ldotp 1 & 0 & 0\ldotp 2\\
% 0\ldotp 1 & 0\ldotp 5 & 0\\
% 0 & 0 & 0\ldotp 5
% \end{array}\right\rbrack
% \end{array}

A = [1.05,0.05,0.05;0.05,1.05,0.05;0,0.05,1.05];
B = [0.1,0,0.2;0.1,0.5,0;0,0,0.5]; 

% The quadratic cost matrices are:
% \begin{array}{l}
% Q=\left\lbrack \begin{array}{ccc}
% 10 & 3 & 1\\
% 3 & 5 & 4\\
% 1 & 4 & 9
% \end{array}\right\rbrack \\
% R=\left\lbrack \begin{array}{ccc}
% 0\ldotp 5 & 0 & 0\\
% 0 & 0\ldotp 5 & 0\\
% 0 & 0 & 0\ldotp 5
% \end{array}\right\rbrack
% \end{array}

Q = [10,3,1;3,5,4;1,4,9]; 
R = 0.5*eye(3);

% For this environment, the reward at time t is given by r_t = -x_t'Qx_t-u_t'Ru_t
% , which is the negative of the quadratic cost. Therefore, maximizing the reward minimizes the cost. The initial conditions are set randomly by the reset function.
% Create the MATLAB environment interface for this linear system and reward. The myDiscreteEnv function creates an environment by defining custom step and reset functions. For more information on creating such a  custom environment, see Create MATLAB Environment Using Custom Functions.

env = myDiscreteEnv(A,B,Q,R);

% Fix the random generator seed for reproducibility.

rng(0)

%
%% Create Custom LQR Agent
% For the LQR problem, the Q-function for a given control gain K can be defined as Q_K \left(x,u\right)={\left\lbrack \begin{array}{c}
% x\\
% u
% \end{array}\right\rbrack }^{\prime } H_K \left\lbrack \begin{array}{c}
% x\\
% u
% \end{array}\right\rbrack, where H_K =\;\left\lbrack \begin{array}{cc}
% H_{\textrm{xx}}  & H_{\textrm{xu}} \\
% H_{\textrm{ux}}  & H_{\textrm{uu}}
% \end{array}\right\rbrack is a symmetric, positive definite matrix.
% The control law to maximize Q_K is u=-{\left(H_{\textrm{uu}} \right)}^{-1} H_{\textrm{ux}} \;x, and the feedback gain is {K=\left(H_{\textrm{uu}} \right)}^{-1} H_{\textrm{ux}} \;.
% The matrix H_K contains m=\frac{1}{2}n\left(n+1\right) distinct element values, where n is the sum of the number of states and number of inputs. Denote \theta as the vector corresponding to these m elements, where the off-diagonal elements in H_K are multiplied by two.
% Represent the Q-function by \theta, where \theta contains the parameters to be learned.
% Q_K \left(x,u\right)=\theta^{\prime } \left(K\right)\phi \left(x,u\right), where \phi \left(x,u\right) is the quadratic basis function in terms of x and u.
% The LQR agent starts with a stabilizing controller K_0. To get an initial stabilizing controller, place the poles of the closed-loop system A-{\textrm{BK}}_0 inside the unit circle.

K0 = place(A,B,[0.4,0.8,0.5]);

% To create a custom agent, you must create a subclass of the rl.agent.CustomAgent abstract class. For the custom LQR agent, the defined custom subclass is LQRCustomAgent. For more information, see Custom Agents. Create the custom LQR agent using Q, R, and K_0. The agent does not require information on the system matrices A and B.

agent = LQRCustomAgent(Q,R,K0);

% For this example, set the agent discount factor to one. To use a discounted future reward, set the discount factor to a value less than one.

agent.Gamma = 1;

% Because the linear system has three states and three inputs, the total number of learnable parameters is m=21. To ensure satisfactory performance of the agent, set the number of parameter estimates N_p to be greater than twice the number of learnable parameters. In this example, the value is N_p =45.

agent.EstimateNum = 45;

% To get good estimation results for \theta, you must apply a persistently excited exploration model to the system. In this example, encourage model exploration by adding white noise to the controller output: u_t =-{\textrm{Kx}}_t +e_t. In general, the exploration model depends on the system models.
%
%% Train Agent
% To train the agent, first specify the training options. For this example, use the following options.
% Run each training episode for at most 10 episodes, with each episode lasting at most 50 time steps.
% Display the command line display (set the Verbose option) and disable the training progress in the Episode Manager dialog box (set the Plots option).
% For more information, see rlTrainingOptions.

trainingOpts = rlTrainingOptions(...
    'MaxEpisodes',10, ...
    'MaxStepsPerEpisode',50, ...
    'Verbose',true, ...
    'Plots','none');

%
% Train the agent using the train function.

trainingStats = train(agent,env,trainingOpts);

%
%% Simulate Agent and Compare with Optimal Solution
% To validate the performance of the trained agent, simulate it within the MATLAB environment. For more information on agent simulation, see rlSimulationOptions and sim.

simOptions = rlSimulationOptions('MaxSteps',20);
experience = sim(env,agent,simOptions);
totalReward = sum(experience.Reward)

% You can compute the optimal solution for the LQR problem using the dlqr function.

[Koptimal,P] = dlqr(A,B,Q,R); 

% The optimal reward is given by J_{\textrm{optimal}} ={{-x}_0 }^{\prime } {\textrm{Px}}_0.

x0 = experience.Observation.obs1.getdatasamples(1);
Joptimal = -x0'*P*x0;

% Compute the error in the reward between the trained LQR agent and the optimal LQR solution.

rewardError = totalReward - Joptimal

% View the history of the 2-norm of error in the gains between the trained LQR agent and the optimal LQR solution.

% number of gain updates
len = agent.KUpdate;
err = zeros(len,1);
for i = 1:len
    % norm of error in the gain
    err(i) = norm(agent.KBuffer{i}-Koptimal);
end  
plot(err,'b*-')

% Compute the norm of final error for the feedback gain.

gainError = norm(agent.K - Koptimal)

% Overall, the trained agent finds an LQR solution that is close to the true optimal LQR solution.
% Copyright 2019 The MathWorks, Inc.