Fix nans in gradient from inverse softplus

[SimonKamuk]

Describe your changes

Fix a bug where the inverse softplus causes gradients to become nans.

I believe the issue was that even though the non_linear_part had clamping inside, and the values were later being discarded if they were above the threshold, the input x values could still cause numerical instabilities before we checked if the input was above the threshold.

What I did to fix it was to clamp the input x values. The limits were previously defined as $exp(x\cdot beta)-1 \ge 10^{-6}$ and $x\cdot beta \le threshold$, so I changed it to the equivalent $log(10^{-6}+1)/beta \le x \le threshold/beta$.

Issue Link

closes #119

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the README to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form (context).
• I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section
  reflecting type of change (add section where missing):
  • added: when you have added new functionality
  • changed: when default behaviour of the code has been changed
  • fixes: when your contribution fixes a bug

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• author has added an entry to the changelog (and designated the change as added, changed or fixed)
• Once the PR is ready to be merged, squash commits and merge the PR.

[SimonKamuk]

Since our tests now use the reduced danra domain and are relatively quick, what do you think about adding detect_anomaly=True to the pl.Trainer in test_training.py ? That would prevent an issue like this from getting through (assuming the training config used for the test includes any such future features). Is it a bit overkill?

[sadamov]

Tested this with two different datasets and ar_steps_train, solves the bug. Thanks Simon

[joeloskarsson]

I think it would be a good idea to add detect_anomaly=True to the tests