diff --git a/test/test_envs.py b/test/test_envs.py index 469cd7f7110..d247dbd7327 100644 --- a/test/test_envs.py +++ b/test/test_envs.py @@ -56,7 +56,12 @@ SerialEnv, set_gym_backend, ) -from torchrl.envs.batched_envs import _stackable +from torchrl.envs.batched_envs import ( + _has_float64_leaf, + _stackable, + _td_to_device_mps_safe, + _to_device_mps_safe, +) from torchrl.envs.gym_like import default_info_dict_reader from torchrl.envs.libs.dm_control import _has_dmc, DMControlEnv from torchrl.envs.libs.gym import _has_gym, gym_backend, GymEnv, GymWrapper @@ -4834,6 +4839,193 @@ def test_async_transformed(self, backend, make_envs): base_env._maybe_shutdown() +def _has_mps(): + if hasattr(torch, "mps") and hasattr(torch.mps, "is_available"): + return torch.mps.is_available() + return ( + getattr(torch.backends, "mps", None) is not None + and torch.backends.mps.is_available() + ) + + +@pytest.mark.skipif(not _has_mps(), reason="MPS device not available") +class TestMPSDeviceCasting: + def test_mps_does_not_support_float64(self): + """Assert that MPS still doesn't support float64. + + If this test fails, MPS has gained float64 support and the downcasts + can be removed. + """ + with pytest.raises(TypeError, match="MPS framework doesn't support float64"): + torch.ones(2, dtype=torch.float64, device="mps") + + def test_to_device_mps_safe_float64_downcast(self): + t = torch.randn(4, dtype=torch.float64, device="cpu") + result = _to_device_mps_safe(t, torch.device("mps")) + assert result.device.type == "mps" + assert result.dtype == torch.float32 + + def test_td_to_device_mps_safe_downcasts_float64(self): + td = TensorDict( + { + "obs": torch.randn(3, dtype=torch.float64), + "flag": torch.ones(1, dtype=torch.bool), + }, + batch_size=[], + ) + result = _td_to_device_mps_safe(td, torch.device("mps")) + assert result["obs"].device.type == "mps" + assert result["obs"].dtype == torch.float32 + assert result["flag"].device.type == "mps" + assert result["flag"].dtype == torch.bool + + def test_has_float64_leaf(self): + from torchrl.data import Unbounded + + spec_f64 = Composite( + obs=Unbounded(shape=(3,), dtype=torch.float64, device="cpu") + ) + assert _has_float64_leaf(spec_f64) is True + + spec_mixed_dtype_with_f64 = Composite( + obs32=Unbounded(shape=(3,), dtype=torch.float32, device="cpu"), + obs64=Unbounded(shape=(3,), dtype=torch.float64, device="cpu"), + ) + assert _has_float64_leaf(spec_mixed_dtype_with_f64) is True + + spec_f32 = Composite( + obs=Unbounded(shape=(3,), dtype=torch.float32, device="cpu") + ) + assert _has_float64_leaf(spec_f32) is False + + inner = Composite(obs=Unbounded(shape=(3,), dtype=torch.float64, device="cpu")) + outer = Composite(next=inner) + assert _has_float64_leaf(outer) is True + + assert _has_float64_leaf(Composite()) is False + + class _Float64ObsEnv(EnvBase): + """Minimal env that produces float64 observations on CPU.""" + + def __init__(self): + super().__init__(device="cpu") + self.observation_spec = Composite( + observation=Unbounded(shape=(4,), dtype=torch.float64), + ) + self.action_spec = Unbounded(shape=(2,)) + self.reward_spec = Unbounded(shape=(1,)) + + def _reset(self, tensordict): + return TensorDict( + { + "observation": torch.randn(4, dtype=torch.float64), + "done": torch.zeros(1, dtype=torch.bool), + "terminated": torch.zeros(1, dtype=torch.bool), + }, + batch_size=[], + ) + + def _step(self, tensordict): + return TensorDict( + { + "observation": torch.randn(4, dtype=torch.float64), + "reward": torch.zeros(1, dtype=torch.float64), + "done": torch.zeros(1, dtype=torch.bool), + "terminated": torch.zeros(1, dtype=torch.bool), + }, + batch_size=[], + ) + + def _set_seed(self, seed): + return seed + + _MPS_FLOAT64_WARNING = r"Sub-environments produce float64 data but the batched env device is 'mps.*' which does not support float64\. All float64 specs and tensors will be downcast to float32\." + + def test_serial_env_mps_parent_cpu_worker_reset(self): + with pytest.warns(UserWarning, match=self._MPS_FLOAT64_WARNING): + env = SerialEnv(2, self._Float64ObsEnv, device="mps") + try: + td = env.reset() + assert td.device.type == "mps" + assert td["observation"].dtype == torch.float32 + finally: + env.close(raise_if_closed=False) + + def test_serial_env_mps_parent_cpu_worker_rollout(self): + with pytest.warns(UserWarning, match=self._MPS_FLOAT64_WARNING): + env = SerialEnv(2, self._Float64ObsEnv, device="mps") + try: + policy = RandomPolicy(env.action_spec) + td = env.rollout(max_steps=3, policy=policy) + assert td.device.type == "mps" + assert td["observation"].dtype == torch.float32 + finally: + env.close(raise_if_closed=False) + + def test_parallel_env_mps_parent_cpu_worker_reset(self): + with pytest.warns(UserWarning, match=self._MPS_FLOAT64_WARNING): + env = ParallelEnv(2, self._Float64ObsEnv, device="mps") + try: + td = env.reset() + assert td.device.type == "mps" + assert td["observation"].dtype == torch.float32 + finally: + env.close(raise_if_closed=False) + + def test_parallel_env_mps_parent_cpu_worker_rollout(self): + with pytest.warns(UserWarning, match=self._MPS_FLOAT64_WARNING): + env = ParallelEnv(2, self._Float64ObsEnv, device="mps") + try: + policy = RandomPolicy(env.action_spec) + td = env.rollout(max_steps=3, policy=policy) + assert td.device.type == "mps" + assert td["observation"].dtype == torch.float32 + finally: + env.close(raise_if_closed=False) + + def test_serial_env_no_buffers_mps_reset(self): + with pytest.warns(UserWarning, match=self._MPS_FLOAT64_WARNING): + env = SerialEnv(2, self._Float64ObsEnv, device="mps", use_buffers=False) + try: + td = env.reset() + assert td.device.type == "mps" + assert td["observation"].dtype == torch.float32 + finally: + env.close(raise_if_closed=False) + + def test_serial_env_no_buffers_mps_rollout(self): + with pytest.warns(UserWarning, match=self._MPS_FLOAT64_WARNING): + env = SerialEnv(2, self._Float64ObsEnv, device="mps", use_buffers=False) + try: + policy = RandomPolicy(env.action_spec) + td = env.rollout(max_steps=3, policy=policy) + assert td.device.type == "mps" + assert td["observation"].dtype == torch.float32 + finally: + env.close(raise_if_closed=False) + + def test_parallel_env_no_buffers_mps_reset(self): + with pytest.warns(UserWarning, match=self._MPS_FLOAT64_WARNING): + env = ParallelEnv(2, self._Float64ObsEnv, device="mps", use_buffers=False) + try: + td = env.reset() + assert td.device.type == "mps" + assert td["observation"].dtype == torch.float32 + finally: + env.close(raise_if_closed=False) + + def test_parallel_env_no_buffers_mps_rollout(self): + with pytest.warns(UserWarning, match=self._MPS_FLOAT64_WARNING): + env = ParallelEnv(2, self._Float64ObsEnv, device="mps", use_buffers=False) + try: + policy = RandomPolicy(env.action_spec) + td = env.rollout(max_steps=3, policy=policy) + assert td.device.type == "mps" + assert td["observation"].dtype == torch.float32 + finally: + env.close(raise_if_closed=False) + + if __name__ == "__main__": args, unknown = argparse.ArgumentParser().parse_known_args() pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown) diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py index 78bdc8aea40..72851f7005f 100644 --- a/torchrl/envs/batched_envs.py +++ b/torchrl/envs/batched_envs.py @@ -61,6 +61,54 @@ ) +def _to_device_mps_safe( + tensor: torch.Tensor, + device: torch.device, + *, + non_blocking: bool = False, +) -> torch.Tensor: + """Move a tensor to the target device, downcasting float64 to float32 for MPS. + + MPS does not support float64. When the target device is MPS and the source + tensor is float64, this automatically downcasts to float32 during the transfer. + For all other devices the call is equivalent to ``tensor.to(device, ...)``. + """ + if device.type == "mps" and tensor.dtype == torch.float64: + return tensor.to(device=device, dtype=torch.float32, non_blocking=non_blocking) + return tensor.to(device=device, non_blocking=non_blocking) + + +def _td_to_device_mps_safe( + td: TensorDictBase, + device: torch.device, + *, + non_blocking: bool = False, +) -> TensorDictBase: + """Move a TensorDict to the target device, handling MPS float64 via :func:`_to_device_mps_safe`. + + For non-MPS devices this falls back to ``td.to(device, ...)``. + For MPS it applies :func:`_to_device_mps_safe` per-tensor to avoid the + float64-on-MPS limitation. + """ + if device.type == "mps": + return td._fast_apply( + functools.partial( + _to_device_mps_safe, device=device, non_blocking=non_blocking + ), + device=device, + filter_empty=True, + ) + return td.to(device, non_blocking=non_blocking) + + +def _has_float64_leaf(spec) -> bool: + """Return True if any non-Composite leaf spec in *spec* has dtype float64.""" + for _, s in spec.items(include_nested=True, leaves_only=True): + if getattr(s, "dtype", None) == torch.float64: + return True + return False + + def _check_start(fun): def decorated_fun(self: BatchedEnvBase, *args, **kwargs): if self.is_closed: @@ -551,7 +599,9 @@ def select_and_clone(self, name, tensor, selected_keys=None): selected_keys = self._selected_step_keys if name in selected_keys: if self.device is not None and tensor.device != self.device: - return tensor.to(self.device, non_blocking=self.non_blocking) + return _to_device_mps_safe( + tensor, self.device, non_blocking=self.non_blocking + ) return tensor.clone() @property @@ -775,11 +825,29 @@ def _check_for_empty_spec(specs: Composite): meta_data.specs["output_spec"].to(device) ) - self.action_spec = input_spec["full_action_spec"] - self.state_spec = input_spec["full_state_spec"] - self.observation_spec = output_spec["full_observation_spec"] - self.reward_spec = output_spec["full_reward_spec"] - self.done_spec = output_spec["full_done_spec"] + if ( + self._device is not None + and torch.device(self._device).type == "mps" + and (_has_float64_leaf(input_spec) or _has_float64_leaf(output_spec)) + ): + warnings.warn( + f"Sub-environments produce float64 data but the batched env device is " + f"'{self._device}' which does not support float64. " + f"All float64 specs and tensors will be downcast to float32.", + UserWarning, + stacklevel=2, + ) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="MPS device does not support float64", + category=UserWarning, + ) + self.action_spec = input_spec["full_action_spec"] + self.state_spec = input_spec["full_state_spec"] + self.observation_spec = output_spec["full_observation_spec"] + self.reward_spec = output_spec["full_reward_spec"] + self.done_spec = output_spec["full_done_spec"] self._dummy_env_str = meta_data.env_str self._env_tensordict = meta_data.tensordict @@ -819,12 +887,34 @@ def map_device(key, value, device_map=device_map): output_spec.append(_check_for_empty_spec(md.specs["output_spec"])) output_spec = torch.stack(output_spec, 0) - self.action_spec = input_spec["full_action_spec"] - self.state_spec = input_spec["full_state_spec"] + if ( + self._device is not None + and torch.device(self._device).type == "mps" + and any( + _has_float64_leaf(md.specs["input_spec"]) + or _has_float64_leaf(md.specs["output_spec"]) + for md in meta_data + ) + ): + warnings.warn( + f"Sub-environments produce float64 data but the batched env device is " + f"'{self._device}' which does not support float64. " + f"All float64 specs and tensors will be downcast to float32.", + UserWarning, + stacklevel=2, + ) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="MPS device does not support float64", + category=UserWarning, + ) + self.action_spec = input_spec["full_action_spec"] + self.state_spec = input_spec["full_state_spec"] - self.observation_spec = output_spec["full_observation_spec"] - self.reward_spec = output_spec["full_reward_spec"] - self.done_spec = output_spec["full_done_spec"] + self.observation_spec = output_spec["full_observation_spec"] + self.reward_spec = output_spec["full_reward_spec"] + self.done_spec = output_spec["full_done_spec"] self._dummy_env_str = str(meta_data[0]) if self.share_individual_td: @@ -1193,7 +1283,7 @@ def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase: if self.device is None: ftd.clear_device_() else: - ftd = ftd.to(self.device) + ftd = _td_to_device_mps_safe(ftd, self.device) out_tds[i] = ftd continue if tensordict is not None: @@ -1244,7 +1334,9 @@ def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase: if device is None: result = result.clear_device_() else: - result = result.to(device, non_blocking=self.non_blocking) + result = _td_to_device_mps_safe( + result, device, non_blocking=self.non_blocking + ) self._sync_w2m() return result @@ -1309,6 +1401,8 @@ def _step( if not self._use_buffers or self._non_tensor_keys: out_tds = [] + device = self.device + if self._use_buffers: next_td = self.shared_tensordict_parent.get("next") for i, _data_in in zip(workers_range, data_in): @@ -1325,7 +1419,6 @@ def _step( # We must pass a clone of the tensordict, as the values of this tensordict # will be modified in-place at further steps - device = self.device selected_keys = self._selected_step_keys @@ -1353,6 +1446,14 @@ def _step( out_td = self._envs[i]._step(_data_in) out_tds.append(out_td) out = LazyStackedTensorDict.maybe_dense_stack(out_tds) + if out.device != device: + if device is None: + out = out.clear_device_() + else: + out = _td_to_device_mps_safe( + out, device, non_blocking=self.non_blocking + ) + self._sync_w2m() if partial_steps is not None and not partial_steps.all(): result = out.new_zeros(tensordict_save.shape) @@ -2115,7 +2216,9 @@ def _step_no_buffers( out = LazyStackedTensorDict.maybe_dense_stack(out_tds) if self.device is not None and out.device != self.device: - out = out.to(self.device, non_blocking=self.non_blocking) + out = _td_to_device_mps_safe( + out, self.device, non_blocking=self.non_blocking + ) if partial_steps is not None: result = out.new_zeros(tensordict_save.shape) @@ -2349,7 +2452,9 @@ def _reset_no_buffers( result = LazyStackedTensorDict.maybe_dense_stack(out_tds) device = self.device if device is not None and result.device != device: - return result.to(self.device, non_blocking=self.non_blocking) + return _td_to_device_mps_safe( + result, self.device, non_blocking=self.non_blocking + ) return result @torch.no_grad()