AI reinforcement learning via asynchronous advantage actor critic?

It has been almost 2 years since a major breakthough in AI research made Reinforcement Learning computationally feasible on laptops even without use of GPU computing:

https://arxiv.org/abs/1602.01783

Searching for "asynchronous advantage actor critic" (A3C in short) in the Mathematica documentation does not seem to find anything.

My questions are:

Has modern reinforcement learning been implemented in Mathematica? If not, does Wolfram plan to do that?

More importantly:

Is it possible to implement A3C in Mathematica with the currently available functionality? Perhaps some one knowledgeable can post a minimal working example as an answer to this question?

• WRI has implemented, and released as open source, a binding for OpenAI, which is a framework for reinforcement learning. Here is another question about reinforcement learning (unanswered.) I don't know the answer to your question though, whether state of the art reinforcement learning has been implemented by anyone. – C. E. Jan 25 '18 at 19:05
• @C.E. can you please add a link to that binding? – M.R. Feb 25 at 8:14
• @M.R. You can get it here. Since writing the previous comment, I have tried it to see if it works, and it did. – C. E. Feb 25 at 8:39

The following is my attempt at writing an A3C system in Mathematica from scratch. The result is not particularly good at learning, but definitely does learn something at least, which suffices as a sanity check and proof of concept.

We will need parallel computations, so we activate some number of Kernels:

numKernels = 8;
LaunchKernels[numKernels]


The following defines an instance of the circular pong game on each parallel Kernel:

ParallelEvaluate[
scores = {};
dt = 1/100.;
r = 0;
gameOver = True;
orient[P1_, P2_, P3_] := Sign[(P2[] - P1[]) (P3[] - P2[]) - (P2[] - P1[]) (P3[] - P2[])];
initialize[] := Block[{},
phi = RandomReal[{0, 2 \[Pi]}]; (*initial player starting angle*)
dir = 0;(*1=move clockwise, 0=stay, -1=move anti-clockwise*)
p1 = ReIm[Exp[I (phi + ArcCsc)]];(*bounding points for*)
p2 = ReIm[Exp[I (phi - ArcCsc)]];(*the player racket*)
theta = Mod[phi + \[Pi] + RandomReal[{-\[Pi]/2, \[Pi]/2}], 2. \[Pi]];
X = ReIm[Exp[I theta]];(*initial ball position*)
P = -X;(*initial ball movement direction*)
score = 0;
gameOver = False;
];
r = 0;
phi = Mod[phi + \[Pi] dir dt, 2. \[Pi]];
p1 = ReIm[Exp[I (phi - ArcCsc)]];
p2 = ReIm[Exp[I (phi + ArcCsc)]];
Xn = X + P dt;
If[(orient[X, Xn, p1] != orient[X, Xn, p2]) && (orient[p1, p2, X] != orient[p1, p2, Xn]),
normal = ReIm[Exp[I (phi)]];
If[normal.P > 0,
rt = ReflectionTransform[normal];
P = rt[P];
r = r + 500;
, X = Xn;
];
, X = Xn;
];
If[X.X > 1, gameOver = True; score = score - 2000; AppendTo[scores, score];];
vec = ReIm[Exp[I (phi)]] - X;
r = r + 1/(1 + Sqrt[vec.vec]); score = score + r;
];
];


Then we pre-define some overall parameters for our neural network. The game has five values describing each state:

1. $x$-coordinate of racket center
2. $y$-coordinate of racket center
3. $x$-coordinate of ball
4. $y$-coordinate of ball
5. current score

Therefore we choose input layer to be $2\times 5=10$ nodes, and pass a history of two consecutive states of the system as input. All layers are chosen to be fully connected with their neighbors. There are 4 hidden rectified linear unit layers (each with 10 nodes). The output layer has three softmax policy $\pi$ fields (to be interpreted as probabilities to choose clockwise, stay and anti-clockwise motion actions) and one linear node for value function $V$ output. Overall, parameters read:

reLuLayers = 4;
inLayerDim = 10;
reLuDim = inLayerDim;
outLayerDim = 4;(*policy: go clockwise, stay, go a-clockwise and value function*)
Tmax = 1000;(*number of model updates*)
Nrep = 20;(*number of steps to accumulate gradient*)
gamma = 0.95;(*discount of rewards*)
rndT = 500;(*number of updates for random choices to fall from 1 to 0.1 probability*)
overallT = 0;(*total number of updates in this session*)
epsilon = 10^-3;(*learning rate*)
rho1 = 0.9;(*parameters for adam learning algorithm*)
rho2 = 0.999;
delta = 10^-8;

