Why is the optimal output out of domain in A2C?

Question

If each state has an optimal action, then the optimal actions distribution vector is a one-hot vector kind of like [0,0,1,0,0,0]. But with algorithms like A2C, we use a softmax to get a distribution, meaning that to get a 0, you would need $e^x = 0$, which is impossible.

So why don't we use another activation function? Or am I simply missing something?

Answer 1

Implementations of RL policies (for example, in Stable Baselines) will often have a boolean argument deterministic. If we set this flag to true, then the policy will return the argmax over the distribution, and we would get the optimal actions for each state, even if the policy's distribution over actions doesn't put 100% of probability mass on the optimal choice.

Two important reasons that we don't take the argmax action during training, and instead sample a random action from the categorical distribution induced by the softmax:

1. Only the softmax distribution is differentiable. You can't differentiate through an argmax.
2. Without some nondeterminism in the policy, we would fail to explore the MDP.