From 0879a8712f6f4c85bdd53e7b5f558205ebddc929 Mon Sep 17 00:00:00 2001 From: Stelios Tymvios Date: Sat, 11 Jul 2020 21:16:12 +0300 Subject: [PATCH] clarifications --- stable_baselines3/common/buffers.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 68973aa..68c449b 100644 --- a/stable_baselines3/common/buffers.py +++ b/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