Avoid torch type-error under torch.compile (#1922)

* Avoid torch type-error under torch.compile * Update changelog and version * Update stable_baselines3/common/buffers.py Co-authored-by: Antonin RAFFIN <[email protected]> --------- Co-authored-by: Antonin Raffin <[email protected]>
DLR-RM · May 13, 2024 · 766b9e9 · 766b9e9
1 parent 285e01f
commit 766b9e9
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
@@ -14,6 +14,7 @@ New Features:
 
 Bug Fixes:
 ^^^^^^^^^^
+- Cast type in compute gae method to avoid error when using torch compile (@amjames)
 
 `SB3-Contrib`_
 ^^^^^^^^^^^^^^

diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py
@@ -424,7 +424,7 @@ def compute_returns_and_advantage(self, last_values: th.Tensor, dones: np.ndarra
         last_gae_lam = 0
         for step in reversed(range(self.buffer_size)):
             if step == self.buffer_size - 1:
-                next_non_terminal = 1.0 - dones
+                next_non_terminal = 1.0 - dones.astype(np.float32)
                 next_values = last_values
             else:
                 next_non_terminal = 1.0 - self.episode_starts[step + 1]

diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt
@@ -1 +1 @@
-2.3.2
+2.4.0a0