[RLLib] Fix OneHotPreprocessor, use gym.spaces.utils.flatten

[olipinski]

Why are these changes needed?

As per #27496, the current implementation of OneHotPreprocessor loses information when one-hot encoding. Using the gym implementation should avoid this issue.

Related issue number

Closes #27496.

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

[gjoliver]

nice! thanks a ton for the fix.
can you help add a quick unit test in rllib/model/tests/test_preprocess.py for this behavior?
the team is really trying to push for better quality and testing in the long terms, fix like this is perfect for adding a unit test so there won't be regression.

[olipinski]

can you help add a quick unit test in rllib/model/tests/test_preprocess.py for this behavior?

I've added a quick test just now, both for 2D and 3D multidiscretes!

[olipinski]

Do we need to do anything else here?

[gjoliver]

thanks a ton for the fix.
I will ping @sven1977 to get this merged.

[olipinski]

Thanks for the approve!

[sven1977]

Thanks for the PR @olipinski , could you make sure the autoregressive_action_dist test case passes? There seems to be a connection to your changes:

https://buildkite.com/ray-project/ray-builders-pr/builds/43759#0182fa77-a584-4bef-a8ef-91f7e0ff4279

The script that fails is rllib/examples/autoregressive_action_dist.py run with the --framework torch flag.

[olipinski]

https://buildkite.com/ray-project/ray-builders-pr/builds/43759#0182fa77-a584-4bef-a8ef-91f7e0ff4279

The script that fails is rllib/examples/autoregressive_action_dist.py run with the --framework torch flag.

This issue appears to be due to torch.nn.linear supporting only float32, which the old code used to cast every one hotted action into. I have added the cast into the preprocessor, though I'm uncertain if that should be done there, or should the dtype be specified in the environment?

[gjoliver]

I think this is fine. thanks 👍

[olipinski]

Perfect, should be all good then!

[sven1977]

Thanks again @olipinski , and @gjoliver for the review!