diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 45b7ee7..d874abc 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 = np.hstack([np.ones((len(batch_inds),1)), not_dones[:, 1:]]) * (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/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 72c175a..baa291f 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -21,6 +21,7 @@ from stable_baselines3.common.save_util import save_to_pkl, load_from_pkl _BufferType = TypeVar("_BufferType", bound=ReplayBuffer) + class OffPolicyAlgorithm(BaseAlgorithm, Generic[_BufferType]): """ The base for Off-Policy algorithms (ex: SAC/TD3) diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index dc5e451..b171cc6 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -11,6 +11,7 @@ from stable_baselines3.common.utils import get_linear_fn from stable_baselines3.dqn.policies import DQNPolicy from stable_baselines3.common.buffers import ReplayBuffer + class DQN(OffPolicyAlgorithm): """ Deep Q-Network (DQN) diff --git a/tests/test_buffers.py b/tests/test_buffers.py index 41b9f2b..95c8df5 100644 --- a/tests/test_buffers.py +++ b/tests/test_buffers.py @@ -65,7 +65,7 @@ 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)