Calculating the Policy Gradient for a Monte Carlo REINFORCE Algorithm

Question

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.

Answer 1

The scheme you describe sounds right. But let me clear up some of the confusion about the policy network, which will hopefully explain why that makes sense:

$\pi(\cdot| s; \theta)$ is a distribution across actions conditioned on a state.

If we call the policy network which outputs logits scores for each action $f(s;\theta)$, then $\pi(\cdot | s;\theta) = \text{Cat}(\text{softmax}(f(s;\theta)))$, assuming this is a discrete action space. However, people often use $\pi$ to both denote the network AND the distribution its output parameterizes.

Finally, $\pi(A_t | S_t, \theta)$ is the probability of sampling a single action $A_t$ from the distribution $\pi(\cdot |s;\theta)$. So $\pi$ is overloaded yet again to denote the probability mass function.

So now we know that actually $\pi(A_t | S_t, \theta) = \text{softmax}(f(s;\theta))_i$, where $i$ is the index of action $A_t$. And this is clearly a scalar value.