From 92f7a6f23b90a9965c2e9cbe53f775055c5e0dc4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Fri, 13 Jan 2023 18:28:22 +0100 Subject: [PATCH 01/12] Fix `test_vec_normalize.py`, `test_tensorboard.py` and `common/monitor.py` type hint (#1194) * Remove from mypy exclude * type hint for metadata * Union[float, int] -> float * Remove useless __init__ * Type hint for model and logger in BaseCallback * Type hint for metric_dict * Update changelog * fix test_tensorboard * ignore gamma type checking * Fix monitor type hint * Update logger type hints * Fix type annotation and bump version * Fix circular import Co-authored-by: Antonin RAFFIN --- docs/guide/tensorboard.rst | 10 ++++------ docs/misc/changelog.rst | 5 ++++- setup.cfg | 3 --- stable_baselines3/common/callbacks.py | 10 +++++++--- .../common/envs/multi_input_envs.py | 2 +- stable_baselines3/common/logger.py | 10 +++++----- stable_baselines3/common/monitor.py | 19 ++++++++++--------- stable_baselines3/common/running_mean_std.py | 4 ++-- stable_baselines3/common/utils.py | 2 +- stable_baselines3/version.txt | 2 +- tests/test_tensorboard.py | 19 ++++++++++--------- tests/test_vec_normalize.py | 3 ++- 12 files changed, 47 insertions(+), 42 deletions(-) diff --git a/docs/guide/tensorboard.rst b/docs/guide/tensorboard.rst index 610afd6..2699d4a 100644 --- a/docs/guide/tensorboard.rst +++ b/docs/guide/tensorboard.rst @@ -268,11 +268,9 @@ Here is an example of how to save hyperparameters in TensorBoard: class HParamCallback(BaseCallback): - def __init__(self): - """ - Saves the hyperparameters and metrics at the start of the training, and logs them to TensorBoard. - """ - super().__init__() + """ + Saves the hyperparameters and metrics at the start of the training, and logs them to TensorBoard. + """ def _on_training_start(self) -> None: hparam_dict = { @@ -284,7 +282,7 @@ Here is an example of how to save hyperparameters in TensorBoard: # Tensorbaord will find & display metrics from the `SCALARS` tab metric_dict = { "rollout/ep_len_mean": 0, - "train/value_loss": 0, + "train/value_loss": 0.0, } self.logger.record( "hparams", diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 402ad59..db9d6c6 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -4,7 +4,7 @@ Changelog ========== -Release 1.8.0a0 (WIP) +Release 1.8.0a1 (WIP) -------------------------- @@ -28,6 +28,9 @@ Deprecations: Others: ^^^^^^^ +- Fixed ``tests/test_tensorboard.py`` type hint +- Fixed ``tests/test_vec_normalize.py`` type hint +- Fixed ``stable_baselines3/common/monitor.py`` type hint Documentation: ^^^^^^^^^^^^^^ diff --git a/setup.cfg b/setup.cfg index ca40fa6..37ffa17 100644 --- a/setup.cfg +++ b/setup.cfg @@ -39,7 +39,6 @@ exclude = (?x)( | stable_baselines3/common/envs/identity_env.py$ | stable_baselines3/common/envs/multi_input_envs.py$ | stable_baselines3/common/logger.py$ - | stable_baselines3/common/monitor.py$ | stable_baselines3/common/off_policy_algorithm.py$ | stable_baselines3/common/on_policy_algorithm.py$ | stable_baselines3/common/policies.py$ @@ -67,9 +66,7 @@ exclude = (?x)( | stable_baselines3/td3/policies.py$ | stable_baselines3/td3/td3.py$ | tests/test_logger.py$ - | tests/test_tensorboard.py$ | tests/test_train_eval_mode.py$ - | tests/test_vec_normalize.py$ ) [flake8] diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py index 0ecd1ad..a96c52c 100644 --- a/stable_baselines3/common/callbacks.py +++ b/stable_baselines3/common/callbacks.py @@ -6,6 +6,8 @@ from typing import Any, Callable, Dict, List, Optional, Union import gym import numpy as np +from stable_baselines3.common.logger import Logger + try: from tqdm import TqdmExperimentalWarning @@ -29,10 +31,13 @@ class BaseCallback(ABC): :param verbose: Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages """ + # The RL model + # Type hint as string to avoid circular import + model: "base_class.BaseAlgorithm" + logger: Logger + def __init__(self, verbose: int = 0): super().__init__() - # The RL model - self.model = None # type: Optional[base_class.BaseAlgorithm] # An alias for self.model.get_env(), the environment used for training self.training_env = None # type: Union[gym.Env, VecEnv, None] # Number of time the callback was called @@ -42,7 +47,6 @@ class BaseCallback(ABC): self.verbose = verbose self.locals: Dict[str, Any] = {} self.globals: Dict[str, Any] = {} - self.logger = None # Sometimes, for event callback, it is useful # to have access to the parent object self.parent = None # type: Optional[BaseCallback] diff --git a/stable_baselines3/common/envs/multi_input_envs.py b/stable_baselines3/common/envs/multi_input_envs.py index 433591d..166c699 100644 --- a/stable_baselines3/common/envs/multi_input_envs.py +++ b/stable_baselines3/common/envs/multi_input_envs.py @@ -121,7 +121,7 @@ class SimpleMultiObsEnv(gym.Env): self.right_possible = [0, 1, 2, 12, 13, 14] self.up_possible = [4, 8, 12, 7, 11, 15] - def step(self, action: Union[int, float, np.ndarray]) -> GymStepReturn: + def step(self, action: Union[float, np.ndarray]) -> GymStepReturn: """ Run one timestep of the environment's dynamics. When end of episode is reached, you are responsible for calling `reset()` diff --git a/stable_baselines3/common/logger.py b/stable_baselines3/common/logger.py index e065992..939d924 100644 --- a/stable_baselines3/common/logger.py +++ b/stable_baselines3/common/logger.py @@ -5,7 +5,7 @@ import sys import tempfile import warnings from collections import defaultdict -from typing import Any, Dict, List, Optional, Sequence, TextIO, Tuple, Union +from typing import Any, Dict, List, Mapping, Optional, Sequence, TextIO, Tuple, Union import numpy as np import pandas @@ -16,7 +16,7 @@ try: from torch.utils.tensorboard import SummaryWriter from torch.utils.tensorboard.summary import hparams except ImportError: - SummaryWriter = None + SummaryWriter = None # type: ignore[misc, assignment] try: from tqdm import tqdm @@ -38,7 +38,7 @@ class Video: :param fps: frames per second """ - def __init__(self, frames: th.Tensor, fps: Union[float, int]): + def __init__(self, frames: th.Tensor, fps: float): self.frames = frames self.fps = fps @@ -80,7 +80,7 @@ class HParam: A non-empty metrics dict is required to display hyperparameters in the corresponding Tensorboard section. """ - def __init__(self, hparam_dict: Dict[str, Union[bool, str, float, int, None]], metric_dict: Dict[str, Union[float, int]]): + def __init__(self, hparam_dict: Mapping[str, Union[bool, str, float, None]], metric_dict: Mapping[str, float]): self.hparam_dict = hparam_dict if not metric_dict: raise Exception("`metric_dict` must not be empty to display hyperparameters to the HPARAMS tensorboard tab.") @@ -329,7 +329,7 @@ class CSVOutputFormat(KVWriter): def __init__(self, filename: str): self.file = open(filename, "w+t") - self.keys = [] + self.keys: List[str] = [] self.separator = "," self.quotechar = '"' diff --git a/stable_baselines3/common/monitor.py b/stable_baselines3/common/monitor.py index 1e56fdb..ffc318d 100644 --- a/stable_baselines3/common/monitor.py +++ b/stable_baselines3/common/monitor.py @@ -5,7 +5,7 @@ import json import os import time from glob import glob -from typing import Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union import gym import numpy as np @@ -41,6 +41,7 @@ class Monitor(gym.Wrapper): ): super().__init__(env=env) self.t_start = time.time() + self.results_writer = None if filename is not None: self.results_writer = ResultsWriter( filename, @@ -48,18 +49,18 @@ class Monitor(gym.Wrapper): extra_keys=reset_keywords + info_keywords, override_existing=override_existing, ) - else: - self.results_writer = None + self.reset_keywords = reset_keywords self.info_keywords = info_keywords self.allow_early_resets = allow_early_resets - self.rewards = None + self.rewards: List[float] = [] self.needs_reset = True - self.episode_returns = [] - self.episode_lengths = [] - self.episode_times = [] + self.episode_returns: List[float] = [] + self.episode_lengths: List[int] = [] + self.episode_times: List[float] = [] self.total_steps = 0 - self.current_reset_info = {} # extra info about the current episode, that was passed in during reset() + # extra info about the current episode, that was passed in during reset() + self.current_reset_info: Dict[str, Any] = {} def reset(self, **kwargs) -> GymObs: """ @@ -200,7 +201,7 @@ class ResultsWriter: self.file_handler.flush() - def write_row(self, epinfo: Dict[str, Union[float, int]]) -> None: + def write_row(self, epinfo: Dict[str, float]) -> None: """ Close the file handler diff --git a/stable_baselines3/common/running_mean_std.py b/stable_baselines3/common/running_mean_std.py index b48f922..9dfa4b8 100644 --- a/stable_baselines3/common/running_mean_std.py +++ b/stable_baselines3/common/running_mean_std.py @@ -1,4 +1,4 @@ -from typing import Tuple, Union +from typing import Tuple import numpy as np @@ -40,7 +40,7 @@ class RunningMeanStd: batch_count = arr.shape[0] self.update_from_moments(batch_mean, batch_var, batch_count) - def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: Union[int, float]) -> None: + def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: float) -> None: delta = batch_mean - self.mean tot_count = self.count + batch_count diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 4dc284e..b5cbff4 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -76,7 +76,7 @@ def update_learning_rate(optimizer: th.optim.Optimizer, learning_rate: float) -> param_group["lr"] = learning_rate -def get_schedule_fn(value_schedule: Union[Schedule, float, int]) -> Schedule: +def get_schedule_fn(value_schedule: Union[Schedule, float]) -> Schedule: """ Transform (if needed) learning rate and clip range (for PPO) to callable. diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 52d893b..0d03ef9 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.8.0a0 +1.8.0a1 diff --git a/tests/test_tensorboard.py b/tests/test_tensorboard.py index 8aa864d..eee0ec0 100644 --- a/tests/test_tensorboard.py +++ b/tests/test_tensorboard.py @@ -1,4 +1,5 @@ import os +from typing import Dict, Union import pytest @@ -18,21 +19,21 @@ N_STEPS = 100 class HParamCallback(BaseCallback): - def __init__(self): - """ - Saves the hyperparameters and metrics at the start of the training, and logs them to TensorBoard. - """ - super().__init__() + """ + Saves the hyperparameters and metrics at the start of the training, and logs them to TensorBoard. + """ def _on_training_start(self) -> None: - hparam_dict = { + hparam_dict: Dict[str, Union[str, float]] = { "algorithm": self.model.__class__.__name__, - "learning rate": self.model.learning_rate, - "gamma": self.model.gamma, + # Ignore type checking for gamma, see https://github.com/DLR-RM/stable-baselines3/pull/1194/files#r1035006458 + "gamma": self.model.gamma, # type: ignore[attr-defined] } + if isinstance(self.model.learning_rate, float): # Can also be Schedule, in that case, we don't report + hparam_dict["learning rate"] = self.model.learning_rate # define the metrics that will appear in the `HPARAMS` Tensorboard tab by referencing their tag # Tensorbaord will find & display metrics from the `SCALARS` tab - metric_dict = { + metric_dict: Dict[str, float] = { "rollout/ep_len_mean": 0, } self.logger.record( diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index 7b443c2..b17d28c 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -1,4 +1,5 @@ import operator +from typing import Any, Dict import gym import numpy as np @@ -20,7 +21,7 @@ ENV_ID = "Pendulum-v1" class DummyRewardEnv(gym.Env): - metadata = {} + metadata: Dict[str, Any] = {} def __init__(self, return_reward_idx=0): self.action_space = spaces.Discrete(2) From 69fdf155e1076730f88f19f39f598a0b418e180a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Mon, 23 Jan 2023 10:56:45 +0100 Subject: [PATCH 02/12] Downgrade `sphinx-autodoc-typehints` (#1291) * Update setup.py * black * hotfix pytype --- setup.py | 2 +- stable_baselines3/common/buffers.py | 10 +++++----- stable_baselines3/common/envs/identity_env.py | 8 ++++---- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/setup.py b/setup.py index df97633..612b2be 100644 --- a/setup.py +++ b/setup.py @@ -117,7 +117,7 @@ setup( # For spelling "sphinxcontrib.spelling", # Type hints support - "sphinx-autodoc-typehints", + "sphinx-autodoc-typehints==1.21.1", # TODO: remove version constraint, see #1290 # Copy button for code snippets "sphinx_copybutton", ], diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 2dafd41..f9f0c72 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -474,7 +474,7 @@ class RolloutBuffer(BaseBuffer): 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: + def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> RolloutBufferSamples: # type: ignore[signature-mismatch] #FIXME data = ( self.observations[batch_inds], self.actions[batch_inds], @@ -603,7 +603,7 @@ class DictReplayBuffer(ReplayBuffer): self.full = True self.pos = 0 - def sample(self, batch_size: int, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples: + def sample(self, batch_size: int, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples: # type: ignore[signature-mismatch] #FIXME: """ Sample elements from the replay buffer. @@ -614,7 +614,7 @@ class DictReplayBuffer(ReplayBuffer): """ return super(ReplayBuffer, self).sample(batch_size=batch_size, env=env) - def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples: + def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples: # type: ignore[signature-mismatch] #FIXME: # Sample randomly the env idx env_indices = np.random.randint(0, high=self.n_envs, size=(len(batch_inds),)) @@ -743,7 +743,7 @@ class DictRolloutBuffer(RolloutBuffer): if self.pos == self.buffer_size: self.full = True - def get(self, batch_size: Optional[int] = None) -> Generator[DictRolloutBufferSamples, None, None]: + def get(self, batch_size: Optional[int] = None) -> Generator[DictRolloutBufferSamples, None, None]: # type: ignore[signature-mismatch] #FIXME assert self.full, "" indices = np.random.permutation(self.buffer_size * self.n_envs) # Prepare the data @@ -767,7 +767,7 @@ class DictRolloutBuffer(RolloutBuffer): 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) -> DictRolloutBufferSamples: + def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> DictRolloutBufferSamples: # type: ignore[signature-mismatch] #FIXME return DictRolloutBufferSamples( observations={key: self.to_torch(obs[batch_inds]) for (key, obs) in self.observations.items()}, diff --git a/stable_baselines3/common/envs/identity_env.py b/stable_baselines3/common/envs/identity_env.py index 9635e53..a8bed17 100644 --- a/stable_baselines3/common/envs/identity_env.py +++ b/stable_baselines3/common/envs/identity_env.py @@ -71,7 +71,7 @@ class IdentityEnvBox(IdentityEnv[np.ndarray]): super().__init__(ep_length=ep_length, space=space) self.eps = eps - def step(self, action: np.ndarray) -> GymStepReturn: + def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, Dict[str, Any]]: reward = self._get_reward(action) self._choose_next_state() self.current_step += 1 @@ -83,7 +83,7 @@ class IdentityEnvBox(IdentityEnv[np.ndarray]): class IdentityEnvMultiDiscrete(IdentityEnv[np.ndarray]): - def __init__(self, dim: int = 1, ep_length: int = 100): + def __init__(self, dim: int = 1, ep_length: int = 100) -> None: """ Identity environment for testing purposes @@ -95,7 +95,7 @@ class IdentityEnvMultiDiscrete(IdentityEnv[np.ndarray]): class IdentityEnvMultiBinary(IdentityEnv[np.ndarray]): - def __init__(self, dim: int = 1, ep_length: int = 100): + def __init__(self, dim: int = 1, ep_length: int = 100) -> None: """ Identity environment for testing purposes @@ -126,7 +126,7 @@ class FakeImageEnv(gym.Env): n_channels: int = 1, discrete: bool = True, channel_first: bool = False, - ): + ) -> None: self.observation_shape = (screen_height, screen_width, n_channels) if channel_first: self.observation_shape = (n_channels, screen_height, screen_width) From b702884c23b6aeaa5d2a830b37d6b15fb1bdf983 Mon Sep 17 00:00:00 2001 From: Alex Pasquali Date: Mon, 23 Jan 2023 14:55:19 +0100 Subject: [PATCH 03/12] Removed shared layers in mlp_extractor (#1292) * Modified actor-critic policies & MlpExtractor class ActorCriticPolicy: - changed type hint of net_arch param: now it's a dict - removed check that if features extractor is not shared: no shared layers are allowed in the mlp_extractor regardless of the features extractor ActorCriticCnnPolicy: - changed type hint of net_arch param: now it's a dict MultiInputActorcriticPolicy: - changed type hint of net_arch param: now it's a dict MlpExtractor: - changed type hint of net_arch param: now it's a dict - adapted networks creation - adapted methods: forward, forward_actor & forward_critic * Removed shared layers in mlp_extractor * Updated docs and changelog + reformat * Updated custom policy tests * Removed test on deprecation warning for share layers in mlp_extractor Now shared layers are removed * Update version * Update RL Zoo doc * Fix linter warnings * Add ruff to Makefile (experimental) * Add backward compat code and minor updates * Update tests * Add backward compatibility * Fix test * Improve compat code Co-authored-by: Antonin RAFFIN --- Makefile | 7 ++ docs/guide/custom_policy.rst | 45 ++++------ docs/guide/rl_zoo.rst | 19 ++-- docs/misc/changelog.rst | 3 +- stable_baselines3/common/base_class.py | 6 +- stable_baselines3/common/buffers.py | 29 ++++-- stable_baselines3/common/policies.py | 37 +++----- stable_baselines3/common/torch_layers.py | 109 +++++++---------------- stable_baselines3/version.txt | 2 +- tests/test_custom_policy.py | 25 ++---- tests/test_identity.py | 2 + 11 files changed, 122 insertions(+), 162 deletions(-) diff --git a/Makefile b/Makefile index c806507..6351162 100644 --- a/Makefile +++ b/Makefile @@ -19,6 +19,13 @@ lint: # exit-zero treats all errors as warnings. flake8 ${LINT_PATHS} --count --exit-zero --statistics +ruff: + # stop the build if there are Python syntax errors or undefined names + # see https://lintlyci.github.io/Flake8Rules/ + ruff ${LINT_PATHS} --select=E9,F63,F7,F82 --show-source + # exit-zero treats all errors as warnings. + ruff ${LINT_PATHS} --exit-zero --line-length 127 + format: # Sort imports isort ${LINT_PATHS} diff --git a/docs/guide/custom_policy.rst b/docs/guide/custom_policy.rst index c9e598e..dae6048 100644 --- a/docs/guide/custom_policy.rst +++ b/docs/guide/custom_policy.rst @@ -117,11 +117,6 @@ that derives from ``BaseFeaturesExtractor`` and then pass it to the model when t ``policy_kwargs`` (both for on-policy and off-policy algorithms). -.. warning:: - If the features extractor is **non-shared**, it is **not** possible to have shared layers in the ``mlp_extractor``. - Please note that this option is **deprecated**, therefore in a future release the layers in the ``mlp_extractor`` will have to be non-shared. - - .. code-block:: python import torch as th @@ -242,41 +237,31 @@ On-Policy Algorithms Custom Networks --------------- -.. warning:: - Shared layers in the the ``mlp_extractor`` are **deprecated**. - In a future release all layers will have to be non-shared. - If needed, you can implement a custom policy network (see `advanced example below <#advanced-example>`_). - -.. warning:: - In the next Stable-Baselines3 release, the behavior of ``net_arch=[128, 128]`` will change - to match the one of off-policy algorithms: it will create **separate** networks (instead of shared currently) - for the actor and the critic, with the same architecture. - - If you need a network architecture that is different for the actor and the critic when using ``PPO``, ``A2C`` or ``TRPO``, you can pass a dictionary of the following structure: ``dict(pi=[], vf=[])``. For example, if you want a different architecture for the actor (aka ``pi``) and the critic ( value-function aka ``vf``) networks, then you can specify ``net_arch=dict(pi=[32, 32], vf=[64, 64])``. -.. Otherwise, to have actor and critic that share the same network architecture, -.. you only need to specify ``net_arch=[128, 128]`` (here, two hidden layers of 128 units each). +Otherwise, to have actor and critic that share the same network architecture, +you only need to specify ``net_arch=[128, 128]`` (here, two hidden layers of 128 units each, this is equivalent to ``net_arch=dict(pi=[128, 128], vf=[128, 128])``). + +If shared layers are needed, you need to implement a custom policy network (see `advanced example below <#advanced-example>`_). Examples ~~~~~~~~ -.. TODO(antonin): uncomment when shared network is removed -.. Same architecture for actor and critic with two layers of size 128: ``net_arch=[128, 128]`` -.. -.. .. code-block:: none -.. -.. obs -.. / \ -.. <128> <128> -.. | | -.. <128> <128> -.. | | -.. action value +Same architecture for actor and critic with two layers of size 128: ``net_arch=[128, 128]`` + +.. code-block:: none + + obs + / \ + <128> <128> + | | + <128> <128> + | | + action value Different architectures for actor and critic: ``net_arch=dict(pi=[32, 32], vf=[64, 64])`` diff --git a/docs/guide/rl_zoo.rst b/docs/guide/rl_zoo.rst index ea15832..8a611d8 100644 --- a/docs/guide/rl_zoo.rst +++ b/docs/guide/rl_zoo.rst @@ -20,6 +20,10 @@ Goals of this repository: Installation ------------ +Option 1: install the python package ``pip install rl_zoo3`` + +or: + 1. Clone the repository: :: @@ -42,7 +46,10 @@ Installation :: apt-get install swig cmake ffmpeg + # full dependencies pip install -r requirements.txt + # minimal dependencies + pip install -e . Train an Agent @@ -56,13 +63,13 @@ using: :: - python train.py --algo algo_name --env env_id + python -m rl_zoo3.train --algo algo_name --env env_id For example (with evaluation and checkpoints): :: - python train.py --algo ppo --env CartPole-v1 --eval-freq 10000 --save-freq 50000 + python -m rl_zoo3.train --algo ppo --env CartPole-v1 --eval-freq 10000 --save-freq 50000 Continue training (here, load pretrained agent for Breakout and continue @@ -70,7 +77,7 @@ training for 5000 steps): :: - python train.py --algo a2c --env BreakoutNoFrameskip-v4 -i trained_agents/a2c/BreakoutNoFrameskip-v4_1/BreakoutNoFrameskip-v4.zip -n 5000 + python -m rl_zoo3.train --algo a2c --env BreakoutNoFrameskip-v4 -i trained_agents/a2c/BreakoutNoFrameskip-v4_1/BreakoutNoFrameskip-v4.zip -n 5000 Enjoy a Trained Agent @@ -80,13 +87,13 @@ If the trained agent exists, then you can see it in action using: :: - python enjoy.py --algo algo_name --env env_id + python -m rl_zoo3.enjoy --algo algo_name --env env_id For example, enjoy A2C on Breakout during 5000 timesteps: :: - python enjoy.py --algo a2c --env BreakoutNoFrameskip-v4 --folder rl-trained-agents/ -n 5000 + python -m rl_zoo3.enjoy --algo a2c --env BreakoutNoFrameskip-v4 --folder rl-trained-agents/ -n 5000 Hyperparameter Optimization @@ -100,7 +107,7 @@ with a budget of 1000 trials and a maximum of 50000 steps: :: - python train.py --algo ppo --env MountainCar-v0 -n 50000 -optimize --n-trials 1000 --n-jobs 2 \ + python -m rl_zoo3.train --algo ppo --env MountainCar-v0 -n 50000 -optimize --n-trials 1000 --n-jobs 2 \ --sampler random --pruner median diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index db9d6c6..7cb344b 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -4,12 +4,13 @@ Changelog ========== -Release 1.8.0a1 (WIP) +Release 1.8.0a2 (WIP) -------------------------- Breaking Changes: ^^^^^^^^^^^^^^^^^ +- Removed shared layers in ``mlp_extractor`` (@AlexPasqua) New Features: ^^^^^^^^^^^^^ diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index a71043d..b6fba85 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -667,6 +667,11 @@ class BaseAlgorithm(ABC): if "policy_kwargs" in data: if "device" in data["policy_kwargs"]: del data["policy_kwargs"]["device"] + # backward compatibility, convert to new format + if "net_arch" in data["policy_kwargs"] and len(data["policy_kwargs"]["net_arch"]) > 0: + saved_net_arch = data["policy_kwargs"]["net_arch"] + if isinstance(saved_net_arch, list) and isinstance(saved_net_arch[0], dict): + data["policy_kwargs"]["net_arch"] = saved_net_arch[0] if "policy_kwargs" in kwargs and kwargs["policy_kwargs"] != data["policy_kwargs"]: raise ValueError( @@ -726,7 +731,6 @@ class BaseAlgorithm(ABC): ) else: raise e - # put other pytorch variables back in place if pytorch_variables is not None: for name in pytorch_variables: diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index f9f0c72..f71dd29 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -474,7 +474,11 @@ class RolloutBuffer(BaseBuffer): 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: # type: ignore[signature-mismatch] #FIXME + def _get_samples( + self, + batch_inds: np.ndarray, + env: Optional[VecNormalize] = None, + ) -> RolloutBufferSamples: # type: ignore[signature-mismatch] #FIXME data = ( self.observations[batch_inds], self.actions[batch_inds], @@ -603,7 +607,11 @@ class DictReplayBuffer(ReplayBuffer): self.full = True self.pos = 0 - def sample(self, batch_size: int, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples: # type: ignore[signature-mismatch] #FIXME: + def sample( + self, + batch_size: int, + env: Optional[VecNormalize] = None, + ) -> DictReplayBufferSamples: # type: ignore[signature-mismatch] #FIXME: """ Sample elements from the replay buffer. @@ -614,7 +622,11 @@ class DictReplayBuffer(ReplayBuffer): """ return super(ReplayBuffer, self).sample(batch_size=batch_size, env=env) - def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples: # type: ignore[signature-mismatch] #FIXME: + def _get_samples( + self, + batch_inds: np.ndarray, + env: Optional[VecNormalize] = None, + ) -> DictReplayBufferSamples: # type: ignore[signature-mismatch] #FIXME: # Sample randomly the env idx env_indices = np.random.randint(0, high=self.n_envs, size=(len(batch_inds),)) @@ -743,7 +755,10 @@ class DictRolloutBuffer(RolloutBuffer): if self.pos == self.buffer_size: self.full = True - def get(self, batch_size: Optional[int] = None) -> Generator[DictRolloutBufferSamples, None, None]: # type: ignore[signature-mismatch] #FIXME + def get( + self, + batch_size: Optional[int] = None, + ) -> Generator[DictRolloutBufferSamples, None, None]: # type: ignore[signature-mismatch] #FIXME assert self.full, "" indices = np.random.permutation(self.buffer_size * self.n_envs) # Prepare the data @@ -767,7 +782,11 @@ class DictRolloutBuffer(RolloutBuffer): 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) -> DictRolloutBufferSamples: # type: ignore[signature-mismatch] #FIXME + def _get_samples( + self, + batch_inds: np.ndarray, + env: Optional[VecNormalize] = None, + ) -> DictRolloutBufferSamples: # type: ignore[signature-mismatch] #FIXME return DictRolloutBufferSamples( observations={key: self.to_torch(obs[batch_inds]) for (key, obs) in self.observations.items()}, diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 0cf4917..793cfc5 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -418,8 +418,7 @@ class ActorCriticPolicy(BasePolicy): observation_space: spaces.Space, action_space: spaces.Space, lr_schedule: Schedule, - # TODO(antonin): update type annotation when we remove shared network support - net_arch: Union[List[int], Dict[str, List[int]], List[Dict[str, List[int]]], None] = None, + net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.Tanh, ortho_init: bool = True, use_sde: bool = False, @@ -452,21 +451,15 @@ class ActorCriticPolicy(BasePolicy): normalize_images=normalize_images, ) - # Convert [dict()] to dict() as shared network are deprecated - if isinstance(net_arch, list) and len(net_arch) > 0: - if isinstance(net_arch[0], dict): - warnings.warn( - ( - "As shared layers in the mlp_extractor are deprecated and will be removed in SB3 v1.8.0, " - "you should now pass directly a dictionary and not a list " - "(net_arch=dict(pi=..., vf=...) instead of net_arch=[dict(pi=..., vf=...)])" - ), - ) - net_arch = net_arch[0] - else: - # Note: deprecation warning will be emitted - # by the MlpExtractor constructor - pass + if isinstance(net_arch, list) and len(net_arch) > 0 and isinstance(net_arch[0], dict): + warnings.warn( + ( + "As shared layers in the mlp_extractor are removed since SB3 v1.8.0, " + "you should now pass directly a dictionary and not a list " + "(net_arch=dict(pi=..., vf=...) instead of net_arch=[dict(pi=..., vf=...)])" + ), + ) + net_arch = net_arch[0] # Default network architecture, from stable-baselines if net_arch is None: @@ -488,12 +481,6 @@ class ActorCriticPolicy(BasePolicy): else: self.pi_features_extractor = self.features_extractor self.vf_features_extractor = self.make_features_extractor() - # if the features extractor is not shared, there cannot be shared layers in the mlp_extractor - # TODO(antonin): update the check once we change net_arch behavior - if isinstance(net_arch, list) and len(net_arch) > 0: - raise ValueError( - "Error: if the features extractor is not shared, there cannot be shared layers in the mlp_extractor" - ) self.log_std_init = log_std_init dist_kwargs = None @@ -770,7 +757,7 @@ class ActorCriticCnnPolicy(ActorCriticPolicy): observation_space: spaces.Space, action_space: spaces.Space, lr_schedule: Schedule, - net_arch: Union[List[int], Dict[str, List[int]], List[Dict[str, List[int]]], None] = None, + net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.Tanh, ortho_init: bool = True, use_sde: bool = False, @@ -843,7 +830,7 @@ class MultiInputActorCriticPolicy(ActorCriticPolicy): observation_space: spaces.Dict, action_space: spaces.Space, lr_schedule: Schedule, - net_arch: Union[List[int], Dict[str, List[int]], List[Dict[str, List[int]]], None] = None, + net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.Tanh, ortho_init: bool = True, use_sde: bool = False, diff --git a/stable_baselines3/common/torch_layers.py b/stable_baselines3/common/torch_layers.py index 302d9b1..44714d6 100644 --- a/stable_baselines3/common/torch_layers.py +++ b/stable_baselines3/common/torch_layers.py @@ -1,5 +1,3 @@ -import warnings -from itertools import zip_longest from typing import Dict, List, Tuple, Type, Union import gym @@ -151,98 +149,57 @@ class MlpExtractor(nn.Module): Constructs an MLP that receives the output from a previous features extractor (i.e. a CNN) or directly the observations (if no features extractor is applied) as an input and outputs a latent representation for the policy and a value network. - The ``net_arch`` parameter allows to specify the amount and size of the hidden layers and how many - of them are shared between the policy network and the value network. It is assumed to be a list with the following - structure: - 1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer. - If the number of ints is zero, there will be no shared layers. - 2. An optional dict, to specify the following non-shared layers for the value network and the policy network. - It is formatted like ``dict(vf=[], pi=[])``. - If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed. + The ``net_arch`` parameter allows to specify the amount and size of the hidden layers. + It can be in either of the following forms: + 1. ``dict(vf=[], pi=[])``: to specify the amount and size of the layers in the + policy and value nets individually. If it is missing any of the keys (pi or vf), + zero layers will be considered for that key. + 2. ``[]``: "shortcut" in case the amount and size of the layers + in the policy and value nets are the same. Same as ``dict(vf=int_list, pi=int_list)`` + where int_list is the same for the actor and critic. - Deprecation note: shared layers in ``net_arch`` are deprecated, please use separate - pi and vf networks (e.g. net_arch=dict(pi=[...], vf=[...])) - - For example to construct a network with one shared layer of size 55 followed by two non-shared layers for the value - network of size 255 and a single non-shared layer of size 128 for the policy network, the following layers_spec - would be used: ``[55, dict(vf=[255, 255], pi=[128])]``. A simple shared network topology with two layers of size 128 - would be specified as [128, 128]. - - Adapted from Stable Baselines. + .. note:: + If a key is not specified or an empty list is passed ``[]``, a linear network will be used. :param feature_dim: Dimension of the feature vector (can be the output of a CNN) :param net_arch: The specification of the policy and value networks. See above for details on its formatting. :param activation_fn: The activation function to use for the networks. - :param device: + :param device: PyTorch device. """ def __init__( self, feature_dim: int, - net_arch: Union[Dict[str, List[int]], List[Union[int, Dict[str, List[int]]]]], + net_arch: Union[List[int], Dict[str, List[int]]], activation_fn: Type[nn.Module], device: Union[th.device, str] = "auto", ) -> None: super().__init__() device = get_device(device) - shared_net: List[nn.Module] = [] policy_net: List[nn.Module] = [] value_net: List[nn.Module] = [] - policy_only_layers: List[int] = [] # Layer sizes of the network that only belongs to the policy network - value_only_layers: List[int] = [] # Layer sizes of the network that only belongs to the value network - last_layer_dim_shared = feature_dim + last_layer_dim_pi = feature_dim + last_layer_dim_vf = feature_dim - if isinstance(net_arch, list) and len(net_arch) > 0 and isinstance(net_arch[0], int): - warnings.warn( - ( - "Shared layers in the mlp_extractor are deprecated and will be removed in SB3 v1.8.0, " - "please use separate pi and vf networks " - "(e.g. net_arch=dict(pi=[...], vf=[...]))" - ), - DeprecationWarning, - ) - - # TODO(antonin): update behavior for net_arch=[64, 64] - # once shared networks are removed + # save dimensions of layers in policy and value nets if isinstance(net_arch, dict): - policy_only_layers = net_arch["pi"] - value_only_layers = net_arch["vf"] + # Note: if key is not specificed, assume linear network + pi_layers_dims = net_arch.get("pi", []) # Layer sizes of the policy network + vf_layers_dims = net_arch.get("vf", []) # Layer sizes of the value network else: - # Iterate through the shared layers and build the shared parts of the network - for layer in net_arch: - if isinstance(layer, int): # Check that this is a shared layer - shared_net.append(nn.Linear(last_layer_dim_shared, layer)) # add linear of size layer - shared_net.append(activation_fn()) - last_layer_dim_shared = layer - else: - assert isinstance(layer, dict), "Error: the net_arch list can only contain ints and dicts" - if "pi" in layer: - assert isinstance(layer["pi"], list), "Error: net_arch[-1]['pi'] must contain a list of integers." - policy_only_layers = layer["pi"] - - if "vf" in layer: - assert isinstance(layer["vf"], list), "Error: net_arch[-1]['vf'] must contain a list of integers." - value_only_layers = layer["vf"] - break # From here on the network splits up in policy and value network - - last_layer_dim_pi = last_layer_dim_shared - last_layer_dim_vf = last_layer_dim_shared - - # Build the non-shared part of the network - for pi_layer_size, vf_layer_size in zip_longest(policy_only_layers, value_only_layers): - if pi_layer_size is not None: - assert isinstance(pi_layer_size, int), "Error: net_arch[-1]['pi'] must only contain integers." - policy_net.append(nn.Linear(last_layer_dim_pi, pi_layer_size)) - policy_net.append(activation_fn()) - last_layer_dim_pi = pi_layer_size - - if vf_layer_size is not None: - assert isinstance(vf_layer_size, int), "Error: net_arch[-1]['vf'] must only contain integers." - value_net.append(nn.Linear(last_layer_dim_vf, vf_layer_size)) - value_net.append(activation_fn()) - last_layer_dim_vf = vf_layer_size + pi_layers_dims = vf_layers_dims = net_arch + # Iterate through the policy layers and build the policy net + for curr_layer_dim in pi_layers_dims: + policy_net.append(nn.Linear(last_layer_dim_pi, curr_layer_dim)) + policy_net.append(activation_fn()) + last_layer_dim_pi = curr_layer_dim + # Iterate through the value layers and build the value net + for curr_layer_dim in vf_layers_dims: + value_net.append(nn.Linear(last_layer_dim_vf, curr_layer_dim)) + value_net.append(activation_fn()) + last_layer_dim_vf = curr_layer_dim # Save dim, used to create the distributions self.latent_dim_pi = last_layer_dim_pi @@ -250,7 +207,6 @@ class MlpExtractor(nn.Module): # Create networks # If the list of layers is empty, the network will just act as an Identity module - self.shared_net = nn.Sequential(*shared_net).to(device) self.policy_net = nn.Sequential(*policy_net).to(device) self.value_net = nn.Sequential(*value_net).to(device) @@ -259,14 +215,13 @@ class MlpExtractor(nn.Module): :return: latent_policy, latent_value of the specified network. If all layers are shared, then ``latent_policy == latent_value`` """ - shared_latent = self.shared_net(features) - return self.policy_net(shared_latent), self.value_net(shared_latent) + return self.forward_actor(features), self.forward_critic(features) def forward_actor(self, features: th.Tensor) -> th.Tensor: - return self.policy_net(self.shared_net(features)) + return self.policy_net(features) def forward_critic(self, features: th.Tensor) -> th.Tensor: - return self.value_net(self.shared_net(features)) + return self.value_net(features) class CombinedExtractor(BaseFeaturesExtractor): diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 0d03ef9..c3d22c0 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.8.0a1 +1.8.0a2 diff --git a/tests/test_custom_policy.py b/tests/test_custom_policy.py index 85c3d37..1f89b23 100644 --- a/tests/test_custom_policy.py +++ b/tests/test_custom_policy.py @@ -9,21 +9,21 @@ from stable_baselines3.common.sb2_compat.rmsprop_tf_like import RMSpropTFLike "net_arch", [ [], - dict(vf=[16], pi=[8]), - # [] behavior will change [4], [4, 4], - # All values below are deprecated - [12, dict(vf=[16], pi=[8])], - [12, dict(vf=[8, 4], pi=[8])], - [12, dict(vf=[8], pi=[8, 4])], - [12, dict(pi=[8])], + dict(vf=[16], pi=[8]), + dict(vf=[8, 4], pi=[8]), + dict(vf=[8], pi=[8, 4]), + dict(pi=[8]), + # Old format, emits a warning + [dict(vf=[8])], + [dict(vf=[8], pi=[4])], ], ) @pytest.mark.parametrize("model_class", [A2C, PPO]) def test_flexible_mlp(model_class, net_arch): - if isinstance(net_arch, list) and len(net_arch) > 0 and isinstance(net_arch[0], int): - with pytest.warns(DeprecationWarning): + if isinstance(net_arch, list) and len(net_arch) > 0 and isinstance(net_arch[0], dict): + with pytest.warns(UserWarning): _ = model_class("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=net_arch), n_steps=64).learn(300) else: _ = model_class("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=net_arch), n_steps=64).learn(300) @@ -62,10 +62,3 @@ def test_tf_like_rmsprop_optimizer(): def test_dqn_custom_policy(): policy_kwargs = dict(optimizer_class=RMSpropTFLike, net_arch=[32]) _ = DQN("MlpPolicy", "CartPole-v1", policy_kwargs=policy_kwargs, learning_starts=100).learn(300) - - -@pytest.mark.parametrize("model_class", [A2C, PPO]) -def test_not_shared_features_extractor(model_class): - policy_kwargs = dict(net_arch=[12, dict(vf=[16], pi=[8])], share_features_extractor=False) - with pytest.raises(ValueError): - model_class("MlpPolicy", "Pendulum-v1", policy_kwargs=policy_kwargs) diff --git a/tests/test_identity.py b/tests/test_identity.py index f5bbc49..cc7746b 100644 --- a/tests/test_identity.py +++ b/tests/test_identity.py @@ -45,6 +45,8 @@ def test_continuous(model_class): n_actions = 1 action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions)) kwargs["action_noise"] = action_noise + elif model_class in [A2C]: + kwargs["policy_kwargs"]["log_std_init"] = -0.5 model = model_class("MlpPolicy", env, **kwargs).learn(n_steps) From 637988c9cc78b0f285074195055ae933f4a6df5c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Thu, 26 Jan 2023 00:31:20 +0100 Subject: [PATCH 04/12] Fix Atari wrapper bug: tried to step environment that needs reset (#1297) * fix 1060 * update changelog --- docs/misc/changelog.rst | 3 ++- stable_baselines3/common/atari_wrappers.py | 11 +++++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 7cb344b..c686539 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -23,6 +23,7 @@ New Features: Bug Fixes: ^^^^^^^^^^ +- Fixed Atari wrapper that missed the reset condition (@luizapozzobon) Deprecations: ^^^^^^^^^^^^^ @@ -1218,4 +1219,4 @@ And all the contributors: @Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede @Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont @honglu2875 @yuanmingqi @anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong -@DavyMorgan +@DavyMorgan @luizapozzobon diff --git a/stable_baselines3/common/atari_wrappers.py b/stable_baselines3/common/atari_wrappers.py index 785d911..32c1bda 100644 --- a/stable_baselines3/common/atari_wrappers.py +++ b/stable_baselines3/common/atari_wrappers.py @@ -106,7 +106,13 @@ class EpisodicLifeEnv(gym.Wrapper): obs = self.env.reset(**kwargs) else: # no-op step to advance from terminal/lost life state - obs, _, _, _ = self.env.step(0) + obs, _, done, _ = self.env.step(0) + + # The no-op step can lead to a game over, so we need to check it again + # to see if we should reset the environment and avoid the + # monitor.py `RuntimeError: Tried to step environment that needs reset` + if done: + obs = self.env.reset(**kwargs) self.lives = self.env.unwrapped.ale.lives() return obs @@ -150,9 +156,6 @@ class MaxAndSkipEnv(gym.Wrapper): return max_frame, total_reward, done, info - def reset(self, **kwargs) -> GymObs: - return self.env.reset(**kwargs) - class ClipRewardEnv(gym.RewardWrapper): """ From 5ee90095352d6e6f4da901cff341e399f5a458c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Thu, 26 Jan 2023 10:32:58 +0100 Subject: [PATCH 05/12] Add sticky actions for Atari games (#1286) * repeat_action_probability * Add test * Undo atari wrapper doc change since CI fails * remove action_repeat_probability from make_atari_env * Add sticky action wrapper and improve documentation * Update changelog * handle the case noop_max=0 * Update tests * Comply to ALE implementation * Reorder doc * Add doc warning and don't wrap with sticky action when not needed * fix docstring and reorder * Move `action_repeat_probability` args at the last position * Add ref * Update doc and wrap with frameskip only if needed * Update changelog Co-authored-by: Antonin RAFFIN --- docs/misc/changelog.rst | 4 +- stable_baselines3/common/atari_wrappers.py | 81 +++++++++++++++++----- stable_baselines3/version.txt | 2 +- tests/test_utils.py | 64 +++++++++++------ 4 files changed, 110 insertions(+), 41 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index c686539..ebbb9f2 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -4,7 +4,7 @@ Changelog ========== -Release 1.8.0a2 (WIP) +Release 1.8.0a3 (WIP) -------------------------- @@ -14,6 +14,8 @@ Breaking Changes: New Features: ^^^^^^^^^^^^^ +- Added ``repeat_action_probability`` argument in ``AtariWrapper``. +- Only use ``NoopResetEnv`` and ``MaxAndSkipEnv`` when needed in ``AtariWrapper`` `SB3-Contrib`_ ^^^^^^^^^^^^^^ diff --git a/stable_baselines3/common/atari_wrappers.py b/stable_baselines3/common/atari_wrappers.py index 32c1bda..1e06a7f 100644 --- a/stable_baselines3/common/atari_wrappers.py +++ b/stable_baselines3/common/atari_wrappers.py @@ -12,13 +12,39 @@ except ImportError: from stable_baselines3.common.type_aliases import GymObs, GymStepReturn +class StickyActionEnv(gym.Wrapper): + """ + Sticky action. + + Paper: https://arxiv.org/abs/1709.06009 + Official implementation: https://github.com/mgbellemare/Arcade-Learning-Environment + + :param env: Environment to wrap + :param action_repeat_probability: Probability of repeating the last action + """ + + def __init__(self, env: gym.Env, action_repeat_probability: float) -> None: + super().__init__(env) + self.action_repeat_probability = action_repeat_probability + assert env.unwrapped.get_action_meanings()[0] == "NOOP" + + def reset(self, **kwargs) -> GymObs: + self._sticky_action = 0 # NOOP + return self.env.reset(**kwargs) + + def step(self, action: int) -> GymStepReturn: + if self.np_random.random() >= self.action_repeat_probability: + self._sticky_action = action + return self.env.step(self._sticky_action) + + class NoopResetEnv(gym.Wrapper): """ Sample initial states by taking random number of no-ops on reset. No-op is assumed to be action 0. - :param env: the environment to wrap - :param noop_max: the maximum value of no-ops to run + :param env: Environment to wrap + :param noop_max: Maximum value of no-ops to run """ def __init__(self, env: gym.Env, noop_max: int = 30) -> None: @@ -47,7 +73,7 @@ class FireResetEnv(gym.Wrapper): """ Take action on reset for environments that are fixed until firing. - :param env: the environment to wrap + :param env: Environment to wrap """ def __init__(self, env: gym.Env) -> None: @@ -71,7 +97,7 @@ class EpisodicLifeEnv(gym.Wrapper): 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: the environment to wrap + :param env: Environment to wrap """ def __init__(self, env: gym.Env) -> None: @@ -120,9 +146,11 @@ class EpisodicLifeEnv(gym.Wrapper): class MaxAndSkipEnv(gym.Wrapper): """ Return only every ``skip``-th frame (frameskipping) + and return the max between the two last frames. - :param env: the environment - :param skip: number of ``skip``-th frame + :param env: Environment to wrap + :param skip: Number of ``skip``-th frame + The same action will be taken ``skip`` times. """ def __init__(self, env: gym.Env, skip: int = 4) -> None: @@ -159,9 +187,9 @@ class MaxAndSkipEnv(gym.Wrapper): class ClipRewardEnv(gym.RewardWrapper): """ - Clips the reward to {+1, 0, -1} by its sign. + Clip the reward to {+1, 0, -1} by its sign. - :param env: the environment + :param env: Environment to wrap """ def __init__(self, env: gym.Env) -> None: @@ -182,9 +210,9 @@ class WarpFrame(gym.ObservationWrapper): Convert to grayscale and warp frames to 84x84 (default) as done in the Nature paper and later work. - :param env: the environment - :param width: - :param height: + :param env: Environment to wrap + :param width: New frame width + :param height: New frame height """ def __init__(self, env: gym.Env, width: int = 84, height: int = 84) -> None: @@ -213,20 +241,29 @@ class AtariWrapper(gym.Wrapper): Specifically: - * NoopReset: obtain initial state by taking random number of no-ops on reset. + * Noop reset: 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. * Resize to a square image: 84x84 by default * Grayscale observation * Clip reward to {-1, 0, 1} + * Sticky actions: disabled by default - :param env: gym environment - :param noop_max: max number of no-ops - :param frame_skip: the frequency at which the agent experiences the game. - :param screen_size: resize Atari frame - :param terminal_on_life_loss: if True, then step() returns done=True whenever a life is lost. + See https://danieltakeshi.github.io/2016/11/25/frame-skipping-and-preprocessing-for-deep-q-networks-on-atari-2600-games/ + for a visual explanation. + + .. warning:: + Use this wrapper only with Atari v4 without frame skip: ``env_id = "*NoFrameskip-v4"``. + + :param env: Environment to wrap + :param noop_max: Max number of no-ops + :param frame_skip: Frequency at which the agent experiences the game. + This correspond to repeating the action ``frame_skip`` times. + :param screen_size: Resize Atari frame + :param terminal_on_life_loss: If True, then step() returns done=True whenever a life is lost. :param clip_reward: If True (default), the reward is clip to {-1, 0, 1} depending on its sign. + :param action_repeat_probability: Probability of repeating the last action """ def __init__( @@ -237,9 +274,15 @@ class AtariWrapper(gym.Wrapper): screen_size: int = 84, terminal_on_life_loss: bool = True, clip_reward: bool = True, + action_repeat_probability: float = 0.0, ) -> None: - env = NoopResetEnv(env, noop_max=noop_max) - env = MaxAndSkipEnv(env, skip=frame_skip) + if action_repeat_probability > 0.0: + env = StickyActionEnv(env, action_repeat_probability) + if noop_max > 0: + env = NoopResetEnv(env, noop_max=noop_max) + # frame_skip=1 is the same as no frame-skip (action repeat) + if frame_skip > 1: + env = MaxAndSkipEnv(env, skip=frame_skip) if terminal_on_life_loss: env = EpisodicLifeEnv(env) if "FIRE" in env.unwrapped.get_action_meanings(): diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index c3d22c0..f5e9264 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.8.0a2 +1.8.0a3 diff --git a/tests/test_utils.py b/tests/test_utils.py index 83d695a..e5236e8 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -9,7 +9,7 @@ from gym import spaces import stable_baselines3 as sb3 from stable_baselines3 import A2C -from stable_baselines3.common.atari_wrappers import ClipRewardEnv, MaxAndSkipEnv +from stable_baselines3.common.atari_wrappers import MaxAndSkipEnv from stable_baselines3.common.env_util import is_wrapped, make_atari_env, make_vec_env, unwrap_wrapper from stable_baselines3.common.evaluation import evaluate_policy from stable_baselines3.common.monitor import Monitor @@ -55,30 +55,54 @@ def test_make_vec_env_func_checker(): env.close() -@pytest.mark.parametrize("env_id", ["BreakoutNoFrameskip-v4"]) -@pytest.mark.parametrize("n_envs", [1, 2]) -@pytest.mark.parametrize("wrapper_kwargs", [None, dict(clip_reward=False, screen_size=60)]) -def test_make_atari_env(env_id, n_envs, wrapper_kwargs): - env = make_atari_env(env_id, n_envs, wrapper_kwargs=wrapper_kwargs, monitor_dir=None, seed=0) +# Use Asterix as it does not requires fire reset +@pytest.mark.parametrize("env_id", ["BreakoutNoFrameskip-v4", "AsterixNoFrameskip-v4"]) +@pytest.mark.parametrize("noop_max", [0, 10]) +@pytest.mark.parametrize("action_repeat_probability", [0.0, 0.25]) +@pytest.mark.parametrize("frame_skip", [1, 4]) +@pytest.mark.parametrize("screen_size", [60]) +@pytest.mark.parametrize("terminal_on_life_loss", [True, False]) +@pytest.mark.parametrize("clip_reward", [True]) +def test_make_atari_env( + env_id, noop_max, action_repeat_probability, frame_skip, screen_size, terminal_on_life_loss, clip_reward +): + n_envs = 2 + wrapper_kwargs = { + "noop_max": noop_max, + "action_repeat_probability": action_repeat_probability, + "frame_skip": frame_skip, + "screen_size": screen_size, + "terminal_on_life_loss": terminal_on_life_loss, + "clip_reward": clip_reward, + } + venv = make_atari_env( + env_id, + n_envs=2, + wrapper_kwargs=wrapper_kwargs, + monitor_dir=None, + seed=0, + ) - assert env.num_envs == n_envs + assert venv.num_envs == n_envs - obs = env.reset() + needs_fire_reset = env_id == "BreakoutNoFrameskip-v4" + expected_frame_number_low = frame_skip * 2 if needs_fire_reset else 0 # FIRE - UP on reset + expected_frame_number_high = expected_frame_number_low + noop_max + expected_shape = (n_envs, screen_size, screen_size, 1) - new_obs, reward, _, _ = env.step([env.action_space.sample() for _ in range(n_envs)]) + obs = venv.reset() + frame_numbers = [env.unwrapped.ale.getEpisodeFrameNumber() for env in venv.envs] + for frame_number in frame_numbers: + assert expected_frame_number_low <= frame_number <= expected_frame_number_high + assert obs.shape == expected_shape - assert obs.shape == new_obs.shape + new_obs, reward, _, _ = venv.step([venv.action_space.sample() for _ in range(n_envs)]) - # Wrapped into DummyVecEnv - wrapped_atari_env = env.envs[0] - 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 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 isinstance(wrapped_atari_env.env, ClipRewardEnv) + new_frame_numbers = [env.unwrapped.ale.getEpisodeFrameNumber() for env in venv.envs] + for frame_number, new_frame_number in zip(frame_numbers, new_frame_numbers): + assert new_frame_number - frame_number == frame_skip + assert new_obs.shape == expected_shape + if clip_reward: assert np.max(np.abs(reward)) < 1.0 From bea3c44ba52278ec755af0179859b04ab80cdcaf Mon Sep 17 00:00:00 2001 From: Alex Pasquali Date: Sat, 28 Jan 2023 12:04:07 +0100 Subject: [PATCH 06/12] Fixed typo in A2C's docstring (#1303) --- docs/misc/changelog.rst | 1 + stable_baselines3/a2c/a2c.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index ebbb9f2..7ff0c0e 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -39,6 +39,7 @@ Others: Documentation: ^^^^^^^^^^^^^^ - Renamed ``load_parameters`` to ``set_parameters`` (@DavyMorgan) +- Fixed typo in ``A2C`` docstring (@AlexPasqua) Release 1.7.0 (2023-01-10) diff --git a/stable_baselines3/a2c/a2c.py b/stable_baselines3/a2c/a2c.py index 972e700..ec4ae2e 100644 --- a/stable_baselines3/a2c/a2c.py +++ b/stable_baselines3/a2c/a2c.py @@ -29,7 +29,7 @@ class A2C(OnPolicyAlgorithm): :param n_steps: The number of steps to run for each environment per update (i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel) :param gamma: Discount factor - :param gae_lambda: Factor for trade-off of bias vs variance for Generalized Advantage Estimator + :param gae_lambda: Factor for trade-off of bias vs variance for Generalized Advantage Estimator. Equivalent to classic advantage when set to 1. :param ent_coef: Entropy coefficient for the loss calculation :param vf_coef: Value function coefficient for the loss calculation From 82bc63fca4273c72deea247faa655e17d47bdd85 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Thu, 2 Feb 2023 11:58:41 +0100 Subject: [PATCH 07/12] Upgrade black formatting (#1310) * apply black * Reformat tests --------- Co-authored-by: Antonin Raffin --- stable_baselines3/a2c/a2c.py | 3 --- stable_baselines3/common/buffers.py | 7 ------- stable_baselines3/common/callbacks.py | 2 -- stable_baselines3/common/evaluation.py | 1 - stable_baselines3/common/logger.py | 5 +---- stable_baselines3/common/off_policy_algorithm.py | 2 -- stable_baselines3/common/on_policy_algorithm.py | 2 -- stable_baselines3/common/policies.py | 1 - stable_baselines3/common/results_plotter.py | 2 +- stable_baselines3/common/save_util.py | 2 +- stable_baselines3/common/vec_env/stacked_observations.py | 1 - stable_baselines3/common/vec_env/vec_frame_stack.py | 1 - stable_baselines3/common/vec_env/vec_video_recorder.py | 1 - stable_baselines3/ddpg/ddpg.py | 2 -- stable_baselines3/dqn/dqn.py | 2 -- stable_baselines3/her/her_replay_buffer.py | 2 -- stable_baselines3/ppo/ppo.py | 2 -- stable_baselines3/sac/sac.py | 2 -- stable_baselines3/td3/td3.py | 3 --- tests/test_run.py | 1 - tests/test_save_load.py | 1 - tests/test_vec_normalize.py | 3 +-- 22 files changed, 4 insertions(+), 44 deletions(-) diff --git a/stable_baselines3/a2c/a2c.py b/stable_baselines3/a2c/a2c.py index ec4ae2e..9e8b40c 100644 --- a/stable_baselines3/a2c/a2c.py +++ b/stable_baselines3/a2c/a2c.py @@ -81,7 +81,6 @@ class A2C(OnPolicyAlgorithm): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy, env, @@ -132,7 +131,6 @@ class A2C(OnPolicyAlgorithm): # This will only loop once (get all data in one go) for rollout_data in self.rollout_buffer.get(batch_size=None): - actions = rollout_data.actions if isinstance(self.action_space, spaces.Discrete): # Convert discrete action from float to long @@ -189,7 +187,6 @@ class A2C(OnPolicyAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfA2C: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index f71dd29..273dba9 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -240,7 +240,6 @@ class ReplayBuffer(BaseBuffer): done: np.ndarray, infos: List[Dict[str, Any]], ) -> None: - # Reshape needed when using multiple envs with discrete observations # as numpy cannot broadcast (n_discrete,) to (n_discrete, 1) if isinstance(self.observation_space, spaces.Discrete): @@ -346,7 +345,6 @@ class RolloutBuffer(BaseBuffer): gamma: float = 0.99, n_envs: int = 1, ): - super().__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs) self.gae_lambda = gae_lambda self.gamma = gamma @@ -356,7 +354,6 @@ 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.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) @@ -451,7 +448,6 @@ class RolloutBuffer(BaseBuffer): indices = np.random.permutation(self.buffer_size * self.n_envs) # Prepare the data if not self.generator_ready: - _tensor_names = [ "observations", "actions", @@ -688,7 +684,6 @@ class DictRolloutBuffer(RolloutBuffer): gamma: float = 0.99, n_envs: int = 1, ): - super(RolloutBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs) assert isinstance(self.obs_shape, dict), "DictRolloutBuffer must be used with Dict obs space only" @@ -763,7 +758,6 @@ class DictRolloutBuffer(RolloutBuffer): indices = np.random.permutation(self.buffer_size * self.n_envs) # Prepare the data if not self.generator_ready: - for key, obs in self.observations.items(): self.observations[key] = self.swap_and_flatten(obs) @@ -787,7 +781,6 @@ class DictRolloutBuffer(RolloutBuffer): batch_inds: np.ndarray, env: Optional[VecNormalize] = None, ) -> DictRolloutBufferSamples: # type: ignore[signature-mismatch] #FIXME - return DictRolloutBufferSamples( observations={key: self.to_torch(obs[batch_inds]) for (key, obs) in self.observations.items()}, actions=self.to_torch(self.actions[batch_inds]), diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py index a96c52c..69a21ab 100644 --- a/stable_baselines3/common/callbacks.py +++ b/stable_baselines3/common/callbacks.py @@ -429,11 +429,9 @@ class EvalCallback(EventCallback): self._is_success_buffer.append(maybe_is_success) def _on_step(self) -> bool: - continue_training = True if self.eval_freq > 0 and self.n_calls % self.eval_freq == 0: - # Sync training and eval env if there is VecNormalize if self.model.get_vec_normalize_env() is not None: try: diff --git a/stable_baselines3/common/evaluation.py b/stable_baselines3/common/evaluation.py index ff18137..b65edf8 100644 --- a/stable_baselines3/common/evaluation.py +++ b/stable_baselines3/common/evaluation.py @@ -91,7 +91,6 @@ def evaluate_policy( current_lengths += 1 for i in range(n_envs): if episode_counts[i] < episode_count_targets[i]: - # unpack values so that the callback can access the local variables reward = rewards[i] done = dones[i] diff --git a/stable_baselines3/common/logger.py b/stable_baselines3/common/logger.py index 939d924..a8aa766 100644 --- a/stable_baselines3/common/logger.py +++ b/stable_baselines3/common/logger.py @@ -173,7 +173,6 @@ class HumanOutputFormat(KVWriter, SeqWriter): key2str = {} tag = None for (key, value), (_, excluded) in zip(sorted(key_values.items()), sorted(key_excluded.items())): - if excluded is not None and ("stdout" in excluded or "log" in excluded): continue @@ -342,7 +341,7 @@ class CSVOutputFormat(KVWriter): self.file.seek(0) lines = self.file.readlines() self.file.seek(0) - for (i, key) in enumerate(self.keys): + for i, key in enumerate(self.keys): if i > 0: self.file.write(",") self.file.write(key) @@ -399,9 +398,7 @@ class TensorBoardOutputFormat(KVWriter): self.writer = SummaryWriter(log_dir=folder) def write(self, key_values: Dict[str, Any], key_excluded: Dict[str, Union[str, Tuple[str, ...]]], step: int = 0) -> None: - for (key, value), (_, excluded) in zip(sorted(key_values.items()), sorted(key_excluded.items())): - if excluded is not None and "tensorboard" in excluded: continue diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 48779ef..c1ab215 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -102,7 +102,6 @@ class OffPolicyAlgorithm(BaseAlgorithm): sde_support: bool = True, supported_action_spaces: Optional[Tuple[spaces.Space, ...]] = None, ): - super().__init__( policy=policy, env=env, @@ -319,7 +318,6 @@ class OffPolicyAlgorithm(BaseAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfOffPolicyAlgorithm: - total_timesteps, callback = self._setup_learn( total_timesteps, callback, diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index bc0dda4..44d8b26 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/stable_baselines3/common/on_policy_algorithm.py @@ -72,7 +72,6 @@ class OnPolicyAlgorithm(BaseAlgorithm): _init_setup_model: bool = True, supported_action_spaces: Optional[Tuple[spaces.Space, ...]] = None, ): - super().__init__( policy=policy, env=env, @@ -244,7 +243,6 @@ class OnPolicyAlgorithm(BaseAlgorithm): callback.on_training_start(locals(), globals()) while self.num_timesteps < total_timesteps: - continue_training = self.collect_rollouts(self.env, callback, self.rollout_buffer, n_rollout_steps=self.n_steps) if continue_training is False: diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 793cfc5..457274a 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -433,7 +433,6 @@ class ActorCriticPolicy(BasePolicy): optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam, optimizer_kwargs: Optional[Dict[str, Any]] = None, ): - if optimizer_kwargs is None: optimizer_kwargs = {} # Small values to avoid NaN in Adam optimizer diff --git a/stable_baselines3/common/results_plotter.py b/stable_baselines3/common/results_plotter.py index 92f67ac..dac2b6c 100644 --- a/stable_baselines3/common/results_plotter.py +++ b/stable_baselines3/common/results_plotter.py @@ -84,7 +84,7 @@ def plot_curves( plt.figure(title, figsize=figsize) max_x = max(xy[0][-1] for xy in xy_list) min_x = 0 - for (_, (x, y)) in enumerate(xy_list): + for _, (x, y) in enumerate(xy_list): plt.scatter(x, y, s=2) # Do not plot the smoothed curve at all if the timeseries is shorter than window size. if x.shape[0] >= EPISODES_WINDOW: diff --git a/stable_baselines3/common/save_util.py b/stable_baselines3/common/save_util.py index facc55a..7ae1e22 100644 --- a/stable_baselines3/common/save_util.py +++ b/stable_baselines3/common/save_util.py @@ -367,7 +367,7 @@ def load_from_zip_file( device: Union[th.device, str] = "auto", verbose: int = 0, print_system_info: bool = False, -) -> (Tuple[Optional[Dict[str, Any]], Optional[TensorDict], Optional[TensorDict]]): +) -> Tuple[Optional[Dict[str, Any]], Optional[TensorDict], Optional[TensorDict]]: """ Load model data from a .zip archive diff --git a/stable_baselines3/common/vec_env/stacked_observations.py b/stable_baselines3/common/vec_env/stacked_observations.py index 8583518..d373b87 100644 --- a/stable_baselines3/common/vec_env/stacked_observations.py +++ b/stable_baselines3/common/vec_env/stacked_observations.py @@ -30,7 +30,6 @@ class StackedObservations: observation_space: spaces.Space, channels_order: Optional[str] = None, ): - self.n_stack = n_stack ( self.channels_first, diff --git a/stable_baselines3/common/vec_env/vec_frame_stack.py b/stable_baselines3/common/vec_env/vec_frame_stack.py index e06d512..d933104 100644 --- a/stable_baselines3/common/vec_env/vec_frame_stack.py +++ b/stable_baselines3/common/vec_env/vec_frame_stack.py @@ -44,7 +44,6 @@ class VecFrameStack(VecEnvWrapper): def step_wait( self, ) -> Tuple[Union[np.ndarray, Dict[str, np.ndarray]], np.ndarray, np.ndarray, List[Dict[str, Any]],]: - observations, rewards, dones, infos = self.venv.step_wait() observations, infos = self.stackedobs.update(observations, dones, infos) diff --git a/stable_baselines3/common/vec_env/vec_video_recorder.py b/stable_baselines3/common/vec_env/vec_video_recorder.py index 70d74eb..83d058a 100644 --- a/stable_baselines3/common/vec_env/vec_video_recorder.py +++ b/stable_baselines3/common/vec_env/vec_video_recorder.py @@ -30,7 +30,6 @@ class VecVideoRecorder(VecEnvWrapper): video_length: int = 200, name_prefix: str = "rl-video", ): - VecEnvWrapper.__init__(self, venv) self.env = venv diff --git a/stable_baselines3/ddpg/ddpg.py b/stable_baselines3/ddpg/ddpg.py index 40d67b5..c311b23 100644 --- a/stable_baselines3/ddpg/ddpg.py +++ b/stable_baselines3/ddpg/ddpg.py @@ -76,7 +76,6 @@ class DDPG(TD3): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy=policy, env=env, @@ -121,7 +120,6 @@ class DDPG(TD3): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfDDPG: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index dd8794e..ea1946a 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -94,7 +94,6 @@ class DQN(OffPolicyAlgorithm): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy, env, @@ -261,7 +260,6 @@ class DQN(OffPolicyAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfDQN: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/stable_baselines3/her/her_replay_buffer.py b/stable_baselines3/her/her_replay_buffer.py index 1518436..0c3da25 100644 --- a/stable_baselines3/her/her_replay_buffer.py +++ b/stable_baselines3/her/her_replay_buffer.py @@ -81,7 +81,6 @@ class HerReplayBuffer(DictReplayBuffer): online_sampling: bool = True, handle_timeout_termination: bool = True, ): - super().__init__(buffer_size, env.observation_space, env.action_space, device, env.num_envs) # convert goal_selection_strategy into GoalSelectionStrategy if string @@ -389,7 +388,6 @@ class HerReplayBuffer(DictReplayBuffer): done: np.ndarray, infos: List[Dict[str, Any]], ) -> None: - if self.current_idx == 0 and self.full: # Clear info buffer self.info_buffer[self.pos] = deque(maxlen=self.max_episode_length) diff --git a/stable_baselines3/ppo/ppo.py b/stable_baselines3/ppo/ppo.py index bd80736..c934527 100644 --- a/stable_baselines3/ppo/ppo.py +++ b/stable_baselines3/ppo/ppo.py @@ -98,7 +98,6 @@ class PPO(OnPolicyAlgorithm): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy, env, @@ -303,7 +302,6 @@ class PPO(OnPolicyAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfPPO: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 74285b6..d1a6610 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -109,7 +109,6 @@ class SAC(OffPolicyAlgorithm): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy, env, @@ -295,7 +294,6 @@ class SAC(OffPolicyAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfSAC: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index ae442e1..c844a99 100644 --- a/stable_baselines3/td3/td3.py +++ b/stable_baselines3/td3/td3.py @@ -94,7 +94,6 @@ class TD3(OffPolicyAlgorithm): device: Union[th.device, str] = "auto", _init_setup_model: bool = True, ): - super().__init__( policy, env, @@ -151,7 +150,6 @@ class TD3(OffPolicyAlgorithm): actor_losses, critic_losses = [], [] for _ in range(gradient_steps): - self._n_updates += 1 # Sample replay buffer replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env) @@ -210,7 +208,6 @@ class TD3(OffPolicyAlgorithm): reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfTD3: - return super().learn( total_timesteps=total_timesteps, callback=callback, diff --git a/tests/test_run.py b/tests/test_run.py index 71236a3..ca7548f 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -115,7 +115,6 @@ def test_dqn(): @pytest.mark.parametrize("train_freq", [4, (4, "step"), (1, "episode")]) def test_train_freq(tmp_path, train_freq): - model = SAC( "MlpPolicy", "Pendulum-v1", diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 2c35e43..9d3d537 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -648,7 +648,6 @@ def test_open_file_str_pathlib(tmp_path, pathtype): def test_open_file(tmp_path): - # path must much the type with pytest.raises(TypeError): open_path(123, None, None, None) diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index b17d28c..fb37fd3 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -178,7 +178,7 @@ def _make_warmstart_dict_env(**kwargs): def test_runningmeanstd(): """Test RunningMeanStd object""" - for (x_1, x_2, x_3) in [ + for x_1, x_2, x_3 in [ (np.random.randn(3), np.random.randn(4), np.random.randn(5)), (np.random.randn(3, 2), np.random.randn(4, 2), np.random.randn(5, 2)), ]: @@ -336,7 +336,6 @@ def test_normalize_dict_selected_keys(): @pytest.mark.parametrize("model_class", [SAC, TD3, HerReplayBuffer]) @pytest.mark.parametrize("online_sampling", [False, True]) def test_offpolicy_normalization(model_class, online_sampling): - if online_sampling and model_class != HerReplayBuffer: pytest.skip() From d0c1a87faf8490a8bdaa1ccd79da15cbf105ee81 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marco=20Tr=C3=B6ster?= Date: Thu, 2 Feb 2023 12:34:38 +0100 Subject: [PATCH 08/12] Add scaling section to A2C documentation (#1250) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add scaling section to A2C documentation * add cross-reference to vectorized envs article * turn it as note * update changelog * add Bonifatius94 to the list of contributors * fix issue number --------- Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com> Co-authored-by: Quentin GALLOUÉDEC Co-authored-by: Antonin RAFFIN --- docs/misc/changelog.rst | 3 ++- docs/modules/a2c.rst | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 7ff0c0e..9dee356 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -39,6 +39,7 @@ Others: Documentation: ^^^^^^^^^^^^^^ - Renamed ``load_parameters`` to ``set_parameters`` (@DavyMorgan) +- Clarified documentation about subproc multiprocessing for A2C (@Bonifatius94) - Fixed typo in ``A2C`` docstring (@AlexPasqua) @@ -1222,4 +1223,4 @@ And all the contributors: @Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede @Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont @honglu2875 @yuanmingqi @anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong -@DavyMorgan @luizapozzobon +@DavyMorgan @luizapozzobon @Bonifatius94 diff --git a/docs/modules/a2c.rst b/docs/modules/a2c.rst index e871424..670da61 100644 --- a/docs/modules/a2c.rst +++ b/docs/modules/a2c.rst @@ -76,6 +76,24 @@ Train a A2C agent on ``CartPole-v1`` using 4 environments. env.render() +.. note:: + + A2C is meant to be run primarily on the CPU, especially when you are not using a CNN. To improve CPU utilization, try turning off the GPU and using ``SubprocVecEnv`` instead of the default ``DummyVecEnv``: + + .. code-block:: + + from stable_baselines3 import A2C + from stable_baselines3.common.env_util import make_vec_env + from stable_baselines3.common.vec_env import SubprocVecEnv + + if __name__=="__main__": + env = make_vec_env("CartPole-v1", n_envs=8, vec_env_cls=SubprocVecEnv) + model = A2C("MlpPolicy", env, device="cpu") + model.learn(total_timesteps=25_000) + + For more information, see :ref:`Vectorized Environments `, `Issue #1245 `_ or the `Multiprocessing notebook `_. + + Results ------- From 411ff697dde31df9ff914ba0c335875d1bd5998d Mon Sep 17 00:00:00 2001 From: adamfrly <45516720+adamfrly@users.noreply.github.com> Date: Mon, 6 Feb 2023 09:48:41 -0500 Subject: [PATCH 09/12] Ensure train/n_updates metric accounts for early stopping of training loop (#1311) * Correct _n_updates when target_kl stops loop early * Update changelog * Simplify code --------- Co-authored-by: Antonin Raffin --- docs/misc/changelog.rst | 1 + stable_baselines3/ppo/ppo.py | 3 +-- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 9dee356..acfa5e3 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -26,6 +26,7 @@ New Features: Bug Fixes: ^^^^^^^^^^ - Fixed Atari wrapper that missed the reset condition (@luizapozzobon) +- Fixed PPO train/n_updates metric not accounting for early stopping (@adamfrly) Deprecations: ^^^^^^^^^^^^^ diff --git a/stable_baselines3/ppo/ppo.py b/stable_baselines3/ppo/ppo.py index c934527..3ea6756 100644 --- a/stable_baselines3/ppo/ppo.py +++ b/stable_baselines3/ppo/ppo.py @@ -189,7 +189,6 @@ class PPO(OnPolicyAlgorithm): clip_fractions = [] continue_training = True - # train for n_epochs epochs for epoch in range(self.n_epochs): approx_kl_divs = [] @@ -271,10 +270,10 @@ class PPO(OnPolicyAlgorithm): th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) self.policy.optimizer.step() + self._n_updates += 1 if not continue_training: break - self._n_updates += self.n_epochs explained_var = explained_variance(self.rollout_buffer.values.flatten(), self.rollout_buffer.returns.flatten()) # Logs From 2e4a45020ec619b09e2b1ccff14fa4f2c291dc77 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Mon, 6 Feb 2023 22:41:59 +0100 Subject: [PATCH 10/12] Refactor observation stacking (#1238) * refactor stacking obs * Improve docstring * remove all StackedDictObservations * Update tests and make stacked obs clearer * Fix type check * fix stacked_observation_space * undo init change, deprecate StackedDictObservations * deprecate stack_observation_space * type hints * ignore pytype errors * undo vecenv doc change * Deprecation warning in StackedDictObs doctstring * Fix vec_env.rst * Fix __all__ sorting * fix pytype ignore statement * Update docstring * stack * Remove n_stack * Update changelog * Simplify code * Rename test file * Re-use variable for shift * Fix doc build * Remove pytype comment * Disable pytype error --------- Co-authored-by: Antonin RAFFIN --- docs/guide/vec_envs.rst | 6 - docs/misc/changelog.rst | 4 +- setup.cfg | 1 - stable_baselines3/common/vec_env/__init__.py | 3 +- .../common/vec_env/stacked_observations.py | 316 +++++++----------- .../common/vec_env/vec_frame_stack.py | 47 +-- stable_baselines3/version.txt | 2 +- tests/test_vec_stacked_obs.py | 314 +++++++++++++++++ 8 files changed, 459 insertions(+), 234 deletions(-) create mode 100644 tests/test_vec_stacked_obs.py diff --git a/docs/guide/vec_envs.rst b/docs/guide/vec_envs.rst index b074dad..d847811 100644 --- a/docs/guide/vec_envs.rst +++ b/docs/guide/vec_envs.rst @@ -122,12 +122,6 @@ StackedObservations .. autoclass:: stable_baselines3.common.vec_env.stacked_observations.StackedObservations :members: -StackedDictObservations -~~~~~~~~~~~~~~~~~~~~~~~ - -.. autoclass:: stable_baselines3.common.vec_env.stacked_observations.StackedDictObservations - :members: - VecNormalize ~~~~~~~~~~~~ diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index acfa5e3..9d24bbe 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -4,13 +4,14 @@ Changelog ========== -Release 1.8.0a3 (WIP) +Release 1.8.0a4 (WIP) -------------------------- Breaking Changes: ^^^^^^^^^^^^^^^^^ - Removed shared layers in ``mlp_extractor`` (@AlexPasqua) +- Refactored ``StackedObservations`` (it now handles dict obs, ``StackedDictObservations`` was removed) New Features: ^^^^^^^^^^^^^ @@ -36,6 +37,7 @@ Others: - Fixed ``tests/test_tensorboard.py`` type hint - Fixed ``tests/test_vec_normalize.py`` type hint - Fixed ``stable_baselines3/common/monitor.py`` type hint +- Added tests for StackedObservations Documentation: ^^^^^^^^^^^^^^ diff --git a/setup.cfg b/setup.cfg index 37ffa17..3698ded 100644 --- a/setup.cfg +++ b/setup.cfg @@ -48,7 +48,6 @@ exclude = (?x)( | stable_baselines3/common/vec_env/__init__.py$ | stable_baselines3/common/vec_env/base_vec_env.py$ | stable_baselines3/common/vec_env/dummy_vec_env.py$ - | stable_baselines3/common/vec_env/stacked_observations.py$ | stable_baselines3/common/vec_env/subproc_vec_env.py$ | stable_baselines3/common/vec_env/util.py$ | stable_baselines3/common/vec_env/vec_extract_dict_obs.py$ diff --git a/stable_baselines3/common/vec_env/__init__.py b/stable_baselines3/common/vec_env/__init__.py index 33a103a..2c03637 100644 --- a/stable_baselines3/common/vec_env/__init__.py +++ b/stable_baselines3/common/vec_env/__init__.py @@ -4,7 +4,7 @@ from typing import Optional, Type, Union from stable_baselines3.common.vec_env.base_vec_env import CloudpickleWrapper, VecEnv, VecEnvWrapper from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv -from stable_baselines3.common.vec_env.stacked_observations import StackedDictObservations, StackedObservations +from stable_baselines3.common.vec_env.stacked_observations import StackedObservations from stable_baselines3.common.vec_env.subproc_vec_env import SubprocVecEnv from stable_baselines3.common.vec_env.vec_check_nan import VecCheckNan from stable_baselines3.common.vec_env.vec_extract_dict_obs import VecExtractDictObs @@ -78,7 +78,6 @@ __all__ = [ "VecEnv", "VecEnvWrapper", "DummyVecEnv", - "StackedDictObservations", "StackedObservations", "SubprocVecEnv", "VecCheckNan", diff --git a/stable_baselines3/common/vec_env/stacked_observations.py b/stable_baselines3/common/vec_env/stacked_observations.py index d373b87..a26812c 100644 --- a/stable_baselines3/common/vec_env/stacked_observations.py +++ b/stable_baselines3/common/vec_env/stacked_observations.py @@ -1,61 +1,80 @@ import warnings -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, Generic, List, Mapping, Optional, Tuple, TypeVar, Union import numpy as np from gym import spaces from stable_baselines3.common.preprocessing import is_image_space, is_image_space_channels_first +TObs = TypeVar("TObs", np.ndarray, Dict[str, np.ndarray]) -class StackedObservations: + +# Disable errors for pytype which doesn't play well with Generic[TypeVar] +# mypy check passes though +# pytype: disable=attribute-error +class StackedObservations(Generic[TObs]): """ Frame stacking wrapper for data. - Dimension to stack over is either first (channels-first) or - last (channels-last), which is detected automatically using - ``common.preprocessing.is_image_space_channels_first`` if - observation is an image space. + Dimension to stack over is either first (channels-first) or last (channels-last), which is detected automatically using + ``common.preprocessing.is_image_space_channels_first`` if observation is an image space. - :param num_envs: number of environments + :param num_envs: Number of environments :param n_stack: Number of frames to stack - :param observation_space: Environment observation space. + :param observation_space: Environment observation space :param channels_order: If "first", stack on first image dimension. If "last", stack on last dimension. - If None, automatically detect channel to stack over in case of image observation or default to "last" (default). + If None, automatically detect channel to stack over in case of image observation or default to "last". + For Dict space, channels_order can also be a dictionary. """ def __init__( self, num_envs: int, n_stack: int, - observation_space: spaces.Space, - channels_order: Optional[str] = None, - ): + observation_space: Union[spaces.Box, spaces.Dict], # Replace by Space[TObs] in gym>=0.26 + channels_order: Optional[Union[str, Mapping[str, Optional[str]]]] = None, + ) -> None: self.n_stack = n_stack - ( - self.channels_first, - self.stack_dimension, - self.stackedobs, - self.repeat_axis, - ) = self.compute_stacking(num_envs, n_stack, observation_space, channels_order) - super().__init__() + self.observation_space = observation_space + if isinstance(observation_space, spaces.Dict): + if not isinstance(channels_order, Mapping): + channels_order = {key: channels_order for key in observation_space.spaces.keys()} + self.sub_stacked_observations = { + key: StackedObservations(num_envs, n_stack, subspace, channels_order[key]) + for key, subspace in observation_space.spaces.items() + } + self.stacked_observation_space = spaces.Dict( + {key: substack_obs.stacked_observation_space for key, substack_obs in self.sub_stacked_observations.items()} + ) # type: spaces.Dict # make mypy happy + elif isinstance(observation_space, spaces.Box): + if isinstance(channels_order, Mapping): + raise TypeError("When the observation space is Box, channels_order can't be a dict.") + + self.channels_first, self.stack_dimension, self.stacked_shape, self.repeat_axis = self.compute_stacking( + n_stack, observation_space, channels_order + ) + low = np.repeat(observation_space.low, n_stack, axis=self.repeat_axis) + high = np.repeat(observation_space.high, n_stack, axis=self.repeat_axis) + self.stacked_observation_space = spaces.Box(low=low, high=high, dtype=observation_space.dtype) + self.stacked_obs = np.zeros((num_envs,) + self.stacked_shape, dtype=observation_space.dtype) + else: + raise TypeError( + f"StackedObservations only supports Box and Dict as observation spaces. {observation_space} was provided." + ) @staticmethod def compute_stacking( - num_envs: int, - n_stack: int, - observation_space: spaces.Box, - channels_order: Optional[str] = None, - ) -> Tuple[bool, int, np.ndarray, int]: + n_stack: int, observation_space: spaces.Box, channels_order: Optional[str] = None + ) -> Tuple[bool, int, Tuple[int, ...], int]: """ Calculates the parameters in order to stack observations - :param num_envs: Number of environments in the stack - :param n_stack: The number of observations to stack - :param observation_space: The observation space - :param channels_order: The order of the channels - :return: tuple of channels_first, stack_dimension, stackedobs, repeat_axis + :param n_stack: Number of observations to stack + :param observation_space: Observation space + :param channels_order: Order of the channels + :return: Tuple of channels_first, stack_dimension, stackedobs, repeat_axis """ - channels_first = False + if channels_order is None: # Detect channel location automatically for images if is_image_space(observation_space): @@ -74,192 +93,113 @@ class StackedObservations: # This includes the vec-env dimension (first) stack_dimension = 1 if channels_first else -1 repeat_axis = 0 if channels_first else -1 - low = np.repeat(observation_space.low, n_stack, axis=repeat_axis) - stackedobs = np.zeros((num_envs,) + low.shape, low.dtype) - return channels_first, stack_dimension, stackedobs, repeat_axis + stacked_shape = list(observation_space.shape) + stacked_shape[repeat_axis] *= n_stack + return channels_first, stack_dimension, tuple(stacked_shape), repeat_axis - def stack_observation_space(self, observation_space: spaces.Box) -> spaces.Box: + def stack_observation_space(self, observation_space: Union[spaces.Box, spaces.Dict]) -> Union[spaces.Box, spaces.Dict]: """ - Given an observation space, returns a new observation space with stacked observations + This function is deprecated. + + As an alternative, use + + .. code-block:: python + + low = np.repeat(observation_space.low, stacked_observation.n_stack, axis=stacked_observation.repeat_axis) + high = np.repeat(observation_space.high, stacked_observation.n_stack, axis=stacked_observation.repeat_axis) + stacked_observation_space = spaces.Box(low=low, high=high, dtype=observation_space.dtype) :return: New observation space with stacked dimensions """ + warnings.warn( + "stack_observation_space is deprecated and will be removed in the next SB3 release. " + "Please refer to the docstring for a workaround.", + DeprecationWarning, + ) + if isinstance(observation_space, spaces.Dict): + return spaces.Dict( + { + key: sub_stacked_observation.stack_observation_space(sub_stacked_observation.observation_space) + for key, sub_stacked_observation in self.sub_stacked_observations.items() + } + ) low = np.repeat(observation_space.low, self.n_stack, axis=self.repeat_axis) high = np.repeat(observation_space.high, self.n_stack, axis=self.repeat_axis) return spaces.Box(low=low, high=high, dtype=observation_space.dtype) - def reset(self, observation: np.ndarray) -> np.ndarray: + def reset(self, observation: TObs) -> TObs: """ - Resets the stackedobs, adds the reset observation to the stack, and returns the stack + Reset the stacked_obs, add the reset observation to the stack, and return the stack. :param observation: Reset observation :return: The stacked reset observation """ - self.stackedobs[...] = 0 + if isinstance(observation, dict): + return {key: self.sub_stacked_observations[key].reset(obs) for key, obs in observation.items()} + + self.stacked_obs[...] = 0 if self.channels_first: - self.stackedobs[:, -observation.shape[self.stack_dimension] :, ...] = observation + self.stacked_obs[:, -observation.shape[self.stack_dimension] :, ...] = observation else: - self.stackedobs[..., -observation.shape[self.stack_dimension] :] = observation - return self.stackedobs + self.stacked_obs[..., -observation.shape[self.stack_dimension] :] = observation + return self.stacked_obs def update( self, - observations: np.ndarray, + observations: TObs, dones: np.ndarray, infos: List[Dict[str, Any]], - ) -> Tuple[np.ndarray, List[Dict[str, Any]]]: + ) -> Tuple[TObs, List[Dict[str, Any]]]: """ - Adds the observations to the stack and uses the dones to update the infos. + Add the observations to the stack and use the dones to update the infos. - :param observations: numpy array of observations - :param dones: numpy array of done info - :param infos: numpy array of info dicts - :return: tuple of the stacked observations and the updated infos + :param observations: Observations + :param dones: Dones + :param infos: Infos + :return: Tuple of the stacked observations and the updated infos """ - stack_ax_size = observations.shape[self.stack_dimension] - self.stackedobs = np.roll(self.stackedobs, shift=-stack_ax_size, axis=self.stack_dimension) - for i, done in enumerate(dones): + if isinstance(observations, dict): + # From [{}, {terminal_obs: {key1: ..., key2: ...}}] + # to {key1: [{}, {terminal_obs: ...}], key2: [{}, {terminal_obs: ...}]} + sub_infos = { + key: [ + {"terminal_observation": info["terminal_observation"][key]} if "terminal_observation" in info else {} + for info in infos + ] + for key in observations.keys() + } + + stacked_obs = {} + stacked_infos = {} + for key, obs in observations.items(): + stacked_obs[key], stacked_infos[key] = self.sub_stacked_observations[key].update(obs, dones, sub_infos[key]) + + # From {key1: [{}, {terminal_obs: ...}], key2: [{}, {terminal_obs: ...}]} + # to [{}, {terminal_obs: {key1: ..., key2: ...}}] + for key in stacked_infos.keys(): + for env_idx in range(len(infos)): + if "terminal_observation" in infos[env_idx]: + infos[env_idx]["terminal_observation"][key] = stacked_infos[key][env_idx]["terminal_observation"] + return stacked_obs, infos + + shift = -observations.shape[self.stack_dimension] + self.stacked_obs = np.roll(self.stacked_obs, shift, axis=self.stack_dimension) + for env_idx, done in enumerate(dones): if done: - if "terminal_observation" in infos[i]: - old_terminal = infos[i]["terminal_observation"] + if "terminal_observation" in infos[env_idx]: + old_terminal = infos[env_idx]["terminal_observation"] if self.channels_first: - new_terminal = np.concatenate( - (self.stackedobs[i, :-stack_ax_size, ...], old_terminal), - axis=0, # self.stack_dimension - 1, as there is not batch dim - ) + previous_stack = self.stacked_obs[env_idx, :shift, ...] else: - new_terminal = np.concatenate( - (self.stackedobs[i, ..., :-stack_ax_size], old_terminal), - axis=self.stack_dimension, - ) - infos[i]["terminal_observation"] = new_terminal + previous_stack = self.stacked_obs[env_idx, ..., :shift] + + new_terminal = np.concatenate((previous_stack, old_terminal), axis=self.repeat_axis) + infos[env_idx]["terminal_observation"] = new_terminal else: warnings.warn("VecFrameStack wrapping a VecEnv without terminal_observation info") - self.stackedobs[i] = 0 + self.stacked_obs[env_idx] = 0 if self.channels_first: - self.stackedobs[:, -observations.shape[self.stack_dimension] :, ...] = observations + self.stacked_obs[:, shift:, ...] = observations else: - self.stackedobs[..., -observations.shape[self.stack_dimension] :] = observations - return self.stackedobs, infos - - -class StackedDictObservations(StackedObservations): - """ - Frame stacking wrapper for dictionary data. - - Dimension to stack over is either first (channels-first) or - last (channels-last), which is detected automatically using - ``common.preprocessing.is_image_space_channels_first`` if - observation is an image space. - - :param num_envs: number of environments - :param n_stack: Number of frames to stack - :param channels_order: If "first", stack on first image dimension. If "last", stack on last dimension. - If None, automatically detect channel to stack over in case of image observation or default to "last" (default). - """ - - def __init__( - self, - num_envs: int, - n_stack: int, - observation_space: spaces.Dict, - channels_order: Optional[Union[str, Dict[str, str]]] = None, - ): - self.n_stack = n_stack - self.channels_first = {} - self.stack_dimension = {} - self.stackedobs = {} - self.repeat_axis = {} - - for key, subspace in observation_space.spaces.items(): - assert isinstance(subspace, spaces.Box), "StackedDictObservations only works with nested gym.spaces.Box" - if isinstance(channels_order, str) or channels_order is None: - subspace_channel_order = channels_order - else: - subspace_channel_order = channels_order[key] - ( - self.channels_first[key], - self.stack_dimension[key], - self.stackedobs[key], - self.repeat_axis[key], - ) = self.compute_stacking(num_envs, n_stack, subspace, subspace_channel_order) - - def stack_observation_space(self, observation_space: spaces.Dict) -> spaces.Dict: - """ - Returns the stacked version of a Dict observation space - - :param observation_space: Dict observation space to stack - :return: stacked observation space - """ - spaces_dict = {} - for key, subspace in observation_space.spaces.items(): - low = np.repeat(subspace.low, self.n_stack, axis=self.repeat_axis[key]) - high = np.repeat(subspace.high, self.n_stack, axis=self.repeat_axis[key]) - spaces_dict[key] = spaces.Box(low=low, high=high, dtype=subspace.dtype) - return spaces.Dict(spaces=spaces_dict) - - def reset(self, observation: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: # pytype: disable=signature-mismatch - """ - Resets the stacked observations, adds the reset observation to the stack, and returns the stack - - :param observation: Reset observation - :return: Stacked reset observations - """ - for key, obs in observation.items(): - self.stackedobs[key][...] = 0 - if self.channels_first[key]: - self.stackedobs[key][:, -obs.shape[self.stack_dimension[key]] :, ...] = obs - else: - self.stackedobs[key][..., -obs.shape[self.stack_dimension[key]] :] = obs - return self.stackedobs - - def update( - self, - observations: Dict[str, np.ndarray], - dones: np.ndarray, - infos: List[Dict[str, Any]], - ) -> Tuple[Dict[str, np.ndarray], List[Dict[str, Any]]]: # pytype: disable=signature-mismatch - """ - Adds the observations to the stack and uses the dones to update the infos. - - :param observations: Dict of numpy arrays of observations - :param dones: numpy array of dones - :param infos: dict of infos - :return: tuple of the stacked observations and the updated infos - """ - for key in self.stackedobs.keys(): - stack_ax_size = observations[key].shape[self.stack_dimension[key]] - self.stackedobs[key] = np.roll( - self.stackedobs[key], - shift=-stack_ax_size, - axis=self.stack_dimension[key], - ) - - for i, done in enumerate(dones): - if done: - if "terminal_observation" in infos[i]: - old_terminal = infos[i]["terminal_observation"][key] - if self.channels_first[key]: - new_terminal = np.vstack( - ( - self.stackedobs[key][i, :-stack_ax_size, ...], - old_terminal, - ) - ) - else: - new_terminal = np.concatenate( - ( - self.stackedobs[key][i, ..., :-stack_ax_size], - old_terminal, - ), - axis=self.stack_dimension[key], - ) - infos[i]["terminal_observation"][key] = new_terminal - else: - warnings.warn("VecFrameStack wrapping a VecEnv without terminal_observation info") - self.stackedobs[key][i] = 0 - if self.channels_first[key]: - self.stackedobs[key][:, -stack_ax_size:, ...] = observations[key] - else: - self.stackedobs[key][..., -stack_ax_size:] = observations[key] - return self.stackedobs, infos + self.stacked_obs[..., shift:] = observations + return self.stacked_obs, infos diff --git a/stable_baselines3/common/vec_env/vec_frame_stack.py b/stable_baselines3/common/vec_env/vec_frame_stack.py index d933104..8a020dd 100644 --- a/stable_baselines3/common/vec_env/vec_frame_stack.py +++ b/stable_baselines3/common/vec_env/vec_frame_stack.py @@ -1,63 +1,40 @@ -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Mapping, Optional, Tuple, Union import numpy as np from gym import spaces from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper -from stable_baselines3.common.vec_env.stacked_observations import StackedDictObservations, StackedObservations +from stable_baselines3.common.vec_env.stacked_observations import StackedObservations class VecFrameStack(VecEnvWrapper): """ Frame stacking wrapper for vectorized environment. Designed for image observations. - Uses the StackedObservations class, or StackedDictObservations depending on the observations space - - :param venv: the vectorized environment to wrap + :param venv: Vectorized environment to wrap :param n_stack: Number of frames to stack :param channels_order: If "first", stack on first image dimension. If "last", stack on last dimension. If None, automatically detect channel to stack over in case of image observation or default to "last" (default). Alternatively channels_order can be a dictionary which can be used with environments with Dict observation spaces """ - def __init__(self, venv: VecEnv, n_stack: int, channels_order: Optional[Union[str, Dict[str, str]]] = None): - self.venv = venv - self.n_stack = n_stack + def __init__(self, venv: VecEnv, n_stack: int, channels_order: Optional[Union[str, Mapping[str, str]]] = None) -> None: + assert isinstance( + venv.observation_space, (spaces.Box, spaces.Dict) + ), "VecFrameStack only works with gym.spaces.Box and gym.spaces.Dict observation spaces" - wrapped_obs_space = venv.observation_space - - if isinstance(wrapped_obs_space, spaces.Box): - assert not isinstance( - channels_order, dict - ), f"Expected None or string for channels_order but received {channels_order}" - self.stackedobs = StackedObservations(venv.num_envs, n_stack, wrapped_obs_space, channels_order) - - elif isinstance(wrapped_obs_space, spaces.Dict): - self.stackedobs = StackedDictObservations(venv.num_envs, n_stack, wrapped_obs_space, channels_order) - - else: - raise Exception("VecFrameStack only works with gym.spaces.Box and gym.spaces.Dict observation spaces") - - observation_space = self.stackedobs.stack_observation_space(wrapped_obs_space) - VecEnvWrapper.__init__(self, venv, observation_space=observation_space) + self.stacked_obs = StackedObservations(venv.num_envs, n_stack, venv.observation_space, channels_order) + observation_space = self.stacked_obs.stacked_observation_space + super().__init__(venv, observation_space=observation_space) def step_wait( self, ) -> Tuple[Union[np.ndarray, Dict[str, np.ndarray]], np.ndarray, np.ndarray, List[Dict[str, Any]],]: observations, rewards, dones, infos = self.venv.step_wait() - - observations, infos = self.stackedobs.update(observations, dones, infos) - + observations, infos = self.stacked_obs.update(observations, dones, infos) return observations, rewards, dones, infos def reset(self) -> Union[np.ndarray, Dict[str, np.ndarray]]: - """ - Reset all environments - """ observation = self.venv.reset() # pytype:disable=annotation-type-mismatch - - observation = self.stackedobs.reset(observation) + observation = self.stacked_obs.reset(observation) return observation - - def close(self) -> None: - self.venv.close() diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index f5e9264..e8175d3 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.8.0a3 +1.8.0a4 diff --git a/tests/test_vec_stacked_obs.py b/tests/test_vec_stacked_obs.py new file mode 100644 index 0000000..0a7aa39 --- /dev/null +++ b/tests/test_vec_stacked_obs.py @@ -0,0 +1,314 @@ +import numpy as np +from gym import spaces + +from stable_baselines3.common.vec_env.stacked_observations import StackedObservations + +compute_stacking = StackedObservations.compute_stacking +NUM_ENVS = 2 +N_STACK = 4 +H, W, C = 16, 24, 3 + + +def test_compute_stacking_box(): + space = spaces.Box(-1, 1, (4,)) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking(N_STACK, observation_space=space) + assert not channels_first # default is channel last + assert stack_dimension == -1 + assert stacked_shape == (N_STACK * 4,) + assert repeat_axis == -1 + + +def test_compute_stacking_multidim_box(): + space = spaces.Box(-1, 1, (4, 5)) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking(N_STACK, observation_space=space) + assert not channels_first # default is channel last + assert stack_dimension == -1 + assert stacked_shape == (4, N_STACK * 5) + assert repeat_axis == -1 + + +def test_compute_stacking_multidim_box_channel_first(): + space = spaces.Box(-1, 1, (4, 5)) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking( + N_STACK, observation_space=space, channels_order="first" + ) + assert channels_first # default is channel last + assert stack_dimension == 1 + assert stacked_shape == (N_STACK * 4, 5) + assert repeat_axis == 0 + + +def test_compute_stacking_image_channel_first(): + """Detect that image is channel first and stack in that dimension.""" + space = spaces.Box(0, 255, (C, H, W), dtype=np.uint8) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking(N_STACK, observation_space=space) + assert channels_first # default is channel last + assert stack_dimension == 1 + assert stacked_shape == (N_STACK * C, H, W) + assert repeat_axis == 0 + + +def test_compute_stacking_image_channel_last(): + """Detect that image is channel last and stack in that dimension.""" + space = spaces.Box(0, 255, (H, W, C), dtype=np.uint8) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking(N_STACK, observation_space=space) + assert not channels_first # default is channel last + assert stack_dimension == -1 + assert stacked_shape == (H, W, N_STACK * C) + assert repeat_axis == -1 + + +def test_compute_stacking_image_channel_first_stack_last(): + """Detect that image is channel first and stack in that dimension.""" + space = spaces.Box(0, 255, (C, H, W), dtype=np.uint8) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking( + N_STACK, observation_space=space, channels_order="last" + ) + assert not channels_first # default is channel last + assert stack_dimension == -1 + assert stacked_shape == (C, H, N_STACK * W) + assert repeat_axis == -1 + + +def test_compute_stacking_image_channel_last_stack_first(): + """Detect that image is channel last and stack in that dimension.""" + space = spaces.Box(0, 255, (H, W, C), dtype=np.uint8) + channels_first, stack_dimension, stacked_shape, repeat_axis = compute_stacking( + N_STACK, observation_space=space, channels_order="first" + ) + assert channels_first # default is channel last + assert stack_dimension == 1 + assert stacked_shape == (N_STACK * H, W, C) + assert repeat_axis == 0 + + +def test_reset_update_box(): + space = spaces.Box(-1, 1, (4,)) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space) + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * 4) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * 4) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate( + (np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=-1 + ), + ) + + +def test_reset_update_multidim_box(): + space = spaces.Box(-1, 1, (4, 5)) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space) + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, 4, N_STACK * 5) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, 4, N_STACK * 5) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate( + (np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=-1 + ), + ) + + +def test_reset_update_multidim_box_channel_first(): + space = spaces.Box(-1, 1, (4, 5)) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space, channels_order="first") + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * 4, 5) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * 4, 5) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate((np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=1), + ) + + +def test_reset_update_image_channel_first(): + space = spaces.Box(0, 255, (C, H, W), dtype=np.uint8) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space) + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * C, H, W) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * C, H, W) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate((np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=1), + ) + + +def test_reset_update_image_channel_last(): + space = spaces.Box(0, 255, (H, W, C), dtype=np.uint8) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space) + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, H, W, N_STACK * C) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, H, W, N_STACK * C) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate( + (np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=-1 + ), + ) + + +def test_reset_update_image_channel_first_stack_last(): + space = spaces.Box(0, 255, (C, H, W), dtype=np.uint8) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space, channels_order="last") + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, C, H, N_STACK * W) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, C, H, N_STACK * W) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate( + (np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=-1 + ), + ) + + +def test_reset_update_image_channel_last_stack_first(): + space = spaces.Box(0, 255, (H, W, C), dtype=np.uint8) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space, channels_order="first") + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs = stacked_observations.reset(observations_1) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * H, W, C) + assert stacked_obs.dtype == space.dtype + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs.shape == (NUM_ENVS, N_STACK * H, W, C) + assert stacked_obs.dtype == space.dtype + assert np.array_equal( + stacked_obs, + np.concatenate((np.zeros_like(observations_1), np.zeros_like(observations_1), observations_1, observations_2), axis=1), + ) + + +def test_reset_update_dict(): + space = spaces.Dict({"key1": spaces.Box(0, 255, (H, W, C), dtype=np.uint8), "key2": spaces.Box(-1, 1, (4, 5))}) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space, channels_order={"key1": "first", "key2": "last"}) + observations_1 = {key: np.stack([subspace.sample() for _ in range(NUM_ENVS)]) for key, subspace in space.spaces.items()} + stacked_obs = stacked_observations.reset(observations_1) + assert isinstance(stacked_obs, dict) + assert stacked_obs["key1"].shape == (NUM_ENVS, N_STACK * H, W, C) + assert stacked_obs["key2"].shape == (NUM_ENVS, 4, N_STACK * 5) + assert stacked_obs["key1"].dtype == space["key1"].dtype + assert stacked_obs["key2"].dtype == space["key2"].dtype + observations_2 = {key: np.stack([subspace.sample() for _ in range(NUM_ENVS)]) for key, subspace in space.spaces.items()} + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_obs, infos = stacked_observations.update(observations_2, dones, infos) + assert stacked_obs["key1"].shape == (NUM_ENVS, N_STACK * H, W, C) + assert stacked_obs["key2"].shape == (NUM_ENVS, 4, N_STACK * 5) + assert stacked_obs["key1"].dtype == space["key1"].dtype + assert stacked_obs["key2"].dtype == space["key2"].dtype + + assert np.array_equal( + stacked_obs["key1"], + np.concatenate( + ( + np.zeros_like(observations_1["key1"]), + np.zeros_like(observations_1["key1"]), + observations_1["key1"], + observations_2["key1"], + ), + axis=1, + ), + ) + assert np.array_equal( + stacked_obs["key2"], + np.concatenate( + ( + np.zeros_like(observations_1["key2"]), + np.zeros_like(observations_1["key2"]), + observations_1["key2"], + observations_2["key2"], + ), + axis=-1, + ), + ) + + +def test_episode_termination_box(): + space = spaces.Box(-1, 1, (4,)) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space) + observations_1 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_observations.reset(observations_1) + observations_2 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_observations.update(observations_2, dones, infos) + terminal_observation = space.sample() + infos[1]["terminal_observation"] = terminal_observation # episode termination in env1 + dones[1] = True + observations_3 = np.stack([space.sample() for _ in range(NUM_ENVS)]) + stacked_obs, infos = stacked_observations.update(observations_3, dones, infos) + zeros = np.zeros_like(observations_1[0]) + true_stacked_obs_env1 = np.concatenate((zeros, observations_1[0], observations_2[0], observations_3[0]), axis=-1) + true_stacked_obs_env2 = np.concatenate((zeros, zeros, zeros, observations_3[1]), axis=-1) + true_stacked_obs = np.stack((true_stacked_obs_env1, true_stacked_obs_env2)) + assert np.array_equal(true_stacked_obs, stacked_obs) + + +def test_episode_termination_dict(): + space = spaces.Dict({"key1": spaces.Box(0, 255, (H, W, 3), dtype=np.uint8), "key2": spaces.Box(-1, 1, (4, 5))}) + stacked_observations = StackedObservations(NUM_ENVS, N_STACK, space, channels_order={"key1": "first", "key2": "last"}) + observations_1 = {key: np.stack([subspace.sample() for _ in range(NUM_ENVS)]) for key, subspace in space.spaces.items()} + stacked_observations.reset(observations_1) + observations_2 = {key: np.stack([subspace.sample() for _ in range(NUM_ENVS)]) for key, subspace in space.spaces.items()} + dones = np.zeros((NUM_ENVS,), dtype=bool) + infos = [{} for _ in range(NUM_ENVS)] + stacked_observations.update(observations_2, dones, infos) + terminal_observation = space.sample() + infos[1]["terminal_observation"] = terminal_observation # episode termination in env1 + dones[1] = True + observations_3 = {key: np.stack([subspace.sample() for _ in range(NUM_ENVS)]) for key, subspace in space.spaces.items()} + stacked_obs, infos = stacked_observations.update(observations_3, dones, infos) + + for key, axis in zip(observations_1.keys(), [0, -1]): + zeros = np.zeros_like(observations_1[key][0]) + true_stacked_obs_env1 = np.concatenate( + (zeros, observations_1[key][0], observations_2[key][0], observations_3[key][0]), axis + ) + true_stacked_obs_env2 = np.concatenate((zeros, zeros, zeros, observations_3[key][1]), axis) + true_stacked_obs = np.stack((true_stacked_obs_env1, true_stacked_obs_env2)) + assert np.array_equal(true_stacked_obs, stacked_obs[key]) From 489b1fdaf2958ba5d63ed5a90f20540cd4107b13 Mon Sep 17 00:00:00 2001 From: Sidney Tio <35787241+sidney-tio@users.noreply.github.com> Date: Tue, 7 Feb 2023 20:42:14 +0800 Subject: [PATCH 11/12] Add the argument `dtype` (default to `float32`) to the noise (#1301) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fixed noise to return float32 * Updated changelog * Fixed test to use numpy arrays instead of python floats * Sorted imports for tests * Added dtype to constructor * Removed dtype parameter for VectorizedActionNoise * __init__ -> None; Capitalize and period in docstring when needed; fix dtype type hint; dtype in docstring * fix dtype type hint * Update version * Clarify changelog [skip ci] * empty commit to run ci * Update docs/misc/changelog.rst --------- Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com> Co-authored-by: Antonin RAFFIN --- docs/misc/changelog.rst | 1 + stable_baselines3/common/noise.py | 44 ++++++++++++++++++------------- tests/test_deterministic.py | 5 +++- 3 files changed, 30 insertions(+), 20 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 9d24bbe..6fafd99 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -27,6 +27,7 @@ New Features: Bug Fixes: ^^^^^^^^^^ - Fixed Atari wrapper that missed the reset condition (@luizapozzobon) +- Added the argument ``dtype`` (default to ``float32``) to the noise for consistency with gym action (@sidney-tio) - Fixed PPO train/n_updates metric not accounting for early stopping (@adamfrly) Deprecations: diff --git a/stable_baselines3/common/noise.py b/stable_baselines3/common/noise.py index 5e8632d..944408f 100644 --- a/stable_baselines3/common/noise.py +++ b/stable_baselines3/common/noise.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from typing import Iterable, List, Optional import numpy as np +from numpy.typing import DTypeLike class ActionNoise(ABC): @@ -15,7 +16,7 @@ class ActionNoise(ABC): def reset(self) -> None: """ - call end of episode reset for the noise + Call end of episode reset for the noise """ pass @@ -26,19 +27,21 @@ class ActionNoise(ABC): class NormalActionNoise(ActionNoise): """ - A Gaussian action noise + A Gaussian action noise. - :param mean: the mean value of the noise - :param sigma: the scale of the noise (std here) + :param mean: Mean value of the noise + :param sigma: Scale of the noise (std here) + :param dtype: Type of the output noise """ - def __init__(self, mean: np.ndarray, sigma: np.ndarray): + def __init__(self, mean: np.ndarray, sigma: np.ndarray, dtype: DTypeLike = np.float32) -> None: self._mu = mean self._sigma = sigma + self._dtype = dtype super().__init__() def __call__(self) -> np.ndarray: - return np.random.normal(self._mu, self._sigma) + return np.random.normal(self._mu, self._sigma).astype(self._dtype) def __repr__(self) -> str: return f"NormalActionNoise(mu={self._mu}, sigma={self._sigma})" @@ -50,11 +53,12 @@ class OrnsteinUhlenbeckActionNoise(ActionNoise): Based on http://math.stackexchange.com/questions/1287634/implementing-ornstein-uhlenbeck-in-matlab - :param mean: the mean of the noise - :param sigma: the scale of the noise - :param theta: the rate of mean reversion - :param dt: the timestep for the noise - :param initial_noise: the initial value for the noise output, (if None: 0) + :param mean: Mean of the noise + :param sigma: Scale of the noise + :param theta: Rate of mean reversion + :param dt: Timestep for the noise + :param initial_noise: Initial value for the noise output, (if None: 0) + :param dtype: Type of the output noise """ def __init__( @@ -64,11 +68,13 @@ class OrnsteinUhlenbeckActionNoise(ActionNoise): theta: float = 0.15, dt: float = 1e-2, initial_noise: Optional[np.ndarray] = None, - ): + dtype: DTypeLike = np.float32, + ) -> None: self._theta = theta self._mu = mean self._sigma = sigma self._dt = dt + self._dtype = dtype self.initial_noise = initial_noise self.noise_prev = np.zeros_like(self._mu) self.reset() @@ -81,7 +87,7 @@ class OrnsteinUhlenbeckActionNoise(ActionNoise): + self._sigma * np.sqrt(self._dt) * np.random.normal(size=self._mu.shape) ) self.noise_prev = noise - return noise + return noise.astype(self._dtype) def reset(self) -> None: """ @@ -97,11 +103,11 @@ class VectorizedActionNoise(ActionNoise): """ A Vectorized action noise for parallel environments. - :param base_noise: ActionNoise The noise generator to use - :param n_envs: The number of parallel environments + :param base_noise: Noise generator to use + :param n_envs: Number of parallel environments """ - def __init__(self, base_noise: ActionNoise, n_envs: int): + def __init__(self, base_noise: ActionNoise, n_envs: int) -> None: try: self.n_envs = int(n_envs) assert self.n_envs > 0 @@ -113,9 +119,9 @@ class VectorizedActionNoise(ActionNoise): def reset(self, indices: Optional[Iterable[int]] = None) -> None: """ - Reset all the noise processes, or those listed in indices + Reset all the noise processes, or those listed in indices. - :param indices: Optional[Iterable[int]] The indices to reset. Default: None. + :param indices: The indices to reset. Default: None. If the parameter is None, then all processes are reset to their initial position. """ if indices is None: @@ -129,7 +135,7 @@ class VectorizedActionNoise(ActionNoise): def __call__(self) -> np.ndarray: """ - Generate and stack the action noise from each noise object + Generate and stack the action noise from each noise object. """ noise = np.stack([noise() for noise in self.noises]) return noise diff --git a/tests/test_deterministic.py b/tests/test_deterministic.py index 4c92d26..c165e48 100644 --- a/tests/test_deterministic.py +++ b/tests/test_deterministic.py @@ -1,3 +1,4 @@ +import numpy as np import pytest from stable_baselines3 import A2C, DQN, PPO, SAC, TD3 @@ -15,7 +16,9 @@ def test_deterministic_training_common(algo): kwargs = {"policy_kwargs": dict(net_arch=[64])} env_id = "Pendulum-v1" if algo in [TD3, SAC]: - kwargs.update({"action_noise": NormalActionNoise(0.0, 0.1), "learning_starts": 100, "train_freq": 4}) + kwargs.update( + {"action_noise": NormalActionNoise(np.zeros(1), 0.1 * np.ones(1)), "learning_starts": 100, "train_freq": 4} + ) else: if algo == DQN: env_id = "CartPole-v1" From 69b94dd6a8f93cf0b9d2201dcae9c146b8a9c75d Mon Sep 17 00:00:00 2001 From: Vikas Kumar <32395222+theSquaredError@users.noreply.github.com> Date: Sat, 11 Feb 2023 01:45:09 +0530 Subject: [PATCH 12/12] Rename "timesteps" to "episodes" in `log_interval` documentation (#1325) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * change timestamp to episode for logging * update changelog * minor format modif * minor format modif --------- Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com> --- docs/misc/changelog.rst | 4 ++-- stable_baselines3/common/base_class.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 6fafd99..f70a81b 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -45,7 +45,7 @@ Documentation: - Renamed ``load_parameters`` to ``set_parameters`` (@DavyMorgan) - Clarified documentation about subproc multiprocessing for A2C (@Bonifatius94) - Fixed typo in ``A2C`` docstring (@AlexPasqua) - +- Renamed timesteps to episodes for ``log_interval`` description (@theSquaredError) Release 1.7.0 (2023-01-10) -------------------------- @@ -1227,4 +1227,4 @@ And all the contributors: @Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede @Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont @honglu2875 @yuanmingqi @anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong -@DavyMorgan @luizapozzobon @Bonifatius94 +@DavyMorgan @luizapozzobon @Bonifatius94 @theSquaredError diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index b6fba85..17be67a 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -505,7 +505,7 @@ class BaseAlgorithm(ABC): :param total_timesteps: The total number of samples (env steps) to train on :param callback: callback(s) called at every step with state of the algorithm. - :param log_interval: The number of timesteps before logging. + :param log_interval: The number of episodes before logging. :param tb_log_name: the name of the run for TensorBoard logging :param reset_num_timesteps: whether or not to reset the current timestep number (used in logging) :param progress_bar: Display a progress bar using tqdm and rich.