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

Merge pull request #44 from Antonin-Raffin/typing

Browse source 
Add typing and update requirement to python 3.6
This commit is contained in:
Raffin, Antonin 2020-01-27 11:57:48 +01:00 committed by GitHub Enterprise
commit cc3b023533
29 changed files with 251 additions and 184 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -12,6 +12,7 @@ __pycache__/
_build/
*.npz
*.pth
.pytype/
# Setuptools distribution and build folders.
/dist/

13
Makefile Normal file
View file

@ -0,0 +1,13 @@
SHELL=/bin/bash
pytest:
./scripts/run_tests.sh
type:
pytype
docs:
cd docs && make html
spelling:
cd docs && make spelling

30
README.md
View file

@ -6,6 +6,8 @@
PyTorch version of [Stable Baselines](https://github.com/hill-a/stable-baselines), a set of improved implementations of reinforcement learning algorithms.
NOTE: Python 3.6 is required!
## Implemented Algorithms
- A2C
@ -22,6 +24,34 @@ PyTorch version of [Stable Baselines](https://github.com/hill-a/stable-baselines
- cf github Roadmap
## Run the Tests
```
pip install -e .[tests]
make pytest
```
## Type check
```
pip install -e .[tests]
make type
```
## Build the Documentation
```
pip install -e .[docs]
make docs
```
Spell check for the documentation:
```
make spelling
```
## Citing the Project
To cite this repository in publications:

1
docs/conf.py
View file

@ -70,6 +70,7 @@ release = torchy_baselines.__version__
# ones.
extensions = [
'sphinx.ext.autodoc',
'sphinx_autodoc_typehints',
'sphinx.ext.autosummary',
'sphinx.ext.mathjax',
'sphinx.ext.ifconfig',

5
docs/misc/changelog.rst
View file

@ -8,9 +8,11 @@ Pre-Release 0.2.0a0 (WIP)
Breaking Changes:
^^^^^^^^^^^^^^^^^
- Python 2 support was dropped, Torchy Baselines now requires Python 3.6 or above
New Features:
^^^^^^^^^^^^^
- Add `seed()` method to `VecEnv` class
Bug Fixes:
^^^^^^^^^^
@ -18,9 +20,10 @@ Bug Fixes:
Deprecations:
^^^^^^^^^^^^^
Others:
^^^^^^^
- Add type check
- Converted all format string to f-strings
Documentation:
^^^^^^^^^^^^^^

1
docs/modules/ppo.rst
View file

@ -24,6 +24,7 @@ Notes
- Original paper: https://arxiv.org/abs/1707.06347
- Clear explanation of PPO on Arxiv Insights channel: https://www.youtube.com/watch?v=5P7I-xPq8u8
- OpenAI blog post: https://blog.openai.com/openai-baselines-ppo/
- Spinning Up guide: https://spinningup.openai.com/en/latest/algorithms/ppo.html
Can I use?

1
docs/spelling_wordlist.txt
View file

@ -113,3 +113,4 @@ optimizers
Deprecations
forkserver
cuda
Polyak

3
setup.cfg
View file

@ -15,3 +15,6 @@ filterwarnings =
# Gym warnings
ignore:Parameters to load are deprecated.:DeprecationWarning
ignore:the imp module is deprecated in favour of importlib:PendingDeprecationWarning
[pytype]
inputs = torchy_baselines

5
setup.py
View file

@ -18,13 +18,16 @@ setup(name='torchy_baselines',
'pytest-cov',
'pytest-env',
'pytest-xdist',
'pytype',
],
'docs': [
'sphinx',
'sphinx-autobuild',
'sphinx-rtd-theme',
# For spelling
'sphinxcontrib.spelling'
'sphinxcontrib.spelling',
# Type hints support
'sphinx-autodoc-typehints'
],
'extra': [
# For render

3
torchy_baselines/cem_rl/cem_rl.py
View file

@ -33,7 +33,7 @@ class CEMRL(TD3):
:param learning_starts: (int) how many steps of the model to collect transitions for before learning starts
:param gamma: (float) the discount factor
:param batch_size: (int) Minibatch size for each gradient update
:param tau: (float) the soft update coefficient ("polyak update" of the target networks, between 0 and 1)
:param tau: (float) the soft update coefficient ("Polyak update" of the target networks, between 0 and 1)
:param action_noise: (ActionNoise) the action noise type. Cf common.noise for the different action noise type.
:param target_policy_noise: (float) Standard deviation of Gaussian noise added to target policy
(smoothing noise)
@ -103,6 +103,7 @@ class CEMRL(TD3):
eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="CEMRL", reset_num_timesteps=True):
timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env)
actor_steps = 0
while self.num_timesteps < total_timesteps:

229
torchy_baselines/common/base_class.py
View file

@ -1,54 +1,68 @@
import time
from abc import ABCMeta, abstractmethod
from collections import deque
import os
import io
import zipfile
import typing
from typing import Union, Type, Optional, Dict, Any, List, Tuple
from abc import ABC, abstractmethod
from collections import deque
import gym
import torch as th
import numpy as np
from torchy_baselines.common import logger
from torchy_baselines.common.policies import get_policy_from_name
from torchy_baselines.common.policies import BasePolicy, get_policy_from_name
from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate
from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, sync_envs_normalization
from torchy_baselines.common.monitor import Monitor
from torchy_baselines.common.evaluation import evaluate_policy
from torchy_baselines.common.save_util import data_to_json, json_to_data
# TODO: define aliases, ex GymEnv = Union[gym.Env, VecEnv]
if typing.TYPE_CHECKING:
from torchy_baselines.common.noise import ActionNoise
class BaseRLModel(object):
class BaseRLModel(ABC):
"""
The base RL model
:param policy: (BasePolicy) Policy object
:param env: (Gym environment) The environment to learn from
:param policy: Policy object
:param env: The environment to learn from
(if registered in Gym, can be str. Can be None for loading trained models)
:param policy_base: (BasePolicy) the base policy used by this method
:param policy_kwargs: (dict) additional arguments to be passed to the policy on creation
:param verbose: (int) the verbosity level: 0 none, 1 training information, 2 debug
:param device: (str or th.device) Device on which the code should run.
:param policy_base: The base policy used by this method
:param policy_kwargs: Additional arguments to be passed to the policy on creation
:param verbose: The verbosity level: 0 none, 1 training information, 2 debug
:param device: Device on which the code should run.
By default, it will try to use a Cuda compatible device and fallback to cpu
if it is not possible.
:param support_multi_env: (bool) Whether the algorithm supports training
:param support_multi_env: Whether the algorithm supports training
with multiple environments (as in A2C)
:param create_eval_env: (bool) Whether to create a second environment that will be
:param create_eval_env: Whether to create a second environment that will be
used for evaluating the agent periodically. (Only available when passing string for the environment)
:param monitor_wrapper: (bool) When creating an environment, whether to wrap it
:param monitor_wrapper: When creating an environment, whether to wrap it
or not in a Monitor wrapper.
:param seed: (int) Seed for the pseudo random generators
:param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
:param seed: Seed for the pseudo random generators
:param use_sde: Whether to use State Dependent Exploration (SDE)
instead of action noise exploration (default: False)
:param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
:param sde_sample_freq: Sample a new noise matrix every n steps when using SDE
Default: -1 (only sample at the beginning of the rollout)
"""
__metaclass__ = ABCMeta
def __init__(self,
policy: Type[BasePolicy],
env: Union[gym.Env, VecEnv, str],
policy_base: Type[BasePolicy],
policy_kwargs : Dict[str, Any] = None,
verbose: int = 0,
device: Union[th.device, str] = 'auto',
support_multi_env: bool = False,
create_eval_env: bool = False,
monitor_wrapper: bool = True,
seed: Optional[int] = None,
use_sde: bool = False,
sde_sample_freq: int = -1):
def __init__(self, policy, env, policy_base, policy_kwargs=None,
verbose=0, device='auto', support_multi_env=False,
create_eval_env=False, monitor_wrapper=True, seed=None,
use_sde=False, sde_sample_freq=-1):
if isinstance(policy, str) and policy_base is not None:
self.policy_class = get_policy_from_name(policy_base, policy)
else:
@ -59,9 +73,9 @@ class BaseRLModel(object):
self.device = th.device(device)
if verbose > 0:
print("Using {} device".format(self.device))
print(f"Using {self.device} device")
self.env = env
self.env = None # type: Union[gym.Env, VecEnv]
# get VecNormalize object if needed
self._vec_normalize_env = unwrap_vec_normalize(env)
self.verbose = verbose
@ -73,7 +87,10 @@ class BaseRLModel(object):
self.eval_env = None
self.replay_buffer = None
self.seed = seed
self.action_noise = None
self.action_noise = None # type: ActionNoise
self.start_time = None
self.policy, self.actor = None, None
self.learning_rate = None
# Used for SDE only
self.rollout_data = None
self.on_policy_exploration = False
@ -112,12 +129,12 @@ class BaseRLModel(object):
raise ValueError("Error: the model does not support multiple envs requires a single vectorized"
" environment.")
def _get_eval_env(self, eval_env):
def _get_eval_env(self, eval_env: Union[gym.Env, VecEnv, None]) -> Union[gym.Env, VecEnv, None]:
"""
Return the environment that will be used for evaluation.
:param eval_env: (gym.Env or VecEnv)
:return: (VecEnv)
:param eval_env:
:return:
"""
if eval_env is None:
eval_env = self.eval_env
@ -128,48 +145,47 @@ class BaseRLModel(object):
assert eval_env.num_envs == 1
return eval_env
def scale_action(self, action):
def scale_action(self, action: np.ndarray) -> np.ndarray:
"""
Rescale the action from [low, high] to [-1, 1]
(no need for symmetric action space)
:param action: (np.ndarray)
:return: (np.ndarray)
:param action:
:return:
"""
low, high = self.action_space.low, self.action_space.high
return 2.0 * ((action - low) / (high - low)) - 1.0
def unscale_action(self, scaled_action):
def unscale_action(self, scaled_action: np.ndarray) -> np.ndarray:
"""
Rescale the action from [-1, 1] to [low, high]
(no need for symmetric action space)
:param scaled_action: (np.ndarray)
:return: (np.ndarray)
:param scaled_action:
:return:
"""
low, high = self.action_space.low, self.action_space.high
return low + (0.5 * (scaled_action + 1.0) * (high - low))
def _setup_learning_rate(self):
def _setup_learning_rate(self) -> None:
"""Transform to callable if needed."""
self.learning_rate = get_schedule_fn(self.learning_rate)
def _update_current_progress(self, num_timesteps, total_timesteps):
def _update_current_progress(self, num_timesteps: int, total_timesteps: int) -> None:
"""
Compute current progress (from 1 to 0)
:param num_timesteps: (int) current number of timesteps
:param total_timesteps: (int)
:param num_timesteps: current number of timesteps
:param total_timesteps:
"""
self._current_progress = 1.0 - float(num_timesteps) / float(total_timesteps)
def _update_learning_rate(self, optimizers):
def _update_learning_rate(self, optimizers: Union[List[th.optim.Optimizer], th.optim.Optimizer]) -> None:
"""
Update the optimizers learning rate using the current learning rate schedule
and the current progress (from 1 to 0).
:param optimizers: ([th.optim.Optimizer] or Optimizer) An optimizer
or a list of optimizer.
:param optimizers: An optimizer or a list of optimizer.
"""
# Log the current learning rate
logger.logkv("learning_rate", self.learning_rate(self._current_progress))
@ -180,32 +196,32 @@ class BaseRLModel(object):
update_learning_rate(optimizer, self.learning_rate(self._current_progress))
@staticmethod
def safe_mean(arr):
def safe_mean(arr: Union[np.ndarray, list]) -> np.ndarray:
"""
Compute the mean of an array if there is at least one element.
For empty array, return NaN. It is used for logging only.
:param arr: (np.ndarray)
:return: (float)
:param arr:
:return:
"""
return np.nan if len(arr) == 0 else np.mean(arr)
def get_env(self):
def get_env(self) -> Union[VecEnv, None]:
"""
Returns the current environment (can be None if not defined).
:return: (gym.Env) The current environment
:return: The current environment
"""
return self.env
@staticmethod
def check_env(env, observation_space, action_space):
def check_env(env, observation_space: gym.spaces.Space, action_space: gym.spaces.Space) -> bool:
"""
Checks the validity of the environment and returns if it is consistent.
Checked parameters:
- observation_space
- action_space
:return: (bool) True if environment seems to be coherent
:return: True if environment seems to be coherent
"""
if observation_space != env.observation_space:
return False
@ -214,7 +230,7 @@ class BaseRLModel(object):
# return true if no check failed
return True
def set_env(self, env):
def set_env(self, env: Union[gym.Env, VecEnv]) -> None:
"""
Checks the validity of the environment, and if it is coherent, set it as the current environment.
Furthermore wrap any non vectorized env into a vectorized
@ -222,7 +238,7 @@ class BaseRLModel(object):
- observation_space
- action_space
:param env: (gym.Env) The environment for learning a policy
:param env: The environment for learning a policy
"""
if self.check_env(env, self.observation_space, self.action_space) is False:
raise ValueError("The given environment is not compatible with model: "
@ -236,23 +252,24 @@ class BaseRLModel(object):
self.n_envs = env.num_envs
self.env = env
def get_parameters(self):
def get_parameters(self) -> Tuple[Dict[str, Any], Dict[str, Any]]:
"""
Returns policy and optimizer parameters as a tuple
:return: (dict,dict) policy_parameters, opt_parameters
:return: policy_parameters, opt_parameters
"""
return self.get_policy_parameters(), self.get_opt_parameters()
def get_policy_parameters(self):
def get_policy_parameters(self) -> Dict[str, Any]:
"""
Get current model policy parameters as dictionary of variable name -> tensors.
:return: (dict) Dictionary of variable name -> tensor of model's policy parameters.
:return: Dictionary of variable name -> tensor of model's policy parameters.
"""
return self.policy.state_dict()
@abstractmethod
def get_opt_parameters(self):
def get_opt_parameters(self)-> Dict[str, Any]:
"""
Get current model optimizer parameters as dictionary of variable names -> tensors
:return: (dict) Dictionary of variable name -> tensor of model's optimizer parameters
@ -260,8 +277,13 @@ class BaseRLModel(object):
raise NotImplementedError()
@abstractmethod
def learn(self, total_timesteps, callback=None, log_interval=100, tb_log_name="run",
eval_env=None, eval_freq=-1, n_eval_episodes=5, reset_num_timesteps=True):
def learn(self, total_timesteps: int,
callback=None, log_interval: int = 100,
tb_log_name: str = "run",
eval_env: Union[gym.Env, VecEnv, None] = None,
eval_freq: int = -1,
n_eval_episodes: int = 5,
reset_num_timesteps: bool = True):
"""
Return a trained model.
@ -279,19 +301,22 @@ class BaseRLModel(object):
raise NotImplementedError()
@abstractmethod
def predict(self, observation, state=None, mask=None, deterministic=False):
def predict(self, observation: np.ndarray,
state: Optional[np.ndarray] = None,
mask: Optional[np.ndarray] = None,
deterministic: bool = False) -> np.ndarray:
"""
Get the model's action from an observation
:param observation: (np.ndarray) the input observation
:param state: (np.ndarray) The last states (can be None, used in recurrent policies)
:param mask: (np.ndarray) The last masks (can be None, used in recurrent policies)
:param deterministic: (bool) Whether or not to return deterministic actions.
:return: (np.ndarray, np.ndarray) the model's action and the next state (used in recurrent policies)
:param observation: the input observation
:param state: The last states (can be None, used in recurrent policies)
:param mask: The last masks (can be None, used in recurrent policies)
:param deterministic: Whether or not to return deterministic actions.
:return: the model's action and the next state (used in recurrent policies)
"""
raise NotImplementedError()
def load_parameters(self, load_dict, opt_params):
def load_parameters(self, load_dict: Dict[str, Any], opt_params: Dict[str, Any]) -> None:
"""
Load model parameters from a dictionary
load_dict should contain all keys from torch.model.state_dict()
@ -299,29 +324,28 @@ class BaseRLModel(object):
but can only be handled in child classes.
:param load_dict: (dict) dict of parameters from model.state_dict()
:param opt_params: (dict of dicts) dict of optimizer state_dicts should be handled in child_class
:param load_dict: dict of parameters from model.state_dict()
:param opt_params: dict of optimizer state_dicts should be handled in child_class
"""
if opt_params is not None:
raise ValueError("Optimizer Parameters where given but no overloaded load function exists for this class")
self.policy.load_state_dict(load_dict)
@classmethod
def load(cls, load_path, env=None, **kwargs):
def load(cls, load_path: str, env: Union[gym.Env, VecEnv, None] = None, **kwargs):
"""
Load the model from a zip-file
:param load_path: (str) the location of the saved data
:param env: (Gym Environment) the new environment to run the loaded model on
:param load_path: the location of the saved data
:param env: the new environment to run the loaded model on
(can be None if you only need prediction from a trained model) has priority over any saved environment
:param kwargs: extra arguments to change the model when loading
"""
data, params, opt_params = cls._load_from_file(load_path)
if 'policy_kwargs' in kwargs and kwargs['policy_kwargs'] != data['policy_kwargs']:
raise ValueError("The specified policy kwargs do not equal the stored policy kwargs."
"Stored kwargs: {}, specified kwargs: {}".format(data['policy_kwargs'],
kwargs['policy_kwargs']))
raise ValueError(f"The specified policy kwargs do not equal the stored policy kwargs."
"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}")
# check if observation space and action space are part of the saved parameters
if ("observation_space" not in data or "action_space" not in data) and "env" not in data:
@ -344,11 +368,11 @@ class BaseRLModel(object):
return model
@staticmethod
def _load_from_file(load_path, load_data=True):
def _load_from_file(load_path: str, load_data: bool = True):
""" Load model data from a .zip archive
:param load_path: (str) Where to load the model from
:param load_data: (bool) Whether we should load and return data
:param load_path: Where to load the model from
:param load_data: Whether we should load and return data
(class parameters). Mainly used by 'load_parameters' to only load model parameters (weights)
:return: (dict),(dict),(dict) Class parameters, model parameters (state_dict)
and dict of optimizer parameters (dict of state_dict)
@ -359,7 +383,7 @@ class BaseRLModel(object):
if os.path.exists(load_path + ".zip"):
load_path += ".zip"
else:
raise ValueError("Error: the file {} could not be found".format(load_path))
raise ValueError(f"Error: the file {load_path} could not be found")
# Open the zip archive and load data
try:
@ -406,11 +430,11 @@ class BaseRLModel(object):
opt_params[os.path.splitext(file_path)[0]] = th.load(file_content)
except zipfile.BadZipFile:
# load_path wasn't a zip file
raise ValueError("Error: the file {} wasn't a zip-file".format(load_path))
raise ValueError(f"Error: the file {load_path} wasn't a zip-file")
return data, params, opt_params
def set_random_seed(self, seed=None):
def set_random_seed(self, seed: Optional[int] = None) -> None:
"""
Set the seed of the pseudo-random generators
(python, numpy, pytorch, gym, action_space)
@ -446,7 +470,7 @@ class BaseRLModel(object):
eval_env.seed(self.seed)
eval_env = self._get_eval_env(eval_env)
obs = self.env.reset()
obs = self.env.reset() # type: Union[gym.Env, VecEnv]
return timesteps_since_eval, episode_num, evaluations, obs, eval_env
def _update_info_buffer(self, infos):
@ -629,13 +653,15 @@ class BaseRLModel(object):
return mean_reward, total_steps, total_episodes, obs
@staticmethod
def _save_to_file_zip(save_path, data=None, params=None, opt_params=None):
"""Save model to a zip archive
def _save_to_file_zip(save_path: str, data: Dict[str, Any] = None,
params: Dict[str, Any] = None, opt_params: Dict[str, Any] = None) -> None:
"""
Save model to a zip archive.
:param save_path: (str) Where to store the model
:param data: (dict) Class parameters being stored
:param params: (dict) Model parameters being stored expected to be state_dict
:param opt_params: (dict) Optimizer parameters being stored expected to contain an entry for every
:param save_path: Where to store the model
:param data: Class parameters being stored
:param params: Model parameters being stored expected to be state_dict
:param opt_params: Optimizer parameters being stored expected to contain an entry for every
optimizer with its name and the state_dict
"""
@ -666,7 +692,7 @@ class BaseRLModel(object):
th.save(dict_, opt_param_file)
@staticmethod
def excluded_save_params():
def excluded_save_params() -> List[str]:
"""
Returns the names of the parameters that should be excluded by default
when saving the model.
@ -675,13 +701,13 @@ class BaseRLModel(object):
"""
return ["env", "eval_env", "replay_buffer", "rollout_buffer", "_vec_normalize_env"]
def save(self, path, exclude=None, include=None):
def save(self, path: str, exclude: Optional[List[str]] = None, include: Optional[List[str]] = None) -> None:
"""
Save all the attributes of the object and the model parameters in a zip-file.
:param path: (str) path to the file where the rl agent should be saved
:param exclude: ([str]) name of parameters that should be excluded in addition to the default one
:param include: ([str]) name of parameters that might be excluded but should be included anyway
:param path: path to the file where the rl agent should be saved
:param exclude: name of parameters that should be excluded in addition to the default one
:param include: name of parameters that might be excluded but should be included anyway
"""
# copy parameter list so we don't mutate the original dict
data = self.__dict__.copy()
@ -704,25 +730,26 @@ class BaseRLModel(object):
opt_params_to_save = self.get_opt_parameters()
self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save)
def _eval_policy(self, eval_freq, eval_env, n_eval_episodes,
timesteps_since_eval, deterministic=True):
def _eval_policy(self, eval_freq: int, eval_env: int, n_eval_episodes: int,
timesteps_since_eval: int, render: bool = False, deterministic: bool = True) -> int:
"""
Evaluate the current policy on a test environment.
:param eval_env: (gym.Env) Environment that will be used to evaluate the agent
:param eval_freq: (int) Evaluate the agent every `eval_freq` timesteps (this may vary a little)
:param n_eval_episodes: (int) Number of episode to evaluate the agent
:parma timesteps_since_eval: (int) Number of timesteps since last evaluation
:param deterministic: (bool) Whether to use deterministic or stochastic actions
:return: (int) Number of timesteps since last evaluation
:param eval_freq: Evaluate the agent every `eval_freq` timesteps (this may vary a little)
:param n_eval_episodes: Number of episode to evaluate the agent
:parma timesteps_since_eval: Number of timesteps since last evaluation
:param deterministic: Whether to use deterministic or stochastic actions
:param render: Whether to render the eval env or not
:return: Number of timesteps since last evaluation
"""
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
timesteps_since_eval %= eval_freq
# Synchronise the normalization stats if needed
sync_envs_normalization(self.env, eval_env)
mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes, deterministic=deterministic)
mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes,
render=render, deterministic=deterministic)
if self.verbose > 0:
print("Eval num_timesteps={}, "
"episode_reward={:.2f} +/- {:.2f}".format(self.num_timesteps, mean_reward, std_reward))
print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
print(f"Eval num_timesteps={self.num_timesteps}, "
f"episode_reward={mean_reward:.2f} +/- {std_reward:.2f}")
print(f"FPS: {self.num_timesteps / (time.time() - self.start_time):.2f}")
return timesteps_since_eval

3
torchy_baselines/common/distributions.py
View file

@ -467,6 +467,5 @@ def make_proba_distribution(action_space, use_sde=False, dist_kwargs=None):
# elif isinstance(action_space, spaces.MultiBinary):
# return BernoulliDistribution(action_space.n, **dist_kwargs)
else:
raise NotImplementedError("Error: probability distribution, not implemented for action space of type {}."
.format(type(action_space)) +
raise NotImplementedError(f"Error: probability distribution, not implemented for action space of type {type(action_space)}."
" Must be of type Gym Spaces: Box, Discrete, MultiDiscrete or MultiBinary.")

4
torchy_baselines/common/evaluation.py
View file

@ -47,8 +47,8 @@ def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True,
mean_reward = np.mean(episode_rewards)
std_reward = np.std(episode_rewards)
if reward_threshold is not None:
assert mean_reward > reward_threshold, 'Mean reward below threshold: '\
'{:.2f} < {:.2f}'.format(mean_reward, reward_threshold)
assert mean_reward > reward_threshold, (f'Mean reward below threshold: '
'{mean_reward:.2f} < {reward_threshold:.2f}')
if return_episode_rewards:
return episode_rewards, n_steps
return mean_reward, std_reward

8
torchy_baselines/common/logger.py
View file

@ -53,7 +53,7 @@ class HumanOutputFormat(KVWriter, SeqWriter):
self.file = open(filename_or_file, 'wt')
self.own_file = True
else:
assert hasattr(filename_or_file, 'write'), 'Expected file or str, got {}'.format(filename_or_file)
assert hasattr(filename_or_file, 'write'), f'Expected file or str, got {filename_or_file}'
self.file = filename_or_file
self.own_file = False
@ -272,7 +272,7 @@ def getkvs():
return Logger.CURRENT.name2val
def log(*args, **kwargs):
def log(*args, level=INFO):
"""
Write the sequence of args, with no separators,
to the console and output files (if you've configured an output file).
@ -283,7 +283,6 @@ def log(*args, **kwargs):
:param args: (list) log the arguments
:param level: (int) the logging level (can be DEBUG=10, INFO=20, WARN=30, ERROR=40, DISABLED=50)
"""
level = kwargs.get('level', INFO)
Logger.CURRENT.log(*args, level=level)
@ -424,7 +423,7 @@ class Logger(object):
self.name2val.clear()
self.name2cnt.clear()
def log(self, *args, **kwargs):
def log(self, *args, level=INFO):
"""
Write the sequence of args, with no separators,
to the console and output files (if you've configured an output file).
@ -435,7 +434,6 @@ class Logger(object):
:param args: (list) log the arguments
:param level: (int) the logging level (can be DEBUG=10, INFO=20, WARN=30, ERROR=40, DISABLED=50)
"""
level = kwargs.get('level', INFO)
if self.level <= level:
self._do_log(args)

4
torchy_baselines/common/noise.py
View file

@ -30,7 +30,7 @@ class NormalActionNoise(ActionNoise):
return np.random.normal(self._mu, self._sigma)
def __repr__(self):
return 'NormalActionNoise(mu={}, sigma={})'.format(self._mu, self._sigma)
return f'NormalActionNoise(mu={self._mu}, sigma={self._sigma})'
class OrnsteinUhlenbeckActionNoise(ActionNoise):
@ -68,4 +68,4 @@ class OrnsteinUhlenbeckActionNoise(ActionNoise):
self.noise_prev = self.initial_noise if self.initial_noise is not None else np.zeros_like(self._mu)
def __repr__(self):
return 'OrnsteinUhlenbeckActionNoise(mu={}, sigma={})'.format(self._mu, self._sigma)
return f'OrnsteinUhlenbeckActionNoise(mu={self._mu}, sigma={self._sigma})'

11
torchy_baselines/common/policies.py
View file

@ -150,10 +150,10 @@ def get_policy_from_name(base_policy_type, name):
:return: (base_policy_type) the policy
"""
if base_policy_type not in _policy_registry:
raise ValueError("Error: the policy type {} is not registered!".format(base_policy_type))
raise ValueError(f"Error: the policy type {base_policy_type} is not registered!")
if name not in _policy_registry[base_policy_type]:
raise ValueError("Error: unknown policy type {}, the only registed policy type are: {}!"
.format(name, list(_policy_registry[base_policy_type].keys())))
raise ValueError(f"Error: unknown policy type {name},"
"the only registed policy type are: {list(_policy_registry[base_policy_type].keys())}!")
return _policy_registry[base_policy_type][name]
@ -174,13 +174,12 @@ def register_policy(name, policy):
except AttributeError:
sub_class = str(th.random.randint(100))
if sub_class is None:
raise ValueError("Error: the policy {} is not of any known subclasses of BasePolicy!".format(policy))
raise ValueError(f"Error: the policy {policy} is not of any known subclasses of BasePolicy!")
if sub_class not in _policy_registry:
_policy_registry[sub_class] = {}
if name in _policy_registry[sub_class]:
raise ValueError("Error: the name {} is alreay registered for a different policy, will not override."
.format(name))
raise ValueError(f"Error: the name {name} is alreay registered for a different policy, will not override.")
_policy_registry[sub_class][name] = policy

2
torchy_baselines/common/save_util.py
View file

@ -123,7 +123,7 @@ def json_to_data(json_string, custom_objects=None):
)
except pickle.UnpicklingError:
raise RuntimeError(
"Could not deserialize object {}. ".format(data_key) +
f"Could not deserialize object {data_key}. " +
"Consider using `custom_objects` argument to replace " +
"this object."
)

35
torchy_baselines/common/vec_env/base_vec_env.py
View file

@ -1,6 +1,6 @@
from abc import ABCMeta, abstractmethod
import inspect
import pickle
from abc import ABC, abstractmethod
import cloudpickle
@ -27,7 +27,7 @@ class NotSteppingError(Exception):
Exception.__init__(self, msg)
class VecEnv(object):
class VecEnv(ABC):
"""
An abstract asynchronous, vectorized environment.
@ -39,8 +39,6 @@ class VecEnv(object):
'render.modes': ['human', 'rgb_array']
}
__metaclass__ = ABCMeta
def __init__(self, num_envs, observation_space, action_space):
self.num_envs = num_envs
self.observation_space = observation_space
@ -112,7 +110,7 @@ class VecEnv(object):
raise NotImplementedError()
@abstractmethod
def env_method(self, method_name, *method_args, **method_kwargs):
def env_method(self, method_name, *method_args, indices=None, **method_kwargs):
"""
Call instance methods of vectorized environments.
@ -148,6 +146,17 @@ class VecEnv(object):
"""
raise NotImplementedError()
def seed(self, seed, indices=None):
"""
:param seed: (int or [int])
:param indices: ([int])
"""
indices = self._get_indices(indices)
if not hasattr(seed, 'len'):
seed = [seed] * len(indices)
assert len(seed) == len(indices)
return [self.env_method('seed', seed[i], indices=i) for i in indices]
@property
def unwrapped(self):
if isinstance(self, VecEnvWrapper):
@ -163,7 +172,7 @@ class VecEnv(object):
:return: (str or None) name of module whose attribute is being shadowed, if any.
"""
if hasattr(self, name) and already_found:
return "{0}.{1}".format(type(self).__module__, type(self).__name__)
return f"{type(self).__module__}.{type(self).__name__}"
else:
return None
@ -222,8 +231,8 @@ class VecEnvWrapper(VecEnv):
def set_attr(self, attr_name, value, indices=None):
return self.venv.set_attr(attr_name, value, indices)
def env_method(self, method_name, *method_args, **method_kwargs):
return self.venv.env_method(method_name, *method_args, **method_kwargs)
def env_method(self, method_name, *method_args, indices=None, **method_kwargs):
return self.venv.env_method(method_name, *method_args, indices=indices, **method_kwargs)
def __getattr__(self, name):
"""Find attribute from wrapped venv(s) if this wrapper does not have it.
@ -232,10 +241,10 @@ class VecEnvWrapper(VecEnv):
"""
blocked_class = self.getattr_depth_check(name, already_found=False)
if blocked_class is not None:
own_class = "{0}.{1}".format(type(self).__module__, type(self).__name__)
format_str = ("Error: Recursive attribute lookup for {0} from {1} is "
"ambiguous and hides attribute from {2}")
raise AttributeError(format_str.format(name, own_class, blocked_class))
own_class = f"{type(self).__module__}.{type(self).__name__}"
error_str = (f"Error: Recursive attribute lookup for {name} from {own_class} is "
"ambiguous and hides attribute from {blocked_class}")
raise AttributeError(error_str)
return self.getattr_recursive(name)
@ -274,7 +283,7 @@ class VecEnvWrapper(VecEnv):
all_attributes = self._get_all_attributes()
if name in all_attributes and already_found:
# this venv's attribute is being hidden because of a higher venv.
shadowed_wrapper_class = "{0}.{1}".format(type(self).__module__, type(self).__name__)
shadowed_wrapper_class = f"{type(self).__module__}.{type(self).__name__}"
elif name in all_attributes and not already_found:
# we have found the first reference to the attribute. Now check for duplicates.
shadowed_wrapper_class = self.venv.getattr_depth_check(name, True)

18
torchy_baselines/common/vec_env/dummy_vec_env.py
View file

@ -3,7 +3,7 @@ from copy import deepcopy
import numpy as np
from torchy_baselines.common.vec_env import VecEnv
from torchy_baselines.common.vec_env.base_vec_env import VecEnv
from torchy_baselines.common.vec_env.util import copy_obs_dict, dict_to_obs, obs_space_info
@ -54,17 +54,6 @@ class DummyVecEnv(VecEnv):
self._save_obs(env_idx, obs)
return self._obs_from_buf()
def seed(self, seed, indices=None):
"""
:param seed: (int or [int])
:param indices: ([int])
"""
indices = self._get_indices(indices)
if not hasattr(seed, 'len'):
seed = [seed] * len(indices)
assert len(seed) == len(indices)
return [self.envs[i].seed(seed[i]) for i in indices]
def close(self):
for env in self.envs:
env.close()
@ -99,11 +88,8 @@ class DummyVecEnv(VecEnv):
for env_i in target_envs:
setattr(env_i, attr_name, value)
def env_method(self, method_name, *method_args, **method_kwargs):
def env_method(self, method_name, *method_args, indices=None, **method_kwargs):
"""Call instance methods of vectorized environments."""
indices = method_kwargs.get('indices')
if 'indices' in method_kwargs:
del method_kwargs['indices']
target_envs = self._get_target_envs(indices)
return [getattr(env_i, method_name)(*method_args, **method_kwargs) for env_i in target_envs]

23
torchy_baselines/common/vec_env/subproc_vec_env.py
View file

@ -4,7 +4,7 @@ from collections import OrderedDict
import gym
import numpy as np
from torchy_baselines.common.vec_env import VecEnv, CloudpickleWrapper
from torchy_baselines.common.vec_env.base_vec_env import VecEnv, CloudpickleWrapper
def _worker(remote, parent_remote, env_fn_wrapper):
@ -111,7 +111,7 @@ class SubprocVecEnv(VecEnv):
for work_remote, remote, env_fn in zip(self.work_remotes, self.remotes, env_fns):
args = (work_remote, remote, CloudpickleWrapper(env_fn))
# daemon=True: if the main process crashes, we should not cause things to hang
process = ctx.Process(target=_worker, args=args, daemon=True)
process = ctx.Process(target=_worker, args=args, daemon=True) # pytype:disable=attribute-error
process.start()
self.processes.append(process)
work_remote.close()
@ -153,8 +153,7 @@ class SubprocVecEnv(VecEnv):
for pipe in self.remotes:
# gather images from subprocesses
# `mode` will be taken into account later
kwargs.update({'mode': 'rgb_array'})
pipe.send(('render', (args, kwargs)))
pipe.send(('render', (args, {'mode': 'rgb_array', **kwargs})))
imgs = [pipe.recv() for pipe in self.remotes]
# Create a big image by tiling images from subprocesses
bigimg = tile_images(imgs)
@ -188,22 +187,8 @@ class SubprocVecEnv(VecEnv):
for remote in target_remotes:
remote.recv()
def seed(self, seed, indices=None):
"""
:param seed: (int or [int])
:param indices: ([int])
"""
indices = self._get_indices(indices)
if not hasattr(seed, 'len'):
seed = [seed] * len(indices)
assert len(seed) == len(indices)
return [self.env_method('seed', seed[i], indices=i) for i in indices]
def env_method(self, method_name, *method_args, **method_kwargs):
def env_method(self, method_name, *method_args, indices=None, **method_kwargs):
"""Call instance methods of vectorized environments."""
indices = method_kwargs.get('indices')
if 'indices' in method_kwargs:
del method_kwargs['indices']
target_remotes = self._get_target_remotes(indices)
for remote in target_remotes:
remote.send(('env_method', (method_name, method_args, method_kwargs)))

4
torchy_baselines/common/vec_env/util.py
View file

@ -15,7 +15,7 @@ def copy_obs_dict(obs):
:param obs: (OrderedDict<ndarray>): a dict of numpy arrays.
:return (OrderedDict<ndarray>) a dict of copied numpy arrays.
"""
assert isinstance(obs, OrderedDict), "unexpected type for observations '{}'".format(type(obs))
assert isinstance(obs, OrderedDict), f"unexpected type for observations '{type(obs)}'"
return OrderedDict([(k, np.copy(v)) for k, v in obs.items()])
@ -61,7 +61,7 @@ def obs_space_info(obs_space):
elif isinstance(obs_space, gym.spaces.Tuple):
subspaces = {i: space for i, space in enumerate(obs_space.spaces)}
else:
assert not hasattr(obs_space, 'spaces'), "Unsupported structured space '{}'".format(type(obs_space))
assert not hasattr(obs_space, 'spaces'), f"Unsupported structured space '{type(obs_space)}'"
subspaces = {None: obs_space}
keys = []
shapes = {}

8
torchy_baselines/common/vec_env/vec_frame_stack.py
View file

@ -3,18 +3,18 @@ import warnings
import numpy as np
from gym import spaces
from torchy_baselines.common.vec_env import VecEnvWrapper
from torchy_baselines.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper
class VecFrameStack(VecEnvWrapper):
"""
Frame stacking wrapper for vectorized environment
:param venv: (VecEnv) the vectorized environment to wrap
:param n_stack: (int) Number of frames to stack
:param venv: the vectorized environment to wrap
:param n_stack: Number of frames to stack
"""
def __init__(self, venv, n_stack):
def __init__(self, venv: VecEnv, n_stack: int):
self.venv = venv
self.n_stack = n_stack
wrapped_obs_space = venv.observation_space

6
torchy_baselines/common/vec_env/vec_normalize.py
View file

@ -2,7 +2,7 @@ import pickle
import numpy as np
from torchy_baselines.common.vec_env import VecEnvWrapper
from torchy_baselines.common.vec_env.base_vec_env import VecEnvWrapper
from torchy_baselines.common.running_mean_std import RunningMeanStd
@ -182,7 +182,7 @@ class VecNormalize(VecEnvWrapper):
:param path: (str) path to log dir
"""
for rms, name in zip([self.obs_rms, self.ret_rms], ['obs_rms', 'ret_rms']):
with open("{}/{}.pkl".format(path, name), 'wb') as file_handler:
with open(f"{path}/{name}.pkl", 'wb') as file_handler:
pickle.dump(rms, file_handler)
def load_running_average(self, path):
@ -190,5 +190,5 @@ class VecNormalize(VecEnvWrapper):
:param path: (str) path to log dir
"""
for name in ['obs_rms', 'ret_rms']:
with open("{}/{}.pkl".format(path, name), 'rb') as file_handler:
with open(f"{path}/{name}.pkl", 'rb') as file_handler:
setattr(self, name, pickle.load(file_handler))

4
torchy_baselines/ppo/policies.py
View file

@ -76,18 +76,20 @@ class PPOPolicy(BasePolicy):
self.sde_feature_extractor = None
self.sde_net_arch = sde_net_arch
self.use_sde = use_sde
# Action distribution
self.action_dist = make_proba_distribution(action_space, use_sde=use_sde, dist_kwargs=dist_kwargs)
self._build(learning_rate)
def reset_noise(self, n_envs=1):
def reset_noise(self, n_envs: int = 1):
"""
Sample new weights for the exploration matrix.
:param n_envs: (int)
"""
assert isinstance(self.action_dist, StateDependentNoiseDistribution), 'reset_noise() is only available when using SDE'
self.action_dist.sample_weights(self.log_std, batch_size=n_envs)
def _build(self, learning_rate):

2
torchy_baselines/ppo/ppo.py
View file

@ -25,7 +25,7 @@ class PPO(BaseRLModel):
Proximal Policy Optimization algorithm (PPO) (clip version)
Paper: https://arxiv.org/abs/1707.06347
Code: This implementation borrows code from OpenAI spinningup (https://github.com/openai/spinningup/)
Code: This implementation borrows code from OpenAI Spinning Up (https://github.com/openai/spinningup/)
https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail and
and Stable Baselines (PPO2 from https://github.com/hill-a/stable-baselines)

0
torchy_baselines/py.typed Normal file
View file

2
torchy_baselines/sac/policies.py
View file

@ -94,6 +94,7 @@ class Actor(BaseNetwork):
:return: (th.Tensor)
"""
assert isinstance(self.action_dist, StateDependentNoiseDistribution), 'get_std() is only available when using SDE'
return self.action_dist.get_std(self.log_std)
def reset_noise(self, batch_size=1):
@ -102,6 +103,7 @@ class Actor(BaseNetwork):
:param batch_size: (int)
"""
assert isinstance(self.action_dist, StateDependentNoiseDistribution), 'reset_noise() is only available when using SDE'
self.action_dist.sample_weights(self.log_std, batch_size=batch_size)
def _get_latent(self, obs):

7
torchy_baselines/sac/sac.py
View file

@ -39,7 +39,7 @@ class SAC(BaseRLModel):
:param gradient_steps: (int) How many gradient update after each step
:param n_episodes_rollout: (int) Update the model every `n_episodes_rollout` episodes.
Note that this cannot be used at the same time as `train_freq`
:param target_entropy: (str or float) target entropy when learning ent_coef (ent_coef = 'auto')
:param target_entropy: (str or float) target entropy when learning `ent_coef` (`ent_coef = 'auto'`)
:param action_noise: (ActionNoise) the action noise type (None by default), this can help
for hard exploration problem. Cf common.noise for the different action noise type.
:param gamma: (float) the discount factor
@ -161,7 +161,7 @@ class SAC(BaseRLModel):
"""
return self.unscale_action(self.select_action(observation))
def train(self, gradient_steps, batch_size=64):
def train(self, gradient_steps: int, batch_size: int = 64):
# Update optimizers learning rate
optimizers = [self.actor.optimizer, self.critic.optimizer]
if self.ent_coef_optimizer is not None:
@ -169,6 +169,9 @@ class SAC(BaseRLModel):
self._update_learning_rate(optimizers)
ent_coef_loss, ent_coef = th.zeros(1), th.zeros(1)
actor_loss, critic_loss = th.zeros(1), th.zeros(1)
for gradient_step in range(gradient_steps):
# Sample replay buffer
replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)

2
torchy_baselines/td3/td3.py
View file

@ -31,7 +31,7 @@ class TD3(BaseRLModel):
:param gradient_steps: (int) How many gradient update after each step
:param n_episodes_rollout: (int) Update the model every `n_episodes_rollout` episodes.
Note that this cannot be used at the same time as `train_freq`
:param tau: (float) the soft update coefficient ("polyak update" of the target networks, between 0 and 1)
:param tau: (float) the soft update coefficient ("Polyak update" of the target networks, between 0 and 1)
:param action_noise: (ActionNoise) the action noise type. Cf common.noise for the different action noise type.
:param target_policy_noise: (float) Standard deviation of Gaussian noise added to target policy
(smoothing noise)