%% Create Actor Network
modifiedActorNet = layerGraph();

% Add layer branches
tempLayers = [
    featureInputLayer(2,"Name","input_1")
    fullyConnectedLayer(256,"Name","fc_1")
    reluLayer("Name","relu_body")
    fullyConnectedLayer(256,"Name","fc_2")
    reluLayer("Name","relu_body_2")
    fullyConnectedLayer(256,"Name","fc_body")
    reluLayer("Name","body_output")];
modifiedActorNet = addLayers(modifiedActorNet,tempLayers);

tempLayers = [
    fullyConnectedLayer(1,"Name","fc_std")
    softplusLayer("Name","std")];
modifiedActorNet = addLayers(modifiedActorNet,tempLayers);

tempLayers = fullyConnectedLayer(1,"Name","fc_mean");
modifiedActorNet = addLayers(modifiedActorNet,tempLayers);

tempLayers = [
    concatenationLayer(1,2,"Name","output")
    regressionLayer("Name","RepresentationLoss")];
modifiedActorNet = addLayers(modifiedActorNet,tempLayers);

% Connect layer branches
modifiedActorNet = connectLayers(modifiedActorNet,"body_output","fc_std");
modifiedActorNet = connectLayers(modifiedActorNet,"body_output","fc_mean");
modifiedActorNet = connectLayers(modifiedActorNet,"fc_mean","output/in1");
modifiedActorNet = connectLayers(modifiedActorNet,"std","output/in2");

% clean up helper variable
clear tempLayers;

%% Create Critic Network
modifiedCriticNet = [
    featureInputLayer(2,"Name","input_1")
    fullyConnectedLayer(256,"Name","fc_1_1")
    reluLayer("Name","relu_body_1")
    fullyConnectedLayer(256,"Name","fc_1_2")
    reluLayer("Name","relu_body_2")
    fullyConnectedLayer(256,"Name","fc_body")
    reluLayer("Name","body_output")
    fullyConnectedLayer(1,"Name","output")
    regressionLayer("Name","RepresentationLoss")];