diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 45d78d3..5bc2527 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -19,6 +19,7 @@ New Features: when ``psutil`` is available - Saving models now automatically creates the necessary folders and raises appropriate warnings (@PartiallyTyped) - Refactored opening paths for saving and loading to use strings, pathlib or io.BufferedIOBase (@PartiallyTyped) +- Implemented N-Step ReplayBuffer for off-policy methods (@PartiallyTyped) Bug Fixes: ^^^^^^^^^^ diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 6a18335..f668503 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -380,7 +380,7 @@ class RolloutBuffer(BaseBuffer): start_idx = 0 while start_idx < self.buffer_size * self.n_envs: - yield self._get_samples(indices[start_idx : start_idx + batch_size]) + yield self._get_samples(indices[start_idx: start_idx + batch_size]) start_idx += batch_size def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> RolloutBufferSamples: @@ -414,34 +414,37 @@ class NstepReplayBuffer(ReplayBuffer): assert 0 <= gamma <= 1 def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples: - assert not np.any(batch_inds==self.pos) + assert not np.any(batch_inds == self.pos) actions = self.actions[batch_inds, 0, :] - + gamma = self.gamma - # two dimensional matrix contains all indices + # Broadcasting turns 1dim arange matrix to 2 dimensional matrix that contains all + # the indices, % buffersize keeps us in buffer range indices = (np.arange(self.n_step) + batch_inds.reshape(-1, 1)) % self.buffer_size - - # two dim matrix of not dones. If done then subsequent dones are 0 - not_dones = np.multiply.accumulate(1-self.dones[indices], 1).reshape(-1, self.n_step) - - # two dim matrix of gammas to the power, will be used with the rewards - gammas = gamma**np.arange(self.n_step) - - # two dim matrix of rewards + + # 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 = np.multiply.accumulate(1 - self.dones[indices], 1).reshape(-1, self.n_step) + + # vector of the discount factors + gammas = gamma ** np.arange(self.n_step) + + # two dim matrix of rewards for the indices rewards = self.rewards[indices].reshape(not_dones.shape) # 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 = not_dones * (indices != self.pos) filt = np.multiply.accumulate(filt, 1).reshape(filt.shape) # weighted rewards - rewards = filt* gammas * rewards + rewards = filt * gammas * rewards rewards = np.add.reduce(rewards, 1) - + 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 obs = self._normalize_obs(self.observations[batch_inds, 0, :], env) @@ -450,6 +453,7 @@ class NstepReplayBuffer(ReplayBuffer): 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)), next_obs_indices]).reshape(len(batch_inds), 1) + data = (obs, actions, next_obs, dones.reshape(len(batch_inds), 1), rewards.reshape(len(batch_inds), 1)) return ReplayBufferSamples(*tuple(map(self.to_torch, data))) diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 05987dc..72c175a 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -2,7 +2,7 @@ import time import warnings import pathlib import io -from typing import Union, Type, Optional, Dict, Any, Callable, List, Tuple, Mapping +from typing import Union, Type, Optional, Dict, Any, Callable, List, Tuple, Mapping, Generic, TypeVar import gym import torch as th @@ -19,8 +19,9 @@ from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.buffers import ReplayBuffer from stable_baselines3.common.save_util import save_to_pkl, load_from_pkl +_BufferType = TypeVar("_BufferType", bound=ReplayBuffer) -class OffPolicyAlgorithm(BaseAlgorithm): +class OffPolicyAlgorithm(BaseAlgorithm, Generic[_BufferType]): """ The base for Off-Policy algorithms (ex: SAC/TD3) @@ -73,6 +74,8 @@ class OffPolicyAlgorithm(BaseAlgorithm): env: Union[GymEnv, str], policy_base: Type[BasePolicy], learning_rate: Union[float, Callable], + replay_buffer_cls: Type[_BufferType] = None, + replay_buffer_kwargs: Optional[Mapping[str, Any]] = None, buffer_size: int = int(1e6), learning_starts: int = 100, batch_size: int = 256, @@ -95,8 +98,6 @@ class OffPolicyAlgorithm(BaseAlgorithm): sde_sample_freq: int = -1, use_sde_at_warmup: bool = False, sde_support: bool = True, - replay_buffer_cls: Type[ReplayBuffer] = ReplayBuffer, - replay_buffer_kwargs: Mapping[str, Any] = None, ): super(OffPolicyAlgorithm, self).__init__( @@ -137,7 +138,6 @@ class OffPolicyAlgorithm(BaseAlgorithm): ) self.actor = None # type: Optional[th.nn.Module] - self.replay_buffer = None # type: Optional[ReplayBuffer] # Update policy keyword arguments if sde_support: self.policy_kwargs["use_sde"] = self.use_sde @@ -145,14 +145,14 @@ class OffPolicyAlgorithm(BaseAlgorithm): # For gSDE only self.use_sde_at_warmup = use_sde_at_warmup - assert issubclass(replay_buffer_cls, ReplayBuffer) - self.replay_buffer_cls = replay_buffer_cls - self.replay_buffer_kwargs = replay_buffer_kwargs or {} + self.replay_buffer = None + self.replay_buffer_cls = replay_buffer_cls if replay_buffer_cls is not None else ReplayBuffer + self.replay_buffer_kwargs = dict(replay_buffer_kwargs or {}) def _setup_model(self) -> None: self._setup_lr_schedule() self.set_random_seed(self.seed) - self.replay_buffer = ReplayBuffer( + self.replay_buffer = self.replay_buffer_cls( self.buffer_size, self.observation_space, self.action_space, @@ -183,18 +183,19 @@ class OffPolicyAlgorithm(BaseAlgorithm): :param path: (Union[str, pathlib.Path, io.BufferedIOBase]) Path to the pickled replay buffer. """ self.replay_buffer = load_from_pkl(path, self.verbose) - assert isinstance(self.replay_buffer, ReplayBuffer), 'The replay buffer must inherit from ReplayBuffer class' + assert isinstance(self.replay_buffer, ReplayBuffer), "The replay buffer must inherit from ReplayBuffer class" - def _setup_learn(self, - total_timesteps: int, - eval_env: Optional[GymEnv], - callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None, - eval_freq: int = 10000, - n_eval_episodes: int = 5, - log_path: Optional[str] = None, - reset_num_timesteps: bool = True, - tb_log_name: str = 'run', - ) -> Tuple[int, BaseCallback]: + def _setup_learn( + self, + total_timesteps: int, + eval_env: Optional[GymEnv], + callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None, + eval_freq: int = 10000, + n_eval_episodes: int = 5, + log_path: Optional[str] = None, + reset_num_timesteps: bool = True, + tb_log_name: str = "run", + ) -> Tuple[int, BaseCallback]: """ cf `BaseAlgorithm`. """ @@ -207,7 +208,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): and self.replay_buffer is not None and (self.replay_buffer.full or self.replay_buffer.pos > 0) ) - + assert self.replay_buffer is not None if truncate_last_traj: warnings.warn( "The last trajectory in the replay buffer will be truncated, " diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 084824e..dc5e451 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -9,7 +9,7 @@ from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback 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): """ @@ -78,7 +78,7 @@ class DQN(OffPolicyAlgorithm): _init_setup_model: bool = True): super(DQN, self).__init__(policy, env, DQNPolicy, learning_rate, - buffer_size, learning_starts, batch_size, + ReplayBuffer, None, buffer_size, learning_starts, batch_size, tau, gamma, train_freq, gradient_steps, n_episodes_rollout, action_noise=None, # No action noise policy_kwargs=policy_kwargs, diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 203abc4..2c5ad26 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -6,6 +6,7 @@ import numpy as np from stable_baselines3.common import logger from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.buffers import ReplayBuffer from stable_baselines3.common.noise import ActionNoise from stable_baselines3.sac.policies import SACPolicy @@ -93,6 +94,7 @@ class SAC(OffPolicyAlgorithm): _init_setup_model: bool = True): super(SAC, self).__init__(policy, env, SACPolicy, learning_rate, + ReplayBuffer, None, buffer_size, learning_starts, batch_size, tau, gamma, train_freq, gradient_steps, n_episodes_rollout, action_noise, diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index 459e240..dfcee09 100644 --- a/stable_baselines3/td3/td3.py +++ b/stable_baselines3/td3/td3.py @@ -6,6 +6,7 @@ from stable_baselines3.common import logger from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.buffers import ReplayBuffer from stable_baselines3.td3.policies import TD3Policy @@ -77,6 +78,7 @@ class TD3(OffPolicyAlgorithm): _init_setup_model: bool = True): super(TD3, self).__init__(policy, env, TD3Policy, learning_rate, + ReplayBuffer, None, buffer_size, learning_starts, batch_size, tau, gamma, train_freq, gradient_steps, n_episodes_rollout, action_noise=action_noise, diff --git a/tests/test_buffers.py b/tests/test_buffers.py index a7f3ece..3dcaa11 100644 --- a/tests/test_buffers.py +++ b/tests/test_buffers.py @@ -3,29 +3,24 @@ import pytest import numpy as np from gym import spaces -from stable_baselines3.common.buffers import NstepReplayBuffer, ReplayBuffer +from stable_baselines3.common.buffers import NstepReplayBuffer + def test_nsteps(): - buffer = NstepReplayBuffer( - 5, - spaces.Discrete(5), - spaces.Discrete(5), - n_step=5, - gamma=1 - ) + buffer = NstepReplayBuffer(5, spaces.Discrete(5), spaces.Discrete(5), n_step=5, gamma=1) buffer.add(0, 1, 10, 1, 0) buffer.add(1, 2, 11, 1, 0) buffer.add(2, 3, 12, 1, 0) buffer.add(3, 4, 13, 1, 0) buffer.add(4, 5, 14, 1, 0) - obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1,2,3,4])) + obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1, 2, 3, 4])) assert obs.shape == (4, 1) assert act.shape == (4, 1) assert next_obs.shape == (4, 1) assert dones.shape == (4, 1) assert rewards.shape == (4, 1) assert np.allclose(dones, np.zeros_like(dones)) - assert np.allclose(next_obs, np.array([[5],[5],[5],[5]])) + assert np.allclose(next_obs, np.array([[5], [5], [5], [5]])) assert np.allclose(rewards, np.array([4, 3, 2, 1]).reshape(4, 1)) assert np.allclose(act, np.array([11, 12, 13, 14]).reshape(4, 1)) @@ -33,50 +28,38 @@ def test_nsteps(): with pytest.raises(AssertionError): buffer._get_samples(np.array([0, 1, 2, 3])) - buffer = NstepReplayBuffer( - 5, - spaces.Discrete(5), - spaces.Discrete(5), - n_step=5, - gamma=0.9 - ) + buffer = NstepReplayBuffer(5, spaces.Discrete(5), spaces.Discrete(5), n_step=5, gamma=0.9) buffer.add(0, 1, 10, 1, 0) buffer.add(1, 2, 11, 1, 0) buffer.add(2, 3, 12, 1, 0) buffer.add(3, 4, 13, 1, 0) buffer.add(4, 5, 14, 1, 0) - obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1,2,3,4])) + obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1, 2, 3, 4])) assert obs.shape == (4, 1) assert act.shape == (4, 1) assert next_obs.shape == (4, 1) assert dones.shape == (4, 1) assert rewards.shape == (4, 1) assert np.allclose(dones, np.zeros_like(dones)) - assert np.allclose(next_obs, np.array([[5],[5],[5],[5]])) - assert np.allclose(rewards, np.array([1+0.9+0.9**2+0.9**3, 1+0.9+0.9**2, 1+0.9, 1]).reshape(4, 1)) + assert np.allclose(next_obs, np.array([[5], [5], [5], [5]])) + assert np.allclose(rewards, np.array([1 + 0.9 + 0.9 ** 2 + 0.9 ** 3, 1 + 0.9 + 0.9 ** 2, 1 + 0.9, 1]).reshape(4, 1)) assert np.allclose(act, np.array([11, 12, 13, 14]).reshape(4, 1)) - buffer = NstepReplayBuffer( - 10, - spaces.Discrete(5), - spaces.Discrete(5), - n_step=5, - gamma=0.9 - ) + buffer = NstepReplayBuffer(10, spaces.Discrete(5), spaces.Discrete(5), n_step=5, gamma=0.9) buffer.add(0, 1, 10, 1, 0) buffer.add(1, 2, 11, 1, 0) buffer.add(2, 3, 12, 1, 0) buffer.add(3, 4, 13, 1, 0) buffer.add(4, 5, 14, 1, 0) - obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1,2,3,4])) + obs, act, next_obs, dones, rewards = buffer._get_samples(np.array([1, 2, 3, 4])) assert obs.shape == (4, 1) assert act.shape == (4, 1) assert next_obs.shape == (4, 1) assert dones.shape == (4, 1) assert rewards.shape == (4, 1) assert np.allclose(dones, np.zeros_like(dones)) - assert np.allclose(next_obs, np.array([[5],[5],[5],[5]])) - assert np.allclose(rewards, np.array([1+0.9+0.9**2+0.9**3, 1+0.9+0.9**2, 1+0.9, 1]).reshape(4, 1)) + assert np.allclose(next_obs, np.array([[5], [5], [5], [5]])) + assert np.allclose(rewards, np.array([1 + 0.9 + 0.9 ** 2 + 0.9 ** 3, 1 + 0.9 + 0.9 ** 2, 1 + 0.9, 1]).reshape(4, 1)) assert np.allclose(act, np.array([11, 12, 13, 14]).reshape(4, 1)) # shouldn't be able to get batch with indice 5 because the pointer is at 5