# Whats the correct loss function to use during deep Q-learning (discrete action space)

After playing around with normal Q-learning I have decided to switch to deep Q-learning and I have encountered this problem.

As I understand, for a task with discrete action space, where there are 4 possible actions (lets say left, right, up, down) my DQN needs to have four outputs. And then argmax of prediction needs to be taken, which will be my predicted action (if argmax(prediction)==2 then I will pick third action, in my case up).

If I use Mean Squared Error loss for my network then the dimension of output needs to be same as dimension of expected target variables. I am calculating target variables using following code:

target = rewards[i] + GAMMA * max(Qs_next_state[i]) which gives me a single number (while predicted output is four dimensional) as I workaround I decided to use custom loss function:

def custom_loss(y_true, y_pred): # where y_true is target, y_pred is output of neural net
return ((max(y_pred) - y_true)**2) / BATCH_SIZE


But I am not sure if it is correct, from what I have read from tutorials/papers loss functions do not have this max() in them. But how else am I going to end up with same dimensionality between targets and NN outputs. What is the correct approach here?