Using "softmax" (non-linear) vs "linear" activation function in Deep Reinforcement Learning

Question

I am following the tutorial in this video: https://youtu.be/cO5g5qLrLSo which implements deep reinforcement learning (DQN) to balance cart pole in OpenAI default environment.

The DQN model looks like as follows:

model = Sequential()
model.add(Flatten(input_shape=(1,states)))
model.add(Dense(24, activation='relu'))
model.add(Dense(24, activation='relu'))
model.add(Dense(actions, activation='linear'))

Full code is available here: https://github.com/nicknochnack/TensorflowKeras-ReinforcementLearning/blob/master/Deep%20Reinforcement%20Learning.ipynb

The implemented code is also available on Google Collabs for you to run/test here: https://colab.research.google.com/drive/1oQILItVu6Y8jOCprzwMGzwlztYmVKK-F?usp=sharing

I do understand the concept and mathematics behind using "linear" vs "non-linear" (softmax) activation function in the output layer.

But, what I am struggling to understand is that why in this target application linear activation function is used in the output layer instead of softmax? Can someone specify how to realize that which type of activation function will be the best for what type of target application using DQN?

P.S. I tried to change the activation function to softmax instead and got completely different result. Hence, I am confused as to why changing the activation function in the output layer could generate completely opposite results.

Answer 1

The normal use case for softmax in the output layer is for a classification problem, where the output is an array of probabilities for each class.

The normal use case for a linear output is for a regression problem, where the output is an array of floating point numbers that are estimates for some measurement.

In a DQN, the desired output for the neural network is the action values (or Q values) for each action. The action value is the expected sum of future rewards given the starting condition of state and action choice, plus following the current target policy from that point onwards (this is a little self-referential in DQN in that the target policy is defined by taking the action with the highest action value in all states).

It is entirely possible for this predicted action value to be greater than $1$ or less than $0$, and there is also no reason to expect a the set of action values for all actions possible in a specific state to sum to $1$. This is a regression problem, and you should generally use a linear output for such problems.

Some Deep RL will learn a policy function, either instead of, or in addition to the value function - for example REINFORCE or A3C. In that case, if the action space is discrete, then the policy function will behave a lot like a classifier (it classifies the preferred action for a given state), and you would use a softmax output for that network. Other action spaces may be described differently though, so you are not guaranteed to see softmax used in output layer of policy function networks.