more tests, corrected a bug incrementing was wrong when done=1

stable_baselines3/common/buffers.py

@@ -438,7 +438,7 @@ class NstepReplayBuffer(ReplayBuffer):
         # basically we need to ignore indices that are equal to self.pos
         # because the indice at self.pos is older
         # not_dones that are from other transitions are filtered out.
-        filt = not_dones * (indices != self.pos)
+        filt = np.hstack([np.ones((len(batch_inds),1)), not_dones[:, 1:]]) * (indices != self.pos)
         filt = np.multiply.accumulate(filt, 1).reshape(filt.shape)
 
         # weighted rewards

tests/test_buffers.py

@@ -65,3 +65,20 @@ def test_nsteps():
     # shouldn't be able to get batch with indice 5 because the pointer is at 5
     with pytest.raises(AssertionError):
         buffer._get_samples(np.array([5]))
+    
+    buffer = NstepReplayBuffer(10, spaces.Discrete(5), spaces.Discrete(5), n_step=5, gamma=0.9)
+    buffer.add(0, 1, 10, 1, 1)
+    buffer.add(1, 2, 11, 1, 1)
+    buffer.add(2, 3, 12, 1, 1)
+    buffer.add(3, 4, 13, 1, 1)
+    buffer.add(4, 5, 14, 1, 1)
+    obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1, 2, 3, 4]))
+    assert obs.shape == (4, 1)
+    assert act.shape == (4, 1)
+    assert next_obs.shape == (4, 1)
+    assert dones.shape == (4, 1)
+    assert rewards.shape == (4, 1)
+    assert np.allclose(dones, np.ones_like(dones))
+    assert np.allclose(rewards, np.array([1, 1, 1, 1]).reshape(4, 1))
+    assert np.allclose(act, np.array([11, 12, 13, 14]).reshape(4, 1))
+    assert np.allclose(next_obs, np.array([[2], [3], [4], [5]]))