From 641a330732b9c6b7e95e1fc5cd5a38cd321988e6 Mon Sep 17 00:00:00 2001
From: Stelios Tymvios <stelios.tymvios@icloud.com>
Date: Sat, 4 Jul 2020 16:11:24 +0300
Subject: [PATCH] fixed bug with done indexing, fixed bug with reward type

---
 stable_baselines3/common/buffers.py | 11 ++++++-----
 tests/test_buffers.py               | 20 ++++++++++++++++++++
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py
index d874abc..6de5219 100644
--- a/stable_baselines3/common/buffers.py
+++ b/stable_baselines3/common/buffers.py
@@ -432,6 +432,7 @@ class NstepReplayBuffer(ReplayBuffer):
 
         # two dim matrix of rewards for the indices
         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
@@ -443,17 +444,17 @@ class NstepReplayBuffer(ReplayBuffer):
 
         # weighted rewards
         rewards = filt * gammas * rewards
-        rewards = np.add.reduce(rewards, 1)
+        rewards = np.add.reduce(rewards, 1).reshape(len(batch_inds), 1).astype(np.float32)
 
-        next_obs_indices = np.add.reduce(filt, 1).astype(np.int).reshape(batch_inds.shape)
-        next_obs_indices = (next_obs_indices + batch_inds - 1) % self.buffer_size
+        increments = np.add.reduce(filt, 1).astype(np.int).reshape(batch_inds.shape)
+        next_obs_indices = (increments + batch_inds - 1) % 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)), next_obs_indices]).reshape(len(batch_inds), 1)
+        dones = 1. - (not_dones[np.arange(len(batch_inds)), increments-1]).reshape(len(batch_inds), 1)
 
-        data = (obs, actions, next_obs, dones.reshape(len(batch_inds), 1), rewards.reshape(len(batch_inds), 1))
+        data = (obs, actions, next_obs, dones, rewards)
         return ReplayBufferSamples(*tuple(map(self.to_torch, data)))
diff --git a/tests/test_buffers.py b/tests/test_buffers.py
index 95c8df5..a441678 100644
--- a/tests/test_buffers.py
+++ b/tests/test_buffers.py
@@ -4,6 +4,8 @@ import numpy as np
 from gym import spaces
 
 from stable_baselines3.common.buffers import NstepReplayBuffer
+from stable_baselines3.common.noise import NormalActionNoise
+from stable_baselines3 import DQN, TD3, SAC
 
 
 def test_nsteps():
@@ -82,3 +84,21 @@ def test_nsteps():
     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]]))
+
+
+@pytest.mark.parametrize("algo", [DQN, TD3])
+def test_with_algo(algo):
+    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})
+    else:
+        env_id = 'CartPole-v1'
+        if algo == DQN:
+            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._setup_model()
+    agent.learn(500)
\ No newline at end of file