[RLlib] Fixed bug in restoring a gpu trained algorithm

[kouroshHakha]

Why are these changes needed?

Apparently when you restore the torch optimizer states, the param_groups should not be converted to tensor and moved to cuda devices. (I don't think it's even necessary for the states to be moved to the device they just have to be converted to a tensor). This bug was that when someone trained an algorithm with GPU and wanted to restore for further training again on a cuda device it would yell at you saying the beta parameter in Adam optimizer should not be on cuda or should be a scalar.

This PR fixes that by separating out restoration process of param_groups vs. state keys in state_dict of torch optimizers. It also adds unittests to ensure this edge case is covered in the unittests.

#closes #34159

Related issue number

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 :(

[avnishn]

ill probably need to make this change myself. adding a todo