If so, then would it be correct to take the gradient of this expression of $\hat{q}(S,A,w)$ with respect to the weights of the network
Yes, your expression for $\hat{q}(S,A,w)$ looks correct for your network, and can be differentiated.
and plug into into the following update equation for back propagation?
Not clear what you intend. If you already have the gradient $\nabla\hat{q}(S,A,w)$ then you have already performed backpropagation from the current output value to all weights. The additional factor of $[R + \gamma \hat{q}(S',A',w) - \hat{q}(S,A,w)]$ is effectively the gradient of the MSE cost function w.r.t. the current output that needs to be multiplied in so that the backprop starts from the MSE loss.
The equations given in Sutton&Barto are not that easy to apply when you already have a neural network library that has everything defined around supervised learning.
You may find it easier to calculate a target value for the whole action value vector and use your library's definitions of loss functions and update steps. That would go something like this:
- Set the loss function of the neural network to MSE.
- Set the optimisation method to SGD (although you now have the option of using others e.g. Adam, momentum, mini-batches etc)
- Calculate $\begin{bmatrix}
\hat{q}(S,a_1,w) & \hat{q}(S,a_2,w) & \hat{q}(S,a_3,w)
\end{bmatrix}$ forward using the neural network.
- Calculate the TD target $G = R + \gamma \hat{q}(S',A',w)$
- Substitute the target value into the forward matrix e.g. $\begin{bmatrix}
\hat{q}(S,a_1,w) & \hat{q}(S,a_2,w) & G
\end{bmatrix}$ assuming it was action 3 that the agent took.
- Use the adjusted forward matrix as the SGD training value.
This will automatically provide a gradient 0 associated with actions that were not taken, and is effectively the same update as per Sutton&Barto, but using the NN library as-is without trying to tweak the gradient calculations directly.
If you are using a library like PyTorch, then it is definitely possible to calculate the pure gradient of the selected action as you are attempting. However, you may find yourself in more familiar and supported territory if you re-frame the updates as supervised learning as suggested above.
I suggest one of the first things you try if this does not converge is to batch up the updates and only perform them every e.g. 100 steps. Again, this is possible using the S&B update rule, and something like PyTorch, but things will be easier if you go with the grain and re-frame the updates as supervised learning "inside" the RL mechanism.
I am not sure how well SARSA will cope with Lunar Lander. I have solved it using DQN, which is similar, but has the advantage of re-using older off-policy experience, sampling from it randomly, which SARSA cannot do. To correct for SARSA's weakness there, you may need quite large batches to avoid too much correlation between experience. A common pattern with "Deep SARSA" is to sample multiple environments at once under the current policy, to build a large batch (you see this in A2C and other on-policy policy gradient methods too).
