[Bugfix] Downcast to float32 when Parallel/SerialEnv is on MPS

[bsprenger]

Description

Fix SerialEnv and ParallelEnv crashing with TypeError: Cannot convert a MPS Tensor to float64 dtype when the parent environment is placed on an MPS device and sub-environments produce float64 observations (the default for many Gym/Gymnasium envs).

Motivation and Context

MPS has no hardware support for 64-bit floating-point. Standard envs (e.g. HalfCheetah-v4) return float64 observations by default. As a result, any user who tried to run a batched environment on MPS with a standard Gym env received a hard crash with an opaque TypeError from PyTorch's MPS backend.

How can this happen: the Parallel/SerialEnv have workers which actually run the rollouts. These workers may run the env on CPU, so it can collect float64 data. The error comes when returning that data to the parent Parallel/SerialEnv. To do so, it has to cast it to the parent's device, which can be a different device from the workers. If the parent is on MPS, then an error is raised as it cannot have float64 tensors.

Even if the workers are on CPU in many cases it is better to have the parent be on MPS, as the parent env can be wrapped in a sequence of transforms which are typically more efficient on MPS. More practically, it is just easier to always keep the parent on GPU as a default, and let the worker envs be cpu/gpu as their defaults specify.

Closes: #3550

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• Bug fix (non-breaking change which fixes an issue)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3551

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 4 Cancelled Jobs, 1 Pending, 1 Unrelated Failure

As of commit e2d0cf7 with merge base b6326f7 ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  Workflow failed! Resource not accessible by integration
• Generate documentation / build-docs (3.12, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 963834c2b7a2fb006062843a7c2e10c1d04b511bff9606f42cf45e2559d7fcb8 /exec failed with exit code 2

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Continuous Benchmark / GPU Pytest benchmark (gh)
  WARNING: The directory '/github/home/.cache/pip' or its parent directory is not owned or is not writable by the current user. The cache has been disabled. Check the permissions and owner of that directory. If executing pip with sudo, you should use sudo's -H flag.
• Libs Tests on Linux / unittests-gym (3.10, 12.8) / linux-job (gh)
• Libs Tests on Linux / unittests-sklearn (3.10, 12.8) / linux-job (gh)

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Continuous Benchmark / CPU Pytest benchmark (gh) (trunk failure)
  TypeError: expected Tensor()

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

LGTM thanks!