diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index f668503..45b7ee7 100644 --- a/stable_baselines3/common/buffers.py +++ b/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 diff --git a/tests/test_buffers.py b/tests/test_buffers.py index 3dcaa11..41b9f2b 100644 --- a/tests/test_buffers.py +++ b/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]]))