Use proper dtype for RolloutBuffer storage (#2163)

* Initial implementation of dtype-decision logic

* Fixed init logic

* Updated changelog

* Added a test

* Reformatted using make format

* Ensure make type passes

* Fixed DictRolloutBuffer dtype assignment

* Updated to create a BufferDTypes dataclass and updated pytests

* Fix type check errors on Github, separate dict_obs and obs, honor _normalize_obs for rollout buffers

* Revert _normalize_obs calls in rollout buffers

* Updated docs

* Updated docs

* Added save / load support with backward compatibility

* Cast sampled actions of rollout buffers to float32 to avoid breaking changes

* Fixed pickle loading of BufferDTypes

* Use default_factory instead of default for BufferDTypes.dict_obs

* Simplified BufferDTypes and reverted changes on replay buffers as requested

* Removed BufferDTypes

* Fixed oversight in dictrolloutbuffer dtype

* Update changelog and version

* Remove cast to float32

* Update tests

* Remove cast to long

* Revert "Remove cast to long"

This reverts commit 216d757.

* Revert "Remove cast to float32"

This reverts commit d1e5221.

* Reapply "Remove cast to float32"

This reverts commit 3511452.

* Reapply "Remove cast to long"

This reverts commit c2d532c.

* Cast int8 to float32 to avoid PyTorch issues (MultiBinary)

* Revert "Reapply "Remove cast to long""

This reverts commit 88e6b68.

* Cast at sample time only

* Update changelog.rst

---------

Co-authored-by: Antonin RAFFIN <[email protected]>

Author: Trenza1ore

docs/misc/changelog.rst

@@ -3,6 +3,38 @@
 Changelog
 ==========
 
+Release 2.7.1a1 (WIP)
+--------------------------
+
+Breaking Changes:
+^^^^^^^^^^^^^^^^^
+
+New Features:
+^^^^^^^^^^^^^
+- ``RolloutBuffer`` and ``DictRolloutBuffer`` now uses the actual observation / action space ``dtype`` (instead of float32), this should save memory (@Trenza1ore)
+
+Bug Fixes:
+^^^^^^^^^^
+
+`SB3-Contrib`_
+^^^^^^^^^^^^^^
+
+`RL Zoo`_
+^^^^^^^^^
+
+`SBX`_ (SB3 + Jax)
+^^^^^^^^^^^^^^^^^^
+
+Deprecations:
+^^^^^^^^^^^^^
+
+Others:
+^^^^^^^
+
+Documentation:
+^^^^^^^^^^^^^^
+
+
 Release 2.7.0 (2025-07-25)
 --------------------------
 
@@ -1857,4 +1889,4 @@ And all the contributors:
 @DavyMorgan @luizapozzobon @Bonifatius94 @theSquaredError @harveybellini @DavyMorgan @FieteO @jonasreiher @npit @WeberSamuel @troiganto
 @lutogniew @lbergmann1 @lukashass @BertrandDecoster @pseudo-rnd-thoughts @stefanbschneider @kyle-he @PatrickHelm @corentinlger
 @marekm4 @stagoverflow @rushitnshah @markscsmith @NickLucche @cschindlbeck @peteole @jak3122 @will-maclean
-@brn-dev @jmacglashan @kplers @MarcDcls @chrisgao99 @pstahlhofen @akanto
+@brn-dev @jmacglashan @kplers @MarcDcls @chrisgao99 @pstahlhofen @akanto @Trenza1ore

stable_baselines3/common/buffers.py

@@ -389,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=np.float32)
-        self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=np.float32)
+        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)
@@ -512,7 +512,8 @@ def _get_samples(
     ) -> RolloutBufferSamples:
         data = (
             self.observations[batch_inds],
-            self.actions[batch_inds],
+            # Cast to float32 (backward compatible), this would lead to RuntimeError for MultiBinary space
+            self.actions[batch_inds].astype(np.float32, copy=False),
             self.values[batch_inds].flatten(),
             self.log_probs[batch_inds].flatten(),
             self.advantages[batch_inds].flatten(),
@@ -745,8 +746,10 @@ def __init__(
     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=np.float32)
-        self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=np.float32)
+            self.observations[key] = np.zeros(
+                (self.buffer_size, self.n_envs, *obs_input_shape), dtype=self.observation_space[key].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)
@@ -832,7 +835,8 @@ def _get_samples(  # type: ignore[override]
     ) -> DictRolloutBufferSamples:
         return DictRolloutBufferSamples(
             observations={key: self.to_torch(obs[batch_inds]) for (key, obs) in self.observations.items()},
-            actions=self.to_torch(self.actions[batch_inds]),
+            # Cast to float32 (backward compatible), this would lead to RuntimeError for MultiBinary space
+            actions=self.to_torch(self.actions[batch_inds].astype(np.float32, copy=False)),
             old_values=self.to_torch(self.values[batch_inds].flatten()),
             old_log_prob=self.to_torch(self.log_probs[batch_inds].flatten()),
             advantages=self.to_torch(self.advantages[batch_inds].flatten()),

stable_baselines3/version.txt

@@ -1 +1 @@
-2.7.0
+2.7.1a0

tests/test_buffers.py

@@ -163,6 +163,75 @@ def test_device_buffer(replay_buffer_cls, device):
                 raise TypeError(f"Unknown value type: {type(value)}")
 
 
+@pytest.mark.parametrize(
+    "obs_dtype",
+    [
+        np.dtype(np.uint8),
+        np.dtype(np.int8),
+        np.dtype(np.uint16),
+        np.dtype(np.int16),
+        np.dtype(np.uint32),
+        np.dtype(np.int32),
+        np.dtype(np.uint64),
+        np.dtype(np.int64),
+        np.dtype(np.float16),
+        np.dtype(np.float32),
+        np.dtype(np.float64),
+    ],
+)
+@pytest.mark.parametrize("use_dict", [False, True])
+@pytest.mark.parametrize(
+    "action_space",
+    [
+        spaces.Discrete(10),
+        spaces.Box(low=-1.0, high=1.0, dtype=np.float32),
+        spaces.Box(low=-1.0, high=1.0, dtype=np.float64),
+    ],
+)
+def test_buffer_dtypes(obs_dtype, use_dict, action_space):
+    obs_space = spaces.Box(0, 100, dtype=obs_dtype)
+    buffer_params = dict(buffer_size=1, action_space=action_space)
+    # For off-policy algorithms, we cast float64 actions to float32, see GH#1145
+    actual_replay_action_dtype = ReplayBuffer._maybe_cast_dtype(action_space.dtype)
+    # For on-policy, we cast at sample time to float32 for backward compat
+    # and to avoid issue computing log prob with multibinary
+    actual_rollout_action_dtype = np.float32
+
+    if use_dict:
+        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)
+        replay_buffer = DictReplayBuffer(**buffer_params)
+        assert rollout_buffer.observations["obs"].dtype == obs_dtype
+        assert replay_buffer.observations["obs"].dtype == obs_dtype
+        assert rollout_buffer.observations["obs_2"].dtype == np.uint8
+        assert replay_buffer.observations["obs_2"].dtype == np.uint8
+    else:
+        buffer_params["observation_space"] = obs_space
+        rollout_buffer = RolloutBuffer(**buffer_params)
+        replay_buffer = ReplayBuffer(**buffer_params)
+        assert rollout_buffer.observations.dtype == obs_dtype
+        assert replay_buffer.observations.dtype == obs_dtype
+
+    assert rollout_buffer.actions.dtype == action_space.dtype
+    assert replay_buffer.actions.dtype == actual_replay_action_dtype
+    # Check that sampled types are corrects
+    rollout_buffer.full = True
+    replay_buffer.full = True
+    rollout_data = next(rollout_buffer.get(batch_size=64))
+    buffer_data = replay_buffer.sample(batch_size=64)
+    assert rollout_data.actions.numpy().dtype == actual_rollout_action_dtype
+    assert buffer_data.actions.numpy().dtype == actual_replay_action_dtype
+    if use_dict:
+        assert buffer_data.observations["obs"].numpy().dtype == obs_dtype
+        assert buffer_data.observations["obs_2"].numpy().dtype == np.uint8
+        assert rollout_data.observations["obs"].numpy().dtype == obs_dtype
+        assert rollout_data.observations["obs_2"].numpy().dtype == np.uint8
+    else:
+        assert buffer_data.observations.numpy().dtype == obs_dtype
+        assert rollout_data.observations.numpy().dtype == obs_dtype
+
+
 def test_custom_rollout_buffer():
     A2C("MlpPolicy", "Pendulum-v1", rollout_buffer_class=RolloutBuffer, rollout_buffer_kwargs=dict())