From 69d3d3d6d597b8284eb371b4cbf710a8dac0d8eb Mon Sep 17 00:00:00 2001
From: Jacob Oursland <jacob.oursland@gmail.com>
Date: Sat, 13 Dec 2025 00:49:33 -0800
Subject: [PATCH 1/2] Add SinglePrecision env for backends that do not support
 64-bit floats.

This env taken from BiggerRegularizedOptimistic by naumix:
https://github.com/naumix/BiggerRegularizedOptimistic/blob/main/jaxrl/envs/single_precision.py
---
 .../common/envs/single_precision.py           | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)
 create mode 100644 stable_baselines3/common/envs/single_precision.py

diff --git a/stable_baselines3/common/envs/single_precision.py b/stable_baselines3/common/envs/single_precision.py
new file mode 100644
index 0000000000..d9dbc7fb42
--- /dev/null
+++ b/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

From d644e4b3eb2616b5b4d40471c468f5f811c4944a Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@ensta.org>
Date: Mon, 4 Jul 2022 14:51:46 +0200
Subject: [PATCH 2/2] Use MPS device when available

---
 docs/misc/changelog.rst           |  2 ++
 stable_baselines3/common/utils.py | 29 ++++++++++++++++++++++-------
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
index 9303d6d36d..10a1c5a156 100644
--- a/docs/misc/changelog.rst
+++ b/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:
 ^^^^^^^^^^
diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py
index 846a954359..cf41cff823 100644
--- a/stable_baselines3/common/utils.py
+++ b/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__,