From 37f9f136842fdf9ffa44e4c743c1f4230cb3afec Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Wed, 22 Jan 2020 16:18:27 +0100 Subject: [PATCH 1/8] Revert all changes for python 2 + Add makefile and pytype --- .gitignore | 1 + Makefile | 13 +++++++++++++ setup.cfg | 3 +++ setup.py | 1 + torchy_baselines/common/base_class.py | 8 +++----- torchy_baselines/common/logger.py | 6 ++---- torchy_baselines/common/vec_env/base_vec_env.py | 12 +++++------- torchy_baselines/common/vec_env/dummy_vec_env.py | 5 +---- torchy_baselines/common/vec_env/subproc_vec_env.py | 8 ++------ torchy_baselines/py.typed | 0 10 files changed, 31 insertions(+), 26 deletions(-) create mode 100644 Makefile create mode 100644 torchy_baselines/py.typed diff --git a/.gitignore b/.gitignore index 4e58c84..1c61a6c 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,7 @@ __pycache__/ _build/ *.npz *.pth +.pytype/ # Setuptools distribution and build folders. /dist/ diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..c3636cf --- /dev/null +++ b/Makefile @@ -0,0 +1,13 @@ +SHELL=/bin/bash + +pytest: + ./scripts/run_tests.sh + +pytype: + pytype + +doc: + cd docs && make html + +spelling: + cd docs && make spelling diff --git a/setup.cfg b/setup.cfg index aca686b..708de33 100644 --- a/setup.cfg +++ b/setup.cfg @@ -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 diff --git a/setup.py b/setup.py index 393174f..343967a 100644 --- a/setup.py +++ b/setup.py @@ -18,6 +18,7 @@ setup(name='torchy_baselines', 'pytest-cov', 'pytest-env', 'pytest-xdist', + 'pytype', ], 'docs': [ 'sphinx', diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index bdbef53..a1f7722 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -1,9 +1,9 @@ import time -from abc import ABCMeta, abstractmethod -from collections import deque import os import io import zipfile +from abc import ABC, abstractmethod +from collections import deque import gym import torch as th @@ -18,7 +18,7 @@ from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.common.save_util import data_to_json, json_to_data -class BaseRLModel(object): +class BaseRLModel(ABC): """ The base RL model @@ -43,8 +43,6 @@ class BaseRLModel(object): :param sde_sample_freq: (int) 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, env, policy_base, policy_kwargs=None, verbose=0, device='auto', support_multi_env=False, create_eval_env=False, monitor_wrapper=True, seed=None, diff --git a/torchy_baselines/common/logger.py b/torchy_baselines/common/logger.py index a909285..78845ee 100644 --- a/torchy_baselines/common/logger.py +++ b/torchy_baselines/common/logger.py @@ -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) diff --git a/torchy_baselines/common/vec_env/base_vec_env.py b/torchy_baselines/common/vec_env/base_vec_env.py index abc36a8..47defa9 100644 --- a/torchy_baselines/common/vec_env/base_vec_env.py +++ b/torchy_baselines/common/vec_env/base_vec_env.py @@ -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. @@ -222,8 +220,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. diff --git a/torchy_baselines/common/vec_env/dummy_vec_env.py b/torchy_baselines/common/vec_env/dummy_vec_env.py index 7bfe1a5..c55ca6a 100644 --- a/torchy_baselines/common/vec_env/dummy_vec_env.py +++ b/torchy_baselines/common/vec_env/dummy_vec_env.py @@ -99,11 +99,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] diff --git a/torchy_baselines/common/vec_env/subproc_vec_env.py b/torchy_baselines/common/vec_env/subproc_vec_env.py index f89cbc6..0d1ecde 100644 --- a/torchy_baselines/common/vec_env/subproc_vec_env.py +++ b/torchy_baselines/common/vec_env/subproc_vec_env.py @@ -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) @@ -199,11 +198,8 @@ class SubprocVecEnv(VecEnv): 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))) diff --git a/torchy_baselines/py.typed b/torchy_baselines/py.typed new file mode 100644 index 0000000..e69de29 From 88f07bafb6cf853c0e1ad2a46da677b209303361 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Wed, 22 Jan 2020 16:39:25 +0100 Subject: [PATCH 2/8] Convert format to f-strings --- torchy_baselines/common/base_class.py | 17 ++++++++--------- torchy_baselines/common/distributions.py | 3 +-- torchy_baselines/common/evaluation.py | 4 ++-- torchy_baselines/common/logger.py | 2 +- torchy_baselines/common/noise.py | 4 ++-- torchy_baselines/common/policies.py | 11 +++++------ torchy_baselines/common/save_util.py | 2 +- torchy_baselines/common/vec_env/base_vec_env.py | 12 ++++++------ torchy_baselines/common/vec_env/util.py | 4 ++-- .../common/vec_env/vec_normalize.py | 4 ++-- 10 files changed, 30 insertions(+), 33 deletions(-) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index a1f7722..5c780c9 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -57,7 +57,7 @@ class BaseRLModel(ABC): self.device = th.device(device) if verbose > 0: - print("Using {} device".format(self.device)) + print(f"Using {self.device} device") self.env = env # get VecNormalize object if needed @@ -317,9 +317,8 @@ class BaseRLModel(ABC): 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: @@ -357,7 +356,7 @@ class BaseRLModel(ABC): 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: @@ -404,7 +403,7 @@ class BaseRLModel(ABC): 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 @@ -720,7 +719,7 @@ class BaseRLModel(ABC): sync_envs_normalization(self.env, eval_env) mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes, 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}, " + "episode_reward={mean_reward:.2f} +/- {std_reward:.2f}") + print(f"FPS: {self.num_timesteps / (time.time() - self.start_time):.2f}") return timesteps_since_eval diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index e4ca9bc..6ede43f 100644 --- a/torchy_baselines/common/distributions.py +++ b/torchy_baselines/common/distributions.py @@ -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.") diff --git a/torchy_baselines/common/evaluation.py b/torchy_baselines/common/evaluation.py index 04ae296..0133621 100644 --- a/torchy_baselines/common/evaluation.py +++ b/torchy_baselines/common/evaluation.py @@ -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 diff --git a/torchy_baselines/common/logger.py b/torchy_baselines/common/logger.py index 78845ee..0da70e6 100644 --- a/torchy_baselines/common/logger.py +++ b/torchy_baselines/common/logger.py @@ -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 diff --git a/torchy_baselines/common/noise.py b/torchy_baselines/common/noise.py index 09e07fb..dd7d8a3 100644 --- a/torchy_baselines/common/noise.py +++ b/torchy_baselines/common/noise.py @@ -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})' diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index f536624..ea7f755 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -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 diff --git a/torchy_baselines/common/save_util.py b/torchy_baselines/common/save_util.py index 7b6a2be..1b5d801 100644 --- a/torchy_baselines/common/save_util.py +++ b/torchy_baselines/common/save_util.py @@ -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." ) diff --git a/torchy_baselines/common/vec_env/base_vec_env.py b/torchy_baselines/common/vec_env/base_vec_env.py index 47defa9..d239ddd 100644 --- a/torchy_baselines/common/vec_env/base_vec_env.py +++ b/torchy_baselines/common/vec_env/base_vec_env.py @@ -161,7 +161,7 @@ class VecEnv(ABC): :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 @@ -230,10 +230,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) @@ -272,7 +272,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) diff --git a/torchy_baselines/common/vec_env/util.py b/torchy_baselines/common/vec_env/util.py index 03ce286..0e8629a 100644 --- a/torchy_baselines/common/vec_env/util.py +++ b/torchy_baselines/common/vec_env/util.py @@ -15,7 +15,7 @@ def copy_obs_dict(obs): :param obs: (OrderedDict): a dict of numpy arrays. :return (OrderedDict) 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 = {} diff --git a/torchy_baselines/common/vec_env/vec_normalize.py b/torchy_baselines/common/vec_env/vec_normalize.py index 0824f6a..4f5cb11 100644 --- a/torchy_baselines/common/vec_env/vec_normalize.py +++ b/torchy_baselines/common/vec_env/vec_normalize.py @@ -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)) From 44fce7c02a08266f147d1800ad1b4167093cf52d Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Wed, 22 Jan 2020 17:17:12 +0100 Subject: [PATCH 3/8] Fix typing errors and typos --- Makefile | 4 ++-- docs/modules/ppo.rst | 1 + docs/spelling_wordlist.txt | 1 + torchy_baselines/cem_rl/cem_rl.py | 3 ++- torchy_baselines/common/base_class.py | 20 +++++++++++++------ .../common/vec_env/base_vec_env.py | 11 ++++++++++ .../common/vec_env/dummy_vec_env.py | 13 +----------- .../common/vec_env/subproc_vec_env.py | 15 ++------------ .../common/vec_env/vec_frame_stack.py | 2 +- .../common/vec_env/vec_normalize.py | 2 +- torchy_baselines/ppo/policies.py | 4 +++- torchy_baselines/ppo/ppo.py | 2 +- torchy_baselines/sac/policies.py | 2 ++ torchy_baselines/sac/sac.py | 7 +++++-- torchy_baselines/td3/td3.py | 2 +- 15 files changed, 48 insertions(+), 41 deletions(-) diff --git a/Makefile b/Makefile index c3636cf..78b043f 100644 --- a/Makefile +++ b/Makefile @@ -3,10 +3,10 @@ SHELL=/bin/bash pytest: ./scripts/run_tests.sh -pytype: +type: pytype -doc: +docs: cd docs && make html spelling: diff --git a/docs/modules/ppo.rst b/docs/modules/ppo.rst index 5531d3a..9c9887d 100644 --- a/docs/modules/ppo.rst +++ b/docs/modules/ppo.rst @@ -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? diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 696ae55..6ac1c76 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -113,3 +113,4 @@ optimizers Deprecations forkserver cuda +Polyak diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 1b84018..8ae7bb4 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -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: diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 5c780c9..a745d4c 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -2,6 +2,8 @@ import time import os import io import zipfile +import typing +from typing import Union, Type, Optional from abc import ABC, abstractmethod from collections import deque @@ -10,13 +12,16 @@ 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 +if typing.TYPE_CHECKING: + from torchy_baselines.common.noise import ActionNoise + class BaseRLModel(ABC): """ @@ -43,7 +48,7 @@ class BaseRLModel(ABC): :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE Default: -1 (only sample at the beginning of the rollout) """ - def __init__(self, policy, env, policy_base, policy_kwargs=None, + def __init__(self, policy: Type[BasePolicy], env: Union[gym.Env, VecEnv, str], 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): @@ -59,7 +64,7 @@ class BaseRLModel(ABC): if verbose > 0: 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 @@ -71,7 +76,10 @@ class BaseRLModel(ABC): 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 @@ -407,7 +415,7 @@ class BaseRLModel(ABC): return data, params, opt_params - def set_random_seed(self, seed=None): + def set_random_seed(self, seed: Optional[int] = None): """ Set the seed of the pseudo-random generators (python, numpy, pytorch, gym, action_space) @@ -443,7 +451,7 @@ class BaseRLModel(ABC): 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): diff --git a/torchy_baselines/common/vec_env/base_vec_env.py b/torchy_baselines/common/vec_env/base_vec_env.py index d239ddd..e29c35e 100644 --- a/torchy_baselines/common/vec_env/base_vec_env.py +++ b/torchy_baselines/common/vec_env/base_vec_env.py @@ -146,6 +146,17 @@ class VecEnv(ABC): """ 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): diff --git a/torchy_baselines/common/vec_env/dummy_vec_env.py b/torchy_baselines/common/vec_env/dummy_vec_env.py index c55ca6a..2d0211e 100644 --- a/torchy_baselines/common/vec_env/dummy_vec_env.py +++ b/torchy_baselines/common/vec_env/dummy_vec_env.py @@ -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() diff --git a/torchy_baselines/common/vec_env/subproc_vec_env.py b/torchy_baselines/common/vec_env/subproc_vec_env.py index 0d1ecde..9205683 100644 --- a/torchy_baselines/common/vec_env/subproc_vec_env.py +++ b/torchy_baselines/common/vec_env/subproc_vec_env.py @@ -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() @@ -187,17 +187,6 @@ 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, indices=None, **method_kwargs): """Call instance methods of vectorized environments.""" target_remotes = self._get_target_remotes(indices) diff --git a/torchy_baselines/common/vec_env/vec_frame_stack.py b/torchy_baselines/common/vec_env/vec_frame_stack.py index 2610f42..562c525 100644 --- a/torchy_baselines/common/vec_env/vec_frame_stack.py +++ b/torchy_baselines/common/vec_env/vec_frame_stack.py @@ -3,7 +3,7 @@ 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 VecEnvWrapper class VecFrameStack(VecEnvWrapper): diff --git a/torchy_baselines/common/vec_env/vec_normalize.py b/torchy_baselines/common/vec_env/vec_normalize.py index 4f5cb11..ea94bc0 100644 --- a/torchy_baselines/common/vec_env/vec_normalize.py +++ b/torchy_baselines/common/vec_env/vec_normalize.py @@ -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 diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 5dd746f..4421d50 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -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): diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 8641552..77a47cf 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -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) diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index f40f298..3fe11c5 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -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): diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 70ca621..1e682cf 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -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) diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 8a19d8f..32cc1b2 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -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) From 9345b85cfc5c46b7b0033398c3374082f92797b3 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Wed, 22 Jan 2020 17:23:42 +0100 Subject: [PATCH 4/8] Update changelog and README --- README.md | 30 ++++++++++++++++++++++++++++++ docs/misc/changelog.rst | 4 +++- 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 018eb35..7ae8981 100644 --- a/README.md +++ b/README.md @@ -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: diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 39e30db..b774e02 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -8,6 +8,7 @@ Pre-Release 0.2.0a0 (WIP) Breaking Changes: ^^^^^^^^^^^^^^^^^ +- Python 2 support was dropped, Torchy Baselines now requires Python 3.6 or above New Features: ^^^^^^^^^^^^^ @@ -18,9 +19,10 @@ Bug Fixes: Deprecations: ^^^^^^^^^^^^^ - Others: ^^^^^^^ +- Add type check +- Converted all format string to f-strings Documentation: ^^^^^^^^^^^^^^ From 0328a39d1b406b4646881699a2770897ead17d87 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Wed, 22 Jan 2020 17:25:08 +0100 Subject: [PATCH 5/8] Update changelog --- docs/misc/changelog.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index b774e02..4f95650 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -12,6 +12,7 @@ Breaking Changes: New Features: ^^^^^^^^^^^^^ +- Add `seed()` method to `VecEnv` class Bug Fixes: ^^^^^^^^^^ From ff0eddfb1796e84dde7560db1f56f2ebefe83737 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Wed, 22 Jan 2020 17:51:27 +0100 Subject: [PATCH 6/8] Partially type base class --- docs/conf.py | 1 + setup.py | 4 +- torchy_baselines/common/base_class.py | 149 +++++++++++++++----------- 3 files changed, 88 insertions(+), 66 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 06c977b..624e305 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -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', diff --git a/setup.py b/setup.py index 343967a..e598108 100644 --- a/setup.py +++ b/setup.py @@ -25,7 +25,9 @@ setup(name='torchy_baselines', 'sphinx-autobuild', 'sphinx-rtd-theme', # For spelling - 'sphinxcontrib.spelling' + 'sphinxcontrib.spelling', + # Type hints support + 'sphinx-autodoc-typehints' ], 'extra': [ # For render diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index a745d4c..d714300 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -3,7 +3,7 @@ import os import io import zipfile import typing -from typing import Union, Type, Optional +from typing import Union, Type, Optional, Dict, Any, List, Tuple from abc import ABC, abstractmethod from collections import deque @@ -19,6 +19,7 @@ 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 @@ -27,31 +28,41 @@ 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) """ - def __init__(self, policy: Type[BasePolicy], env: Union[gym.Env, VecEnv, str], 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): + 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): + if isinstance(policy, str) and policy_base is not None: self.policy_class = get_policy_from_name(policy_base, policy) else: @@ -118,12 +129,12 @@ class BaseRLModel(ABC): 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 @@ -134,48 +145,47 @@ class BaseRLModel(ABC): 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)) @@ -186,32 +196,32 @@ class BaseRLModel(ABC): 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 @@ -220,7 +230,7 @@ class BaseRLModel(ABC): # 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 @@ -228,7 +238,7 @@ class BaseRLModel(ABC): - 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: " @@ -242,23 +252,24 @@ class BaseRLModel(ABC): 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 @@ -266,8 +277,13 @@ class BaseRLModel(ABC): 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. @@ -285,19 +301,22 @@ class BaseRLModel(ABC): 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() @@ -305,20 +324,20 @@ class BaseRLModel(ABC): 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 """ @@ -349,11 +368,11 @@ class BaseRLModel(ABC): 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) @@ -415,7 +434,7 @@ class BaseRLModel(ABC): return data, params, opt_params - def set_random_seed(self, seed: Optional[int] = 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) From 7265d9e3529916b8455fd8831fbb6775c91993ec Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Thu, 23 Jan 2020 10:56:53 +0100 Subject: [PATCH 7/8] Fix multiline f-string --- torchy_baselines/common/base_class.py | 45 ++++++++++++++------------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index d714300..563e419 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -653,13 +653,15 @@ class BaseRLModel(ABC): 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 """ @@ -690,7 +692,7 @@ class BaseRLModel(ABC): 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. @@ -699,13 +701,13 @@ class BaseRLModel(ABC): """ 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() @@ -728,25 +730,26 @@ class BaseRLModel(ABC): 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(f"Eval num_timesteps={self.num_timesteps}, " - "episode_reward={mean_reward:.2f} +/- {std_reward:.2f}") + 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 From fb57a6b80c1936b5f9fc40e2ce2ea2c075d325b1 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Thu, 23 Jan 2020 11:20:12 +0100 Subject: [PATCH 8/8] Update docstring --- torchy_baselines/common/vec_env/vec_frame_stack.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/torchy_baselines/common/vec_env/vec_frame_stack.py b/torchy_baselines/common/vec_env/vec_frame_stack.py index 562c525..6676162 100644 --- a/torchy_baselines/common/vec_env/vec_frame_stack.py +++ b/torchy_baselines/common/vec_env/vec_frame_stack.py @@ -3,18 +3,18 @@ import warnings import numpy as np from gym import spaces -from torchy_baselines.common.vec_env.base_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