hNodes = Prepend[ Append[Table[0, {i, 1, reLuLayers}, {j, 1, reLuDim}], Table[0, {j, 1, outLayerDim}]], Table[0, {j, 1, inLayerDim}]];
activations = Append[Table[0, {i, 1, reLuLayers}, {j, 1, reLuDim}], Table[0, {j, 1, outLayerDim}]];


Randomly initialize a master neural network model and share parameters with each parallel Kernel:

rdm := RandomReal[{0, 1}];
cnst = 1/3;
masterModel = Prepend[Append[ Table[{Table[rdm, {k, 1, reLuDim}], cnst}, {i, 2, reLuLayers}, {j, 1, reLuDim}], Table[{Table[rdm, {k, 1, reLuDim}], cnst}, {j, 1, outLayerDim}]], Table[{Table[rdm, {k, 1, inLayerDim}], cnst}, {j, 1, reLuDim}]];
SetSharedVariable[masterModel];


Initialize local neural network model versions and quantities that are needed to perform back-propagation on each parallel Kernel:

ParallelEvaluate[
(*---repeat above parameter definitions - local versions yield speedup---*)
reLuLayers = 4;
inLayerDim = 10;
reLuDim = inLayerDim;
outLayerDim = 4;(*policy: go clockwise, stay, go a-clockwise and value function*)
W = 1;(*label shortcut*)
b = 2;(*label shortcut*)
Tmax = 1000;(*number of model updates*)
Nrep = 20;(*number of steps to accumulate gradient*)
gamma = 0.95;(*discount of rewards*)
epsilon = 10^-3;(*learning rate*)
rndT = 500;(*number of updates for random choices to fall from 1 to 0.1 probability*)
overallT = 0;(*total number of updates in this session*)
rndProb[t_] := Exp[-23/(10 rndT) t];(*probability to perform random action*)
(*----------------------------------------------------------------*)
model = masterModel;
activations = Append[Table[0, {i, 1, reLuLayers}, {j, 1, reLuDim}], Table[0, {j, 1, outLayerDim}]];
hNodes = Prepend[Append[Table[0, {i, 1, reLuLayers}, {j, 1, reLuDim}], Table[0, {j, 1, outLayerDim}]], Table[0, {j, 1, inLayerDim}]];
R = 0.;
];


We need a forward propagation routine to evaluate outputs from the neural networks. We use the algorithm presented in the deep learning book: ParallelEvaluate[
policyAndV[state_] := Block[{exps},
(* write input state to input layer output *)
hNodes[] = state;
(* propagate the input state through hidden layers, get activations and outputs *)
Do[
activations[[k, j]] = model[[k, j, W]].hNodes[[k]] + model[[k, j, b]];
hNodes[[k + 1, j]] = Max[0, activations[[k, j]] ];
, {k, 1, reLuLayers}, {j, 1, reLuDim}];
(* go from last hidden layer to output layer, get activations *)
activations[[-1]] = Table[model[[-1, j, W]].hNodes[[-2]] + model[[-1, j, b]], {j, 1, outLayerDim}];
(* get output-layer output *)
exps = Total[ Exp[ activations[[-1, 1 ;; -2]] ] ];
hNodes[[-1]] = Exp[ activations[[-1]] ]/exps;
hNodes[[-1, -1]] = activations[[-1, -1]];
(* return output-layer results *)
hNodes[[-1]]
];
];


The back-propagation algorithm that gives a gradient is also from this book. My implementation is: ParallelEvaluate[
backProp[] := Block[{},
(* output layer Log[policy] gradient for activations (ignoring (R-V) factor for now) *)
activationsGradient[[-1]] = 1 -(outLayerDim - 1) hNodes[[-1]];
(* output layer gradient of (R- V)^2 for the activation (ignoring (R-V) factor for now) *)
(* reLuLayers gradient for activations *)
Do[
activationsGradient[[k, j]] = HeavisideTheta[activations[[k, j]]] Total[ model[[k + 1, ;; , W, j]] activationsGradient[[k + 1, ;;]] ];
, {k, reLuLayers, 1, -1}, {j, 1, Length[activationsGradient[[k]] ]}];
(* gradient of model parameters *)
Do[
, {k, 1, reLuLayers + 1}, {i, 1, Length[gradTable[[k]] ]}];
(* output result *)
];
];


where we made use of the fact that logged softmax node derivatives are given by 1 minus (outLayerDim-1) of the respective softmax itself:

num = 3;
Dt[Sum[Log[Exp[a[i]]/Sum[Exp[a[j]], {j, 1, num}]], {i, 1, num}] /. a[x_] :> ToExpression["a" <> ToString[x]]] // Simplify Additionally, we will make use of another back propagation routine, which will generate the gradient of a policy entropy $\nabla_\theta\pi\log\pi$ (improves system exploration by seeking higher entropy), and gradient of negative value function $-\nabla_\theta V$ (incentivises the model to seek higher scores). The code is the same as above, except for these changed lines:

