From 2057a1082c59f8030c3e793c90ce90b0e63131fd Mon Sep 17 00:00:00 2001 From: Stelios Tymvios Date: Sat, 11 Jul 2020 21:07:34 +0300 Subject: [PATCH] corrected miscalculation of filter --- stable_baselines3/common/buffers.py | 35 ++++++++++++++++++----------- tests/test_buffers.py | 33 ++++++++++++++++++++------- 2 files changed, 47 insertions(+), 21 deletions(-) diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 6de5219..68973aa 100644 --- a/stable_baselines3/common/buffers.py +++ b/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))) diff --git a/tests/test_buffers.py b/tests/test_buffers.py index 1bb7f9c..5c5b8b8 100644 --- a/tests/test_buffers.py +++ b/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)