Yes it does, thanks!
One last question: when using DQN, how does RLlib handles the custom model if it extends TFModelV2 versus extending DistributionalQTFModel?
I see that in parametric_actions_model.py they extend the latter but the implementation is the same (without considering the action embedding).