Removed BufferDTypes

Author: Trenza1ore

docs/misc/changelog.rst

@@ -13,7 +13,6 @@ New Features:
 
 Others:
 ^^^^^^^
-- Added an additional ``BufferDTypes`` dataclass to ``stable_baselines3/common/buffers.py`` for representing buffer datatypes (@Trenza1ore)
 - Added an additional ``test_buffers.py::test_buffer_dtypes`` which tests the `dtype` of `RolloutBuffer` and `DictRolloutBuffer` (@Trenza1ore)
 
 Release 2.7.0 (2025-07-25)

stable_baselines3/common/buffers.py

@@ -1,9 +1,7 @@
 import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Generator, Mapping
-from dataclasses import InitVar, dataclass, field
-from types import MappingProxyType
-from typing import Any, ClassVar, Optional, Union
+from collections.abc import Generator
+from typing import Any, Optional, Union
 
 import numpy as np
 import torch as th
@@ -13,7 +11,6 @@
 from stable_baselines3.common.type_aliases import (
     DictReplayBufferSamples,
     DictRolloutBufferSamples,
-    DTypeLike,
     ReplayBufferSamples,
     RolloutBufferSamples,
 )
@@ -27,55 +24,6 @@
     psutil = None
 
 
-@dataclass
-class BufferDTypes:
-    """
-    Data class representing the data types used by a buffer.
-
-    :param observations: Datatype of observation space
-    :param actions: Datatype of action space
-    """
-
-    MAP_TORCH_DTYPES: ClassVar[dict] = dict(complex32="complex64", float="float32", bfloat16="float32", bool="bool_")
-
-    observations: InitVar[Union[DTypeLike, Mapping[str, DTypeLike]]]
-    actions: InitVar[DTypeLike]
-
-    dict_obs: MappingProxyType[str, np.dtype] = field(default_factory=lambda: MappingProxyType({}), init=False)
-    obs: Optional[np.dtype] = field(default=None, init=False)
-    act: Optional[np.dtype] = field(default=None, init=False)
-
-    def __post_init__(self, observations: Union[DTypeLike, Mapping[str, DTypeLike]], actions: DTypeLike):
-        if isinstance(observations, Mapping):
-            self.dict_obs = MappingProxyType({k: self.to_numpy_dtype(v) for k, v in observations.items()})
-        else:
-            self.obs = self.to_numpy_dtype(observations)
-        self.act = self.to_numpy_dtype(actions)
-
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        if isinstance(self.dict_obs, MappingProxyType):
-            state["dict_obs"] = dict(self.dict_obs)
-        return state
-
-    def __setstate__(self, state: Mapping[str, Any]):
-        state = dict(state)
-        if state.get("dict_obs"):
-            state["dict_obs"] = MappingProxyType(state["dict_obs"].copy())
-        self.__dict__.update(state)
-
-    @classmethod
-    def to_numpy_dtype(cls, dtype_like: DTypeLike) -> np.dtype:
-        if isinstance(dtype_like, th.dtype):
-            torch_dtype_name = repr(dtype_like).removeprefix("torch.")
-            numpy_dtype_name = cls.MAP_TORCH_DTYPES.get(torch_dtype_name, torch_dtype_name)
-            try:
-                return np.dtype(getattr(np, numpy_dtype_name))
-            except AttributeError as e:
-                raise TypeError(f"Cannot cast torch dtype '{torch_dtype_name}' to numpy.dtype implicitly.") from e
-        return np.dtype(dtype_like)
-
-
 class BaseBuffer(ABC):
     """
     Base class that represent a buffer (rollout or replay)
@@ -111,16 +59,6 @@ def __init__(
         self.device = get_device(device)
         self.n_envs = n_envs
 
-        # unify the dtype decision logic for all buffer classes
-        # see https://github.com/DLR-RM/stable-baselines3/issues/2162
-        if isinstance(observation_space, spaces.Dict):
-            self.dtypes = BufferDTypes(
-                observations={key: space.dtype for (key, space) in observation_space.spaces.items()},
-                actions=action_space.dtype,
-            )
-        else:
-            self.dtypes = BufferDTypes(observations=observation_space.dtype, actions=action_space.dtype)
-
     @staticmethod
     def swap_and_flatten(arr: np.ndarray) -> np.ndarray:
         """
@@ -451,8 +389,8 @@ def __init__(
         self.reset()
 
     def reset(self) -> None:
-        self.observations = np.zeros((self.buffer_size, self.n_envs, *self.obs_shape), dtype=self.dtypes.obs)
-        self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=self.dtypes.act)
+        self.observations = np.zeros((self.buffer_size, self.n_envs, *self.obs_shape), dtype=self.observation_space.dtype)
+        self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=self.action_space.dtype)
         self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
         self.returns = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
         self.episode_starts = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
@@ -808,9 +746,9 @@ def reset(self) -> None:
         self.observations = {}
         for key, obs_input_shape in self.obs_shape.items():
             self.observations[key] = np.zeros(
-                (self.buffer_size, self.n_envs, *obs_input_shape), dtype=self.dtypes.dict_obs[key]
+                (self.buffer_size, self.n_envs, *obs_input_shape), dtype=self.observation_space.dtype
             )
