Fixed log term when there are illegal actions in JAX implementation of DQN algorithm.
Fixed log term when there are illegal actions in JAX implementation of DQN
algorithm.

PiperOrigin-RevId: 430153241
Change-Id: Ie01e36d89cbbe388380533e71c4049ef77fe1588

deepmind/open_spiel,deepmind/open_spiel,deepmind/open_spiel,deepmind/open_spiel,deepmind/open_spiel,deepmind/open_spiel