ParallelEvaluate[
backPropEntropy[] := Block[{},
...
(* output layer policy*Log[policy] gradient for activations  *)
activationsGradient[[-1, 1]] = hNodes[[-1, 1]] (hNodes[[-1, 2]] Log[ hNodes[[-1, 1]]/(hNodes[[-1, 2]] + delta) + delta] + hNodes[[-1, 3]] Log[ hNodes[[-1, 1]]/(hNodes[[-1, 3]] + delta) + delta]);
activationsGradient[[-1, 2]] = hNodes[[-1, 2]] (hNodes[[-1, 1]] Log[ hNodes[[-1, 2]]/(hNodes[[-1, 1]] + delta) + delta] + hNodes[[-1, 3]] Log[ hNodes[[-1, 2]]/(hNodes[[-1, 3]] + delta) + delta]);
activationsGradient[[-1, 3]] = hNodes[[-1, 3]] (hNodes[[-1, 2]] Log[ hNodes[[-1, 3]]/(hNodes[[-1, 2]] + delta) + delta] + hNodes[[-1, 1]] Log[ hNodes[[-1, 3]]/(hNodes[[-1, 1]] + delta) + delta]);
(* output layer gradient of -V for the activation *)
(* reLuLayers gradient for activations *)
...
];
];


Finally, we implement the A3C algorithm from the original paper: T = 1;
Dynamic[{T, Tmax}]
ParallelEvaluate[gameOver = True;];
While[T < Tmax,
model = masterModel;
If[gameOver,
initialize[];
state = {{ReIm[Exp[I phi]], X, score}, {ReIm[Exp[I phi]], X, score}};
Do[stepAdvance[]; state[[i]] = {ReIm[Exp[I phi]], X, score};, {i, 2, 2}];
state = state // Flatten;(*prepare initial 4 steps of the game *)
gameOver = False;];
t = 1;
revs = Table[0, Nrep];(* store intermediate rewards *)
entrp = Table[0, Nrep];(* store intermediate entropy gradients *)
Vals = Table[0, Nrep];(* store intermediate value functions  *)
While[t < Nrep + 1,
policyAndV[state];(* get policy for current state *)
Vals[[t]] = hNodes[[-1, -1]];
entrp[[t]] = backPropEntropy[];(* accumulate entropy gradients *)
If[rndProb[overallT] > RandomReal[{0, 1}],
dir = RandomInteger[{-1, 1}];
,
dir = Ordering[hNodes[[-1, ;; -2]], -1][] - 2;(*which way? - change for other games/tasks*)
];
state = Append[state[[inLayerDim/2 + 1 ;;]], {ReIm[Exp[I phi]], X,
score}] // Flatten;(* append new state - change for other games/tasks*)
revs[[t]] = r;(* record reward of this time step*)
t++;
If[gameOver, Break[];](* exit While loop if terminal state *)
];
If[gameOver, R = 0;, R = hNodes[[-1, -1]];];(* projected reward in current episode *)
return = 0;
Do[
R = revs[[i]] + gamma R;
return = return + gradi[[i]] (R - Vals[[i]]) + entrp[[i]];
, {i, t - 1, 1, -1}];
overallT = overallT + 1;
return
];
overallT = overallT + 1;
masterModel = masterModel - epsilon (ADAMs/(1 - rho1^overallT)/( delta + Sqrt[ADAMr/(1 - rho2^overallT)]));
T++;];


where in the end we used the Adam algorithm to perform gradient descent.

The above code concludes the reinforcement learning algorithm. But we also would like to see how the trained model performs. So we should make sure to evaluate the circular pong game code above and the policyAndV[] function on the controlling kernel, to make these definitions available there as well. Then we can let the model play the game as follows:

Dynamic[{N[score, 2], hNodes[[-1, -1]]}]
Dynamic[ Show[ Graphics[Circle[{0, 0}, 1]] , Graphics[Line[{p1, p2}]] , Graphics[{PointSize[Large], Point[X]}] ] ]
gameOver = True;
model = masterModel;
While[True,
If[gameOver,
initialize[];
state = {{ReIm[Exp[I phi]], X, score}, {ReIm[Exp[I phi]], X, score}};
Do[stepAdvance[]; state[[i]] = {ReIm[Exp[I phi]], X, score};, {i, 2, 2}];
state = state // Flatten;(*prepare initial 2 steps of the game *)
gameOver = False;];
Pause[2 10^-2];
policyAndV[state];(* get policy for current state *)
dir = (Ordering[hNodes[[-1, ;; -2]], -1][] - 2);(*which way? - change for other games/tasks*)  