Get rid of NaNs in kl penalty losses in the presence of masked networks.
Get rid of NaNs in kl penalty losses in the presence of masked networks.

This is to support masked networks (still wip), and have no impact on existing flows.

PiperOrigin-RevId: 321480566
Change-Id: I28012a25e0a3a3258ac6adfa7eb6d25190606cb9

tensorflow/agents,tensorflow/agents