Make MSE the default loss type for JAX DQN. This is based on the suggestions from https://arxiv.org/abs/2011.14826, and consistent with the baselines page (https://google.github.io/dopamine/baselines/atari/plots.html) where we have DQN+RMSProp+Huber (for TF) and DQN+Adam+MSE (for JAX).
Make MSE the default loss type for JAX DQN. This is based on the suggestions from https://arxiv.org/abs/2011.14826, and consistent with the baselines page (https://google.github.io/dopamine/baselines/atari/plots.html) where we have DQN+RMSProp+Huber (for TF) and DQN+Adam+MSE (for JAX).

PiperOrigin-RevId: 444822901

google/dopamine,google/dopamine,google/dopamine