saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-30 20:18:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improve typing coverage (#175)

Browse source 
* Improve typing coverage

* Even more types

* Fixes

* Update changelog

* Unified docstrings

* Improve error messages for unsupported spaces
This commit is contained in:
Antonin RAFFIN 2020-10-07 10:51:49 +02:00 committed by GitHub
parent a10e3ae587
commit a1e055695c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
25 changed files with 206 additions and 138 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
docs/misc/changelog.rst
View file

@ -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)
------------------------------

4
stable_baselines3/common/atari_wrappers.py
View file

@ -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)

4
stable_baselines3/common/base_class.py
View file

@ -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).

8
stable_baselines3/common/buffers.py
View file

@ -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:

11
stable_baselines3/common/callbacks.py
View file

@ -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

6
stable_baselines3/common/cmd_util.py
View file

@ -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.

4
stable_baselines3/common/distributions.py
View file

@ -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

2
stable_baselines3/common/logger.py
View file

@ -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

10
stable_baselines3/common/monitor.py
View file

@ -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
"""

2
stable_baselines3/common/noise.py
View file

@ -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):

2
stable_baselines3/common/policies.py
View file

@ -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.

6
stable_baselines3/common/preprocessing.py
View file

@ -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")

17
stable_baselines3/common/save_util.py
View file

@ -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

44
stable_baselines3/common/sb2_compat/rmsprop_tf_like.py
View file

@ -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:

2
stable_baselines3/common/utils.py
View file

@ -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.

26
stable_baselines3/common/vec_env/base_vec_env.py
View file

@ -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.

28
stable_baselines3/common/vec_env/dummy_vec_env.py
View file

@ -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]

46
stable_baselines3/common/vec_env/subproc_vec_env.py
View file

@ -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<X> or tuple<X> where X is dict<ndarray>, tuple<ndarray> 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<ndarray>, tuple<ndarray> 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.
"""

8
stable_baselines3/common/vec_env/util.py
View file

@ -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.

18
stable_baselines3/common/vec_env/vec_check_nan.py
View file

@ -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)

37
stable_baselines3/common/vec_env/vec_normalize.py
View file

@ -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
"""

28
stable_baselines3/common/vec_env/vec_video_recorder.py
View file

@ -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()

2
stable_baselines3/dqn/dqn.py
View file

@ -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.

2
stable_baselines3/td3/policies.py
View file

@ -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:

2
stable_baselines3/version.txt
View file

@ -1 +1 @@
0.9.0
0.10.0a0