-        self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=self.dtypes.act)
+        self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=self.action_space.dtype)
         self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
         self.returns = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
         self.episode_starts = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)

stable_baselines3/common/off_policy_algorithm.py

@@ -11,7 +11,7 @@
 from gymnasium import spaces
 
 from stable_baselines3.common.base_class import BaseAlgorithm
-from stable_baselines3.common.buffers import BufferDTypes, DictReplayBuffer, NStepReplayBuffer, ReplayBuffer
+from stable_baselines3.common.buffers import DictReplayBuffer, NStepReplayBuffer, ReplayBuffer
 from stable_baselines3.common.callbacks import BaseCallback
 from stable_baselines3.common.noise import ActionNoise, VectorizedActionNoise
 from stable_baselines3.common.policies import BasePolicy
@@ -247,19 +247,6 @@ def load_replay_buffer(
             self.replay_buffer.handle_timeout_termination = False
             self.replay_buffer.timeouts = np.zeros_like(self.replay_buffer.dones)
 
-        # Backward compatibility with SB3 < 2.7.1 replay buffer
-        if not hasattr(self.replay_buffer, "dtypes"):
-            if isinstance(self.replay_buffer, DictReplayBuffer):
-                self.replay_buffer.dtypes = BufferDTypes(
-                    observations={key: obs.dtype for (key, obs) in self.replay_buffer.observations.items()},
-                    actions=self.replay_buffer.actions.dtype,
-                )
-            else:
-                self.replay_buffer.dtypes = BufferDTypes(
-                    observations=self.replay_buffer.observations.dtype,
-                    actions=self.replay_buffer.actions.dtype,
-                )
-
         if isinstance(self.replay_buffer, HerReplayBuffer):
             assert self.env is not None, "You must pass an environment at load time when using `HerReplayBuffer`"
             self.replay_buffer.set_env(self.env)

stable_baselines3/common/type_aliases.py

@@ -22,7 +22,6 @@
 OptimizerStateDict = dict[str, Any]
 MaybeCallback = Union[None, Callable, list["BaseCallback"], "BaseCallback"]
 PyTorchObs = Union[th.Tensor, TensorDict]
-DTypeLike = Union[None, np.dtype, th.dtype, type, str]
 
 # A schedule takes the remaining progress as input
 # and outputs a scalar (e.g. learning rate, clip range, ...)

tests/test_buffers.py

@@ -184,20 +184,26 @@ def test_device_buffer(replay_buffer_cls, device):
 @pytest.mark.parametrize("use_dict", [False, True])
 def test_buffer_dtypes(obs_dtype: Union[type[np.integer], type[np.floating]], use_dict: bool):
     rollout_buffer: Union[RolloutBuffer, DictRolloutBuffer]
+    replay_buffer: Union[ReplayBuffer, DictReplayBuffer]
     obs_space = spaces.Box(0, 100, dtype=obs_dtype)
-    action_space = spaces.Discrete(10)
+    act_space = spaces.Discrete(10)
+    buffer_params = dict(buffer_size=1, action_space=act_space)
 
     if use_dict:
-        obs_space_2 = spaces.Box(0, 100, dtype=np.uint8)
-        observation_space = spaces.Dict({"obs": obs_space, "obs_2": obs_space_2})
-        rollout_buffer = DictRolloutBuffer(buffer_size=1, observation_space=observation_space, action_space=action_space)
-        assert rollout_buffer.observations["obs"].dtype == obs_dtype
-        assert rollout_buffer.observations["obs_2"].dtype == np.uint8
+        dict_obs_space = spaces.Dict({"obs": obs_space, "obs_2": spaces.Box(0, 100, dtype=np.uint8)})
+        buffer_params["observation_space"] = dict_obs_space
+        rollout_buffer = DictRolloutBuffer(**buffer_params)  # type: ignore[arg-type]
+        replay_buffer = DictReplayBuffer(**buffer_params)  # type: ignore[arg-type]
+        assert rollout_buffer.observations["obs"].dtype == replay_buffer.observations["obs"].dtype == obs_dtype
+        assert rollout_buffer.observations["obs_2"].dtype == replay_buffer.observations["obs_2"].dtype == np.uint8
     else:
-        rollout_buffer = RolloutBuffer(buffer_size=1, observation_space=obs_space, action_space=action_space)
-        assert rollout_buffer.observations.dtype == obs_dtype
+        buffer_params["observation_space"] = obs_space
+        rollout_buffer = RolloutBuffer(**buffer_params)  # type: ignore[arg-type]
+        replay_buffer = ReplayBuffer(**buffer_params)  # type: ignore[arg-type]
+        assert rollout_buffer.observations.dtype == replay_buffer.observations.dtype == obs_dtype
 
-    assert rollout_buffer.actions.dtype == np.int64
+    assert rollout_buffer.actions.dtype == np.float32, "RolloutBuffer action dtype must be np.float32"
+    assert replay_buffer.actions.dtype == act_space.dtype
 
 
 def test_custom_rollout_buffer():