I am currently trying to implement the Monte Carlo REINFORCE algorithm, as described in Sutton and Barto's book Reinforcement Learning (p. 328, Second Edition). If $\theta$ denotes the parameter for the policy function $\pi(a | s, \theta)$, then the general update rule goes as follows
$$\theta_{t+1} = \theta_t + \alpha G_t \frac{\nabla \pi(A_t|S_t,\theta)}{\pi(A_t|S_t,\theta)}$$
where $G_t$ denotes the expected return at timestep $t$ and $\alpha$ is the learning rate.
So far I programmed a general framework for neural networks of variable size, which use the sigmoid activation function on hidden layers, as well as a final softmax layer, since I want to learn action probabilities for a discrete set of actions. In this case $$\pi(.|S_t,\theta)$$ would just represent the output of the network, when given the input $S_t$ and the $k$ output neurons each correspond to the probabilities of taking action $A_j$ in state $S_t$.
The problem I ran into is, that I am not quite sure how to compute $\nabla\pi(A_t|S_t,\theta)$ in this case. Viewing the network as its own function then $\pi$ would be a vector valued function, taking a state and outputting a vector. Clearly, in this case the gradient would not make too much sense.
I know how backpropagation works, so I thought maybe calculating the policy gradient at $A_t, S_t$ would simply correspond to doing backpropagation with error function $a_j^{(L)}$, where $a_j^{(L)}$ is the output of the $j$-th neuron in the output layer (and $A_t$ being the $j$-th action). The errors of the output neurons would then be $$\delta_i^{(L)} = \frac{\partial a_j^{(L)}}{\partial z_i^{(L)}}$$ (L denoting the index of the highest layer, $z_i ^ {(L)}$ denoting the weighted input into neuron $i$ in layer $L$). This is just the derivative of the softmax function, which can be easily computed. Then through a backwards pass, I would calculate the errors of the neurons below.
Is this the correct way of calculating $\nabla \pi(A_t | S_t, \theta)$ ? I would appreciate any help.