clarifications

stable_baselines3/common/buffers.py

@@ -409,9 +409,9 @@ class NstepReplayBuffer(ReplayBuffer):
     ):
         super().__init__(buffer_size, observation_space, action_space, device, n_envs, optimize_memory_usage)
         self.n_step = int(n_step)
-        assert self.n_step <= buffer_size
-        self.gamma = float(gamma)
-        assert 0 <= gamma <= 1
+        if not 0 < n_step <= buffer_size:
+            raise ValueError("n_step needs to be strictly smaller than buffer_size, and strictly larger than 0")
+        self.gamma = gamma
 
     def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples:
         assert not np.any(batch_inds == self.pos)
@@ -421,23 +421,28 @@ class NstepReplayBuffer(ReplayBuffer):
 
         # Broadcasting turns 1dim arange matrix to 2 dimensional matrix that contains all
         # the indices, % buffersize keeps us in buffer range
+        # indices is a [B x n_step ] matrix
         indices = (np.arange(self.n_step) + batch_inds.reshape(-1, 1)) % self.buffer_size
 
         # two dim matrix of not dones. If done is true, then subsequent dones are turned to 0
         # using accumulate. This ensures that we don't use invalid transitions
+        # not_dones is a [B x n_step] matrix
         not_dones = np.multiply.accumulate(1 - self.dones[indices], 1).reshape(-1, self.n_step)
 
         # vector of the discount factors
+        # [n_step] vector
         gammas = gamma ** np.arange(self.n_step)
 
         # two dim matrix of rewards for the indices
+        # using indices we select the current transition, plus the next n_step ones
         rewards = self.rewards[indices].reshape(not_dones.shape)
         rewards = self._normalize_reward(rewards, env)
 
         # 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
+        # The immediate indice, i.e. col 0 needs to be 1, so we ensure that it is here using np.ones
+        # If the jth transition is terminal, we need to ignore the j+1 but keep the reward of the jth
         # we do this by "shifting" the not_dones one step to the right
+        # so a terminal transition has a 1, and the next has a 0
         filt = np.hstack([np.ones((len(batch_inds), 1)), not_dones[:, :-1]])
         
         # We ignore self.pos indice since it points to older transitions.
@@ -452,7 +457,7 @@ class NstepReplayBuffer(ReplayBuffer):
         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
+        # filt always sums up to 1 + k 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