Add MPS Backend

docs/misc/changelog.rst

@@ -15,6 +15,8 @@ Breaking Changes:
 New Features:
 ^^^^^^^^^^^^^
 - Added official support for Python 3.13
+- Use MacOS Metal "mps" device when available
+- Save cloudpickle version
 
 Bug Fixes:
 ^^^^^^^^^^

stable_baselines3/common/envs/single_precision.py

@@ -0,0 +1,30 @@
+import copy
+
+import gym
+import numpy as np
+from gym.spaces import Box, Dict
+
+class SinglePrecision(gym.ObservationWrapper):
+    def __init__(self, env):
+        super().__init__(env)
+
+        if isinstance(self.observation_space, Box):
+            obs_space = self.observation_space
+            self.observation_space = Box(obs_space.low, obs_space.high,
+                                         obs_space.shape)
+        elif isinstance(self.observation_space, Dict):
+            obs_spaces = copy.copy(self.observation_space.spaces)
+            for k, v in obs_spaces.items():
+                obs_spaces[k] = Box(v.low, v.high, v.shape)
+            self.observation_space = Dict(obs_spaces)
+        else:
+            raise NotImplementedError
+
+    def observation(self, observation: np.ndarray) -> np.ndarray:
+        if isinstance(observation, np.ndarray):
+            return observation.astype(np.float32)
+        elif isinstance(observation, dict):
+            observation = copy.copy(observation)
+            for k, v in observation.items():
+                observation[k] = v.astype(np.float32)
+            return observation

stable_baselines3/common/utils.py

@@ -224,19 +224,20 @@ def get_device(device: th.device | str = "auto") -> th.device:
     """
     Retrieve PyTorch device.
     It checks that the requested device is available first.
-    For now, it supports only cpu and cuda.
-    By default, it tries to use the gpu.
+    For now, it supports only CPU and CUDA.
+    By default, it tries to use the GPU.
 
-    :param device: One for 'auto', 'cuda', 'cpu'
+    :param device: One of "auto", "cuda", "cpu",
+        or any PyTorch supported device (for instance "mps")
     :return: Supported Pytorch device
     """
-    # Cuda by default
+    # MPS/CUDA by default
     if device == "auto":
-        device = "cuda"
+        device = get_available_accelerator()
     # Force conversion to th.device
     device = th.device(device)
 
-    # Cuda not available
+    # CUDA not available
     if device.type == th.device("cuda").type and not th.cuda.is_available():
         return th.device("cpu")
 
@@ -597,6 +598,20 @@ def should_collect_more_steps(
         )
 
 
+def get_available_accelerator() -> str:
+    """
+    Return the available accelerator
+    (currently checking only for CUDA and MPS device)
+    """
+    if hasattr(th, "has_mps") and th.backends.mps.is_available():
+        # MacOS Metal GPU
+        return "mps"
+    elif th.cuda.is_available():
+        return "cuda"
+    else:
+        return "cpu"
+
+
 def get_system_info(print_info: bool = True) -> tuple[dict[str, str], str]:
     """
     Retrieve system and python env info for the current system.
@@ -612,7 +627,7 @@ def get_system_info(print_info: bool = True) -> tuple[dict[str, str], str]:
         "Python": platform.python_version(),
         "Stable-Baselines3": sb3.__version__,
         "PyTorch": th.__version__,
-        "GPU Enabled": str(th.cuda.is_available()),
+        "Accelerator": get_available_accelerator(),
         "Numpy": np.__version__,
         "Cloudpickle": cloudpickle.__version__,
         "Gymnasium": gym.__version__,