corrected miscalculation of filter

stable_baselines3/common/buffers.py

@@ -434,27 +434,36 @@ class NstepReplayBuffer(ReplayBuffer):
         rewards = self.rewards[indices].reshape(not_dones.shape)
         rewards = self._normalize_reward(rewards, env)
 
-        # filter to ignore loops around the buffer, see:
-        # https://github.com/DLR-RM/stable-baselines3/pull/81#issuecomment-653741592
-        # 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 = np.hstack([np.ones((len(batch_inds), 1)), not_dones[:, 1:]]) * (indices != self.pos)
-        filt = np.multiply.accumulate(filt, 1).reshape(filt.shape)
+        # we filter through the indices.
+        # The immediate indice, i.e. col 0 is 1, so we ensure that it is here  using np.ones
+        # If the transition is terminal, we need to ignore the next but keep the reward
+        # we do this by "shifting" the not_dones one step to the right
+        filt = np.hstack([np.ones((len(batch_inds), 1)), not_dones[:, :-1]])
+        
+        # We ignore self.pos indice since it points to older transitions.
+        # we then accumulate to prevent continuing to the wrong transitions.
+        current_episode = np.multiply.accumulate(indices != self.pos, 1).reshape(filt.shape)
+        
+        # combine the filters
+        filt = filt * current_episode
 
-        # weighted rewards
-        rewards = filt * gammas * rewards
-        rewards = np.add.reduce(rewards, 1).reshape(len(batch_inds), 1).astype(np.float32)
+        # discount the rewards
+        rewards = (rewards * filt) @ gammas.T
+        rewards = rewards.reshape(len(batch_inds), 1).astype(np.float32)
+        
+        # Increments counts how many transitions we need to skip
+        # filt always sums up to 1 + n non terminal transitions due to hstack above
+        # so we subtract 1.
+        increments = np.add.reduce(filt, 1).astype(np.int).reshape(batch_inds.shape) - 1
 
-        increments = np.add.reduce(filt, 1).astype(np.int).reshape(batch_inds.shape)
-        next_obs_indices = (increments + batch_inds - 1) % self.buffer_size
+        next_obs_indices = (increments + batch_inds) % self.buffer_size
         obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
         if self.optimize_memory_usage:
             next_obs = self._normalize_obs(self.observations[(next_obs_indices + 1) % self.buffer_size, 0, :], env)
         else:
             next_obs = self._normalize_obs(self.next_observations[next_obs_indices, 0, :], env)
 
-        dones = 1. - (not_dones[np.arange(len(batch_inds)), increments-1]).reshape(len(batch_inds), 1)
+        dones = 1. - (not_dones[np.arange(len(batch_inds)), increments]).reshape(len(batch_inds), 1)
 
         data = (obs, actions, next_obs, dones, rewards)
         return ReplayBufferSamples(*tuple(map(self.to_torch, data)))

tests/test_buffers.py

@@ -85,20 +85,37 @@ def test_nsteps():
     assert np.allclose(act, np.array([11, 12, 13, 14]).reshape(4, 1))
     assert np.allclose(next_obs, np.array([[2], [3], [4], [5]]))
 
+    buffer = NstepReplayBuffer(5, spaces.Discrete(5), spaces.Discrete(5), n_step=2, gamma=0.99)
+    buffer.add(0, 1, 10, 2, 0)
+    buffer.add(1, 2, 11, 1, 0)
+    buffer.add(2, 3, 12, 2, 0)
+    buffer.add(3, 4, 13, 2, 1)
+    buffer.add(4, 5, 14, 2, 0)
+    obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1, 1, 2, 3, 4]))
+    assert obs.shape == (5, 1)
+    assert act.shape == (5, 1)
+    assert next_obs.shape == (5, 1)
+    assert dones.shape == (5, 1)
+    assert rewards.shape == (5, 1)
+    assert np.allclose(dones, np.array([0, 0, 1, 1, 0]).reshape(5, 1))
+    assert np.allclose(rewards, np.array([2.98, 2.98, 3.98, 2, 2]).reshape(5, 1))
+    assert np.allclose(act, np.array([11, 11, 12, 13, 14]).reshape(5, 1))
+    assert np.allclose(next_obs, np.array([[3], [3], [4], [4], [5]]))
+
 
 @pytest.mark.parametrize("algo", [DQN, TD3])
 def test_with_algo(algo):
-    kwargs = {'policy_kwargs': dict(net_arch=[64]), '_init_setup_model':False}
+    # Integration test
+    kwargs = {"policy_kwargs": dict(net_arch=[64]), "_init_setup_model": False}
     if algo in [TD3, SAC]:
-        env_id = 'Pendulum-v0'
-        kwargs.update({'action_noise': NormalActionNoise(0.0, 0.1),
-                       'learning_starts': 100})
+        env_id = "Pendulum-v0"
+        kwargs.update({"action_noise": NormalActionNoise(0.0, 0.1), "learning_starts": 100})
     else:
-        env_id = 'CartPole-v1'
+        env_id = "CartPole-v1"
         if algo == DQN:
-            kwargs.update({'learning_starts': 100})
-    agent = algo('MlpPolicy', env_id, **kwargs)
+            kwargs.update({"learning_starts": 100})
+    agent = algo("MlpPolicy", env_id, **kwargs)
     agent.replay_buffer_cls = NstepReplayBuffer
-    agent.replay_buffer_kwargs.update({"gamma":agent.gamma, "n_step": 10})
+    agent.replay_buffer_kwargs.update({"gamma": agent.gamma, "n_step": 10})
     agent._setup_model()
     agent.learn(500)