diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 0b89d41..8eeaf41 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -1,4 +1,4 @@ -image: stablebaselines/stable-baselines3-cpu:0.6.0 +image: stablebaselines/stable-baselines3-cpu:0.8.0a1 type-check: script: diff --git a/README.md b/README.md index a14fee6..048f2e0 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,6 @@ These algorithms will make it easier for the research community and industry to Please look at the issue for more details. Planned features: -- [ ] DQN (almost ready, currently in testing phase) - [ ] DDPG (you can use its successor TD3 for now) - [ ] HER diff --git a/docs/guide/algos.rst b/docs/guide/algos.rst index 94dc43b..821fee9 100644 --- a/docs/guide/algos.rst +++ b/docs/guide/algos.rst @@ -12,6 +12,7 @@ A2C ✔️ ✔️ ✔️ ✔️ PPO ✔️ ✔️ ✔️ ✔️ ✔️ SAC ✔️ ❌ ❌ ❌ ❌ TD3 ✔️ ❌ ❌ ❌ ❌ +DQN ❌ ✔️ ❌ ❌ ❌ ============ =========== ============ ================= =============== ================ diff --git a/docs/guide/examples.rst b/docs/guide/examples.rst index 0f98de7..0e3242b 100644 --- a/docs/guide/examples.rst +++ b/docs/guide/examples.rst @@ -33,7 +33,7 @@ notebooks: Basic Usage: Training, Saving, Loading -------------------------------------- -In the following example, we will train, save and load a A2C model on the Lunar Lander environment. +In the following example, we will train, save and load a DQN model on the Lunar Lander environment. .. image:: ../_static/img/colab-badge.svg :target: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/saving_loading_dqn.ipynb @@ -57,7 +57,7 @@ In the following example, we will train, save and load a A2C model on the Lunar import gym - from stable_baselines3 import A2C + from stable_baselines3 import DQN from stable_baselines3.common.evaluation import evaluate_policy @@ -65,15 +65,15 @@ In the following example, we will train, save and load a A2C model on the Lunar env = gym.make('LunarLander-v2') # Instantiate the agent - model = A2C('MlpPolicy', env, verbose=1) + model = DQN('MlpPolicy', env, verbose=1) # Train the agent model.learn(total_timesteps=int(2e5)) # Save the agent - model.save("a2c_lunar") + model.save("dqn_lunar") del model # delete trained model to demonstrate loading # Load the trained agent - model = A2C.load("a2c_lunar") + model = DQN.load("dqn_lunar") # Evaluate the agent mean_reward, std_reward = evaluate_policy(model, model.get_env(), n_eval_episodes=10) @@ -81,7 +81,7 @@ In the following example, we will train, save and load a A2C model on the Lunar # Enjoy trained agent obs = env.reset() for i in range(1000): - action, _states = model.predict(obs) + action, _states = model.predict(obs, deterministic=True) obs, rewards, dones, info = env.step(action) env.render() diff --git a/docs/index.rst b/docs/index.rst index 2c531c9..2eea1bf 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -58,6 +58,7 @@ Main Features modules/ppo modules/sac modules/td3 + modules/dqn .. toctree:: :maxdepth: 1 diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index ec43194..a2c9645 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,15 +3,20 @@ Changelog ========== -Pre-Release 0.8.0a0 (WIP) +Pre-Release 0.8.0a1 (WIP) ------------------------------ Breaking Changes: ^^^^^^^^^^^^^^^^^ +- ``AtariWrapper`` and other Atari wrappers were updated to match SB2 ones - ``save_replay_buffer`` now receives as argument the file path instead of the folder path (@tirafesi) New Features: ^^^^^^^^^^^^^ +- Added ``DQN`` Algorithm (@Artemis-Skade) +- Buffer dtype is now set according to action and observation spaces for ``ReplayBuffer`` +- Added warning when allocation of a buffer may exceed the available memory of the system + when ``psutil`` is available Bug Fixes: ^^^^^^^^^^ @@ -22,6 +27,10 @@ Deprecations: Others: ^^^^^^^ +- Refactored off-policy algorithm to share the same ``.learn()`` method +- Split the ``collect_rollout()`` method for off-policy algorithms +- Added ``_on_step()`` for off-policy base class +- Optimized replay buffer size by removing the need of ``next_observations`` numpy array Documentation: ^^^^^^^^^^^^^^ @@ -29,6 +38,7 @@ Documentation: - Fixed a typo in the section of Enjoy a Trained Agent, in RL Baselines3 Zoo README. (@blurLake) + Pre-Release 0.7.0 (2020-06-10) ------------------------------ diff --git a/docs/modules/dqn.rst b/docs/modules/dqn.rst new file mode 100644 index 0000000..034e5b8 --- /dev/null +++ b/docs/modules/dqn.rst @@ -0,0 +1,94 @@ +.. _dqn: + +.. automodule:: stable_baselines3.dqn + + +DQN +=== + +`Deep Q Network (DQN) `_ + +.. rubric:: Available Policies + +.. autosummary:: + :nosignatures: + + MlpPolicy + CnnPolicy + + +Notes +----- + +- Original paper: https://arxiv.org/abs/1312.5602 +- Further reference: https://www.nature.com/articles/nature14236 + +.. note:: + This implementation provides only vanilla Deep Q-Learning and has no extensions such as Double-DQN, Dueling-DQN and Prioritized Experience Replay. + + +Can I use? +---------- + +- Recurrent policies: ❌ +- Multi processing: ❌ +- Gym spaces: + + +============= ====== =========== +Space Action Observation +============= ====== =========== +Discrete ✔ ✔ +Box ❌ ✔ +MultiDiscrete ❌ ✔ +MultiBinary ❌ ✔ +============= ====== =========== + + +Example +------- + +.. code-block:: python + + import gym + import numpy as np + + from stable_baselines3 import DQN + from stable_baselines3.dqn import MlpPolicy + + env = gym.make('Pendulum-v0') + + model = DQN(MlpPolicy, env, verbose=1) + model.learn(total_timesteps=10000, log_interval=4) + model.save("dqn_pendulum") + + del model # remove to demonstrate saving and loading + + model = DQN.load("dqn_pendulum") + + obs = env.reset() + while True: + action, _states = model.predict(obs, deterministic=True) + obs, reward, done, info = env.step(action) + env.render() + if done: + obs = env.reset() + +Parameters +---------- + +.. autoclass:: DQN + :members: + :inherited-members: + +.. _dqn_policies: + +DQN Policies +------------- + +.. autoclass:: MlpPolicy + :members: + :inherited-members: + +.. autoclass:: CnnPolicy + :members: diff --git a/setup.cfg b/setup.cfg index 41e1c32..c20a6ad 100644 --- a/setup.cfg +++ b/setup.cfg @@ -27,6 +27,7 @@ per-file-ignores = ./stable_baselines3/__init__.py:F401 ./stable_baselines3/common/__init__.py:F401 ./stable_baselines3/a2c/__init__.py:F401 + ./stable_baselines3/dqn/__init__.py:F401 ./stable_baselines3/ppo/__init__.py:F401 ./stable_baselines3/sac/__init__.py:F401 ./stable_baselines3/td3/__init__.py:F401 diff --git a/setup.py b/setup.py index 87c737c..5733441 100644 --- a/setup.py +++ b/setup.py @@ -108,7 +108,9 @@ setup(name='stable_baselines3', # For atari games, 'atari_py~=0.2.0', 'pillow', # Tensorboard support - 'tensorboard' + 'tensorboard', + # Checking memory taken by replay buffer + 'psutil' ] }, description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.', diff --git a/stable_baselines3/__init__.py b/stable_baselines3/__init__.py index d921eeb..47a9bed 100644 --- a/stable_baselines3/__init__.py +++ b/stable_baselines3/__init__.py @@ -4,6 +4,7 @@ from stable_baselines3.a2c import A2C from stable_baselines3.ppo import PPO from stable_baselines3.sac import SAC from stable_baselines3.td3 import TD3 +from stable_baselines3.dqn import DQN # Read version from file version_file = os.path.join(os.path.dirname(__file__), 'version.txt') diff --git a/stable_baselines3/common/atari_wrappers.py b/stable_baselines3/common/atari_wrappers.py index c25dc6d..7af95fa 100644 --- a/stable_baselines3/common/atari_wrappers.py +++ b/stable_baselines3/common/atari_wrappers.py @@ -1,31 +1,213 @@ import gym -from gym.wrappers import AtariPreprocessing +from gym import spaces import numpy as np +try: + import cv2 # pytype:disable=import-error + cv2.ocl.setUseOpenCL(False) +except ImportError: + cv2 = None from stable_baselines3.common.type_aliases import GymStepReturn +class NoopResetEnv(gym.Wrapper): + def __init__(self, env: gym.Env, noop_max: int = 30): + """ + Sample initial states by taking random number of no-ops on reset. + No-op is assumed to be action 0. + + :param env: (gym.Env) the environment to wrap + :param noop_max: (int) the maximum value of no-ops to run + """ + gym.Wrapper.__init__(self, env) + self.noop_max = noop_max + self.override_num_noops = None + self.noop_action = 0 + assert env.unwrapped.get_action_meanings()[0] == 'NOOP' + + def reset(self, **kwargs) -> np.ndarray: + self.env.reset(**kwargs) + if self.override_num_noops is not None: + noops = self.override_num_noops + else: + noops = self.unwrapped.np_random.randint(1, self.noop_max + 1) + assert noops > 0 + obs = np.empty(0) + for _ in range(noops): + obs, _, done, _ = self.env.step(self.noop_action) + if done: + obs = self.env.reset(**kwargs) + return obs + + +class FireResetEnv(gym.Wrapper): + def __init__(self, env: gym.Env): + """ + Take action on reset for environments that are fixed until firing. + + :param env: (gym.Env) the environment to wrap + """ + gym.Wrapper.__init__(self, env) + assert env.unwrapped.get_action_meanings()[1] == 'FIRE' + assert len(env.unwrapped.get_action_meanings()) >= 3 + + def reset(self, **kwargs) -> np.ndarray: + self.env.reset(**kwargs) + obs, _, done, _ = self.env.step(1) + if done: + self.env.reset(**kwargs) + obs, _, done, _ = self.env.step(2) + if done: + self.env.reset(**kwargs) + return obs + + +class EpisodicLifeEnv(gym.Wrapper): + def __init__(self, env: gym.Env): + """ + Make end-of-life == end-of-episode, but only reset on true game over. + Done by DeepMind for the DQN and co. since it helps value estimation. + + :param env: (gym.Env) the environment to wrap + """ + gym.Wrapper.__init__(self, env) + self.lives = 0 + self.was_real_done = True + + def step(self, action: int) -> GymStepReturn: + obs, reward, done, info = self.env.step(action) + self.was_real_done = done + # check current lives, make loss of life terminal, + # then update lives to handle bonus lives + lives = self.env.unwrapped.ale.lives() + if 0 < lives < self.lives: + # for Qbert sometimes we stay in lives == 0 condtion for a few frames + # so its important to keep lives > 0, so that we only reset once + # the environment advertises done. + done = True + self.lives = lives + return obs, reward, done, info + + def reset(self, **kwargs) -> np.ndarray: + """ + Calls the Gym environment reset, only when lives are exhausted. + This way all states are still reachable even though lives are episodic, + and the learner need not know about any of this behind-the-scenes. + + :param kwargs: Extra keywords passed to env.reset() call + :return: (np.ndarray) the first observation of the environment + """ + if self.was_real_done: + obs = self.env.reset(**kwargs) + else: + # no-op step to advance from terminal/lost life state + obs, _, _, _ = self.env.step(0) + self.lives = self.env.unwrapped.ale.lives() + return obs + + +class MaxAndSkipEnv(gym.Wrapper): + def __init__(self, env: gym.Env, skip: int = 4): + """ + Return only every ``skip``-th frame (frameskipping) + + :param env: (gym.Env) the environment + :param skip: (int) number of ``skip``-th frame + """ + gym.Wrapper.__init__(self, env) + # most recent raw observations (for max pooling across time steps) + self._obs_buffer = np.zeros((2,) + env.observation_space.shape, dtype=env.observation_space.dtype) + self._skip = skip + + def step(self, action: int) -> GymStepReturn: + """ + Step the environment with the given action + Repeat action, sum reward, and max over last observations. + + :param action: ([int] or [float]) the action + :return: ([int] or [float], [float], [bool], dict) observation, reward, done, information + """ + total_reward = 0.0 + done = None + for i in range(self._skip): + obs, reward, done, info = self.env.step(action) + if i == self._skip - 2: + self._obs_buffer[0] = obs + if i == self._skip - 1: + self._obs_buffer[1] = obs + total_reward += reward + if done: + break + # Note that the observation on the done=True frame + # doesn't matter + max_frame = self._obs_buffer.max(axis=0) + + return max_frame, total_reward, done, info + + def reset(self, **kwargs): + return self.env.reset(**kwargs) + + +class ClipRewardEnv(gym.RewardWrapper): + def __init__(self, env: gym.Env): + """ + Clips the reward to {+1, 0, -1} by its sign. + + :param env: (gym.Env) the environment + """ + gym.RewardWrapper.__init__(self, env) + + def reward(self, reward: float) -> float: + """ + Bin reward to {+1, 0, -1} by its sign. + + :param reward: (float) + :return: (float) + """ + return np.sign(reward) + + +class WarpFrame(gym.ObservationWrapper): + def __init__(self, env: gym.Env, width: int = 84, height: int = 84): + """ + Convert to grayscale and warp frames to 84x84 (default) + as done in the Nature paper and later work. + + :param env: (gym.Env) the environment + :param width: (int) + :param height: (int) + """ + gym.ObservationWrapper.__init__(self, env) + self.width = width + self.height = height + self.observation_space = spaces.Box(low=0, high=255, shape=(self.height, self.width, 1), + dtype=env.observation_space.dtype) + + def observation(self, frame: np.ndarray) -> np.ndarray: + """ + returns the current observation from a frame + + :param frame: (np.ndarray) environment frame + :return: (np.ndarray) the observation + """ + frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) + frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_AREA) + return frame[:, :, None] + + class AtariWrapper(gym.Wrapper): """ Atari 2600 preprocessings - It is a wrapper around the one found in gym. - It reshapes the observation to have an additional dimension and clip the reward. - See https://github.com/openai/gym/blob/master/gym/wrappers/atari_preprocessing.py - . - This class follows the guidelines in - Machado et al. (2018), "Revisiting the Arcade Learning Environment: - Evaluation Protocols and Open Problems for General Agents". - Specifically: * NoopReset: obtain initial state by taking random number of no-ops on reset. * Frame skipping: 4 by default * Max-pooling: most recent two observations - * Termination signal when a life is lost: turned off by default. Not recommended by Machado et al. (2018). + * Termination signal when a life is lost. * Resize to a square image: 84x84 by default - * Grayscale observation: by default - * Scale observation: optional + * Grayscale observation + * Clip reward to {-1, 0, 1} :param env: (gym.Env) gym environment :param noop_max: (int): max number of no-ops @@ -33,44 +215,22 @@ class AtariWrapper(gym.Wrapper): :param screen_size: (int): resize Atari frame :param terminal_on_life_loss: (bool): if True, then step() returns done=True whenever a life is lost. - :param grayscale_obs: (bool): if True (default), then gray scale observation is returned, otherwise, RGB observation - is returned. - :param scale_obs: (bool): if True, then observation normalized in range [0,1] is returned. It also limits memory - optimization benefits of FrameStack Wrapper. - :param scale_obs: (bool) If True (default), the reward is clip to {-1, 0, 1} depending on its sign. + :param clip_reward: (bool) If True (default), the reward is clip to {-1, 0, 1} depending on its sign. """ def __init__(self, env: gym.Env, noop_max: int = 30, frame_skip: int = 4, screen_size: int = 84, - terminal_on_life_loss: bool = False, - grayscale_obs: bool = True, - scale_obs: bool = False, + terminal_on_life_loss: bool = True, clip_reward: bool = True): - env = AtariPreprocessing(env, noop_max=noop_max, frame_skip=frame_skip, screen_size=screen_size, - terminal_on_life_loss=terminal_on_life_loss, grayscale_obs=grayscale_obs, - scale_obs=scale_obs) - # Add channel dimension - if grayscale_obs: - obs_space = env.observation_space - _low, _high, _obs_dtype = (0, 255, np.uint8) if not scale_obs else (0, 1, np.float32) - env.observation_space = gym.spaces.Box(low=_low, high=_high, shape=obs_space.shape + (1,), - dtype=_obs_dtype) + env = NoopResetEnv(env, noop_max=noop_max) + env = MaxAndSkipEnv(env, skip=frame_skip) + if terminal_on_life_loss: + env = EpisodicLifeEnv(env) + if 'FIRE' in env.unwrapped.get_action_meanings(): + env = FireResetEnv(env) + env = WarpFrame(env, width=screen_size, height=screen_size) + if clip_reward: + env = ClipRewardEnv(env) super(AtariWrapper, self).__init__(env) - self.clip_reward = clip_reward - - def _add_axis(self, obs: np.ndarray) -> np.ndarray: - if self.env.grayscale_obs: - return obs[..., np.newaxis] - return obs - - def reset(self) -> np.ndarray: - return self._add_axis(self.env.reset()) - - def step(self, action: int) -> GymStepReturn: - obs, reward, done, info = self.env.step(action) - # Bin reward to {+1, 0, -1} by its sign. - if self.clip_reward: - reward = np.sign(reward) - return self._add_axis(obs), reward, done, info diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index e899a66..9b52ec4 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -84,6 +84,8 @@ class BaseAlgorithm(ABC): self.action_space = None # type: Optional[gym.spaces.Space] self.n_envs = None self.num_timesteps = 0 + # Used for updating schedules + self._total_timesteps = 0 self.eval_env = None self.seed = seed self.action_noise = None # type: Optional[ActionNoise] @@ -398,7 +400,7 @@ class BaseAlgorithm(ABC): log_path: Optional[str] = None, reset_num_timesteps: bool = True, tb_log_name: str = 'run', - ) -> Tuple[int, 'BaseCallback']: + ) -> Tuple[int, BaseCallback]: """ Initialize different variables needed for training. @@ -410,7 +412,7 @@ class BaseAlgorithm(ABC): :param log_path (Optional[str]): Path to a log folder :param reset_num_timesteps: (bool) Whether to reset or not the ``num_timesteps`` attribute :param tb_log_name: (str) the name of the run for tensorboard log - :return: (int, Tuple[BaseCallback]) + :return: (Tuple[int, BaseCallback]) """ self.start_time = time.time() self.ep_info_buffer = deque(maxlen=100) @@ -425,6 +427,7 @@ class BaseAlgorithm(ABC): else: # Make sure training timesteps are ahead of the internal counter total_timesteps += self.num_timesteps + self._total_timesteps = total_timesteps # Avoid resetting the environment when calling ``.learn()`` consecutive times if reset_num_timesteps or self._last_obs is None: diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 4382bd3..24ed8ae 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -1,9 +1,16 @@ from typing import Union, Optional, Generator +import warnings import numpy as np import torch as th from gym import spaces +try: + # Check memory used by replay buffer when possible + import psutil +except ImportError: + psutil = None + from stable_baselines3.common.vec_env import VecNormalize from stable_baselines3.common.type_aliases import RolloutBufferSamples, ReplayBufferSamples from stable_baselines3.common.preprocessing import get_action_dim, get_obs_shape @@ -145,25 +152,52 @@ class ReplayBuffer(BaseBuffer): :param action_space: (spaces.Space) Action space :param device: (th.device) :param n_envs: (int) Number of parallel environments + :param optimize_memory_usage: (bool) Enable a memory efficient variant + of the replay buffer which reduces by almost a factor two the memory used, + at a cost of more complexity. + See https://github.com/DLR-RM/stable-baselines3/issues/37#issuecomment-637501195 + and https://github.com/DLR-RM/stable-baselines3/pull/28#issuecomment-637559274 """ - def __init__(self, buffer_size: int, observation_space: spaces.Space, action_space: spaces.Space, device: Union[th.device, str] = 'cpu', - n_envs: int = 1): + n_envs: int = 1, + optimize_memory_usage: bool = False): super(ReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs) assert n_envs == 1, "Replay buffer only support single environment for now" - self.observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, dtype=np.float32) - self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=np.float32) - self.next_observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, dtype=np.float32) + # Check that the replay buffer can fit into the memory + if psutil is not None: + mem_available = psutil.virtual_memory().available + + self.optimize_memory_usage = optimize_memory_usage + self.observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, dtype=observation_space.dtype) + if optimize_memory_usage: + # `observations` contains also the next observation + self.next_observations = None + else: + self.next_observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, dtype=observation_space.dtype) + self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=action_space.dtype) self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.dones = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) + if psutil is not None: + total_memory_usage = (self.observations.nbytes + self.actions.nbytes + + self.rewards.nbytes + self.dones.nbytes) + if self.next_observations is not None: + total_memory_usage += self.next_observations.nbytes + + if total_memory_usage > mem_available: + # Convert to GB + total_memory_usage /= 1e9 + mem_available /= 1e9 + warnings.warn("This system does not have apparently enough memory to store the complete " + f"replay buffer {total_memory_usage:.2f}GB > {mem_available:.2f}GB") + def add(self, obs: np.ndarray, next_obs: np.ndarray, @@ -172,7 +206,11 @@ class ReplayBuffer(BaseBuffer): done: np.ndarray) -> None: # Copy to avoid modification by reference self.observations[self.pos] = np.array(obs).copy() - self.next_observations[self.pos] = np.array(next_obs).copy() + if self.optimize_memory_usage: + self.observations[(self.pos + 1) % self.buffer_size] = np.array(next_obs).copy() + else: + self.next_observations[self.pos] = np.array(next_obs).copy() + self.actions[self.pos] = np.array(action).copy() self.rewards[self.pos] = np.array(reward).copy() self.dones[self.pos] = np.array(done).copy() @@ -182,13 +220,43 @@ class ReplayBuffer(BaseBuffer): self.full = True self.pos = 0 + def sample(self, + batch_size: int, + env: Optional[VecNormalize] = None + ) -> ReplayBufferSamples: + """ + Sample elements from the replay buffer. + Custom sampling when using memory efficient variant, + as we should not sample the element with index `self.pos` + See https://github.com/DLR-RM/stable-baselines3/pull/28#issuecomment-637559274 + + :param batch_size: (int) Number of element to sample + :param env: (Optional[VecNormalize]) associated gym VecEnv + to normalize the observations/rewards when sampling + :return: (Union[RolloutBufferSamples, ReplayBufferSamples]) + """ + if not self.optimize_memory_usage: + return super().sample(batch_size=batch_size, env=env) + # Do not sample the element with index `self.pos` as the transitions is invalid + # (we use only one array to store `obs` and `next_obs`) + if self.full: + batch_inds = (np.random.randint(1, self.buffer_size, size=batch_size) + self.pos) % self.buffer_size + else: + batch_inds = np.random.randint(0, self.pos, size=batch_size) + return self._get_samples(batch_inds, env=env) + def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None ) -> ReplayBufferSamples: + if self.optimize_memory_usage: + next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env) + else: + next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env) + data = (self._normalize_obs(self.observations[batch_inds, 0, :], env), self.actions[batch_inds, 0, :], - self._normalize_obs(self.next_observations[batch_inds, 0, :], env), + next_obs, self.dones[batch_inds], self._normalize_reward(self.rewards[batch_inds], env)) return ReplayBufferSamples(*tuple(map(self.to_torch, data))) @@ -227,7 +295,8 @@ class RolloutBuffer(BaseBuffer): self.reset() def reset(self) -> None: - self.observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, dtype=np.float32) + self.observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, + dtype=np.float32) self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=np.float32) self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.returns = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) diff --git a/stable_baselines3/common/cmd_util.py b/stable_baselines3/common/cmd_util.py index 5155104..ceeced7 100644 --- a/stable_baselines3/common/cmd_util.py +++ b/stable_baselines3/common/cmd_util.py @@ -94,8 +94,8 @@ def make_atari_env(env_id: Union[str, Type[gym.Env]], in a Monitor wrapper to provide additional information about training. :param wrapper_kwargs: (Dict[str, Any]) Optional keyword argument to pass to the ``AtariWrapper`` :param env_kwargs: (Dict[str, Any]) Optional keyword argument to pass to the env constructor - :param vec_env_cls: (Type[VecEnv]) A custom `VecEnv` class constructor. Default: None. - :param vec_env_kwargs: (Dict[str, Any]) Keyword arguments to pass to the `VecEnv` class constructor. + :param vec_env_cls: (Type[VecEnv]) A custom ``VecEnv`` class constructor. Default: None. + :param vec_env_kwargs: (Dict[str, Any]) Keyword arguments to pass to the ``VecEnv`` class constructor. :return: (VecEnv) The wrapped environment """ if wrapper_kwargs is None: diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 6794971..5b75e44 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -1,8 +1,7 @@ import time -import os import pickle import warnings -from typing import Union, Type, Optional, Dict, Any, Callable +from typing import Union, Type, Optional, Dict, Any, Callable, List, Tuple import gym import torch as th @@ -13,7 +12,7 @@ from stable_baselines3.common.base_class import BaseAlgorithm from stable_baselines3.common.policies import BasePolicy from stable_baselines3.common.utils import safe_mean from stable_baselines3.common.vec_env import VecEnv -from stable_baselines3.common.type_aliases import GymEnv, RolloutReturn +from stable_baselines3.common.type_aliases import GymEnv, RolloutReturn, MaybeCallback from stable_baselines3.common.callbacks import BaseCallback from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.buffers import ReplayBuffer @@ -32,6 +31,17 @@ class OffPolicyAlgorithm(BaseAlgorithm): :param buffer_size: (int) size of the replay buffer :param learning_starts: (int) how many steps of the model to collect transitions for before learning starts :param batch_size: (int) Minibatch size for each gradient update + :param tau: (float) the soft update coefficient ("Polyak update", between 0 and 1) + :param gamma: (float) the discount factor + :param train_freq: (int) Update the model every ``train_freq`` steps. + :param gradient_steps: (int) How many gradient update after each step + :param n_episodes_rollout: (int) Update the model every ``n_episodes_rollout`` episodes. + Note that this cannot be used at the same time as ``train_freq`` + :param action_noise: (ActionNoise) the action noise type (None by default), this can help + for hard exploration problem. Cf common.noise for the different action noise type. + :param optimize_memory_usage: (bool) Enable a memory efficient variant of the replay buffer + at a cost of more complexity. + See https://github.com/DLR-RM/stable-baselines3/issues/37#issuecomment-637501195 :param policy_kwargs: Additional arguments to be passed to the policy on creation :param tensorboard_log: (str) the log location for tensorboard (if None, no logging) :param verbose: The verbosity level: 0 none, 1 training information, 2 debug @@ -62,6 +72,13 @@ class OffPolicyAlgorithm(BaseAlgorithm): buffer_size: int = int(1e6), learning_starts: int = 100, batch_size: int = 256, + tau: float = 0.005, + gamma: float = 0.99, + train_freq: int = 1, + gradient_steps: int = 1, + n_episodes_rollout: int = -1, + action_noise: Optional[ActionNoise] = None, + optimize_memory_usage: bool = False, policy_kwargs: Dict[str, Any] = None, tensorboard_log: Optional[str] = None, verbose: int = 0, @@ -84,6 +101,22 @@ class OffPolicyAlgorithm(BaseAlgorithm): self.buffer_size = buffer_size self.batch_size = batch_size self.learning_starts = learning_starts + self.tau = tau + self.gamma = gamma + self.train_freq = train_freq + self.gradient_steps = gradient_steps + self.n_episodes_rollout = n_episodes_rollout + self.action_noise = action_noise + self.optimize_memory_usage = optimize_memory_usage + + if train_freq > 0 and n_episodes_rollout > 0: + warnings.warn("You passed a positive value for `train_freq` and `n_episodes_rollout`." + "Please make sure this is intended. " + "The agent will collect data by stepping in the environment " + "until both conditions are true: " + "`number of steps in the env` >= `train_freq` and " + "`number of episodes` > `n_episodes_rollout`") + self.actor = None # type: Optional[th.nn.Module] self.replay_buffer = None # type: Optional[ReplayBuffer] # Update policy keyword arguments @@ -97,7 +130,8 @@ class OffPolicyAlgorithm(BaseAlgorithm): self._setup_lr_schedule() self.set_random_seed(self.seed) self.replay_buffer = ReplayBuffer(self.buffer_size, self.observation_space, - self.action_space, self.device) + self.action_space, self.device, + optimize_memory_usage=self.optimize_memory_usage) self.policy = self.policy_class(self.observation_space, self.action_space, self.lr_schedule, **self.policy_kwargs) self.policy = self.policy.to(self.device) @@ -122,10 +156,158 @@ class OffPolicyAlgorithm(BaseAlgorithm): self.replay_buffer = pickle.load(file_handler) assert isinstance(self.replay_buffer, ReplayBuffer), 'The replay buffer must inherit from ReplayBuffer class' - def collect_rollouts(self, # noqa: C901 + 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`. + """ + # Prevent continuity issue by truncating trajectory + # when using memory efficient replay buffer + # see https://github.com/DLR-RM/stable-baselines3/issues/46 + truncate_last_traj = (self.optimize_memory_usage and reset_num_timesteps + and self.replay_buffer is not None + and (self.replay_buffer.full or self.replay_buffer.pos > 0)) + + if truncate_last_traj: + warnings.warn("The last trajectory in the replay buffer will be truncated, " + "see https://github.com/DLR-RM/stable-baselines3/issues/46." + "You should use `reset_num_timesteps=False` or `optimize_memory_usage=False`" + "to avoid that issue.") + # Go to the previous index + pos = (self.replay_buffer.pos - 1) % self.replay_buffer.buffer_size + self.replay_buffer.dones[pos] = True + + return super()._setup_learn(total_timesteps, eval_env, callback, eval_freq, + n_eval_episodes, log_path, reset_num_timesteps, tb_log_name) + + def learn(self, + total_timesteps: int, + callback: MaybeCallback = None, + log_interval: int = 4, + eval_env: Optional[GymEnv] = None, + eval_freq: int = -1, + n_eval_episodes: int = 5, + tb_log_name: str = "run", + eval_log_path: Optional[str] = None, + reset_num_timesteps: bool = True) -> 'OffPolicyAlgorithm': + + total_timesteps, callback = self._setup_learn(total_timesteps, eval_env, callback, eval_freq, + n_eval_episodes, eval_log_path, reset_num_timesteps, + tb_log_name) + + callback.on_training_start(locals(), globals()) + + while self.num_timesteps < total_timesteps: + + rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout, + n_steps=self.train_freq, action_noise=self.action_noise, + callback=callback, + learning_starts=self.learning_starts, + replay_buffer=self.replay_buffer, + log_interval=log_interval) + + if rollout.continue_training is False: + break + + if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: + # If no `gradient_steps` is specified, + # do as many gradients steps as steps performed during the rollout + gradient_steps = self.gradient_steps if self.gradient_steps > 0 else rollout.episode_timesteps + self.train(batch_size=self.batch_size, gradient_steps=gradient_steps) + + callback.on_training_end() + + return self + + def train(self, gradient_steps: int, batch_size: int) -> None: + """ + Sample the replay buffer and do the updates + (gradient descent and update target networks) + """ + raise NotImplementedError() + + def _sample_action(self, learning_starts: int, + action_noise: Optional[ActionNoise] = None) -> Tuple[np.ndarray, np.ndarray]: + """ + Sample an action according to the exploration policy. + This is either done by sampling the probability distribution of the policy, + or sampling a random action (from a uniform distribution over the action space) + or by adding noise to the deterministic output. + + :param action_noise: (Optional[ActionNoise]) Action noise that will be used for exploration + Required for deterministic policy (e.g. TD3). This can also be used + in addition to the stochastic policy for SAC. + :param learning_starts: (int) Number of steps before learning for the warm-up phase. + :return: (Tuple[np.ndarray, np.ndarray]) action to take in the environment + and scaled action that will be stored in the replay buffer. + The two differs when the action space is not normalized (bounds are not [-1, 1]). + """ + # Select action randomly or according to policy + if self.num_timesteps < learning_starts and not (self.use_sde and self.use_sde_at_warmup): + # Warmup phase + unscaled_action = np.array([self.action_space.sample()]) + else: + # Note: when using continuous actions, + # we assume that the policy uses tanh to scale the action + # We use non-deterministic action in the case of SAC, for TD3, it does not matter + unscaled_action, _ = self.predict(self._last_obs, deterministic=False) + + # Rescale the action from [low, high] to [-1, 1] + if isinstance(self.action_space, gym.spaces.Box): + scaled_action = self.policy.scale_action(unscaled_action) + + # Add noise to the action (improve exploration) + if action_noise is not None: + scaled_action = np.clip(scaled_action + action_noise(), -1, 1) + + # We store the scaled action in the buffer + buffer_action = scaled_action + action = self.policy.unscale_action(scaled_action) + else: + # Discrete case, no need to normalize or clip + buffer_action = unscaled_action + action = buffer_action + return action, buffer_action + + def _dump_logs(self) -> None: + """ + Write log. + """ + fps = int(self.num_timesteps / (time.time() - self.start_time)) + logger.record("time/episodes", self._episode_num, exclude="tensorboard") + if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0: + logger.record('rollout/ep_rew_mean', safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer])) + logger.record('rollout/ep_len_mean', safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer])) + logger.record("time/fps", fps) + logger.record('time/time_elapsed', int(time.time() - self.start_time), exclude="tensorboard") + logger.record("time/total timesteps", self.num_timesteps, exclude="tensorboard") + if self.use_sde: + logger.record("train/std", (self.actor.get_std()).mean().item()) + + if len(self.ep_success_buffer) > 0: + logger.record('rollout/success rate', safe_mean(self.ep_success_buffer)) + # Pass the number of timesteps for tensorboard + logger.dump(step=self.num_timesteps) + + def _on_step(self) -> None: + """ + Method called after each step in the environment. + It is meant to trigger DQN target network update + but can be used for other purposes + """ + pass + + def collect_rollouts(self, env: VecEnv, - # Type hint as string to avoid circular import - callback: 'BaseCallback', + callback: BaseCallback, n_episodes: int = 1, n_steps: int = -1, action_noise: Optional[ActionNoise] = None, @@ -156,17 +338,6 @@ class OffPolicyAlgorithm(BaseAlgorithm): assert isinstance(env, VecEnv), "You must pass a VecEnv" assert env.num_envs == 1, "OffPolicyAlgorithm only support single environment" - if n_episodes > 0 and n_steps > 0: - # Note we are refering to the constructor arguments - # that are named `train_freq` and `n_episodes_rollout` - # but correspond to `n_steps` and `n_episodes` here - warnings.warn("You passed a positive value for `train_freq` and `n_episodes_rollout`." - "Please make sure this is intended. " - "The agent will collect data by stepping in the environment " - "until both conditions are true: " - "`number of steps in the env` >= `train_freq` and " - "`number of episodes` > `n_episodes_rollout`") - if self.use_sde: self.actor.reset_noise() @@ -184,31 +355,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): self.actor.reset_noise() # Select action randomly or according to policy - if self.num_timesteps < learning_starts and not (self.use_sde and self.use_sde_at_warmup): - # Warmup phase - unscaled_action = np.array([self.action_space.sample()]) - else: - # Note: we assume that the policy uses tanh to scale the action - # We use non-deterministic action in the case of SAC, for TD3, it does not matter - unscaled_action, _ = self.predict(self._last_obs, deterministic=False) - - # Rescale the action from [low, high] to [-1, 1] - if isinstance(self.action_space, gym.spaces.Box): - scaled_action = self.policy.scale_action(unscaled_action) - - # Add noise to the action (improve exploration) - if action_noise is not None: - # NOTE: in the original implementation of TD3, the noise was applied to the unscaled action - # Update(October 2019): Not anymore - scaled_action = np.clip(scaled_action + action_noise(), -1, 1) - - # We store the scaled action in the buffer - buffer_action = scaled_action - action = self.policy.unscale_action(scaled_action) - else: - # Discrete case, no need to normalize or clip - buffer_action = unscaled_action - action = buffer_action + action, buffer_action = self._sample_action(learning_starts, action_noise) # Rescale and perform action new_obs, reward, done, infos = env.step(action) @@ -242,6 +389,14 @@ class OffPolicyAlgorithm(BaseAlgorithm): self.num_timesteps += 1 episode_timesteps += 1 total_steps += 1 + self._update_current_progress_remaining(self.num_timesteps, self._total_timesteps) + + # For DQN, check if the target network should be updated + # and update the exploration schedule + # For SAC/TD3, the update is done as the same time as the gradient update + # see https://github.com/hill-a/stable-baselines/issues/900 + self._on_step() + if 0 < n_steps <= total_steps: break @@ -256,21 +411,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): # Log training infos if log_interval is not None and self._episode_num % log_interval == 0: - fps = int(self.num_timesteps / (time.time() - self.start_time)) - logger.record("time/episodes", self._episode_num, exclude="tensorboard") - if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0: - logger.record('rollout/ep_rew_mean', safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer])) - logger.record('rollout/ep_len_mean', safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer])) - logger.record("time/fps", fps) - logger.record('time/time_elapsed', int(time.time() - self.start_time), exclude="tensorboard") - logger.record("time/total timesteps", self.num_timesteps, exclude="tensorboard") - if self.use_sde: - logger.record("train/std", (self.actor.get_std()).mean().item()) - - if len(self.ep_success_buffer) > 0: - logger.record('rollout/success rate', safe_mean(self.ep_success_buffer)) - # Pass the number of timesteps for tensorboard - logger.dump(step=self.num_timesteps) + self._dump_logs() mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0 diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 1bf975c..fbe3572 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -89,6 +89,30 @@ def get_schedule_fn(value_schedule: Union[Callable, float]) -> Callable: return value_schedule +def get_linear_fn(start: float, end: float, end_fraction: float) -> Callable: + """ + Create a function that interpolates linearly between start and end + between ``progress_remaining`` = 1 and ``progress_remaining`` = ``end_fraction``. + This is used in DQN for linearly annealing the exploration fraction + (epsilon for the epsilon-greedy strategy). + + :params start: (float) value to start with if ``progress_remaining`` = 1 + :params end: (float) value to end with if ``progress_remaining`` = 0 + :params end_fraction: (float) fraction of ``progress_remaining`` + where end is reached e.g 0.1 then end is reached after 10% + of the complete training process. + :return: (Callable) + """ + + def func(progress_remaining: float) -> float: + if (1 - progress_remaining) > end_fraction: + return end + else: + return start + (1 - progress_remaining) * (end - start) / end_fraction + + return func + + def constant_fn(val: float) -> Callable: """ Create a function that returns a constant diff --git a/stable_baselines3/dqn/__init__.py b/stable_baselines3/dqn/__init__.py new file mode 100644 index 0000000..8f69683 --- /dev/null +++ b/stable_baselines3/dqn/__init__.py @@ -0,0 +1,3 @@ +from stable_baselines3.dqn.dqn import DQN +from stable_baselines3.dqn.policies import MlpPolicy +from stable_baselines3.dqn.policies import CnnPolicy diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py new file mode 100644 index 0000000..084824e --- /dev/null +++ b/stable_baselines3/dqn/dqn.py @@ -0,0 +1,219 @@ +from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any + +import numpy as np +import torch as th +import torch.nn.functional as F + +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.utils import get_linear_fn +from stable_baselines3.dqn.policies import DQNPolicy + + +class DQN(OffPolicyAlgorithm): + """ + Deep Q-Network (DQN) + + Paper: https://arxiv.org/abs/1312.5602, https://www.nature.com/articles/nature14236 + Default hyperparameters are taken from the nature paper, + except for the optimizer and learning rate that were taken from Stable Baselines defaults. + + :param policy: (DQNPolicy or str) The policy model to use (MlpPolicy, CnnPolicy, ...) + :param env: (Gym environment or str) The environment to learn from (if registered in Gym, can be str) + :param learning_rate: (float or callable) The learning rate, it can be a function + of the current progress (from 1 to 0) + :param buffer_size: (int) size of the replay buffer + :param learning_starts: (int) how many steps of the model to collect transitions for before learning starts + :param batch_size: (int) Minibatch size for each gradient update + :param tau: (float) the soft update coefficient ("Polyak update", between 0 and 1) default 1 for hard update + :param gamma: (float) the discount factor + :param train_freq: (int) Update the model every ``train_freq`` steps. + :param gradient_steps: (int) How many gradient update after each step + :param n_episodes_rollout: (int) Update the model every ``n_episodes_rollout`` episodes. + Note that this cannot be used at the same time as ``train_freq`` + :param optimize_memory_usage: (bool) Enable a memory efficient variant of the replay buffer + at a cost of more complexity. + See https://github.com/DLR-RM/stable-baselines3/issues/37#issuecomment-637501195 + :param target_update_interval: (int) update the target network every ``target_update_interval`` + environment steps. + :param exploration_fraction: (float) fraction of entire training period over which the exploration rate is reduced + :param exploration_initial_eps: (float) initial value of random action probability + :param exploration_final_eps: (float) final value of random action probability + :param max_grad_norm: (float) The maximum value for the gradient clipping + :param tensorboard_log: (str) the log location for tensorboard (if None, no logging) + :param create_eval_env: (bool) Whether to create a second environment that will be + used for evaluating the agent periodically. (Only available when passing string for the environment) + :param policy_kwargs: (dict) additional arguments to be passed to the policy on creation + :param verbose: (int) the verbosity level: 0 no output, 1 info, 2 debug + :param seed: (int) Seed for the pseudo random generators + :param device: (str or th.device) Device (cpu, cuda, ...) on which the code should be run. + Setting it to auto, the code will be run on the GPU if possible. + :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance + """ + + def __init__(self, policy: Union[str, Type[DQNPolicy]], + env: Union[GymEnv, str], + learning_rate: Union[float, Callable] = 1e-4, + buffer_size: int = 1000000, + learning_starts: int = 50000, + batch_size: Optional[int] = 32, + tau: float = 1.0, + gamma: float = 0.99, + train_freq: int = 4, + gradient_steps: int = 1, + n_episodes_rollout: int = -1, + optimize_memory_usage: bool = False, + target_update_interval: int = 10000, + exploration_fraction: float = 0.1, + exploration_initial_eps: float = 1.0, + exploration_final_eps: float = 0.05, + max_grad_norm: float = 10, + tensorboard_log: Optional[str] = None, + create_eval_env: bool = False, + policy_kwargs: Optional[Dict[str, Any]] = None, + verbose: int = 0, + seed: Optional[int] = None, + device: Union[th.device, str] = 'auto', + _init_setup_model: bool = True): + + super(DQN, self).__init__(policy, env, DQNPolicy, learning_rate, + 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, + tensorboard_log=tensorboard_log, + verbose=verbose, device=device, + create_eval_env=create_eval_env, + seed=seed, sde_support=False, + optimize_memory_usage=optimize_memory_usage) + + self.exploration_initial_eps = exploration_initial_eps + self.exploration_final_eps = exploration_final_eps + self.exploration_fraction = exploration_fraction + self.target_update_interval = target_update_interval + self.max_grad_norm = max_grad_norm + # "epsilon" for the epsilon-greedy exploration + self.exploration_rate = 0.0 + # Linear schedule will be defined in `_setup_model()` + self.exploration_schedule = None + self.q_net, self.q_net_target = None, None + + if _init_setup_model: + self._setup_model() + + def _setup_model(self) -> None: + super(DQN, self)._setup_model() + self._create_aliases() + self.exploration_schedule = get_linear_fn(self.exploration_initial_eps, self.exploration_final_eps, + self.exploration_fraction) + + def _create_aliases(self) -> None: + self.q_net = self.policy.q_net + self.q_net_target = self.policy.q_net_target + + def _on_step(self): + """ + Update the exploration rate and target network if needed. + This method is called in ``collect_rollout()`` after each step in the environment. + """ + if self.num_timesteps % self.target_update_interval == 0: + for param, target_param in zip(self.q_net.parameters(), self.q_net_target.parameters()): + target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) + + self.exploration_rate = self.exploration_schedule(self._current_progress_remaining) + logger.record("rollout/exploration rate", self.exploration_rate) + + def train(self, gradient_steps: int, batch_size: int = 100) -> None: + # Update learning rate according to schedule + self._update_learning_rate(self.policy.optimizer) + + for gradient_step in range(gradient_steps): + # Sample replay buffer + replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env) + + with th.no_grad(): + # Compute the target Q values + target_q = self.q_net_target(replay_data.next_observations) + # Follow greedy policy: use the one with the highest value + target_q, _ = target_q.max(dim=1) + # Avoid potential broadcast issue + target_q = target_q.reshape(-1, 1) + # 1-step TD target + target_q = replay_data.rewards + (1 - replay_data.dones) * self.gamma * target_q + + # Get current Q estimates + current_q = self.q_net(replay_data.observations) + + # Retrieve the q-values for the actions from the replay buffer + current_q = th.gather(current_q, dim=1, index=replay_data.actions.long()) + + # Compute Huber loss (less sensitive to outliers) + loss = F.smooth_l1_loss(current_q, target_q) + + # Optimize the policy + self.policy.optimizer.zero_grad() + loss.backward() + # Clip gradient norm + th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) + self.policy.optimizer.step() + + # Increase update counter + self._n_updates += gradient_steps + + logger.record("train/n_updates", self._n_updates, exclude='tensorboard') + + def predict(self, observation: np.ndarray, + state: Optional[np.ndarray] = None, + mask: Optional[np.ndarray] = None, + deterministic: bool = False) -> Tuple[np.ndarray, Optional[np.ndarray]]: + """ + Overrides the base_class predict function to include epsilon-greedy exploration. + + :param observation: (np.ndarray) the input observation + :param state: (Optional[np.ndarray]) The last states (can be None, used in recurrent policies) + :param mask: (Optional[np.ndarray]) The last masks (can be None, used in recurrent policies) + :param deterministic: (bool) Whether or not to return deterministic actions. + :return: (Tuple[np.ndarray, Optional[np.ndarray]]) the model's action and the next state + (used in recurrent policies) + """ + if not deterministic and np.random.rand() < self.exploration_rate: + n_batch = observation.shape[0] + action = np.array([self.action_space.sample() for _ in range(n_batch)]) + else: + action, state = self.policy.predict(observation, state, mask, deterministic) + return action, state + + def learn(self, + total_timesteps: int, + callback: MaybeCallback = None, + log_interval: int = 4, + eval_env: Optional[GymEnv] = None, + eval_freq: int = -1, + n_eval_episodes: int = 5, + tb_log_name: str = "DQN", + eval_log_path: Optional[str] = None, + reset_num_timesteps: bool = True) -> OffPolicyAlgorithm: + + return super(DQN, self).learn(total_timesteps=total_timesteps, callback=callback, log_interval=log_interval, + eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes, + tb_log_name=tb_log_name, eval_log_path=eval_log_path, + reset_num_timesteps=reset_num_timesteps) + + def excluded_save_params(self) -> List[str]: + """ + Returns the names of the parameters that should be excluded by default + when saving the model. + + :return: (List[str]) List of parameters that should be excluded from save + """ + # Exclude aliases + return super(DQN, self).excluded_save_params() + ["q_net", "q_net_target"] + + def get_torch_variables(self) -> Tuple[List[str], List[str]]: + """ + cf base class + """ + state_dicts = ["policy", "policy.optimizer"] + + return state_dicts, [] diff --git a/stable_baselines3/dqn/policies.py b/stable_baselines3/dqn/policies.py new file mode 100644 index 0000000..7184522 --- /dev/null +++ b/stable_baselines3/dqn/policies.py @@ -0,0 +1,229 @@ +from typing import Optional, List, Callable, Union, Type, Any, Dict + +import gym +import torch as th +import torch.nn as nn +from stable_baselines3.common.policies import BasePolicy, register_policy +from stable_baselines3.common.torch_layers import create_mlp, NatureCNN, BaseFeaturesExtractor, FlattenExtractor + + +class QNetwork(BasePolicy): + """ + Action-Value (Q-Value) network for DQN + + :param observation_space: (gym.spaces.Space) Observation space + :param action_space: (gym.spaces.Space) Action space + :param net_arch: (Optional[List[int]]) The specification of the policy and value networks. + :param device: (str or th.device) Device on which the code should run. + :param activation_fn: (Type[nn.Module]) Activation function + :param normalize_images: (bool) Whether to normalize images or not, + dividing by 255.0 (True by default) + """ + + def __init__(self, observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + features_extractor: nn.Module, + features_dim: int, + net_arch: Optional[List[int]] = None, + device: Union[th.device, str] = 'auto', + activation_fn: Type[nn.Module] = nn.ReLU, + normalize_images: bool = True): + super(QNetwork, self).__init__(observation_space, action_space, + features_extractor=features_extractor, + normalize_images=normalize_images, + device=device) + + if net_arch is None: + net_arch = [64, 64] + + self.net_arch = net_arch + self.activation_fn = activation_fn + self.features_extractor = features_extractor + self.features_dim = features_dim + self.normalize_images = normalize_images + action_dim = self.action_space.n # number of actions + q_net = create_mlp(self.features_dim, action_dim, self.net_arch, self.activation_fn) + self.q_net = nn.Sequential(*q_net) + + def forward(self, obs: th.Tensor) -> th.Tensor: + """ + Predict the q-values. + + :param obs: (th.Tensor) Observation + :return: (th.Tensor) The estimated Q-Value for each action. + """ + return self.q_net(self.extract_features(obs)) + + def _predict(self, observation: th.Tensor, deterministic: bool = True) -> th.Tensor: + q_values = self.forward(observation) + # Greedy action + action = q_values.argmax(dim=1).reshape(-1) + return action + + def _get_data(self) -> Dict[str, Any]: + data = super()._get_data() + + data.update(dict( + net_arch=self.net_arch, + features_dim=self.features_dim, + activation_fn=self.activation_fn, + features_extractor=self.features_extractor, + epsilon=self.epsilon, + )) + return data + + +class DQNPolicy(BasePolicy): + """ + Policy class with Q-Value Net and target net for DQN + + :param observation_space: (gym.spaces.Space) Observation space + :param action_space: (gym.spaces.Space) Action space + :param lr_schedule: (callable) Learning rate schedule (could be constant) + :param net_arch: (Optional[List[int]]) The specification of the policy and value networks. + :param device: (str or th.device) Device on which the code should run. + :param activation_fn: (Type[nn.Module]) Activation function + :param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use. + :param features_extractor_kwargs: (Optional[Dict[str, Any]]) Keyword arguments + to pass to the feature extractor. + :param normalize_images: (bool) Whether to normalize images or not, + dividing by 255.0 (True by default) + :param optimizer_class: (Type[th.optim.Optimizer]) The optimizer to use, + ``th.optim.Adam`` by default + :param optimizer_kwargs: (Optional[Dict[str, Any]]) Additional keyword arguments, + excluding the learning rate, to pass to the optimizer + """ + + def __init__(self, observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + lr_schedule: Callable, + net_arch: Optional[List[int]] = None, + device: Union[th.device, str] = 'auto', + activation_fn: Type[nn.Module] = nn.ReLU, + features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor, + features_extractor_kwargs: Optional[Dict[str, Any]] = None, + normalize_images: bool = True, + optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam, + optimizer_kwargs: Optional[Dict[str, Any]] = None): + super(DQNPolicy, self).__init__(observation_space, action_space, + device, + features_extractor_class, + features_extractor_kwargs, + optimizer_class=optimizer_class, + optimizer_kwargs=optimizer_kwargs) + + if net_arch is None: + if features_extractor_class == FlattenExtractor: + net_arch = [64, 64] + else: + net_arch = [] + + self.features_extractor = features_extractor_class(self.observation_space, + **self.features_extractor_kwargs) + self.features_dim = self.features_extractor.features_dim + self.net_arch = net_arch + self.activation_fn = activation_fn + self.normalize_images = normalize_images + + self.net_args = { + 'observation_space': self.observation_space, + 'action_space': self.action_space, + 'features_extractor': self.features_extractor, + 'features_dim': self.features_dim, + 'net_arch': self.net_arch, + 'activation_fn': self.activation_fn, + 'normalize_images': normalize_images, + 'device': device + } + + self.q_net, self.q_net_target = None, None + self._build(lr_schedule) + + def _build(self, lr_schedule: Callable) -> None: + """ + Create the network and the optimizer. + + :param lr_schedule: (Callable) Learning rate schedule + lr_schedule(1) is the initial learning rate + """ + + self.q_net = self.make_q_net() + self.q_net_target = self.make_q_net() + self.q_net_target.load_state_dict(self.q_net.state_dict()) + + # Setup optimizer with initial learning rate + self.optimizer = self.optimizer_class(self.parameters(), lr=lr_schedule(1), + **self.optimizer_kwargs) + + def make_q_net(self) -> QNetwork: + return QNetwork(**self.net_args).to(self.device) + + def forward(self, obs: th.Tensor, deterministic: bool = True) -> th.Tensor: + return self._predict(obs, deterministic=deterministic) + + def _predict(self, obs: th.Tensor, deterministic: bool = True) -> th.Tensor: + return self.q_net._predict(obs, deterministic=deterministic) + + def _get_data(self) -> Dict[str, Any]: + data = super()._get_data() + + data.update(dict( + net_arch=self.net_args['net_arch'], + activation_fn=self.net_args['activation_fn'], + lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone + optimizer_class=self.optimizer_class, + optimizer_kwargs=self.optimizer_kwargs, + features_extractor_class=self.features_extractor_class, + features_extractor_kwargs=self.features_extractor_kwargs + )) + return data + + +MlpPolicy = DQNPolicy + + +class CnnPolicy(DQNPolicy): + """ + Policy class for DQN when using images as input. + + :param observation_space: (gym.spaces.Space) Observation space + :param action_space: (gym.spaces.Space) Action space + :param lr_schedule: (callable) Learning rate schedule (could be constant) + :param net_arch: (Optional[List[int]]) The specification of the policy and value networks. + :param device: (str or th.device) Device on which the code should run. + :param activation_fn: (Type[nn.Module]) Activation function + :param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use. + :param normalize_images: (bool) Whether to normalize images or not, + dividing by 255.0 (True by default) + :param optimizer_class: (Type[th.optim.Optimizer]) The optimizer to use, + ``th.optim.Adam`` by default + :param optimizer_kwargs: (Optional[Dict[str, Any]]) Additional keyword arguments, + excluding the learning rate, to pass to the optimizer + """ + + def __init__(self, observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + lr_schedule: Callable, + net_arch: Optional[List[int]] = None, + device: Union[th.device, str] = 'auto', + activation_fn: Type[nn.Module] = nn.ReLU, + features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN, + features_extractor_kwargs: Optional[Dict[str, Any]] = None, + normalize_images: bool = True, + optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam, + optimizer_kwargs: Optional[Dict[str, Any]] = None): + super(CnnPolicy, self).__init__(observation_space, + action_space, + lr_schedule, + net_arch, + device, + activation_fn, + features_extractor_class, + features_extractor_kwargs, + normalize_images, + optimizer_class, + optimizer_kwargs) + + +register_policy("MlpPolicy", MlpPolicy) +register_policy("CnnPolicy", CnnPolicy) diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 511e5a4..99f8888 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -40,10 +40,14 @@ class SAC(OffPolicyAlgorithm): Note that this cannot be used at the same time as ``train_freq`` :param action_noise: (ActionNoise) the action noise type (None by default), this can help for hard exploration problem. Cf common.noise for the different action noise type. + :param optimize_memory_usage: (bool) Enable a memory efficient variant of the replay buffer + at a cost of more complexity. + See https://github.com/DLR-RM/stable-baselines3/issues/37#issuecomment-637501195 :param ent_coef: (str or float) Entropy regularization coefficient. (Equivalent to inverse of reward scale in the original SAC paper.) Controlling exploration/exploitation trade-off. Set it to 'auto' to learn it automatically (and 'auto_0.1' for using 0.1 as initial value) - :param target_update_interval: (int) update the target network every ``target_network_update_freq`` steps. + :param target_update_interval: (int) update the target network every ``target_network_update_freq`` + gradient steps. :param target_entropy: (str or float) target entropy when learning ``ent_coef`` (``ent_coef = 'auto'``) :param use_sde: (bool) Whether to use generalized State Dependent Exploration (gSDE) instead of action noise exploration (default: False) @@ -73,6 +77,7 @@ class SAC(OffPolicyAlgorithm): gradient_steps: int = 1, n_episodes_rollout: int = -1, action_noise: Optional[ActionNoise] = None, + optimize_memory_usage: bool = False, ent_coef: Union[str, float] = 'auto', target_update_interval: int = 1, target_entropy: Union[str, float] = 'auto', @@ -89,24 +94,22 @@ class SAC(OffPolicyAlgorithm): super(SAC, self).__init__(policy, env, SACPolicy, learning_rate, buffer_size, learning_starts, batch_size, - policy_kwargs, tensorboard_log, verbose, device, + tau, gamma, train_freq, gradient_steps, + n_episodes_rollout, action_noise, + policy_kwargs=policy_kwargs, + tensorboard_log=tensorboard_log, + verbose=verbose, device=device, create_eval_env=create_eval_env, seed=seed, use_sde=use_sde, sde_sample_freq=sde_sample_freq, - use_sde_at_warmup=use_sde_at_warmup) + use_sde_at_warmup=use_sde_at_warmup, + optimize_memory_usage=optimize_memory_usage) self.target_entropy = target_entropy self.log_ent_coef = None # type: Optional[th.Tensor] - self.target_update_interval = target_update_interval - self.tau = tau # Entropy coefficient / Entropy temperature # Inverse of the reward scale self.ent_coef = ent_coef self.target_update_interval = target_update_interval - self.train_freq = train_freq - self.gradient_steps = gradient_steps - self.n_episodes_rollout = n_episodes_rollout - self.action_noise = action_noise - self.gamma = gamma self.ent_coef_optimizer = None if _init_setup_model: @@ -254,30 +257,10 @@ class SAC(OffPolicyAlgorithm): eval_log_path: Optional[str] = None, reset_num_timesteps: bool = True) -> OffPolicyAlgorithm: - total_timesteps, callback = self._setup_learn(total_timesteps, eval_env, callback, eval_freq, - n_eval_episodes, eval_log_path, reset_num_timesteps, - tb_log_name) - callback.on_training_start(locals(), globals()) - - while self.num_timesteps < total_timesteps: - rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout, - n_steps=self.train_freq, action_noise=self.action_noise, - callback=callback, - learning_starts=self.learning_starts, - replay_buffer=self.replay_buffer, - log_interval=log_interval) - - if rollout.continue_training is False: - break - - self._update_current_progress_remaining(self.num_timesteps, total_timesteps) - - if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: - gradient_steps = self.gradient_steps if self.gradient_steps > 0 else rollout.episode_timesteps - self.train(gradient_steps, batch_size=self.batch_size) - - callback.on_training_end() - return self + return super(SAC, self).learn(total_timesteps=total_timesteps, callback=callback, log_interval=log_interval, + eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes, + tb_log_name=tb_log_name, eval_log_path=eval_log_path, + reset_num_timesteps=reset_num_timesteps) def excluded_save_params(self) -> List[str]: """ diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index 2ee57b7..459e240 100644 --- a/stable_baselines3/td3/td3.py +++ b/stable_baselines3/td3/td3.py @@ -34,6 +34,9 @@ class TD3(OffPolicyAlgorithm): Note that this cannot be used at the same time as ``train_freq`` :param action_noise: (ActionNoise) the action noise type (None by default), this can help for hard exploration problem. Cf common.noise for the different action noise type. + :param optimize_memory_usage: (bool) Enable a memory efficient variant of the replay buffer + at a cost of more complexity. + See https://github.com/DLR-RM/stable-baselines3/issues/37#issuecomment-637501195 :param policy_delay: (int) Policy and target networks will only be updated once every policy_delay steps per training steps. The Q values will be updated policy_delay more often (update every training step). :param target_policy_noise: (float) Standard deviation of Gaussian noise added to target policy @@ -61,6 +64,7 @@ class TD3(OffPolicyAlgorithm): gradient_steps: int = -1, n_episodes_rollout: int = 1, action_noise: Optional[ActionNoise] = None, + optimize_memory_usage: bool = False, policy_delay: int = 2, target_policy_noise: float = 0.2, target_noise_clip: float = 0.5, @@ -74,16 +78,14 @@ class TD3(OffPolicyAlgorithm): super(TD3, self).__init__(policy, env, TD3Policy, learning_rate, buffer_size, learning_starts, batch_size, - policy_kwargs, tensorboard_log, verbose, device, + tau, gamma, train_freq, gradient_steps, + n_episodes_rollout, action_noise=action_noise, + policy_kwargs=policy_kwargs, + tensorboard_log=tensorboard_log, + verbose=verbose, device=device, create_eval_env=create_eval_env, seed=seed, - sde_support=False) + sde_support=False, optimize_memory_usage=optimize_memory_usage) - self.train_freq = train_freq - self.gradient_steps = gradient_steps - self.n_episodes_rollout = n_episodes_rollout - self.tau = tau - self.gamma = gamma - self.action_noise = action_noise self.policy_delay = policy_delay self.target_noise_clip = target_noise_clip self.target_policy_noise = target_policy_noise @@ -101,7 +103,7 @@ class TD3(OffPolicyAlgorithm): self.critic = self.policy.critic self.critic_target = self.policy.critic_target - def train(self, gradient_steps: int, batch_size: int = 100, policy_delay: int = 2) -> None: + def train(self, gradient_steps: int, batch_size: int = 100) -> None: # Update learning rate according to lr schedule self._update_learning_rate([self.actor.optimizer, self.critic.optimizer]) @@ -134,7 +136,7 @@ class TD3(OffPolicyAlgorithm): self.critic.optimizer.step() # Delayed policy updates - if gradient_step % policy_delay == 0: + if gradient_step % self.policy_delay == 0: # Compute actor loss actor_loss = -self.critic.q1_forward(replay_data.observations, self.actor(replay_data.observations)).mean() @@ -165,32 +167,10 @@ class TD3(OffPolicyAlgorithm): eval_log_path: Optional[str] = None, reset_num_timesteps: bool = True) -> OffPolicyAlgorithm: - total_timesteps, callback = self._setup_learn(total_timesteps, eval_env, callback, eval_freq, - n_eval_episodes, eval_log_path, reset_num_timesteps, - tb_log_name) - callback.on_training_start(locals(), globals()) - - while self.num_timesteps < total_timesteps: - - rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout, - n_steps=self.train_freq, action_noise=self.action_noise, - callback=callback, - learning_starts=self.learning_starts, - replay_buffer=self.replay_buffer, - log_interval=log_interval) - - if rollout.continue_training is False: - break - - self._update_current_progress_remaining(self.num_timesteps, total_timesteps) - - if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: - gradient_steps = self.gradient_steps if self.gradient_steps > 0 else rollout.episode_timesteps - self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay) - - callback.on_training_end() - - return self + return super(TD3, self).learn(total_timesteps=total_timesteps, callback=callback, log_interval=log_interval, + eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes, + tb_log_name=tb_log_name, eval_log_path=eval_log_path, + reset_num_timesteps=reset_num_timesteps) def excluded_save_params(self) -> List[str]: """ diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 4c46bdd..6685a73 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -0.8.0a0 +0.8.0a1 diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index 2b9a39c..8395364 100644 --- a/tests/test_callbacks.py +++ b/tests/test_callbacks.py @@ -4,21 +4,24 @@ import shutil import pytest import gym -from stable_baselines3 import A2C, PPO, SAC, TD3 +from stable_baselines3 import A2C, PPO, SAC, TD3, DQN from stable_baselines3.common.callbacks import (CallbackList, CheckpointCallback, EvalCallback, EveryNTimesteps, StopTrainingOnRewardThreshold) -@pytest.mark.parametrize("model_class", [A2C, PPO, SAC, TD3]) -def test_callbacks(model_class): - log_folder = './logs/callbacks/' +@pytest.mark.parametrize("model_class", [A2C, PPO, SAC, TD3, DQN]) +def test_callbacks(tmp_path, model_class): + log_folder = tmp_path / 'logs/callbacks/' + + # Dyn only support discrete actions + env_name = select_env(model_class) # Create RL model # Small network for fast test - model = model_class('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[32])) + model = model_class('MlpPolicy', env_name, policy_kwargs=dict(net_arch=[32])) checkpoint_callback = CheckpointCallback(save_freq=1000, save_path=log_folder) - eval_env = gym.make('Pendulum-v0') + eval_env = gym.make(env_name) # Stop training if the performance is good enough callback_on_best = StopTrainingOnRewardThreshold(reward_threshold=-1200, verbose=1) @@ -42,3 +45,10 @@ def test_callbacks(model_class): model.learn(500, callback=lambda _locals, _globals: True) if os.path.exists(log_folder): shutil.rmtree(log_folder) + + +def select_env(model_class) -> str: + if model_class is DQN: + return 'CartPole-v0' + else: + return 'Pendulum-v0' diff --git a/tests/test_cnn.py b/tests/test_cnn.py index 9dc342d..8c17276 100644 --- a/tests/test_cnn.py +++ b/tests/test_cnn.py @@ -3,14 +3,13 @@ import os import numpy as np import pytest -from stable_baselines3 import A2C, PPO, SAC, TD3 +from stable_baselines3 import A2C, PPO, SAC, TD3, DQN from stable_baselines3.common.identity_env import FakeImageEnv -SAVE_PATH = './cnn_model.zip' - -@pytest.mark.parametrize('model_class', [A2C, PPO, SAC, TD3]) -def test_cnn(model_class): +@pytest.mark.parametrize('model_class', [A2C, PPO, SAC, TD3, DQN]) +def test_cnn(tmp_path, model_class): + SAVE_NAME = 'cnn_model.zip' # Fake grayscale with frameskip # Atari after preprocessing: 84x84x1, here we are using lower resolution # to check that the network handle it automatically @@ -29,12 +28,12 @@ def test_cnn(model_class): action, _ = model.predict(obs, deterministic=True) - model.save(SAVE_PATH) + model.save(tmp_path / SAVE_NAME) del model - model = model_class.load(SAVE_PATH) + model = model_class.load(tmp_path / SAVE_NAME) # Check that the prediction is the same assert np.allclose(action, model.predict(obs, deterministic=True)[0]) - os.remove(SAVE_PATH) + os.remove(str(tmp_path / SAVE_NAME)) diff --git a/tests/test_deterministic.py b/tests/test_deterministic.py index bf86de9..645c4d3 100644 --- a/tests/test_deterministic.py +++ b/tests/test_deterministic.py @@ -1,13 +1,13 @@ import pytest -from stable_baselines3 import A2C, PPO, SAC, TD3 +from stable_baselines3 import A2C, PPO, SAC, TD3, DQN from stable_baselines3.common.noise import NormalActionNoise N_STEPS_TRAINING = 3000 SEED = 0 -@pytest.mark.parametrize("algo", [A2C, PPO, SAC, TD3]) +@pytest.mark.parametrize("algo", [A2C, DQN, PPO, SAC, TD3]) def test_deterministic_training_common(algo): results = [[], []] rewards = [[], []] @@ -19,8 +19,8 @@ def test_deterministic_training_common(algo): 'learning_starts': 100}) else: env_id = 'CartPole-v1' - # if algo == DQN: - # kwargs.update({'learning_starts': 100}) + if algo == DQN: + kwargs.update({'learning_starts': 100}) for i in range(2): model = algo('MlpPolicy', env_id, seed=SEED, **kwargs) diff --git a/tests/test_identity.py b/tests/test_identity.py index b41b70c..f56622a 100644 --- a/tests/test_identity.py +++ b/tests/test_identity.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from stable_baselines3 import A2C, PPO, SAC, TD3 +from stable_baselines3 import A2C, PPO, SAC, TD3, DQN from stable_baselines3.common.identity_env import (IdentityEnvBox, IdentityEnv, IdentityEnvMultiBinary, IdentityEnvMultiDiscrete) @@ -9,17 +9,25 @@ from stable_baselines3.common.vec_env import DummyVecEnv from stable_baselines3.common.evaluation import evaluate_policy from stable_baselines3.common.noise import NormalActionNoise - DIM = 4 -@pytest.mark.parametrize("model_class", [A2C, PPO]) +@pytest.mark.parametrize("model_class", [A2C, PPO, DQN]) @pytest.mark.parametrize("env", [IdentityEnv(DIM), IdentityEnvMultiDiscrete(DIM), IdentityEnvMultiBinary(DIM)]) def test_discrete(model_class, env): - env = DummyVecEnv([lambda: env]) - model = model_class('MlpPolicy', env, gamma=0.5, seed=1).learn(3000) + env_ = DummyVecEnv([lambda: env]) + kwargs = {} + n_steps = 3000 + if model_class == DQN: + kwargs = dict(learning_starts=0) + n_steps = 4000 + # DQN only support discrete actions + if isinstance(env, (IdentityEnvMultiDiscrete, IdentityEnvMultiBinary)): + return - evaluate_policy(model, env, n_eval_episodes=20, reward_threshold=90) + model = model_class('MlpPolicy', env_, gamma=0.5, seed=1, **kwargs).learn(n_steps) + + evaluate_policy(model, env_, n_eval_episodes=20, reward_threshold=90) obs = env.reset() assert np.shape(model.predict(obs)[0]) == np.shape(obs) diff --git a/tests/test_predict.py b/tests/test_predict.py index fad35c7..c95e19f 100644 --- a/tests/test_predict.py +++ b/tests/test_predict.py @@ -1,7 +1,7 @@ import gym import pytest -from stable_baselines3 import A2C, PPO, SAC, TD3 +from stable_baselines3 import A2C, PPO, SAC, TD3, DQN from stable_baselines3.common.vec_env import DummyVecEnv MODEL_LIST = [ @@ -9,14 +9,21 @@ MODEL_LIST = [ A2C, TD3, SAC, + DQN, ] @pytest.mark.parametrize("model_class", MODEL_LIST) def test_auto_wrap(model_class): # test auto wrapping of env into a VecEnv - env = gym.make('Pendulum-v0') - eval_env = gym.make('Pendulum-v0') + + # Use different environment for DQN + if model_class is DQN: + env_name = 'CartPole-v0' + else: + env_name = 'Pendulum-v0' + env = gym.make(env_name) + eval_env = gym.make(env_name) model = model_class('MlpPolicy', env) model.learn(100, eval_env=eval_env) @@ -24,7 +31,10 @@ def test_auto_wrap(model_class): @pytest.mark.parametrize("model_class", MODEL_LIST) @pytest.mark.parametrize("env_id", ['Pendulum-v0', 'CartPole-v1']) def test_predict(model_class, env_id): - if env_id == 'CartPole-v1' and model_class not in [PPO, A2C]: + if env_id == 'CartPole-v1': + if model_class in [SAC, TD3]: + return + elif model_class in [DQN]: return # test detection of different shapes by the predict method diff --git a/tests/test_run.py b/tests/test_run.py index fc05988..5aa936a 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from stable_baselines3 import A2C, PPO, SAC, TD3 +from stable_baselines3 import A2C, PPO, SAC, TD3, DQN from stable_baselines3.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise normal_action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1)) @@ -40,3 +40,9 @@ def test_sac(ent_coef): learning_starts=100, verbose=1, create_eval_env=True, ent_coef=ent_coef, action_noise=NormalActionNoise(np.zeros(1), np.zeros(1))) model.learn(total_timesteps=1000, eval_freq=500) + + +def test_dqn(): + model = DQN('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=[64, 64]), + learning_starts=500, buffer_size=500, learning_rate=3e-4, verbose=1, create_eval_env=True) + model.learn(total_timesteps=1000, eval_freq=500) diff --git a/tests/test_save_load.py b/tests/test_save_load.py index dec3ce6..d107115 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -1,26 +1,39 @@ import os +import warnings from copy import deepcopy import pytest +import gym import numpy as np import torch as th -from stable_baselines3 import A2C, PPO, SAC, TD3 -from stable_baselines3.common.identity_env import IdentityEnvBox +from stable_baselines3 import A2C, PPO, SAC, TD3, DQN +from stable_baselines3.common.base_class import BaseAlgorithm +from stable_baselines3.common.identity_env import IdentityEnvBox, IdentityEnv from stable_baselines3.common.vec_env import DummyVecEnv from stable_baselines3.common.identity_env import FakeImageEnv - MODEL_LIST = [ PPO, A2C, TD3, SAC, + DQN, ] +def select_env(model_class: BaseAlgorithm) -> gym.Env: + """ + Selects an environment with the correct action space as DQN only supports discrete action space + """ + if model_class == DQN: + return IdentityEnv(10) + else: + return IdentityEnvBox(10) + + @pytest.mark.parametrize("model_class", MODEL_LIST) -def test_save_load(model_class): +def test_save_load(tmp_path, model_class): """ Test if 'save' and 'load' saves and loads model correctly and if 'load_parameters' and 'get_policy_parameters' work correctly @@ -29,14 +42,15 @@ def test_save_load(model_class): :param model_class: (BaseAlgorithm) A RL model """ - env = DummyVecEnv([lambda: IdentityEnvBox(10)]) + + env = DummyVecEnv([lambda: select_env(model_class)]) # create model model = model_class('MlpPolicy', env, policy_kwargs=dict(net_arch=[16]), verbose=1) model.learn(total_timesteps=500, eval_freq=250) env.reset() - observations = np.concatenate([env.step(env.action_space.sample())[0] for _ in range(10)], axis=0) + observations = np.concatenate([env.step([env.action_space.sample()])[0] for _ in range(10)], axis=0) # Get dictionary of current parameters params = deepcopy(model.policy.state_dict()) @@ -58,9 +72,9 @@ def test_save_load(model_class): selected_actions, _ = model.predict(observations, deterministic=True) # Check - model.save("test_save.zip") + model.save(tmp_path / "test_save.zip") del model - model = model_class.load("test_save", env=env) + model = model_class.load(str(tmp_path / "test_save"), env=env) # check if params are still the same after load new_params = model.policy.state_dict() @@ -77,7 +91,7 @@ def test_save_load(model_class): model.learn(total_timesteps=1000, eval_freq=500) # clear file from os - os.remove("test_save.zip") + os.remove(tmp_path / "test_save.zip") @pytest.mark.parametrize("model_class", MODEL_LIST) @@ -86,9 +100,11 @@ def test_set_env(model_class): Test if set_env function does work correct :param model_class: (BaseAlgorithm) A RL model """ - env = DummyVecEnv([lambda: IdentityEnvBox(10)]) - env2 = DummyVecEnv([lambda: IdentityEnvBox(10)]) - env3 = IdentityEnvBox(10) + + # use discrete for DQN + env = DummyVecEnv([lambda: select_env(model_class)]) + env2 = DummyVecEnv([lambda: select_env(model_class)]) + env3 = select_env(model_class) # create model model = model_class('MlpPolicy', env, policy_kwargs=dict(net_arch=[16])) @@ -107,42 +123,40 @@ def test_set_env(model_class): @pytest.mark.parametrize("model_class", MODEL_LIST) -def test_exclude_include_saved_params(model_class): +def test_exclude_include_saved_params(tmp_path, model_class): """ Test if exclude and include parameters of save() work :param model_class: (BaseAlgorithm) A RL model """ - env = DummyVecEnv([lambda: IdentityEnvBox(10)]) + env = DummyVecEnv([lambda: select_env(model_class)]) # create model, set verbose as 2, which is not standard model = model_class('MlpPolicy', env, policy_kwargs=dict(net_arch=[16]), verbose=2) # Check if exclude works - model.save("test_save.zip", exclude=["verbose"]) + model.save(tmp_path / "test_save.zip", exclude=["verbose"]) del model - model = model_class.load("test_save") + model = model_class.load(str(tmp_path / "test_save")) # check if verbose was not saved assert model.verbose != 2 # set verbose as something different then standard settings model.verbose = 2 # Check if include works - model.save("test_save.zip", exclude=["verbose"], include=["verbose"]) + model.save(tmp_path / "test_save.zip", exclude=["verbose"], include=["verbose"]) del model - model = model_class.load("test_save") + model = model_class.load(str(tmp_path / "test_save")) assert model.verbose == 2 # clear file from os - os.remove("test_save.zip") + os.remove(tmp_path / "test_save.zip") -@pytest.mark.parametrize("model_class", [SAC, TD3]) -def test_save_load_replay_buffer(model_class): - log_folder = 'logs' - replay_path = os.path.join(log_folder, 'replay_buffer.pkl') - os.makedirs(log_folder, exist_ok=True) - model = model_class('MlpPolicy', 'Pendulum-v0', buffer_size=1000) +@pytest.mark.parametrize("model_class", [SAC, TD3, DQN]) +def test_save_load_replay_buffer(tmp_path, model_class): + replay_path = tmp_path / 'replay_buffer.pkl' + model = model_class('MlpPolicy', select_env(model_class), buffer_size=1000) model.learn(500) old_replay_buffer = deepcopy(model.replay_buffer) model.save_replay_buffer(replay_path) @@ -151,21 +165,54 @@ def test_save_load_replay_buffer(model_class): assert np.allclose(old_replay_buffer.observations, model.replay_buffer.observations) assert np.allclose(old_replay_buffer.actions, model.replay_buffer.actions) - assert np.allclose(old_replay_buffer.next_observations, model.replay_buffer.next_observations) assert np.allclose(old_replay_buffer.rewards, model.replay_buffer.rewards) assert np.allclose(old_replay_buffer.dones, model.replay_buffer.dones) # test extending replay buffer - model.replay_buffer.extend(old_replay_buffer.observations, old_replay_buffer.next_observations, + model.replay_buffer.extend(old_replay_buffer.observations, old_replay_buffer.observations, old_replay_buffer.actions, old_replay_buffer.rewards, old_replay_buffer.dones) # clear file from os os.remove(replay_path) +@pytest.mark.parametrize("model_class", [DQN, SAC, TD3]) +@pytest.mark.parametrize("optimize_memory_usage", [False, True]) +def test_warn_buffer(recwarn, model_class, optimize_memory_usage): + """ + When using memory efficient replay buffer, + a warning must be emitted when calling `.learn()` + multiple times. + See https://github.com/DLR-RM/stable-baselines3/issues/46 + """ + # remove gym warnings + warnings.filterwarnings(action='ignore', category=DeprecationWarning) + warnings.filterwarnings(action='ignore', category=UserWarning, module='gym') + + model = model_class('MlpPolicy', select_env(model_class), buffer_size=100, + optimize_memory_usage=optimize_memory_usage, policy_kwargs=dict(net_arch=[64]), + learning_starts=10) + + model.learn(150) + + model.learn(150, reset_num_timesteps=False) + + # Check that there is no warning + assert len(recwarn) == 0 + + model.learn(150) + + if optimize_memory_usage: + assert len(recwarn) == 1 + warning = recwarn.pop(UserWarning) + assert "The last trajectory in the replay buffer will be truncated" in str(warning.message) + else: + assert len(recwarn) == 0 + + @pytest.mark.parametrize("model_class", MODEL_LIST) @pytest.mark.parametrize("policy_str", ['MlpPolicy', 'CnnPolicy']) -def test_save_load_policy(model_class, policy_str): +def test_save_load_policy(tmp_path, model_class, policy_str): """ Test saving and loading policy only. @@ -174,14 +221,14 @@ def test_save_load_policy(model_class, policy_str): """ kwargs = {} if policy_str == 'MlpPolicy': - env = IdentityEnvBox(10) + env = select_env(model_class) else: - if model_class in [SAC, TD3]: + if model_class in [SAC, TD3, DQN]: # Avoid memory error when using replay buffer # Reduce the size of the features kwargs = dict(buffer_size=250) env = FakeImageEnv(screen_height=40, screen_width=40, n_channels=2, - discrete=False) + discrete=model_class == DQN) env = DummyVecEnv([lambda: env]) @@ -191,7 +238,7 @@ def test_save_load_policy(model_class, policy_str): model.learn(total_timesteps=500, eval_freq=250) env.reset() - observations = np.concatenate([env.step(env.action_space.sample())[0] for _ in range(10)], axis=0) + observations = np.concatenate([env.step([env.action_space.sample()])[0] for _ in range(10)], axis=0) policy = model.policy policy_class = policy.__class__ @@ -223,16 +270,16 @@ def test_save_load_policy(model_class, policy_str): selected_actions_actor, _ = actor.predict(observations, deterministic=True) # Save and load policy - policy.save("./logs/policy.pkl") + policy.save(tmp_path / "policy.pkl") # Save and load actor if actor is not None: - actor.save("./logs/actor.pkl") + actor.save(tmp_path / "actor.pkl") del policy, actor - policy = policy_class.load("./logs/policy.pkl") + policy = policy_class.load(tmp_path / "policy.pkl") if actor_class is not None: - actor = actor_class.load("./logs/actor.pkl") + actor = actor_class.load(tmp_path / "actor.pkl") # check if params are still the same after load new_params = policy.state_dict() @@ -251,6 +298,6 @@ def test_save_load_policy(model_class, policy_str): assert np.allclose(selected_actions_actor, new_selected_actions, 1e-4) # clear file from os - os.remove("./logs/policy.pkl") + os.remove(tmp_path / "policy.pkl") if actor_class is not None: - os.remove("./logs/actor.pkl") + os.remove(tmp_path / "actor.pkl") diff --git a/tests/test_spaces.py b/tests/test_spaces.py index dfd4a60..e057f97 100644 --- a/tests/test_spaces.py +++ b/tests/test_spaces.py @@ -2,7 +2,7 @@ import numpy as np import pytest import gym -from stable_baselines3 import SAC, TD3 +from stable_baselines3 import DQN, SAC, TD3 from stable_baselines3.common.evaluation import evaluate_policy @@ -32,13 +32,17 @@ class DummyMultiBinary(gym.Env): return self.observation_space.sample(), 0.0, False, {} -@pytest.mark.parametrize("model_class", [SAC, TD3]) +@pytest.mark.parametrize("model_class", [SAC, TD3, DQN]) @pytest.mark.parametrize("env", [DummyMultiDiscreteSpace([4, 3]), DummyMultiBinary(8)]) def test_identity_spaces(model_class, env): """ - Additional tests for SAC/TD3 to check observation space support + Additional tests for DQ/SAC/TD3 to check observation space support for MultiDiscrete and MultiBinary. """ + # DQN only support discrete actions + if model_class == DQN: + env.action_space = gym.spaces.Discrete(4) + env = gym.wrappers.TimeLimit(env, max_episode_steps=100) model = model_class("MlpPolicy", env, gamma=0.5, seed=1, policy_kwargs=dict(net_arch=[64])) diff --git a/tests/test_utils.py b/tests/test_utils.py index 1c82782..4f04738 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -7,11 +7,12 @@ import numpy as np from stable_baselines3 import A2C from stable_baselines3.common.monitor import Monitor +from stable_baselines3.common.atari_wrappers import ClipRewardEnv from stable_baselines3.common.evaluation import evaluate_policy from stable_baselines3.common.cmd_util import make_vec_env, make_atari_env from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv -from stable_baselines3.common.noise import ( - VectorizedActionNoise, OrnsteinUhlenbeckActionNoise, ActionNoise) +from stable_baselines3.common.noise import (VectorizedActionNoise, + OrnsteinUhlenbeckActionNoise, ActionNoise) @pytest.mark.parametrize("env_id", ['CartPole-v1', lambda: gym.make('CartPole-v1')]) @@ -57,11 +58,11 @@ def test_make_atari_env(env_id, n_envs, wrapper_kwargs): if wrapper_kwargs is not None: assert obs.shape == (n_envs, 60, 60, 1) assert wrapped_atari_env.observation_space.shape == (60, 60, 1) - assert wrapped_atari_env.clip_reward is False + assert not isinstance(wrapped_atari_env.env, ClipRewardEnv) else: assert obs.shape == (n_envs, 84, 84, 1) assert wrapped_atari_env.observation_space.shape == (84, 84, 1) - assert wrapped_atari_env.clip_reward is True + assert isinstance(wrapped_atari_env.env, ClipRewardEnv) assert np.max(np.abs(reward)) < 1.0 diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index 124c5e5..979ea67 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -68,7 +68,7 @@ def test_runningmeanstd(): assert np.allclose(moments_1, moments_2) -def test_vec_env(tmpdir): +def test_vec_env(tmp_path): """Test VecNormalize Object""" clip_obs = 0.5 clip_reward = 5.0 @@ -82,7 +82,7 @@ def test_vec_env(tmpdir): assert np.max(np.abs(obs)) <= clip_obs assert np.max(np.abs(rew)) <= clip_reward - path = str(tmpdir.join("vec_normalize")) + path = tmp_path / "vec_normalize" norm_venv.save(path) deserialized = VecNormalize.load(path, venv=orig_venv) check_vec_norm_equal(norm_venv, deserialized) @@ -125,7 +125,7 @@ def test_offpolicy_normalization(model_class): eval_env = DummyVecEnv([make_env]) eval_env = VecNormalize(eval_env, training=False, norm_obs=True, norm_reward=False, clip_obs=10., clip_reward=10.) - model = model_class('MlpPolicy', env, verbose=1) + model = model_class('MlpPolicy', env, verbose=1, policy_kwargs=dict(net_arch=[64])) model.learn(total_timesteps=1000, eval_env=eval_env, eval_freq=500) # Check getter assert isinstance(model.get_vec_normalize_env(), VecNormalize)