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