When using experience replay in reinforcement learning, which state is used for training?

Question

I'm slightly confused about the experience replay process. I understand why we use batch processing in reinforcement learning, and from my understanding, a batch of states is input into the neural network model.

Suppose there are 2 valid moves in the action space (UP or DOWN)

Suppose the batch size is 5, and the 5 states are this:

$$[s_1, s_2, s_3, s_4, s_5]$$

We put this batch into the neural network model and output Q values. Then we put $[s_1', s_2', s_3', s_4', s_5']$ into a target network.

What I'm confused about is this:

Each state in $[s_1, s_2, s_3, s_4, s_5]$ is different.

Are we computing Q values for UP and DOWN for ALL 5 states after they go through the neural network?

For example, $$[Q_{s_1}(\text{UP}), Q_{s_1}(\text{DOWN})], \\ [Q_{s_2} (\text{UP}), Q_{s_2}(\text{DOWN})], \\ [Q_{s_3}(\text{UP}), Q_{s_3}(\text{DOWN})], \\ [Q_{s_4}(\text{UP}), Q_{s_4}(\text{DOWN})], \\ [Q_{s_5}(\text{UP}), Q_{s_5}(\text{DOWN})]$$

Answer 1

The way the states are used is as follows:

Typically your $Q$-network will state a state as input and output scores over the action space. I.e. $Q : \mathcal{S} \rightarrow \mathbb{R}^{|\mathcal{A}|}$. So, in your replay buffer you should store $s_t, a_t, r_{t+1}, s_{t+1}, \mbox{done}$ (note that done just represents where the episode ended on this transition and I add for completeness.

Now, when you are doing your batch updates you sample uniformly at random from this replay buffer. This means you get $B$ tuples of $s_t, a_t, r_{t+1}, s_{t+1}, \mbox{done}$. Now, I will assume $B=1$ as it is easier to explain and the extension to $B > 1$ should be easy to see.

For our state-action tuple $s_t, a_t$ we want to shift what the network predicts for this pair to be closer to $r_{t+1} + \gamma \arg\max_a Q(s,a)$. However, our neural network only takes the state as input, and outputs a vector of scores for each action. That means we want to shift the output of our network for the state $s_t$ towards the target I just mentioned, but only for the action $a_t$ that we took. To do this we just calculate the target, i.e. we calculate $r_{t+1} + \gamma \arg\max_a Q(s,a)$, and then we do gradient ascent like we would a normal neural network where the target vector is the same as the predicted vector everywhere except the $a_t$th element, which we will change to $r_{t+1} + \gamma \arg\max_a Q(s,a)$. This way, our network moves closer to our Q-learning update for only the action we want, in line with how Q-learning works.

It is also worth nothing that you can parameterise your Neural Network to be a function $Q: \mathcal{S} \times \mathcal{A} \rightarrow \mathbb{R}$ which would make training more in line with how tabular Q-learning but is seldom used in practice as it becomes much more expensive to compute (you have to do a forward pass for each action, rather than one forward pass per state).