Numerical instability in C51

C51 does a cross-entropy loss which could have numerical instability depending on the implementation. See link for an overview. Usually calculating the cross-entropy loss directly from the logits is more numerically stable. However, I am not sure how to do it exactly.

Deepmind's dqn_zoo has an implementation that seems to use the logits directly:

https://github.com/deepmind/dqn_zoo/blob/f011d683529d8d23b017a95194ebbb41a4962fe8/dqn_zoo/c51/agent.py#L35 https://github.com/deepmind/rlax/blob/42bbcf97a69ef9b21cb88322b83169ade7930363/rlax/_src/value_learning.py#L703 https://github.com/deepmind/rlax/blob/42bbcf97a69ef9b21cb88322b83169ade7930363/rlax/_src/value_learning.py#L543

Personally, I am recording this issue but in practice often it's enough to do

loss = (-(target_pmfs * old_pmfs.clamp(min=1e-5, max=1-1e-5).log()).sum(-1)).mean()
# instead of 
# unstable_loss = (-(target_pmfs * old_pmfs.log()).sum(-1)).mean()

The more stable loss results in

and unstable_loss results in

See more at #102

If anyone is interested in digging into this, that will be fantastic.