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[[2]] - P1[[2]]) (P3[[1]] - P2[[1]]) - (P2[[1]] - P1[[1]]) (P3[[2]] - P2[[2]])];
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[3])]];(*bounding points for*)
p2 = ReIm[Exp[I (phi - ArcCsc[3])]];(*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;
];
stepAdvance[] := Block[{},
r = 0;
phi = Mod[phi + \[Pi] dir dt, 2. \[Pi]];
p1 = ReIm[Exp[I (phi - ArcCsc[3])]];
p2 = ReIm[Exp[I (phi + ArcCsc[3])]];
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:
- $x$-coordinate of racket center
- $y$-coordinate of racket center
- $x$-coordinate of ball
- $y$-coordinate of ball
- 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;
ADAMs = 0;(*first moment*)
ADAMr = 0;(*second moment*)
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;
gradTable = model/Infinity;
activations = Append[Table[0, {i, 1, reLuLayers}, {j, 1, reLuDim}], Table[0, {j, 1, outLayerDim}]];
activationsGradient = activations/Infinity;
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[[1]] = 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[{},
(* blank gradTable and activationsGradient *)
gradTable = gradTable/Infinity;
activationsGradient = activations/Infinity;
(* 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) *)
activationsGradient[[-1, -1]] = -2;
(* 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[
gradTable[[k, i, W]] = hNodes[[k]] activationsGradient[[k, i]];
gradTable[[k, i, b]] = activationsGradient[[k, i]];
, {k, 1, reLuLayers + 1}, {i, 1, Length[gradTable[[k]] ]}];
(* output result *)
gradTable
];
];
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 *)
activationsGradient[[-1, -1]] = -1 ;
(* 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,
grads = ParallelEvaluate[
gradTable = model/Infinity;
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 *)
gradi = Table[0, Nrep];(* store intermediate gradients *)
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]];
gradi[[t]] = backProp[];(* accumulate gradients *)
entrp[[t]] = backPropEntropy[];(* accumulate entropy gradients *)
If[rndProb[overallT] > RandomReal[{0, 1}],
dir = RandomInteger[{-1, 1}];
,
dir = Ordering[hNodes[[-1, ;; -2]], -1][[1]] - 2;(*which way? - change for other games/tasks*)
];
stepAdvance[];(* advance system by a time step *)
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
];
gradTable = Sum[grads[[i]], {i, 1, numKernels}]/numKernels;
(* adam learning algorithm: *)
overallT = overallT + 1;
ADAMs = rho1 ADAMs + (1 - rho1) gradTable;
ADAMr = rho2 ADAMr + (1 - rho2) gradTable gradTable;
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][[1]] - 2);(*which way? - change for other games/tasks*)
stepAdvance[];(* advance system by a time step *)
state = Append[state[[inLayerDim/2 + 1 ;;]], {ReIm[Exp[I phi]], X, score}] // Flatten;(* append new state - change for other games/ tasks*)
];
This is it! Training the system a couple of times, the neural network sometimes converges to a state that occasionally performs a somewhat sensible action:
However, most of the time it performs rather useless actions:
Nevertheless, a randomly initialized neural network almost certainly defaults to a steady rotation without ever changing pace or direction. This demonstrates that the system actually learned something during the reinforcement learning process.
