saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-28 20:11:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fixed bug with done indexing, fixed bug with reward type

Browse source
This commit is contained in:
Stelios Tymvios 2020-07-04 16:11:24 +03:00
parent 77d842b74d
commit 641a330732
2 changed files with 26 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
stable_baselines3/common/buffers.py
View file

@ -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)))

20
tests/test_buffers.py
View file

@ -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)