The update rule for the DQN get gradients flowing only for the index that was selected.
The backward pass of the gather function will do just that: pass the gradient of the output to the values that were gathered and 0 for everything else. These indices are saved during the forward pass as a “buffer” needed to compute the backward. We actually save a lot of a them during the forward.