You can find an implementation of the REINFORCE algorithm (as defined in your question) in PyTorch at the following URL: https://github.com/JamesChuanggg/pytorch-REINFORCE/. First of all, I would like to note that a policy can be represented or implemented as a neural network, where the input is the state (you are currently in) and the output is a "probability distribution over the actions you can take from that state received as input".
In the Python module https://github.com/JamesChuanggg/pytorch-REINFORCE/blob/master/reinforce_discrete.py, the policy is defined as a neural network with 2 linear layers, where the first linear layer is followed by a ReLU activation function, whereas the second is followed by a soft-max. In that same Python module, the author also defines another class called
REINFORCE, which creates a
Policy object (in the
__init__ method) and defines it as property of that class. The class
REINFORCE also defines two methods
update_parameters. These two methods are called from the
main.py module, where the main loop of the
REINFORCE algorithm is implemented. In that same main loop, the author declares lists
rewards. Note that these lists are re-initialized at ever episode. A "log_prob" and an "entropy" is returned from the
select_action method, whereas a "reward" is returned from the environment after having executed one environment step. The environment is provided by the OpenAI's Gym library. The lists
rewards are then used to update the parameters, i.e. they are used by the method
update_parameters defined in the class
Let's see better now what these methods,
update_parameters, actually do.
select_action first calls the
forward method of the class
Policy, which returns the output of the forward pass of the NN (i.e. the output of the soft-max layer), so it returns the probabilities of selecting each of the available actions (from the state given as input). It then selects the probability associated with the first action (I guess, it picks the probabilities associated with the action with the highest probabilities), denoted by
prob (in the source code). Essentially, what I've described so far regarding this
select_action method is the computation of $\pi(A_t \mid S_t, \theta)$ (as shown in the pseudocode of your question). Afterwards, in the same method
select_action, the author also computes the log of that probability I've just mentioned above (i.e. the one associated with the action with the highest probability, i.e. the log of
prob), denoted by
log_prob. In that same method, the entropy (as defined in this answer) is calculated. In reality, the author calculates the entropy using only one distribution (instead of two): more specifically, the entropy is calculated as follows
entropy = -(probs*probs.log()).sum(). In fact, the entropy loss function usually requires the ground-truth labels (as explained in the answer I linked you to above), but, in this case, we do not have ground-truth labels (given that we are performing RL and not supervised learning). Nonetheless, I can't really tell you why the entropy is calculated like this, in this case. Finally, the method
select_action then return
action, log_prob, entropy.
First of all, I would like to note that the method
update_parameters is called only at the end of each episode (in the
main.py module). In that same method, a variable called
loss is first initialized to zero. In that method, we then iterate the list of rewards for the current episode. Inside that loop of the
update_parameters method, the return,
R is calculated.
R is also multiplied by $\gamma$. On each time step, the loss is then calculated as follows
loss = loss - (log_probs[i]*(Variable(R).expand_as(log_probs[i])).cuda()).sum() - (0.0001*entropies[i].cuda()).sum()
loss is calculated by subtracting the previous
(log_probs[i]*(Variable(R).expand_as(log_probs[i])).cuda()).sum() - (0.0001*entropies[i].cuda()).sum()
log_probs are the log probabilities calculated in the
log_probs is the part $\log \pi(A_t \mid S_t, \theta)$ of the update rule of your pseudocode.
log_probs are then multiplied by the return
R. We then sum the result of this multiplication over all elements of the vector. We then subtract this just obtained result by the entropies multiplied by 0.0001. I can't really tell you why the author decided to implement the loss in this way. I would need to think about it a little more.
The following article may also be useful: https://pytorch.org/docs/stable/distributions.html.