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
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?