[RLlib] Add gradient checks to avoid nan gradients in TorchLearner.

[simonsays1980]

Why are these changes needed?

If any gradients turn nan in TorchLearner these gradients get added to the network's weights and in turn weights become nan and all network outputs as well. As a result the training errors out and stops. This PR proposes a gradient check to only add gradients if they are sane. It switches nan values to zeros in the gradients or skips en update entirely.
The latter can be of advantage, if training phase ecnounters highly unstable policy updates (e.g. with highly explorative policies or during early stages of training). In such phases many gradients could turn nan and this may lead to corrupted internal optimizer states (e.g. Adam).

Related issue number

#47451

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[sven1977]

Wait, I'm confused. We have this block here further below, which I think does the exact same thing: skips the entire optim.step in case any gradient is non-finite (inf or nan).

                # `step` the optimizer (default), but only if all gradients are finite.
                elif all(
                    param.grad is None or torch.isfinite(param.grad).all()
                    for group in optim.param_groups
                    for param in group["params"]
                ):

Can you check and see whether these two logics can be consolidated?

Kind of like this:

• If user sets this flag (default=False), the optimizer will skip the update step entirely (+ warning raised by RLlib).
• If user does NOT set this flag (default behavior), grads that are non-finite will be set to 0.0 (+ warning raised by RLlib).

[sven1977]

LGTM! Thanks @simonsays1980