From a10e3ae5870836ed38e0b65c4176bbec3f22a1cb Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Sun, 4 Oct 2020 17:12:35 +0200 Subject: [PATCH 1/5] Release v0.9.0 (#174) --- docs/misc/changelog.rst | 5 ++++- stable_baselines3/dqn/dqn.py | 3 +++ stable_baselines3/version.txt | 2 +- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index d74b402..3ee445b 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,9 +3,11 @@ Changelog ========== -Pre-Release 0.9.0a2 (WIP) +Pre-Release 0.9.0 (2020-10-03) ------------------------------ +**Bug fixes, get/set parameters and improved docs** + Breaking Changes: ^^^^^^^^^^^^^^^^^ - Removed ``device`` keyword argument of policies; use ``policy.to(device)`` instead. (@qxcv) @@ -50,6 +52,7 @@ Others: - Clarified docstrings on what is saved and loaded to/from files - Simplified ``save_to_zip_file`` function by removing duplicate code - Store library version along with the saved models +- DQN loss is now logged Documentation: ^^^^^^^^^^^^^^ diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 29323c0..55fb835 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -147,6 +147,7 @@ class DQN(OffPolicyAlgorithm): # Update learning rate according to schedule self._update_learning_rate(self.policy.optimizer) + losses = [] for gradient_step in range(gradient_steps): # Sample replay buffer replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env) @@ -169,6 +170,7 @@ class DQN(OffPolicyAlgorithm): # Compute Huber loss (less sensitive to outliers) loss = F.smooth_l1_loss(current_q, target_q) + losses.append(loss.item()) # Optimize the policy self.policy.optimizer.zero_grad() @@ -181,6 +183,7 @@ class DQN(OffPolicyAlgorithm): self._n_updates += gradient_steps logger.record("train/n_updates", self._n_updates, exclude="tensorboard") + logger.record("train/loss", np.mean(losses)) def predict( self, diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 63aaf3d..ac39a10 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -0.9.0a2 +0.9.0 From a1e055695c3638f9f15de0cb805b8fcbb5c02764 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Wed, 7 Oct 2020 10:51:49 +0200 Subject: [PATCH 2/5] Improve typing coverage (#175) * Improve typing coverage * Even more types * Fixes * Update changelog * Unified docstrings * Improve error messages for unsupported spaces --- docs/misc/changelog.rst | 25 ++++++++++ stable_baselines3/common/atari_wrappers.py | 4 +- stable_baselines3/common/base_class.py | 4 +- stable_baselines3/common/buffers.py | 8 +++- stable_baselines3/common/callbacks.py | 11 +++-- stable_baselines3/common/cmd_util.py | 6 +-- stable_baselines3/common/distributions.py | 4 +- stable_baselines3/common/logger.py | 2 +- stable_baselines3/common/monitor.py | 10 ++-- stable_baselines3/common/noise.py | 2 +- stable_baselines3/common/policies.py | 2 +- stable_baselines3/common/preprocessing.py | 6 +-- stable_baselines3/common/save_util.py | 17 ++++--- .../common/sb2_compat/rmsprop_tf_like.py | 44 +++++++++++------- stable_baselines3/common/utils.py | 2 +- .../common/vec_env/base_vec_env.py | 26 +++++------ .../common/vec_env/dummy_vec_env.py | 28 +++++------ .../common/vec_env/subproc_vec_env.py | 46 +++++++++++-------- stable_baselines3/common/vec_env/util.py | 8 ++-- .../common/vec_env/vec_check_nan.py | 18 ++++---- .../common/vec_env/vec_normalize.py | 37 +++++++++------ .../common/vec_env/vec_video_recorder.py | 28 ++++++----- stable_baselines3/dqn/dqn.py | 2 +- stable_baselines3/td3/policies.py | 2 +- stable_baselines3/version.txt | 2 +- 25 files changed, 206 insertions(+), 138 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 3ee445b..90b5322 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,6 +3,31 @@ Changelog ========== + +Pre-Release 0.10.0a0 (WIP) +------------------------------ + +Breaking Changes: +^^^^^^^^^^^^^^^^^ + +New Features: +^^^^^^^^^^^^^ + +Bug Fixes: +^^^^^^^^^^ + +Deprecations: +^^^^^^^^^^^^^ + +Others: +^^^^^^^ +- Improved typing coverage +- Improved error messages for unsupported spaces + +Documentation: +^^^^^^^^^^^^^^ + + Pre-Release 0.9.0 (2020-10-03) ------------------------------ diff --git a/stable_baselines3/common/atari_wrappers.py b/stable_baselines3/common/atari_wrappers.py index 977f6d3..7cc6836 100644 --- a/stable_baselines3/common/atari_wrappers.py +++ b/stable_baselines3/common/atari_wrappers.py @@ -9,7 +9,7 @@ try: except ImportError: cv2 = None -from stable_baselines3.common.type_aliases import GymStepReturn +from stable_baselines3.common.type_aliases import GymObs, GymStepReturn class NoopResetEnv(gym.Wrapper): @@ -146,7 +146,7 @@ class MaxAndSkipEnv(gym.Wrapper): return max_frame, total_reward, done, info - def reset(self, **kwargs): + def reset(self, **kwargs) -> GymObs: return self.env.reset(**kwargs) diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index d30830b..c904925 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -478,7 +478,7 @@ class BaseAlgorithm(ABC): load_path_or_dict: Union[str, Dict[str, Dict]], exact_match: bool = True, device: Union[th.device, str] = "auto", - ): + ) -> None: """ Load parameters from a given zip-file or a nested dictionary containing parameters for different modules (see ``get_parameters``). @@ -610,7 +610,7 @@ class BaseAlgorithm(ABC): model.policy.reset_noise() # pytype: disable=attribute-error return model - def get_parameters(self): + def get_parameters(self) -> Dict[str, Dict]: """ Return the parameters of the agent. This includes parameters from different networks, e.g. critics (value functions) and policies (pi functions). diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 12f4481..2063d6c 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -1,4 +1,5 @@ import warnings +from abc import ABC, abstractmethod from typing import Generator, Optional, Union import numpy as np @@ -16,7 +17,7 @@ from stable_baselines3.common.type_aliases import ReplayBufferSamples, RolloutBu from stable_baselines3.common.vec_env import VecNormalize -class BaseBuffer(object): +class BaseBuffer(ABC): """ Base class that represent a buffer (rollout or replay) @@ -102,7 +103,10 @@ class BaseBuffer(object): batch_inds = np.random.randint(0, upper_bound, size=batch_size) return self._get_samples(batch_inds, env=env) - def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None): + @abstractmethod + def _get_samples( + self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None + ) -> Union[ReplayBufferSamples, RolloutBufferSamples]: """ :param batch_inds: :param env: diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py index 05f4d14..872aed1 100644 --- a/stable_baselines3/common/callbacks.py +++ b/stable_baselines3/common/callbacks.py @@ -1,7 +1,7 @@ import os import warnings from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Optional, Union import gym import numpy as np @@ -217,9 +217,10 @@ class CheckpointCallback(BaseCallback): :param save_freq: :param save_path: Path to the folder where the model will be saved. :param name_prefix: Common prefix to the saved models + :param verbose: """ - def __init__(self, save_freq: int, save_path: str, name_prefix="rl_model", verbose=0): + def __init__(self, save_freq: int, save_path: str, name_prefix: str = "rl_model", verbose: int = 0): super(CheckpointCallback, self).__init__(verbose) self.save_freq = save_freq self.save_path = save_path @@ -247,7 +248,7 @@ class ConvertCallback(BaseCallback): :param verbose: """ - def __init__(self, callback, verbose=0): + def __init__(self, callback: Callable, verbose: int = 0): super(ConvertCallback, self).__init__(verbose) self.callback = callback @@ -314,7 +315,7 @@ class EvalCallback(EventCallback): self.evaluations_timesteps = [] self.evaluations_length = [] - def _init_callback(self): + def _init_callback(self) -> None: # Does not work in some corner cases, where the wrapper is not the same if not isinstance(self.training_env, type(self.eval_env)): warnings.warn("Training and eval env are not of the same type" f"{self.training_env} != {self.eval_env}") @@ -450,7 +451,7 @@ class StopTrainingOnMaxEpisodes(BaseCallback): self._total_max_episodes = max_episodes self.n_episodes = 0 - def _init_callback(self): + def _init_callback(self) -> None: # At start set total max according to number of envirnments self._total_max_episodes = self.max_episodes * self.training_env.num_envs diff --git a/stable_baselines3/common/cmd_util.py b/stable_baselines3/common/cmd_util.py index 4aede71..4169ae6 100644 --- a/stable_baselines3/common/cmd_util.py +++ b/stable_baselines3/common/cmd_util.py @@ -6,7 +6,7 @@ import gym from stable_baselines3.common.atari_wrappers import AtariWrapper from stable_baselines3.common.monitor import Monitor -from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv +from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv, VecEnv def make_vec_env( @@ -19,7 +19,7 @@ def make_vec_env( env_kwargs: Optional[Dict[str, Any]] = None, vec_env_cls: Optional[Type[Union[DummyVecEnv, SubprocVecEnv]]] = None, vec_env_kwargs: Optional[Dict[str, Any]] = None, -): +) -> VecEnv: """ Create a wrapped, monitored ``VecEnv``. By default it uses a ``DummyVecEnv`` which is usually faster @@ -85,7 +85,7 @@ def make_atari_env( env_kwargs: Optional[Dict[str, Any]] = None, vec_env_cls: Optional[Union[DummyVecEnv, SubprocVecEnv]] = None, vec_env_kwargs: Optional[Dict[str, Any]] = None, -): +) -> VecEnv: """ Create a wrapped, monitored VecEnv for Atari. It is a wrapper around ``make_vec_env`` that includes common preprocessing for Atari games. diff --git a/stable_baselines3/common/distributions.py b/stable_baselines3/common/distributions.py index b8d7837..0f43ca8 100644 --- a/stable_baselines3/common/distributions.py +++ b/stable_baselines3/common/distributions.py @@ -1,7 +1,7 @@ """Probability distributions.""" from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Union import gym import torch as th @@ -19,7 +19,7 @@ class Distribution(ABC): super(Distribution, self).__init__() @abstractmethod - def proba_distribution_net(self, *args, **kwargs): + def proba_distribution_net(self, *args, **kwargs) -> Union[nn.Module, Tuple[nn.Module, nn.Parameter]]: """Create the layers and parameters that represent the distribution. Subclasses must define this, but the arguments and return type vary between diff --git a/stable_baselines3/common/logger.py b/stable_baselines3/common/logger.py index 01e7709..0b201a7 100644 --- a/stable_baselines3/common/logger.py +++ b/stable_baselines3/common/logger.py @@ -50,7 +50,7 @@ class SeqWriter(object): sequence writer """ - def write_sequence(self, sequence: List): + def write_sequence(self, sequence: List) -> None: """ write_sequence an array to file diff --git a/stable_baselines3/common/monitor.py b/stable_baselines3/common/monitor.py index 17c0587..bb50fe4 100644 --- a/stable_baselines3/common/monitor.py +++ b/stable_baselines3/common/monitor.py @@ -5,12 +5,14 @@ import json import os import time from glob import glob -from typing import Any, Dict, List, Optional, Tuple +from typing import List, Optional, Tuple, Union import gym import numpy as np import pandas +from stable_baselines3.common.type_aliases import GymObs, GymStepReturn + class Monitor(gym.Wrapper): """ @@ -62,7 +64,7 @@ class Monitor(gym.Wrapper): self.total_steps = 0 self.current_reset_info = {} # extra info about the current episode, that was passed in during reset() - def reset(self, **kwargs) -> np.ndarray: + def reset(self, **kwargs) -> GymObs: """ Calls the Gym environment reset. Can only be called if the environment is over, or if allow_early_resets is True @@ -83,7 +85,7 @@ class Monitor(gym.Wrapper): self.current_reset_info[key] = value return self.env.reset(**kwargs) - def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, Dict[Any, Any]]: + def step(self, action: Union[np.ndarray, int]) -> GymStepReturn: """ Step the environment with the given action @@ -112,7 +114,7 @@ class Monitor(gym.Wrapper): self.total_steps += 1 return observation, reward, done, info - def close(self): + def close(self) -> None: """ Closes the environment """ diff --git a/stable_baselines3/common/noise.py b/stable_baselines3/common/noise.py index a3b66db..b1db6f4 100644 --- a/stable_baselines3/common/noise.py +++ b/stable_baselines3/common/noise.py @@ -139,7 +139,7 @@ class VectorizedActionNoise(ActionNoise): return self._base_noise @base_noise.setter - def base_noise(self, base_noise: ActionNoise): + def base_noise(self, base_noise: ActionNoise) -> None: if base_noise is None: raise ValueError("Expected base_noise to be an instance of ActionNoise, not None", ActionNoise) if not isinstance(base_noise, ActionNoise): diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index a3f503b..8a0fbbb 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -145,7 +145,7 @@ class BaseModel(nn.Module, ABC): model.to(device) return model - def load_from_vector(self, vector: np.ndarray): + def load_from_vector(self, vector: np.ndarray) -> None: """ Load parameters from a 1D vector. diff --git a/stable_baselines3/common/preprocessing.py b/stable_baselines3/common/preprocessing.py index d96117a..0b84683 100644 --- a/stable_baselines3/common/preprocessing.py +++ b/stable_baselines3/common/preprocessing.py @@ -78,7 +78,7 @@ def preprocess_obs(obs: th.Tensor, observation_space: spaces.Space, normalize_im return obs.float() else: - raise NotImplementedError() + raise NotImplementedError(f"Preprocessing not implemented for {observation_space}") def get_obs_shape(observation_space: spaces.Space) -> Tuple[int, ...]: @@ -100,7 +100,7 @@ def get_obs_shape(observation_space: spaces.Space) -> Tuple[int, ...]: # Number of binary features return (int(observation_space.n),) else: - raise NotImplementedError() + raise NotImplementedError(f"{observation_space} observation space is not supported") def get_flattened_obs_dim(observation_space: spaces.Space) -> int: @@ -139,4 +139,4 @@ def get_action_dim(action_space: spaces.Space) -> int: # Number of binary actions return int(action_space.n) else: - raise NotImplementedError() + raise NotImplementedError(f"{action_space} action space is not supported") diff --git a/stable_baselines3/common/save_util.py b/stable_baselines3/common/save_util.py index f59cdd6..4d9ed8b 100644 --- a/stable_baselines3/common/save_util.py +++ b/stable_baselines3/common/save_util.py @@ -176,7 +176,7 @@ def json_to_data(json_string: str, custom_objects: Optional[Dict[str, Any]] = No @functools.singledispatch -def open_path(path: Union[str, pathlib.Path, io.BufferedIOBase], mode: str, verbose=0, suffix=None): +def open_path(path: Union[str, pathlib.Path, io.BufferedIOBase], mode: str, verbose: int = 0, suffix: Optional[str] = None): """ Opens a path for reading or writing with a preferred suffix and raises debug information. If the provided path is a derivative of io.BufferedIOBase it ensures that the file @@ -197,6 +197,7 @@ def open_path(path: Union[str, pathlib.Path, io.BufferedIOBase], mode: str, verb :param suffix: The preferred suffix. If mode is "w" then the opened file has the suffix. If mode is "r" then we attempt to open the path. If an error is raised and the suffix is not None, we attempt to open the path with the suffix. + :return: """ if not isinstance(path, io.BufferedIOBase): raise TypeError("Path parameter has invalid type.", io.BufferedIOBase) @@ -214,7 +215,7 @@ def open_path(path: Union[str, pathlib.Path, io.BufferedIOBase], mode: str, verb @open_path.register(str) -def open_path_str(path: str, mode: str, verbose=0, suffix=None) -> io.BufferedIOBase: +def open_path_str(path: str, mode: str, verbose: int = 0, suffix: Optional[str] = None) -> io.BufferedIOBase: """ Open a path given by a string. If writing to the path, the function ensures that the path exists. @@ -226,12 +227,13 @@ def open_path_str(path: str, mode: str, verbose=0, suffix=None) -> io.BufferedIO :param suffix: The preferred suffix. If mode is "w" then the opened file has the suffix. If mode is "r" then we attempt to open the path. If an error is raised and the suffix is not None, we attempt to open the path with the suffix. + :return: """ return open_path(pathlib.Path(path), mode, verbose, suffix) @open_path.register(pathlib.Path) -def open_path_pathlib(path: pathlib.Path, mode: str, verbose=0, suffix=None) -> io.BufferedIOBase: +def open_path_pathlib(path: pathlib.Path, mode: str, verbose: int = 0, suffix: Optional[str] = None) -> io.BufferedIOBase: """ Open a path given by a string. If writing to the path, the function ensures that the path exists. @@ -244,6 +246,7 @@ def open_path_pathlib(path: pathlib.Path, mode: str, verbose=0, suffix=None) -> :param suffix: The preferred suffix. If mode is "w" then the opened file has the suffix. If mode is "r" then we attempt to open the path. If an error is raised and the suffix is not None, we attempt to open the path with the suffix. + :return: """ if mode not in ("w", "r"): raise ValueError("Expected mode to be either 'w' or 'r'.") @@ -286,7 +289,7 @@ def save_to_zip_file( data: Dict[str, Any] = None, params: Dict[str, Any] = None, pytorch_variables: Dict[str, Any] = None, - verbose=0, + verbose: int = 0, ) -> None: """ Save model data to a zip archive. @@ -321,7 +324,7 @@ def save_to_zip_file( archive.writestr("_stable_baselines3_version", stable_baselines3.__version__) -def save_to_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], obj, verbose=0) -> None: +def save_to_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], obj: Any, verbose: int = 0) -> None: """ Save an object to path creating the necessary folders along the way. If the path exists and is a directory, it will raise a warning and rename the path. @@ -337,7 +340,7 @@ def save_to_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], obj, verbose= pickle.dump(obj, file_handler) -def load_from_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], verbose=0) -> Any: +def load_from_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], verbose: int = 0) -> Any: """ Load an object from the path. If a suffix is provided in the path, it will use that suffix. If the path does not exist, it will attempt to load using the .pkl suffix. @@ -355,7 +358,7 @@ def load_from_zip_file( load_path: Union[str, pathlib.Path, io.BufferedIOBase], load_data: bool = True, device: Union[th.device, str] = "auto", - verbose=0, + verbose: int = 0, ) -> (Tuple[Optional[Dict[str, Any]], Optional[TensorDict], Optional[TensorDict]]): """ Load model data from a .zip archive diff --git a/stable_baselines3/common/sb2_compat/rmsprop_tf_like.py b/stable_baselines3/common/sb2_compat/rmsprop_tf_like.py index 46ef4b0..758079f 100644 --- a/stable_baselines3/common/sb2_compat/rmsprop_tf_like.py +++ b/stable_baselines3/common/sb2_compat/rmsprop_tf_like.py @@ -1,3 +1,5 @@ +from typing import Any, Callable, Dict, Iterable, Optional + import torch from torch.optim import Optimizer @@ -28,21 +30,29 @@ class RMSpropTFLike(Optimizer): is the scheduled learning rate and :math:`v` is the weighted moving average of the squared gradient. - Arguments: - params (iterable): iterable of parameters to optimize or dicts defining - parameter groups - lr (float, optional): learning rate (default: 1e-2) - momentum (float, optional): momentum factor (default: 0) - alpha (float, optional): smoothing constant (default: 0.99) - eps (float, optional): term added to the denominator to improve - numerical stability (default: 1e-8) - centered (bool, optional) : if ``True``, compute the centered RMSProp, - the gradient is normalized by an estimation of its variance - weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + :params: iterable of parameters to optimize or dicts defining + parameter groups + :param lr: learning rate (default: 1e-2) + :param momentum: momentum factor (default: 0) + :param alpha: smoothing constant (default: 0.99) + :param eps: term added to the denominator to improve + numerical stability (default: 1e-8) + :param centered: if ``True``, compute the centered RMSProp, + the gradient is normalized by an estimation of its variance + :param weight_decay: weight decay (L2 penalty) (default: 0) """ - def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False): + def __init__( + self, + params: Iterable[torch.nn.Parameter], + lr: float = 1e-2, + alpha: float = 0.99, + eps: float = 1e-8, + weight_decay: float = 0, + momentum: float = 0, + centered: bool = False, + ): if not 0.0 <= lr: raise ValueError("Invalid learning rate: {}".format(lr)) if not 0.0 <= eps: @@ -57,19 +67,19 @@ class RMSpropTFLike(Optimizer): defaults = dict(lr=lr, momentum=momentum, alpha=alpha, eps=eps, centered=centered, weight_decay=weight_decay) super(RMSpropTFLike, self).__init__(params, defaults) - def __setstate__(self, state): + def __setstate__(self, state: Dict[str, Any]) -> None: super(RMSpropTFLike, self).__setstate__(state) for group in self.param_groups: group.setdefault("momentum", 0) group.setdefault("centered", False) @torch.no_grad() - def step(self, closure=None): + def step(self, closure: Optional[Callable] = None) -> Optional[torch.Tensor]: """Performs a single optimization step. - Arguments: - closure (callable, optional): A closure that reevaluates the model - and returns the loss. + :param closure: A closure that reevaluates the model + and returns the loss. + :return: loss """ loss = None if closure is not None: diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index deb7bdb..8177842 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -191,7 +191,7 @@ def configure_logger( logger.configure(format_strings=[""]) -def check_for_correct_spaces(env: GymEnv, observation_space: gym.spaces.Space, action_space: gym.spaces.Space): +def check_for_correct_spaces(env: GymEnv, observation_space: gym.spaces.Space, action_space: gym.spaces.Space) -> None: """ Checks that the environment has same spaces as provided ones. Used by BaseAlgorithm to check if spaces match after loading the model with given env. diff --git a/stable_baselines3/common/vec_env/base_vec_env.py b/stable_baselines3/common/vec_env/base_vec_env.py index 9b54245..d208e12 100644 --- a/stable_baselines3/common/vec_env/base_vec_env.py +++ b/stable_baselines3/common/vec_env/base_vec_env.py @@ -13,7 +13,7 @@ from stable_baselines3.common import logger VecEnvIndices = Union[None, int, Iterable[int]] # VecEnvObs is what is returned by the reset() method # it contains the observation for each env -VecEnvObs = Union[np.ndarray, Dict[str, Any]] +VecEnvObs = Union[np.ndarray, Dict[str, np.ndarray], Tuple[np.ndarray, ...]] # VecEnvStepReturn is what is returned by the step() method # it contains the observation, reward, done, info for each env VecEnvStepReturn = Tuple[VecEnvObs, np.ndarray, np.ndarray, List[Dict]] @@ -76,7 +76,7 @@ class VecEnv(ABC): raise NotImplementedError() @abstractmethod - def step_async(self, actions: np.ndarray): + def step_async(self, actions: np.ndarray) -> None: """ Tell all the environments to start taking a step with the given actions. @@ -104,7 +104,7 @@ class VecEnv(ABC): raise NotImplementedError() @abstractmethod - def get_attr(self, attr_name: str, indices: "VecEnvIndices" = None) -> List[Any]: + def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]: """ Return attribute from vectorized environment. @@ -115,7 +115,7 @@ class VecEnv(ABC): raise NotImplementedError() @abstractmethod - def set_attr(self, attr_name: str, value: Any, indices: "VecEnvIndices" = None) -> None: + def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None: """ Set attribute inside vectorized environments. @@ -127,7 +127,7 @@ class VecEnv(ABC): raise NotImplementedError() @abstractmethod - def env_method(self, method_name: str, *method_args, indices: "VecEnvIndices" = None, **method_kwargs) -> List[Any]: + def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]: """ Call instance methods of vectorized environments. @@ -210,7 +210,7 @@ class VecEnv(ABC): else: return None - def _get_indices(self, indices: "VecEnvIndices") -> Iterable[int]: + def _get_indices(self, indices: VecEnvIndices) -> Iterable[int]: """ Convert a flexibly-typed reference to environment indices to an implied list of indices. @@ -248,7 +248,7 @@ class VecEnvWrapper(VecEnv): ) self.class_attributes = dict(inspect.getmembers(self.__class__)) - def step_async(self, actions: np.ndarray): + def step_async(self, actions: np.ndarray) -> None: self.venv.step_async(actions) @abstractmethod @@ -259,7 +259,7 @@ class VecEnvWrapper(VecEnv): def step_wait(self) -> VecEnvStepReturn: pass - def seed(self, seed: Optional[int] = None): + def seed(self, seed: Optional[int] = None) -> List[Union[None, int]]: return self.venv.seed(seed) def close(self) -> None: @@ -271,13 +271,13 @@ class VecEnvWrapper(VecEnv): def get_images(self) -> Sequence[np.ndarray]: return self.venv.get_images() - def get_attr(self, attr_name, indices=None): + def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]: return self.venv.get_attr(attr_name, indices) - def set_attr(self, attr_name, value, indices=None): + def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None: return self.venv.set_attr(attr_name, value, indices) - def env_method(self, method_name, *method_args, indices=None, **method_kwargs): + def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]: return self.venv.env_method(method_name, *method_args, indices=indices, **method_kwargs) def __getattr__(self, name: str) -> Any: @@ -305,7 +305,7 @@ class VecEnvWrapper(VecEnv): all_attributes.update(self.class_attributes) return all_attributes - def getattr_recursive(self, name: str): + def getattr_recursive(self, name: str) -> Any: """Recursively check wrappers to find attribute. :param name: name of attribute to look for @@ -323,7 +323,7 @@ class VecEnvWrapper(VecEnv): return attr - def getattr_depth_check(self, name: str, already_found: bool): + def getattr_depth_check(self, name: str, already_found: bool) -> str: """See base class. :return: name of module whose attribute is being shadowed, if any. diff --git a/stable_baselines3/common/vec_env/dummy_vec_env.py b/stable_baselines3/common/vec_env/dummy_vec_env.py index 1b34d2d..bd656e1 100644 --- a/stable_baselines3/common/vec_env/dummy_vec_env.py +++ b/stable_baselines3/common/vec_env/dummy_vec_env.py @@ -1,11 +1,11 @@ from collections import OrderedDict from copy import deepcopy -from typing import Callable, List, Optional, Sequence +from typing import Any, Callable, List, Optional, Sequence, Union import gym import numpy as np -from stable_baselines3.common.vec_env.base_vec_env import VecEnv +from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvIndices, VecEnvObs, VecEnvStepReturn from stable_baselines3.common.vec_env.util import copy_obs_dict, dict_to_obs, obs_space_info @@ -35,10 +35,10 @@ class DummyVecEnv(VecEnv): self.actions = None self.metadata = env.metadata - def step_async(self, actions: np.ndarray): + def step_async(self, actions: np.ndarray) -> None: self.actions = actions - def step_wait(self): + def step_wait(self) -> VecEnvStepReturn: for env_idx in range(self.num_envs): obs, self.buf_rews[env_idx], self.buf_dones[env_idx], self.buf_infos[env_idx] = self.envs[env_idx].step( self.actions[env_idx] @@ -50,26 +50,26 @@ class DummyVecEnv(VecEnv): self._save_obs(env_idx, obs) return (self._obs_from_buf(), np.copy(self.buf_rews), np.copy(self.buf_dones), deepcopy(self.buf_infos)) - def seed(self, seed: Optional[int] = None) -> List[int]: + def seed(self, seed: Optional[int] = None) -> List[Union[None, int]]: seeds = list() for idx, env in enumerate(self.envs): seeds.append(env.seed(seed + idx)) return seeds - def reset(self): + def reset(self) -> VecEnvObs: for env_idx in range(self.num_envs): obs = self.envs[env_idx].reset() self._save_obs(env_idx, obs) return self._obs_from_buf() - def close(self): + def close(self) -> None: for env in self.envs: env.close() def get_images(self) -> Sequence[np.ndarray]: return [env.render(mode="rgb_array") for env in self.envs] - def render(self, mode: str = "human"): + def render(self, mode: str = "human") -> Optional[np.ndarray]: """ Gym environment rendering. If there are multiple environments then they are tiled together in one image via ``BaseVecEnv.render()``. @@ -86,32 +86,32 @@ class DummyVecEnv(VecEnv): else: return super().render(mode=mode) - def _save_obs(self, env_idx, obs): + def _save_obs(self, env_idx: int, obs: VecEnvObs) -> None: for key in self.keys: if key is None: self.buf_obs[key][env_idx] = obs else: self.buf_obs[key][env_idx] = obs[key] - def _obs_from_buf(self): + def _obs_from_buf(self) -> VecEnvObs: return dict_to_obs(self.observation_space, copy_obs_dict(self.buf_obs)) - def get_attr(self, attr_name, indices=None): + def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]: """Return attribute from vectorized environment (see base class).""" target_envs = self._get_target_envs(indices) return [getattr(env_i, attr_name) for env_i in target_envs] - def set_attr(self, attr_name, value, indices=None): + def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None: """Set attribute inside vectorized environments (see base class).""" target_envs = self._get_target_envs(indices) for env_i in target_envs: setattr(env_i, attr_name, value) - def env_method(self, method_name, *method_args, indices=None, **method_kwargs): + def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]: """Call instance methods of vectorized environments.""" target_envs = self._get_target_envs(indices) return [getattr(env_i, method_name)(*method_args, **method_kwargs) for env_i in target_envs] - def _get_target_envs(self, indices): + def _get_target_envs(self, indices: VecEnvIndices) -> List[gym.Env]: indices = self._get_indices(indices) return [self.envs[i] for i in indices] diff --git a/stable_baselines3/common/vec_env/subproc_vec_env.py b/stable_baselines3/common/vec_env/subproc_vec_env.py index 0f5d8bb..4a9cdc5 100644 --- a/stable_baselines3/common/vec_env/subproc_vec_env.py +++ b/stable_baselines3/common/vec_env/subproc_vec_env.py @@ -1,14 +1,22 @@ -import multiprocessing +import multiprocessing as mp from collections import OrderedDict -from typing import Sequence +from typing import Any, Callable, List, Optional, Sequence, Tuple, Union import gym import numpy as np -from stable_baselines3.common.vec_env.base_vec_env import CloudpickleWrapper, VecEnv +from stable_baselines3.common.vec_env.base_vec_env import ( + CloudpickleWrapper, + VecEnv, + VecEnvIndices, + VecEnvObs, + VecEnvStepReturn, +) -def _worker(remote, parent_remote, env_fn_wrapper): +def _worker( + remote: mp.connection.Connection, parent_remote: mp.connection.Connection, env_fn_wrapper: CloudpickleWrapper +) -> None: parent_remote.close() env = env_fn_wrapper.var() while True: @@ -71,7 +79,7 @@ class SubprocVecEnv(VecEnv): Defaults to 'forkserver' on available platforms, and 'spawn' otherwise. """ - def __init__(self, env_fns, start_method=None): + def __init__(self, env_fns: List[Callable[[], gym.Env]], start_method: Optional[str] = None): self.waiting = False self.closed = False n_envs = len(env_fns) @@ -80,9 +88,9 @@ class SubprocVecEnv(VecEnv): # Fork is not a thread safe method (see issue #217) # but is more user friendly (does not require to wrap the code in # a `if __name__ == "__main__":`) - forkserver_available = "forkserver" in multiprocessing.get_all_start_methods() + forkserver_available = "forkserver" in mp.get_all_start_methods() start_method = "forkserver" if forkserver_available else "spawn" - ctx = multiprocessing.get_context(start_method) + ctx = mp.get_context(start_method) self.remotes, self.work_remotes = zip(*[ctx.Pipe() for _ in range(n_envs)]) self.processes = [] @@ -98,29 +106,29 @@ class SubprocVecEnv(VecEnv): observation_space, action_space = self.remotes[0].recv() VecEnv.__init__(self, len(env_fns), observation_space, action_space) - def step_async(self, actions): + def step_async(self, actions: np.ndarray) -> None: for remote, action in zip(self.remotes, actions): remote.send(("step", action)) self.waiting = True - def step_wait(self): + def step_wait(self) -> VecEnvStepReturn: results = [remote.recv() for remote in self.remotes] self.waiting = False obs, rews, dones, infos = zip(*results) return _flatten_obs(obs, self.observation_space), np.stack(rews), np.stack(dones), infos - def seed(self, seed=None): + def seed(self, seed: Optional[int] = None) -> List[Union[None, int]]: for idx, remote in enumerate(self.remotes): remote.send(("seed", seed + idx)) return [remote.recv() for remote in self.remotes] - def reset(self): + def reset(self) -> VecEnvObs: for remote in self.remotes: remote.send(("reset", None)) obs = [remote.recv() for remote in self.remotes] return _flatten_obs(obs, self.observation_space) - def close(self): + def close(self) -> None: if self.closed: return if self.waiting: @@ -140,14 +148,14 @@ class SubprocVecEnv(VecEnv): imgs = [pipe.recv() for pipe in self.remotes] return imgs - def get_attr(self, attr_name, indices=None): + def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]: """Return attribute from vectorized environment (see base class).""" target_remotes = self._get_target_remotes(indices) for remote in target_remotes: remote.send(("get_attr", attr_name)) return [remote.recv() for remote in target_remotes] - def set_attr(self, attr_name, value, indices=None): + def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None: """Set attribute inside vectorized environments (see base class).""" target_remotes = self._get_target_remotes(indices) for remote in target_remotes: @@ -155,14 +163,14 @@ class SubprocVecEnv(VecEnv): for remote in target_remotes: remote.recv() - def env_method(self, method_name, *method_args, indices=None, **method_kwargs): + def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]: """Call instance methods of vectorized environments.""" target_remotes = self._get_target_remotes(indices) for remote in target_remotes: remote.send(("env_method", (method_name, method_args, method_kwargs))) return [remote.recv() for remote in target_remotes] - def _get_target_remotes(self, indices): + def _get_target_remotes(self, indices: VecEnvIndices) -> List[Any]: """ Get the connection object needed to communicate with the wanted envs that are in subprocesses. @@ -174,14 +182,14 @@ class SubprocVecEnv(VecEnv): return [self.remotes[i] for i in indices] -def _flatten_obs(obs, space): +def _flatten_obs(obs: Union[List[VecEnvObs], Tuple[VecEnvObs]], space: gym.spaces.Space) -> VecEnvObs: """ Flatten observations, depending on the observation space. - :param obs: (list or tuple where X is dict, tuple or ndarray) observations. + :param obs: observations. A list or tuple of observations, one per environment. Each environment observation may be a NumPy array, or a dict or tuple of NumPy arrays. - :return (OrderedDict, tuple or ndarray) flattened observations. + :return: flattened observations. A flattened NumPy array or an OrderedDict or tuple of flattened numpy arrays. Each NumPy array has the environment index as its first axis. """ diff --git a/stable_baselines3/common/vec_env/util.py b/stable_baselines3/common/vec_env/util.py index c38c964..8bc7ee4 100644 --- a/stable_baselines3/common/vec_env/util.py +++ b/stable_baselines3/common/vec_env/util.py @@ -2,11 +2,13 @@ Helpers for dealing with vectorized environments. """ from collections import OrderedDict -from typing import Any, Dict, List, Tuple, Union +from typing import Any, Dict, List, Tuple import gym import numpy as np +from stable_baselines3.common.vec_env.base_vec_env import VecEnvObs + def copy_obs_dict(obs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: """ @@ -19,9 +21,7 @@ def copy_obs_dict(obs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: return OrderedDict([(k, np.copy(v)) for k, v in obs.items()]) -def dict_to_obs( - space: gym.spaces.Space, obs_dict: Dict[Any, np.ndarray] -) -> Union[Dict[Any, np.ndarray], Tuple[np.ndarray, ...], np.ndarray]: +def dict_to_obs(space: gym.spaces.Space, obs_dict: Dict[Any, np.ndarray]) -> VecEnvObs: """ Convert an internal representation raw_obs into the appropriate type specified by space. diff --git a/stable_baselines3/common/vec_env/vec_check_nan.py b/stable_baselines3/common/vec_env/vec_check_nan.py index c943cb2..258f9c2 100644 --- a/stable_baselines3/common/vec_env/vec_check_nan.py +++ b/stable_baselines3/common/vec_env/vec_check_nan.py @@ -2,7 +2,7 @@ import warnings import numpy as np -from stable_baselines3.common.vec_env.base_vec_env import VecEnvWrapper +from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvObs, VecEnvStepReturn, VecEnvWrapper class VecCheckNan(VecEnvWrapper): @@ -16,7 +16,7 @@ class VecCheckNan(VecEnvWrapper): :param check_inf: Whether or not to check for +inf or -inf as well """ - def __init__(self, venv, raise_exception=False, warn_once=True, check_inf=True): + def __init__(self, venv: VecEnv, raise_exception: bool = False, warn_once: bool = True, check_inf: bool = True): VecEnvWrapper.__init__(self, venv) self.raise_exception = raise_exception self.warn_once = warn_once @@ -25,13 +25,13 @@ class VecCheckNan(VecEnvWrapper): self._observations = None self._user_warned = False - def step_async(self, actions): + def step_async(self, actions: np.ndarray) -> None: self._check_val(async_step=True, actions=actions) self._actions = actions self.venv.step_async(actions) - def step_wait(self): + def step_wait(self) -> VecEnvStepReturn: observations, rewards, news, infos = self.venv.step_wait() self._check_val(async_step=False, observations=observations, rewards=rewards, news=news) @@ -39,7 +39,7 @@ class VecCheckNan(VecEnvWrapper): self._observations = observations return observations, rewards, news, infos - def reset(self): + def reset(self) -> VecEnvObs: observations = self.venv.reset() self._actions = None @@ -48,7 +48,7 @@ class VecCheckNan(VecEnvWrapper): self._observations = observations return observations - def _check_val(self, *, async_step, **kwargs): + def _check_val(self, *, async_step: bool, **kwargs) -> None: # if warn and warn once and have warned once: then stop checking if not self.raise_exception and self.warn_once and self._user_warned: return @@ -66,7 +66,7 @@ class VecCheckNan(VecEnvWrapper): self._user_warned = True msg = "" for i, (name, type_val) in enumerate(found): - msg += "found {} in {}".format(type_val, name) + msg += f"found {type_val} in {name}" if i != len(found) - 1: msg += ", " @@ -76,9 +76,9 @@ class VecCheckNan(VecEnvWrapper): if self._actions is None: msg += "environment observation (at reset)" else: - msg += "environment, Last given value was: \r\n\taction={}".format(self._actions) + msg += f"environment, Last given value was: \r\n\taction={self._actions}" else: - msg += "RL model, Last given value was: \r\n\tobservations={}".format(self._observations) + msg += f"RL model, Last given value was: \r\n\tobservations={self._observations}" if self.raise_exception: raise ValueError(msg) diff --git a/stable_baselines3/common/vec_env/vec_normalize.py b/stable_baselines3/common/vec_env/vec_normalize.py index 0c40748..39a5d11 100644 --- a/stable_baselines3/common/vec_env/vec_normalize.py +++ b/stable_baselines3/common/vec_env/vec_normalize.py @@ -1,9 +1,10 @@ import pickle +from typing import Any, Dict import numpy as np from stable_baselines3.common.running_mean_std import RunningMeanStd -from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper +from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvStepReturn, VecEnvWrapper class VecNormalize(VecEnvWrapper): @@ -22,7 +23,15 @@ class VecNormalize(VecEnvWrapper): """ def __init__( - self, venv, training=True, norm_obs=True, norm_reward=True, clip_obs=10.0, clip_reward=10.0, gamma=0.99, epsilon=1e-8 + self, + venv: VecEnv, + training: bool = True, + norm_obs: bool = True, + norm_reward: bool = True, + clip_obs: float = 10.0, + clip_reward: float = 10.0, + gamma: float = 0.99, + epsilon: float = 1e-8, ): VecEnvWrapper.__init__(self, venv) self.obs_rms = RunningMeanStd(shape=self.observation_space.shape) @@ -39,7 +48,7 @@ class VecNormalize(VecEnvWrapper): self.old_obs = np.array([]) self.old_reward = np.array([]) - def __getstate__(self): + def __getstate__(self) -> Dict[str, Any]: """ Gets state for pickling. @@ -52,7 +61,7 @@ class VecNormalize(VecEnvWrapper): del state["ret"] return state - def __setstate__(self, state): + def __setstate__(self, state: Dict[str, Any]) -> None: """ Restores pickled state. @@ -63,7 +72,7 @@ class VecNormalize(VecEnvWrapper): assert "venv" not in state self.venv = None - def set_venv(self, venv): + def set_venv(self, venv: VecEnv) -> None: """ Sets the vector environment to wrap to venv. @@ -78,7 +87,7 @@ class VecNormalize(VecEnvWrapper): raise ValueError("venv is incompatible with current statistics.") self.ret = np.zeros(self.num_envs) - def step_wait(self): + def step_wait(self) -> VecEnvStepReturn: """ Apply sequence of actions to sequence of environments actions -> (observations, rewards, news) @@ -100,12 +109,12 @@ class VecNormalize(VecEnvWrapper): self.ret[news] = 0 return obs, rews, news, infos - def _update_reward(self, reward): + def _update_reward(self, reward: np.ndarray) -> None: """Update reward normalization statistics.""" self.ret = self.ret * self.gamma + reward self.ret_rms.update(self.ret) - def normalize_obs(self, obs): + def normalize_obs(self, obs: np.ndarray) -> np.ndarray: """ Normalize observations using this VecNormalize's observations statistics. Calling this method does not update statistics. @@ -114,7 +123,7 @@ class VecNormalize(VecEnvWrapper): obs = np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs, self.clip_obs) return obs - def normalize_reward(self, reward): + def normalize_reward(self, reward: np.ndarray) -> np.ndarray: """ Normalize rewards using this VecNormalize's rewards statistics. Calling this method does not update statistics. @@ -123,30 +132,30 @@ class VecNormalize(VecEnvWrapper): reward = np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward) return reward - def unnormalize_obs(self, obs): + def unnormalize_obs(self, obs: np.ndarray) -> np.ndarray: if self.norm_obs: return (obs * np.sqrt(self.obs_rms.var + self.epsilon)) + self.obs_rms.mean return obs - def unnormalize_reward(self, reward): + def unnormalize_reward(self, reward: np.ndarray) -> np.ndarray: if self.norm_reward: return reward * np.sqrt(self.ret_rms.var + self.epsilon) return reward - def get_original_obs(self): + def get_original_obs(self) -> np.ndarray: """ Returns an unnormalized version of the observations from the most recent step or reset. """ return self.old_obs.copy() - def get_original_reward(self): + def get_original_reward(self) -> np.ndarray: """ Returns an unnormalized version of the rewards from the most recent step. """ return self.old_reward.copy() - def reset(self): + def reset(self) -> np.ndarray: """ Reset all environments """ diff --git a/stable_baselines3/common/vec_env/vec_video_recorder.py b/stable_baselines3/common/vec_env/vec_video_recorder.py index fc3be1b..08a2950 100644 --- a/stable_baselines3/common/vec_env/vec_video_recorder.py +++ b/stable_baselines3/common/vec_env/vec_video_recorder.py @@ -1,13 +1,12 @@ import os +from typing import Callable from gym.wrappers.monitoring import video_recorder from stable_baselines3.common import logger -from stable_baselines3.common.vec_env.base_vec_env import VecEnvWrapper +from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvObs, VecEnvStepReturn, VecEnvWrapper from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv from stable_baselines3.common.vec_env.subproc_vec_env import SubprocVecEnv -from stable_baselines3.common.vec_env.vec_frame_stack import VecFrameStack -from stable_baselines3.common.vec_env.vec_normalize import VecNormalize class VecVideoRecorder(VecEnvWrapper): @@ -24,7 +23,14 @@ class VecVideoRecorder(VecEnvWrapper): :param name_prefix: Prefix to the video name """ - def __init__(self, venv, video_folder, record_video_trigger, video_length=200, name_prefix="rl-video"): + def __init__( + self, + venv: VecEnv, + video_folder: str, + record_video_trigger: Callable[[int], bool], + video_length: int = 200, + name_prefix: str = "rl-video", + ): VecEnvWrapper.__init__(self, venv) @@ -34,7 +40,7 @@ class VecVideoRecorder(VecEnvWrapper): # Unwrap to retrieve metadata dict # that will be used by gym recorder - while isinstance(temp_env, VecNormalize) or isinstance(temp_env, VecFrameStack): + while isinstance(temp_env, VecEnvWrapper): temp_env = temp_env.venv if isinstance(temp_env, DummyVecEnv) or isinstance(temp_env, SubprocVecEnv): @@ -58,12 +64,12 @@ class VecVideoRecorder(VecEnvWrapper): self.recording = False self.recorded_frames = 0 - def reset(self): + def reset(self) -> VecEnvObs: obs = self.venv.reset() self.start_video_recorder() return obs - def start_video_recorder(self): + def start_video_recorder(self) -> None: self.close_video_recorder() video_name = f"{self.name_prefix}-step-{self.step_id}-to-step-{self.step_id + self.video_length}" @@ -76,10 +82,10 @@ class VecVideoRecorder(VecEnvWrapper): self.recorded_frames = 1 self.recording = True - def _video_enabled(self): + def _video_enabled(self) -> bool: return self.record_video_trigger(self.step_id) - def step_wait(self): + def step_wait(self) -> VecEnvStepReturn: obs, rews, dones, infos = self.venv.step_wait() self.step_id += 1 @@ -94,13 +100,13 @@ class VecVideoRecorder(VecEnvWrapper): return obs, rews, dones, infos - def close_video_recorder(self): + def close_video_recorder(self) -> None: if self.recording: self.video_recorder.close() self.recording = False self.recorded_frames = 1 - def close(self): + def close(self) -> None: VecEnvWrapper.close(self) self.close_video_recorder() diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 55fb835..6d7a455 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -132,7 +132,7 @@ class DQN(OffPolicyAlgorithm): self.q_net = self.policy.q_net self.q_net_target = self.policy.q_net_target - def _on_step(self): + def _on_step(self) -> None: """ Update the exploration rate and target network if needed. This method is called in ``collect_rollout()`` after each step in the environment. diff --git a/stable_baselines3/td3/policies.py b/stable_baselines3/td3/policies.py index 401a541..62601c4 100644 --- a/stable_baselines3/td3/policies.py +++ b/stable_baselines3/td3/policies.py @@ -181,7 +181,7 @@ class TD3Policy(BasePolicy): def make_critic(self) -> ContinuousCritic: return ContinuousCritic(**self.critic_kwargs).to(self.device) - def forward(self, observation: th.Tensor, deterministic: bool = False): + def forward(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor: return self._predict(observation, deterministic=deterministic) def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor: diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index ac39a10..37f1777 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -0.9.0 +0.10.0a0 From fc6c5d3daa6a84c1bf4cf67edd676e164a7b88b2 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Sun, 11 Oct 2020 23:22:12 +0200 Subject: [PATCH 3/5] Migration Guide (#123) * Start migration guide * Update guide * Add comment on RMSpropTFLike plus PPO/A2C migrations * Add note about set/get-parameters * Update migration guide * Update changelog and readme * Update doc + clean changelog * Address comments Co-authored-by: Anssi "Miffyli" Kanervisto --- README.md | 4 +- docs/guide/install.rst | 2 +- docs/guide/migration.rst | 190 ++++++++++++++++++++++++++++++++++++++- docs/misc/changelog.rst | 8 +- 4 files changed, 195 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 95e7c15..94c005d 100644 --- a/README.md +++ b/README.md @@ -50,9 +50,9 @@ Planned features: - [ ] TRPO -## Migration guide +## Migration guide: from Stable-Baselines (SB2) to Stable-Baselines3 (SB3) -**TODO: migration guide from Stable-Baselines in the documentation** +A migration guide from SB2 to SB3 can be found in the [documentation](https://stable-baselines3.readthedocs.io/en/master/guide/migration.html). ## Documentation diff --git a/docs/guide/install.rst b/docs/guide/install.rst index 27c431b..e5ec926 100644 --- a/docs/guide/install.rst +++ b/docs/guide/install.rst @@ -29,7 +29,7 @@ To install Stable Baselines3 with pip, execute: pip install stable-baselines3[extra] -This includes an optional dependencies like Tensorboard, OpenCV or ```atari-py``` to train on atari games. If you do not need those, you can use: +This includes an optional dependencies like Tensorboard, OpenCV or ``atari-py`` to train on atari games. If you do not need those, you can use: .. code-block:: bash diff --git a/docs/guide/migration.rst b/docs/guide/migration.rst index 263c939..82fbc69 100644 --- a/docs/guide/migration.rst +++ b/docs/guide/migration.rst @@ -5,8 +5,194 @@ Migrating from Stable-Baselines ================================ -This is a guide to migrate from Stable-Baselines to Stable-Baselines3. +This is a guide to migrate from Stable-Baselines (SB2) to Stable-Baselines3 (SB3). It also references the main changes. -**TODO** + +Overview +======== + +Overall Stable-Baselines3 (SB3) keeps the high-level API of Stable-Baselines (SB2). +Most of the changes are to ensure more consistency and are internal ones. +Because of the backend change, from Tensorflow to PyTorch, the internal code is much much readable and easy to debug +at the cost of some speed (dynamic graph vs static graph., see `Issue #90 `_) +However, the algorithms were extensively benchmarked on Atari games and continuous control PyBullet envs +(see `Issue #48 `_ and `Issue #49 `_) +so you should not expect performance drop when switching from SB2 to SB3. + + +How to migrate? +=============== + +In most cases, replacing ``from stable_baselines`` by ``from stable_baselines3`` will be sufficient. +Some files were moved to the common folder (cf below) and could result to import errors. +Some algorithms were removed because of their complexity to improve the maintainability of the project. +We recommend reading this guide carefully to understand all the changes that were made. +You can also take a look at the `rl-zoo3 `_ and compare the imports +to the `rl-zoo `_ of SB2 to have a concrete example of successful migration. + + +Breaking Changes +================ + + +- SB3 requires python 3.6+ (instead of python 3.5+ for SB2) +- Dropped MPI support +- Dropped layer normalized policies (e.g. ``LnMlpPolicy``) +- Dropped parameter noise for DDPG and DQN +- PPO is now closer to the original implementation (no clipping of the value function by default), cf PPO section below +- Orthogonal initialization is only used by A2C/PPO +- The features extractor (CNN extractor) is shared between policy and q-networks for DDPG/SAC/TD3 and only the policy loss used to update it (much faster) +- Tensorboard legacy logging was dropped in favor of having one logger for the terminal and Tensorboard (cf :ref:`Tensorboard integration `) +- We dropped ACKTR/ACER support because of their complexity compared to simpler alternatives (PPO, SAC, TD3) performing as good. +- We dropped GAIL support as we are focusing on model-free RL only, you can however take a look at the `Imitation Learning Baseline Implementations `_ + which are based on SB3. +- ``action_probability`` is currently not implemented in the base class + +You can take a look at the `issue about SB3 implementation design `_ for more details. + + +Moved Files +----------- + +- ``bench/monitor.py`` -> ``common/monitor.py`` +- ``logger.py`` -> ``common/logger.py`` +- ``results_plotter.py`` -> ``common/results_plotter.py`` + +Utility functions are no longer exported from ``common`` module, you should import them with their absolute path, e.g.: + +.. code-block:: python + + from stable_baselines3.common.cmd_util import make_atari_env, make_vec_env + from stable_baselines3.common.utils import set_random_seed + +instead of ``from stable_baselines3.common import make_atari_env`` + + + +Changes and renaming in parameters +---------------------------------- + +Base-class (all algorithms) +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +- ``load_parameters`` -> ``set_parameters`` + + - ``get/set_parameters`` return a dictionary mapping object names + to their respective PyTorch tensors and other objects representing + their parameters, instead of simpler mapping of parameter name to + a NumPy array. These functions also return PyTorch tensors rather + than NumPy arrays. + + +Policies +^^^^^^^^ + +- ``cnn_extractor`` -> ``feature_extractor``, as ``feature_extractor`` in now used with ``MlpPolicy`` too + +A2C +^^^ + +- ``epsilon`` -> ``rms_prop_eps`` +- ``lr_schedule`` is part of ``learning_rate`` (it can be a callable). +- ``alpha``, ``momentum`` are modifiable through ``policy_kwargs`` key ``optimizer_kwargs``. + +.. warning:: + + PyTorch implementation of RMSprop `differs from Tensorflow's `_, + which leads to `different and potentially more unstable results `_. + Use ``stable_baselines3.common.sb2_compat.rmsprop_tf_like.RMSpropTFLike`` optimizer to match the results + with Tensorflow's implementation. This can be done through ``policy_kwargs``: ``A2C(policy_kwargs=dict(optimizer_class=RMSpropTFLike))`` + + +PPO +^^^ + +- ``cliprange`` -> ``clip_range`` +- ``cliprange_vf`` -> ``clip_range_vf`` +- ``nminibatches`` -> ``batch_size`` + +.. warning:: + + ``nminibatches`` gave different batch size depending on the number of environments: ``batch_size = (n_steps * n_envs) // nminibatches`` + + +- ``clip_range_vf`` behavior for PPO is slightly different: Set it to ``None`` (default) to deactivate clipping (in SB2, you had to pass ``-1``, ``None`` meant to use ``clip_range`` for the clipping) +- ``lam`` -> ``gae_lambda`` +- ``noptepochs`` -> ``n_epochs`` + +PPO default hyperparameters are the one tuned for continuous control environment. +We recommend taking a look at the :ref:`RL Zoo ` for hyperparameters tuned for Atari games. + + +DQN +^^^ + +Only the vanilla DQN is implemented right now but extensions will follow. +Default hyperparameters are taken from the nature paper, except for the optimizer and learning rate that were taken from Stable Baselines defaults. + +DDPG +^^^^ + +DDPG now follows the same interface as SAC/TD3. +For state/reward normalization, you should use ``VecNormalize`` as for all other algorithms. + +SAC/TD3 +^^^^^^^ + +SAC/TD3 now accept any number of critics, e.g. ``policy_kwargs=dict(n_critics=3)``, instead of only two before. + + +.. note:: + + SAC/TD3 default hyperparameters (including network architecture) now match the ones from the original papers. + DDPG is using TD3 defaults. + + +SAC +^^^ + +SAC implementation matches the latest version of the original implementation: it uses two Q function networks and two target Q function networks +instead of two Q function networks and one Value function network (SB2 implementation, first version of the original implementation). +Despite this change, no change in performance should be expected. + +.. note:: + + SAC ``predict()`` method has now ``deterministic=False`` by default for consistency. + To match SB2 behavior, you need to explicitly pass ``deterministic=True`` + + + +New logger API +-------------- + +- Methods were renamed in the logger: + + - ``logkv`` -> ``record``, ``writekvs`` -> ``write``, ``writeseq`` -> ``write_sequence``, + - ``logkvs`` -> ``record_dict``, ``dumpkvs`` -> ``dump``, + - ``getkvs`` -> ``get_log_dict``, ``logkv_mean`` -> ``record_mean``, + + +Internal Changes +---------------- + +Please read the :ref:`Developer Guide ` section. + + +New Features (SB3 vs SB2) +========================= + +- Much cleaner and consistent base code (and no more warnings =D!) and static type checks +- Independent saving/loading/predict for policies +- A2C now supports Generalized Advantage Estimation (GAE) and advantage normalization (both are deactivated by default) +- Generalized State-Dependent Exploration (gSDE) exploration is available for A2C/PPO/SAC. It allows to use RL directly on real robots (cf https://arxiv.org/abs/2005.05719) +- Proper evaluation (using separate env) is included in the base class (using ``EvalCallback``), + if you pass the environment as a string, you can pass ``create_eval_env=True`` to the algorithm constructor. +- Better saving/loading: optimizers are now included in the saved parameters and there is two new methods ``save_replay_buffer`` and ``load_replay_buffer`` for the replay buffer when using off-policy algorithms (DQN/DDPG/SAC/TD3) +- You can pass ``optimizer_class`` and ``optimizer_kwargs`` to ``policy_kwargs`` in order to easily + customize optimizers +- Seeding now works properly to have deterministic results +- Replay buffer does not grow, allocate everything at build time (faster) +- We added a memory efficient replay buffer variant (pass ``optimize_memory_usage=True`` to the constructor), it reduces drastically the memory used especially when using images +- You can specify an arbitrary number of critics for SAC/TD3 (e.g. ``policy_kwargs=dict(n_critics=3)``) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 90b5322..da32a8e 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -26,6 +26,7 @@ Others: Documentation: ^^^^^^^^^^^^^^ +- Added first draft of migration guide Pre-Release 0.9.0 (2020-10-03) @@ -36,8 +37,7 @@ Pre-Release 0.9.0 (2020-10-03) Breaking Changes: ^^^^^^^^^^^^^^^^^ - Removed ``device`` keyword argument of policies; use ``policy.to(device)`` instead. (@qxcv) -- Rename ``BaseClass.get_torch_variables`` -> ``BaseClass._get_torch_save_params`` and - ``BaseClass.excluded_save_params`` -> ``BaseClass._excluded_save_params`` +- Rename ``BaseClass.get_torch_variables`` -> ``BaseClass._get_torch_save_params`` and ``BaseClass.excluded_save_params`` -> ``BaseClass._excluded_save_params`` - Renamed saved items ``tensors`` to ``pytorch_variables`` for clarity - ``make_atari_env``, ``make_vec_env`` and ``set_random_seed`` must be imported with (and not directly from ``stable_baselines3.common``): @@ -73,7 +73,7 @@ Others: - Removed ``AlreadySteppingError`` and ``NotSteppingError`` that were not used - Fixed typos in SAC and TD3 - Reorganized functions for clarity in ``BaseClass`` (save/load functions close to each other, private - functions at top) + functions at top) - Clarified docstrings on what is saved and loaded to/from files - Simplified ``save_to_zip_file`` function by removing duplicate code - Store library version along with the saved models @@ -100,7 +100,7 @@ Breaking Changes: - Refactored ``Critic`` class for ``TD3`` and ``SAC``, it is now called ``ContinuousCritic`` and has an additional parameter ``n_critics`` - ``SAC`` and ``TD3`` now accept an arbitrary number of critics (e.g. ``policy_kwargs=dict(n_critics=3)``) - instead of only 2 previously + instead of only 2 previously New Features: ^^^^^^^^^^^^^ From fc9527157a6728075043a772a22cfb1da2783efb Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Mon, 12 Oct 2020 23:10:54 +0200 Subject: [PATCH 4/5] Fix off-by-one GAE computation (#185) * Fix off-by-one GAE computation * Fix identity test * Revert gae loop --- docs/misc/changelog.rst | 1 + stable_baselines3/common/buffers.py | 12 ++++++------ stable_baselines3/common/on_policy_algorithm.py | 7 ++++++- tests/test_identity.py | 2 +- 4 files changed, 14 insertions(+), 8 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index da32a8e..7a4bb69 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -15,6 +15,7 @@ New Features: Bug Fixes: ^^^^^^^^^^ +- Fix GAE computation for on-policy algorithms (off-by one for the last value) (thanks @Wovchena) Deprecations: ^^^^^^^^^^^^^ diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 2063d6c..563ddea 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -299,7 +299,7 @@ class RolloutBuffer(BaseBuffer): self.generator_ready = False super(RolloutBuffer, self).reset() - def compute_returns_and_advantage(self, last_value: th.Tensor, dones: np.ndarray) -> None: + def compute_returns_and_advantage(self, last_values: th.Tensor, dones: np.ndarray) -> None: """ Post-processing step: compute the returns (sum of discounted rewards) and GAE advantage. @@ -310,22 +310,22 @@ class RolloutBuffer(BaseBuffer): where R is the discounted reward with value bootstrap, set ``gae_lambda=1.0`` during initialization. - :param last_value: + :param last_values: :param dones: """ # convert to numpy - last_value = last_value.clone().cpu().numpy().flatten() + last_values = last_values.clone().cpu().numpy().flatten() last_gae_lam = 0 for step in reversed(range(self.buffer_size)): if step == self.buffer_size - 1: next_non_terminal = 1.0 - dones - next_value = last_value + next_values = last_values else: next_non_terminal = 1.0 - self.dones[step + 1] - next_value = self.values[step + 1] - delta = self.rewards[step] + self.gamma * next_value * next_non_terminal - self.values[step] + next_values = self.values[step + 1] + delta = self.rewards[step] + self.gamma * next_values * next_non_terminal - self.values[step] last_gae_lam = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae_lam self.advantages[step] = last_gae_lam self.returns = self.advantages + self.values diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 85d08c8..4f2ef63 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/stable_baselines3/common/on_policy_algorithm.py @@ -179,7 +179,12 @@ class OnPolicyAlgorithm(BaseAlgorithm): self._last_obs = new_obs self._last_dones = dones - rollout_buffer.compute_returns_and_advantage(values, dones=dones) + with th.no_grad(): + # Compute value for the last timestep + obs_tensor = th.as_tensor(new_obs).to(self.device) + _, values, _ = self.policy.forward(obs_tensor) + + rollout_buffer.compute_returns_and_advantage(last_values=values, dones=dones) callback.on_rollout_end() diff --git a/tests/test_identity.py b/tests/test_identity.py index 38c6570..678f63c 100644 --- a/tests/test_identity.py +++ b/tests/test_identity.py @@ -23,7 +23,7 @@ def test_discrete(model_class, env): if isinstance(env, (IdentityEnvMultiDiscrete, IdentityEnvMultiBinary)): return - model = model_class("MlpPolicy", env_, gamma=0.5, seed=1, **kwargs).learn(n_steps) + model = model_class("MlpPolicy", env_, gamma=0.4, seed=1, **kwargs).learn(n_steps) evaluate_policy(model, env_, n_eval_episodes=20, reward_threshold=90) obs = env.reset() From 2599f04940f14436290273caa2df9366a54b375c Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Tue, 13 Oct 2020 12:01:33 +0200 Subject: [PATCH 5/5] Add custom arch for off-policy actor/critic networks (#182) * Add custom arch for off-policy actor/critic networks * Fix type hints * Address comments * Make sure number of updated parameters match in polyak * Add zip_strict for strict-length zipping * Fix building docs * Add test for zip strict * Faster tests Co-authored-by: Anssi "Miffyli" Kanervisto --- docs/guide/custom_policy.rst | 34 ++++++++++++++++---- docs/misc/changelog.rst | 1 + stable_baselines3/common/torch_layers.py | 40 ++++++++++++++++++++++++ stable_baselines3/common/utils.py | 22 ++++++++++++- stable_baselines3/sac/policies.py | 22 ++++++++----- stable_baselines3/td3/policies.py | 25 ++++++++++----- tests/test_custom_policy.py | 23 ++++++++++---- tests/test_run.py | 6 ++-- tests/test_save_load.py | 20 +++++++----- tests/test_sde.py | 3 ++ tests/test_utils.py | 20 +++++++++++- 11 files changed, 177 insertions(+), 39 deletions(-) diff --git a/docs/guide/custom_policy.rst b/docs/guide/custom_policy.rst index 85ae801..b490bd7 100644 --- a/docs/guide/custom_policy.rst +++ b/docs/guide/custom_policy.rst @@ -258,9 +258,31 @@ If your task requires even more granular control over the policy/value architect -.. TODO (see https://github.com/DLR-RM/stable-baselines3/issues/113) -.. Off-Policy Algorithms -.. ^^^^^^^^^^^^^^^^^^^^^ -.. -.. If you need a network architecture that is different for the actor and the critic when using ``SAC``, ``DDPG`` or ``TD3``, -.. you can easily redefine the actor class for instance. +Off-Policy Algorithms +^^^^^^^^^^^^^^^^^^^^^ + +If you need a network architecture that is different for the actor and the critic when using ``SAC``, ``DDPG`` or ``TD3``, +you can pass a dictionary of the following structure: ``dict(qf=[], pi=[])``. + +For example, if you want a different architecture for the actor (aka ``pi``) and the critic (Q-function aka ``qf``) networks, +then you can specify ``net_arch=dict(qf=[400, 300], pi=[64, 64])``. + +Otherwise, to have actor and critic that share the same network architecture, +you only need to specify ``net_arch=[256, 256]`` (here, two hidden layers of 256 units each). + + +.. note:: + Compared to their on-policy counterparts, no shared layers (other than the feature extractor) + between the actor and the critic are allowed (to prevent issues with target networks). + + +.. code-block:: python + + from stable_baselines3 import SAC + + # Custom actor architecture with two layers of 64 units each + # Custom critic architecture with two layers of 400 and 300 units + policy_kwargs = dict(net_arch=dict(pi=[64, 64], qf=[400, 300])) + # Create the agent + model = SAC("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs, verbose=1) + model.learn(5000) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 7a4bb69..f883c1e 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -12,6 +12,7 @@ Breaking Changes: New Features: ^^^^^^^^^^^^^ +- Allow custom actor/critic network architectures using ``net_arch=dict(qf=[400, 300], pi=[64, 64])`` for off-policy algorithms (SAC, TD3, DDPG) Bug Fixes: ^^^^^^^^^^ diff --git a/stable_baselines3/common/torch_layers.py b/stable_baselines3/common/torch_layers.py index 43e7c5d..62539d8 100644 --- a/stable_baselines3/common/torch_layers.py +++ b/stable_baselines3/common/torch_layers.py @@ -219,3 +219,43 @@ class MlpExtractor(nn.Module): """ shared_latent = self.shared_net(features) return self.policy_net(shared_latent), self.value_net(shared_latent) + + +def get_actor_critic_arch(net_arch: Union[List[int], Dict[str, List[int]]]) -> Tuple[List[int], List[int]]: + """ + Get the actor and critic network architectures for off-policy actor-critic algorithms (SAC, TD3, DDPG). + + The ``net_arch`` parameter allows to specify the amount and size of the hidden layers, + which can be different for the actor and the critic. + It is assumed to be a list of ints or a dict. + + 1. If it is a list, actor and critic networks will have the same architecture. + The architecture is represented by a list of integers (of arbitrary length (zero allowed)) + each specifying the number of units per layer. + If the number of ints is zero, the network will be linear. + 2. If it is a dict, it should have the following structure: + ``dict(qf=[], pi=[])``. + where the network architecture is a list as described in 1. + + For example, to have actor and critic that share the same network architecture, + you only need to specify ``net_arch=[256, 256]`` (here, two hidden layers of 256 units each). + + If you want a different architecture for the actor and the critic, + then you can specify ``net_arch=dict(qf=[400, 300], pi=[64, 64])``. + + .. note:: + Compared to their on-policy counterparts, no shared layers (other than the feature extractor) + between the actor and the critic are allowed (to prevent issues with target networks). + + :param net_arch: The specification of the actor and critic networks. + See above for details on its formatting. + :return: The network architectures for the actor and the critic + """ + if isinstance(net_arch, list): + actor_arch, critic_arch = net_arch, net_arch + else: + assert isinstance(net_arch, dict), "Error: the net_arch can only contain be a list of ints or a dict" + assert "pi" in net_arch, "Error: no key 'pi' was provided in net_arch for the actor network" + assert "qf" in net_arch, "Error: no key 'qf' was provided in net_arch for the critic network" + actor_arch, critic_arch = net_arch["pi"], net_arch["qf"] + return actor_arch, critic_arch diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 8177842..5787265 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -2,6 +2,7 @@ import glob import os import random from collections import deque +from itertools import zip_longest from typing import Callable, Iterable, Optional, Union import gym @@ -286,6 +287,24 @@ def safe_mean(arr: Union[np.ndarray, list, deque]) -> np.ndarray: return np.nan if len(arr) == 0 else np.mean(arr) +def zip_strict(*iterables: Iterable) -> Iterable: + r""" + ``zip()`` function but enforces that iterables are of equal length. + Raises ``ValueError`` if iterables not of equal length. + Code inspired by Stackoverflow answer for question #32954486. + + :param \*iterables: iterables to ``zip()`` + """ + # As in Stackoverflow #32954486, use + # new object for "empty" in case we have + # Nones in iterable. + sentinel = object() + for combo in zip_longest(*iterables, fillvalue=sentinel): + if sentinel in combo: + raise ValueError("Iterables have different lengths") + yield combo + + def polyak_update(params: Iterable[th.nn.Parameter], target_params: Iterable[th.nn.Parameter], tau: float) -> None: """ Perform a Polyak average update on ``target_params`` using ``params``: @@ -303,6 +322,7 @@ def polyak_update(params: Iterable[th.nn.Parameter], target_params: Iterable[th. :param tau: the soft update coefficient ("Polyak update", between 0 and 1) """ with th.no_grad(): - for param, target_param in zip(params, target_params): + # zip does not raise an exception if length of parameters does not match. + for param, target_param in zip_strict(params, target_params): target_param.data.mul_(1 - tau) th.add(target_param.data, param.data, alpha=tau, out=target_param.data) diff --git a/stable_baselines3/sac/policies.py b/stable_baselines3/sac/policies.py index c11c437..744b0d0 100644 --- a/stable_baselines3/sac/policies.py +++ b/stable_baselines3/sac/policies.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Optional, Tuple, Type +from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union import gym import torch as th @@ -7,7 +7,13 @@ from torch import nn from stable_baselines3.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution from stable_baselines3.common.policies import BasePolicy, ContinuousCritic, create_sde_features_extractor, register_policy from stable_baselines3.common.preprocessing import get_action_dim -from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, NatureCNN, create_mlp +from stable_baselines3.common.torch_layers import ( + BaseFeaturesExtractor, + FlattenExtractor, + NatureCNN, + create_mlp, + get_actor_critic_arch, +) # CAP the standard deviation of the actor LOG_STD_MAX = 2 @@ -220,7 +226,7 @@ class SACPolicy(BasePolicy): observation_space: gym.spaces.Space, action_space: gym.spaces.Space, lr_schedule: Callable, - net_arch: Optional[List[int]] = None, + net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, @@ -250,6 +256,8 @@ class SACPolicy(BasePolicy): else: net_arch = [] + actor_arch, critic_arch = get_actor_critic_arch(net_arch) + # Create shared features extractor self.features_extractor = features_extractor_class(self.observation_space, **self.features_extractor_kwargs) self.features_dim = self.features_extractor.features_dim @@ -261,7 +269,7 @@ class SACPolicy(BasePolicy): "action_space": self.action_space, "features_extractor": self.features_extractor, "features_dim": self.features_dim, - "net_arch": self.net_arch, + "net_arch": actor_arch, "activation_fn": self.activation_fn, "normalize_images": normalize_images, } @@ -275,7 +283,7 @@ class SACPolicy(BasePolicy): } self.actor_kwargs.update(sde_kwargs) self.critic_kwargs = self.net_args.copy() - self.critic_kwargs.update({"n_critics": n_critics}) + self.critic_kwargs.update({"n_critics": n_critics, "net_arch": critic_arch}) self.actor, self.actor_target = None, None self.critic, self.critic_target = None, None @@ -300,7 +308,7 @@ class SACPolicy(BasePolicy): data.update( dict( - net_arch=self.net_args["net_arch"], + net_arch=self.net_arch, activation_fn=self.net_args["activation_fn"], use_sde=self.actor_kwargs["use_sde"], log_std_init=self.actor_kwargs["log_std_init"], @@ -374,7 +382,7 @@ class CnnPolicy(SACPolicy): observation_space: gym.spaces.Space, action_space: gym.spaces.Space, lr_schedule: Callable, - net_arch: Optional[List[int]] = None, + net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, diff --git a/stable_baselines3/td3/policies.py b/stable_baselines3/td3/policies.py index 62601c4..a4d5cea 100644 --- a/stable_baselines3/td3/policies.py +++ b/stable_baselines3/td3/policies.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Optional, Type +from typing import Any, Callable, Dict, List, Optional, Type, Union import gym import torch as th @@ -6,7 +6,13 @@ from torch import nn from stable_baselines3.common.policies import BasePolicy, ContinuousCritic, register_policy from stable_baselines3.common.preprocessing import get_action_dim -from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, NatureCNN, create_mlp +from stable_baselines3.common.torch_layers import ( + BaseFeaturesExtractor, + FlattenExtractor, + NatureCNN, + create_mlp, + get_actor_critic_arch, +) class Actor(BasePolicy): @@ -101,7 +107,7 @@ class TD3Policy(BasePolicy): observation_space: gym.spaces.Space, action_space: gym.spaces.Space, lr_schedule: Callable, - net_arch: Optional[List[int]] = None, + net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.ReLU, features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor, features_extractor_kwargs: Optional[Dict[str, Any]] = None, @@ -127,6 +133,8 @@ class TD3Policy(BasePolicy): else: net_arch = [] + actor_arch, critic_arch = get_actor_critic_arch(net_arch) + self.features_extractor = features_extractor_class(self.observation_space, **self.features_extractor_kwargs) self.features_dim = self.features_extractor.features_dim @@ -137,12 +145,13 @@ class TD3Policy(BasePolicy): "action_space": self.action_space, "features_extractor": self.features_extractor, "features_dim": self.features_dim, - "net_arch": self.net_arch, + "net_arch": actor_arch, "activation_fn": self.activation_fn, "normalize_images": normalize_images, } + self.actor_kwargs = self.net_args.copy() self.critic_kwargs = self.net_args.copy() - self.critic_kwargs.update({"n_critics": n_critics}) + self.critic_kwargs.update({"n_critics": n_critics, "net_arch": critic_arch}) self.actor, self.actor_target = None, None self.critic, self.critic_target = None, None @@ -163,7 +172,7 @@ class TD3Policy(BasePolicy): data.update( dict( - net_arch=self.net_args["net_arch"], + net_arch=self.net_arch, activation_fn=self.net_args["activation_fn"], n_critics=self.critic_kwargs["n_critics"], lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone @@ -176,7 +185,7 @@ class TD3Policy(BasePolicy): return data def make_actor(self) -> Actor: - return Actor(**self.net_args).to(self.device) + return Actor(**self.actor_kwargs).to(self.device) def make_critic(self) -> ContinuousCritic: return ContinuousCritic(**self.critic_kwargs).to(self.device) @@ -217,7 +226,7 @@ class CnnPolicy(TD3Policy): observation_space: gym.spaces.Space, action_space: gym.spaces.Space, lr_schedule: Callable, - net_arch: Optional[List[int]] = None, + net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None, activation_fn: Type[nn.Module] = nn.ReLU, features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN, features_extractor_kwargs: Optional[Dict[str, Any]] = None, diff --git a/tests/test_custom_policy.py b/tests/test_custom_policy.py index 95f4a7c..6ebeda6 100644 --- a/tests/test_custom_policy.py +++ b/tests/test_custom_policy.py @@ -1,7 +1,7 @@ import pytest import torch as th -from stable_baselines3 import A2C, PPO, SAC, TD3 +from stable_baselines3 import A2C, DQN, PPO, SAC, TD3 from stable_baselines3.common.sb2_compat.rmsprop_tf_like import RMSpropTFLike @@ -19,22 +19,33 @@ from stable_baselines3.common.sb2_compat.rmsprop_tf_like import RMSpropTFLike ) @pytest.mark.parametrize("model_class", [A2C, PPO]) def test_flexible_mlp(model_class, net_arch): - _ = model_class("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000) + _ = model_class("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(300) -@pytest.mark.parametrize("net_arch", [[4], [4, 4]]) +@pytest.mark.parametrize("net_arch", [[], [4], [4, 4], dict(qf=[8], pi=[8, 4])]) @pytest.mark.parametrize("model_class", [SAC, TD3]) def test_custom_offpolicy(model_class, net_arch): - _ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=dict(net_arch=net_arch)).learn(1000) + _ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=dict(net_arch=net_arch), learning_starts=100).learn(300) @pytest.mark.parametrize("model_class", [A2C, PPO, SAC, TD3]) @pytest.mark.parametrize("optimizer_kwargs", [None, dict(weight_decay=0.0)]) def test_custom_optimizer(model_class, optimizer_kwargs): + kwargs = {} + if model_class in {DQN, SAC, TD3}: + kwargs = dict(learning_starts=100) + elif model_class in {A2C, PPO}: + kwargs = dict(n_steps=100) + policy_kwargs = dict(optimizer_class=th.optim.AdamW, optimizer_kwargs=optimizer_kwargs, net_arch=[32]) - _ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs).learn(1000) + _ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs, **kwargs).learn(300) def test_tf_like_rmsprop_optimizer(): policy_kwargs = dict(optimizer_class=RMSpropTFLike, net_arch=[32]) - _ = A2C("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs).learn(1000) + _ = A2C("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs).learn(500) + + +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) diff --git a/tests/test_run.py b/tests/test_run.py index 368381e..5f5f13f 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -80,7 +80,7 @@ def test_n_critics(n_critics): model = SAC( "MlpPolicy", "Pendulum-v0", policy_kwargs=dict(net_arch=[64, 64], n_critics=n_critics), learning_starts=100, verbose=1 ) - model.learn(total_timesteps=1000) + model.learn(total_timesteps=500) def test_dqn(): @@ -88,10 +88,10 @@ def test_dqn(): "MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=[64, 64]), - learning_starts=500, + learning_starts=100, buffer_size=500, learning_rate=3e-4, verbose=1, create_eval_env=True, ) - model.learn(total_timesteps=1000, eval_freq=500) + model.learn(total_timesteps=500, eval_freq=250) diff --git a/tests/test_save_load.py b/tests/test_save_load.py index dde9d2b..18fba37 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -45,7 +45,7 @@ def test_save_load(tmp_path, model_class): # create model model = model_class("MlpPolicy", env, policy_kwargs=dict(net_arch=[16]), verbose=1) - model.learn(total_timesteps=500, eval_freq=250) + model.learn(total_timesteps=500) env.reset() observations = np.concatenate([env.step([env.action_space.sample()])[0] for _ in range(10)], axis=0) @@ -154,7 +154,7 @@ def test_save_load(tmp_path, model_class): assert np.allclose(selected_actions, new_selected_actions, 1e-4) # check if learn still works - model.learn(total_timesteps=1000, eval_freq=500) + model.learn(total_timesteps=500) del model @@ -174,20 +174,26 @@ def test_set_env(model_class): env2 = DummyVecEnv([lambda: select_env(model_class)]) env3 = select_env(model_class) + kwargs = {} + if model_class in {DQN, DDPG, SAC, TD3}: + kwargs = dict(learning_starts=100) + elif model_class in {A2C, PPO}: + kwargs = dict(n_steps=100) + # create model - model = model_class("MlpPolicy", env, policy_kwargs=dict(net_arch=[16])) + model = model_class("MlpPolicy", env, policy_kwargs=dict(net_arch=[16]), **kwargs) # learn - model.learn(total_timesteps=1000, eval_freq=500) + model.learn(total_timesteps=300) # change env model.set_env(env2) # learn again - model.learn(total_timesteps=1000, eval_freq=500) + model.learn(total_timesteps=300) # change env test wrapping model.set_env(env3) # learn again - model.learn(total_timesteps=1000, eval_freq=500) + model.learn(total_timesteps=300) @pytest.mark.parametrize("model_class", MODEL_LIST) @@ -309,7 +315,7 @@ def test_save_load_policy(tmp_path, model_class, policy_str): # create model model = model_class(policy_str, env, policy_kwargs=dict(net_arch=[16]), verbose=1, **kwargs) - model.learn(total_timesteps=500, eval_freq=250) + model.learn(total_timesteps=500) env.reset() observations = np.concatenate([env.step([env.action_space.sample()])[0] for _ in range(10)], axis=0) diff --git a/tests/test_sde.py b/tests/test_sde.py index d26e06c..3b010d3 100644 --- a/tests/test_sde.py +++ b/tests/test_sde.py @@ -68,3 +68,6 @@ def test_state_dependent_offpolicy_noise(model_class, sde_net_arch, use_expln): policy_kwargs=dict(log_std_init=-2, sde_net_arch=sde_net_arch, use_expln=use_expln), ) model.learn(total_timesteps=int(500), eval_freq=250) + model.policy.reset_noise() + if model_class == SAC: + model.policy.actor.get_std() diff --git a/tests/test_utils.py b/tests/test_utils.py index 4e1899f..9450651 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -12,7 +12,7 @@ from stable_baselines3.common.cmd_util import make_atari_env, make_vec_env from stable_baselines3.common.evaluation import evaluate_policy from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.noise import ActionNoise, OrnsteinUhlenbeckActionNoise, VectorizedActionNoise -from stable_baselines3.common.utils import polyak_update +from stable_baselines3.common.utils import polyak_update, zip_strict from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv @@ -167,3 +167,21 @@ def test_polyak(): assert th.allclose(param1, target1) assert th.allclose(param2, target2) + + +def test_zip_strict(): + # Iterables with different lengths + list_a = [0, 1] + list_b = [1, 2, 3] + # zip does not raise any error + for _, _ in zip(list_a, list_b): + pass + + # zip_strict does raise an error + with pytest.raises(ValueError): + for _, _ in zip_strict(list_a, list_b): + pass + + # same length, should not raise an error + for _, _ in zip_strict(list_a, list_b[: len(list_a)]): + pass