classdef CustomReinforceAgent < rl.agent.CustomAgent
    %CUSTOMREINFORCEAGENT Custom REINFORCE agent
    
    %% Properties (set properties attributes accordingly)
    properties
        % Actor representation
        Actor
        
        % Agent options
        Options
        
        % Experience buffer
        ObservationBuffer
        ActionBuffer
        RewardBuffer
    end
    
    properties (Access = private)
        % Training utilities
        Counter
        NumObservation
        NumAction
    end
    
    
    %% Necessary Functions
    %======================================================================
    % Implementation of public function
    %======================================================================
    methods
        function obj = CustomReinforceAgent(Actor,Options)
            %CUSTOMREINFORCEAGENT Construct custom agent
            %   AGENT = CUSTOMREINFORCEAGENT(ACTOR,OPTIONS) creates custom
            %   REINFORCE AGENT from rlStochasticActorRepresentation ACTOR
            %   and structure OPTIONS. OPTIONS has fields:
            %       - DiscountFactor
            %       - MaxStepsPerEpisode
            
            % (required) Call the abstract class constructor.
            obj = obj@rl.agent.CustomAgent();
            obj.ObservationInfo = Actor.ObservationInfo;
            obj.ActionInfo = Actor.ActionInfo;
            
            % (required for Simulink environment) Register sample time. 
            % For MATLAB environment, use -1.
            obj.SampleTime = -1;
            
            % (optional) Register actor and agent options.
            Actor = setLoss(Actor,@lossFunction);
            obj.Actor = Actor;
            obj.Options = Options;
            
            % (optional) Cache the number of observations and actions.
            obj.NumObservation = prod(obj.ObservationInfo.Dimension);
            obj.NumAction = prod(obj.ActionInfo.Dimension);
            
            % (optional) Initialize buffer and counter.
            reset(obj);
        end
    end
    
    %======================================================================
    % Implementation of abstract function
    %======================================================================
    methods (Access = protected)
        function Action = getActionImpl(obj,Observation)
            % Compute an action using the policy given the current 
            % observation.
            
            Action = getAction(obj.Actor,Observation);
        end
        
        function Action = getActionWithExplorationImpl(obj,Observation)
            % Compute an action using the exploration policy given the  
            % current observation.
            
            % REINFORCE: Stochastic actors always explore by default
            % (sample from a probability distribution)
            Action = getAction(obj.Actor,Observation);
        end
        
        function Action = learnImpl(obj,Experience)
            % Define how the agent learns from an Experience, which is a
            % cell array with the following format.
            %   Experience = {observation,action,reward,nextObservation,isDone}
            
            % Reset buffer at the beginning of the episode.
            if obj.Counter < 2
                resetBuffer(obj);
            end
            
            % Extract data from experience.
            Obs = Experience{1};
            Action = Experience{2};
            Reward = Experience{3};
            NextObs = Experience{4};
            IsDone = Experience{5};
            
            % Save data to buffer.
            obj.ObservationBuffer(:,:,obj.Counter) = Obs{1};
            obj.ActionBuffer(:,:,obj.Counter) = Action{1};
            obj.RewardBuffer(:,obj.Counter) = Reward;
            
            if ~IsDone
                % Choose an action for the next state.
                
                Action = getActionWithExplorationImpl(obj, NextObs);
                obj.Counter = obj.Counter + 1;
            else
                % Learn from episodic data.
                
                % Collect data from the buffer.
                BatchSize = min(obj.Counter,obj.Options.MaxStepsPerEpisode);
                ObservationBatch = obj.ObservationBuffer(:,:,1:BatchSize);
                ActionBatch = obj.ActionBuffer(:,:,1:BatchSize);
                RewardBatch = obj.RewardBuffer(:,1:BatchSize);
                
                % Compute the discounted future reward.
                DiscountedReturn = zeros(1,BatchSize);
                for t = 1:BatchSize
                    G = 0;
                    for k = t:BatchSize
                        G = G + obj.Options.DiscountFactor ^ (k-t) * RewardBatch(k);
                    end
                    DiscountedReturn(t) = G;
                end
                
                % Organize data to pass to the loss function.
                LossData.batchSize = BatchSize;
                LossData.actInfo = obj.ActionInfo;
                LossData.actionBatch = ActionBatch;
                LossData.discountedReturn = DiscountedReturn;
                
                % Compute the gradient of the loss with respect to the
                % actor parameters.
                ActorGradient = gradient(obj.Actor,'loss-parameters',...
                    {ObservationBatch},LossData);
                
                % Update the actor parameters using the computed gradients.
                obj.Actor = optimize(obj.Actor,ActorGradient);
                
                % Reset the counter.
                obj.Counter = 1;
            end
        end
    end
    
    %% Optional Functions
    %======================================================================
    % Implementation of optional function
    %======================================================================
    methods (Access = protected)
        function resetImpl(obj)
            % (Optional) Define how the agent is reset before training.
            
            resetBuffer(obj);
            obj.Counter = 1;
        end
    end
    
    methods (Access = private)
        function resetBuffer(obj)
            % Reinitialize all experience buffers.
            
            obj.ObservationBuffer = zeros(obj.NumObservation,1,obj.Options.MaxStepsPerEpisode);
            obj.ActionBuffer = zeros(obj.NumAction,1,obj.Options.MaxStepsPerEpisode);
            obj.RewardBuffer = zeros(1,obj.Options.MaxStepsPerEpisode);
        end
    end
end

function loss = lossFunction(policy,lossData)

    % Create the action indication matrix.
    batchSize = lossData.batchSize;
    Z = repmat(lossData.actInfo.Elements',1,batchSize);
    actionIndicationMatrix = lossData.actionBatch(:,:) == Z;
    
    % Resize the discounted return to the size of policy.
    G = actionIndicationMatrix .* lossData.discountedReturn;
    G = reshape(G,size(policy));
    
    % Round any policy values less than eps to eps.
    policy(policy < eps) = eps;
    
    % Compute the loss.
    loss = -sum(G .* log(policy),'all');
    
end