Construct tensors directly on GPU (#1218)

* Replace .to(device) when possible

* fix numpy dep

* black

* Add warning for device != cpu and copy=False

* Update changelog

* Remove warning

* Update buffers.py

docs/misc/changelog.rst

@@ -49,6 +49,7 @@ Others:
 - Fixed ``stable_baselines3/common/atari_wrappers.py`` type hints
 - Exposed modules in ``__init__.py`` with the ``__all__`` attribute (@ZikangXiong)
 - Upgraded GitHub CI/setup-python to v4 and checkout to v3
+- Set tensors construction directly on the device
 
 Documentation:
 ^^^^^^^^^^^^^^

setup.py

@@ -77,7 +77,7 @@ setup(
     package_data={"stable_baselines3": ["py.typed", "version.txt"]},
     install_requires=[
         "gym==0.21",  # Fixed version due to breaking changes in 0.22
-        "numpy",
+        "numpy<1.24",  # Required for gym==0.21
         "torch>=1.11",
         'typing_extensions>=4.0,<5; python_version < "3.8.0"',
         # For saving models

stable_baselines3/common/buffers.py

@@ -127,13 +127,13 @@ class BaseBuffer(ABC):
         Note: it copies the data by default
 
         :param array:
-        :param copy: Whether to copy or not the data
-            (may be useful to avoid changing things be reference)
+        :param copy: Whether to copy or not the data (may be useful to avoid changing things
+            by reference). This argument is inoperative if the device is not the CPU.
         :return:
         """
         if copy:
-            return th.tensor(array).to(self.device)
-        return th.as_tensor(array).to(self.device)
+            return th.tensor(array, device=self.device)
+        return th.as_tensor(array, device=self.device)
 
     @staticmethod
     def _normalize_obs(

stable_baselines3/common/policies.py

@@ -183,7 +183,7 @@ class BaseModel(nn.Module):
 
         :param vector:
         """
-        th.nn.utils.vector_to_parameters(th.FloatTensor(vector).to(self.device), self.parameters())
+        th.nn.utils.vector_to_parameters(th.FloatTensor(vector, device=self.device), self.parameters())
 
     def parameters_to_vector(self) -> np.ndarray:
         """

stable_baselines3/common/utils.py

@@ -461,9 +461,9 @@ def obs_as_tensor(
     :return: PyTorch tensor of the observation on a desired device.
     """
     if isinstance(obs, np.ndarray):
-        return th.as_tensor(obs).to(device)
+        return th.as_tensor(obs, device=device)
     elif isinstance(obs, dict):
-        return {key: th.as_tensor(_obs).to(device) for (key, _obs) in obs.items()}
+        return {key: th.as_tensor(_obs, device=device) for (key, _obs) in obs.items()}
     else:
         raise Exception(f"Unrecognized type of observation {type(obs)}")
 

stable_baselines3/sac/sac.py

@@ -181,7 +181,7 @@ class SAC(OffPolicyAlgorithm):
             # Force conversion to float
             # this will throw an error if a malformed string (different from 'auto')
             # is passed
-            self.ent_coef_tensor = th.tensor(float(self.ent_coef)).to(self.device)
+            self.ent_coef_tensor = th.tensor(float(self.ent_coef), device=self.device)
 
     def _create_aliases(self) -> None:
         self.actor = self.policy.actor