diff --git a/.coveragerc b/.coveragerc index 4e5d7bd..a8fc2af 100644 --- a/.coveragerc +++ b/.coveragerc @@ -8,3 +8,4 @@ omit = exclude_lines = pragma: no cover raise NotImplementedError() + if typing.TYPE_CHECKING: diff --git a/.github/ISSUE_TEMPLATE/issue-template.md b/.github/ISSUE_TEMPLATE/issue-template.md new file mode 100644 index 0000000..2e2e61b --- /dev/null +++ b/.github/ISSUE_TEMPLATE/issue-template.md @@ -0,0 +1,52 @@ +--- +name: Issue Template +about: How to create an issue for this repository + +--- + +**Important Note: We do not do technical support, nor consulting** and don't answer personal questions per email. + +If you have any questions, feel free to create an issue with the tag [question]. +If you wish to suggest an enhancement or feature request, add the tag [feature request]. +If you are submitting a bug report, please fill in the following details. + +If your issue is related to a custom gym environment, please check it first using: + +```python +from torchy_baselines.common.env_checker import check_env + +env = CustomEnv(arg1, ...) +# It will check your custom environment and output additional warnings if needed +check_env(env) +``` + +**Describe the bug** +A clear and concise description of what the bug is. + +**Code example** +Please try to provide a minimal example to reproduce the bug. Error messages and stack traces are also helpful. + +Please use the [markdown code blocks](https://help.github.com/en/articles/creating-and-highlighting-code-blocks) +for both code and stack traces. + +```python +from torchy_baselines import ... + +``` + +```bash +Traceback (most recent call last): File ... + +``` + +**System Info** +Describe the characteristic of your environment: + * Describe how the library was installed (pip, docker, source, ...) + * GPU models and configuration + * Python version + * PyTorch version + * Gym version + * Versions of any other relevant libraries + +**Additional context** +Add any other context about the problem here. diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000..0d40764 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,29 @@ + + +## Description + + +## Motivation and Context + + + +- [ ] I have raised an issue to propose this change ([required](https://github.com/hill-a/stable-baselines/blob/master/CONTRIBUTING.md) for new features and bug fixes) + +## Types of changes + +- [ ] Bug fix (non-breaking change which fixes an issue) +- [ ] New feature (non-breaking change which adds functionality) +- [ ] Breaking change (fix or feature that would cause existing functionality to change) +- [ ] Documentation (update in the documentation) + +## Checklist: + + +- [ ] I've read the [CONTRIBUTION](https://github.com/hill-a/stable-baselines/blob/master/CONTRIBUTING.md) guide (**required**) +- [ ] I have updated the changelog accordingly (**required**). +- [ ] My change requires a change to the documentation. +- [ ] I have updated the tests accordingly (*required for a bug fix or a new feature*). +- [ ] I have updated the documentation accordingly. +- [ ] I have ensured `pytest` and `pytype` both pass. + + diff --git a/Makefile b/Makefile index 78b043f..9c41b52 100644 --- a/Makefile +++ b/Makefile @@ -6,8 +6,11 @@ pytest: type: pytype -docs: +doc: cd docs && make html spelling: cd docs && make spelling + +clean: + cd docs && make clean diff --git a/docs/conf.py b/docs/conf.py index 624e305..ab7fbbc 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -75,6 +75,7 @@ extensions = [ 'sphinx.ext.mathjax', 'sphinx.ext.ifconfig', 'sphinx.ext.viewcode', + # 'sphinx.ext.intersphinx', ] if enable_spell_check: @@ -206,3 +207,10 @@ texinfo_documents = [ # -- Extension configuration ------------------------------------------------- + +# Example configuration for intersphinx: refer to the Python standard library. +# intersphinx_mapping = { +# 'python': ('https://docs.python.org/3/', None), +# 'numpy': ('http://docs.scipy.org/doc/numpy/', None), +# 'torch': ('http://pytorch.org/docs/master/', None), +# } diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 958c78a..d0d69c1 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -9,14 +9,17 @@ Pre-Release 0.2.0a0 (WIP) Breaking Changes: ^^^^^^^^^^^^^^^^^ - Python 2 support was dropped, Torchy Baselines now requires Python 3.6 or above +- Return type of `evaluation.evaluate_policy()` has been changed New Features: ^^^^^^^^^^^^^ - Add `seed()` method to `VecEnv` class +- Add support for Callback (cf https://github.com/hill-a/stable-baselines/pull/644) Bug Fixes: ^^^^^^^^^^ - Fix loading model on CPU that were trained on GPU +- Fix `reset_num_timesteps` that was not used Deprecations: ^^^^^^^^^^^^^ @@ -25,6 +28,8 @@ Others: ^^^^^^^ - Add type check - Converted all format string to f-strings +- Add test for `OrnsteinUhlenbeckActionNoise` +- Add type aliases in `common.type_aliases` Documentation: ^^^^^^^^^^^^^^ diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py new file mode 100644 index 0000000..96db45a --- /dev/null +++ b/tests/test_callbacks.py @@ -0,0 +1,49 @@ +import os +import shutil + +import pytest +import gym + +from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3 +from torchy_baselines.common.callbacks import (CallbackList, CheckpointCallback, EvalCallback, + EveryNTimesteps, StopTrainingOnRewardThreshold) + + +@pytest.mark.parametrize("model_class", [A2C, CEMRL, PPO, SAC, TD3]) +def test_callbacks(model_class): + log_folder = './logs/callbacks/' + kwargs = {} + if model_class == CEMRL: + kwargs['pop_size'] = 2 + kwargs['n_grad'] = 1 + + # Create RL model + # Small network for fast test + model = model_class('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[32]), **kwargs) + + checkpoint_callback = CheckpointCallback(save_freq=1000, save_path=log_folder) + + eval_env = gym.make('Pendulum-v0') + # Stop training if the performance is good enough + callback_on_best = StopTrainingOnRewardThreshold(reward_threshold=-1200, verbose=1) + + eval_callback = EvalCallback(eval_env, callback_on_new_best=callback_on_best, + best_model_save_path=log_folder, + log_path=log_folder, eval_freq=100) + + # Equivalent to the `checkpoint_callback` + # but here in an event-driven manner + checkpoint_on_event = CheckpointCallback(save_freq=1, save_path=log_folder, + name_prefix='event') + event_callback = EveryNTimesteps(n_steps=500, callback=checkpoint_on_event) + + callback = CallbackList([checkpoint_callback, eval_callback, event_callback]) + + model.learn(500, callback=callback) + model.learn(500, callback=None) + # Transform callback into a callback list automatically + model.learn(500, callback=[checkpoint_callback, eval_callback]) + # Automatic wrapping, old way of doing callbacks + model.learn(500, callback=lambda _locals, _globals : True) + if os.path.exists(log_folder): + shutil.rmtree(log_folder) diff --git a/tests/test_run.py b/tests/test_run.py index e551a8d..1d206c9 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -4,12 +4,13 @@ import pytest import numpy as np from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3 -from torchy_baselines.common.noise import NormalActionNoise +from torchy_baselines.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1)) -def test_td3(): +@pytest.mark.parametrize('action_noise', [action_noise, OrnsteinUhlenbeckActionNoise(np.zeros(1), 0.1 * np.ones(1))]) +def test_td3(action_noise): model = TD3('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]), learning_starts=100, verbose=1, create_eval_env=True, action_noise=action_noise) model.learn(total_timesteps=1000, eval_freq=500) diff --git a/tests/test_save_load.py b/tests/test_save_load.py index b4dd7b7..5171042 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -69,7 +69,7 @@ def test_save_load(model_class): # check if model still selects the same actions new_selected_actions = [model.predict(observation, deterministic=True) for observation in observations] - assert np.allclose(selected_actions, new_selected_actions) + assert np.allclose(selected_actions, new_selected_actions, 1e-4) # check if learn still works model.learn(total_timesteps=1000, eval_freq=500) diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index 6fa667c..a4e83be 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/torchy_baselines/a2c/a2c.py @@ -130,8 +130,10 @@ class A2C(PPO): logger.logkv("std", th.exp(self.policy.log_std).mean().item()) def learn(self, total_timesteps, callback=None, log_interval=100, - eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="A2C", reset_num_timesteps=True): + eval_env=None, eval_freq=-1, n_eval_episodes=5, + tb_log_name="A2C", eval_log_path=None, reset_num_timesteps=True): return super(A2C, self).learn(total_timesteps=total_timesteps, callback=callback, log_interval=log_interval, eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes, - tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps) + tb_log_name=tb_log_name, eval_log_path=eval_log_path, + reset_num_timesteps=reset_num_timesteps) diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 8ae7bb4..653b3e2 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -100,21 +100,21 @@ class CEMRL(TD3): elitism=self.elitism) def learn(self, total_timesteps, callback=None, log_interval=4, - eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="CEMRL", reset_num_timesteps=True): + eval_env=None, eval_freq=-1, n_eval_episodes=5, + tb_log_name="CEMRL", eval_log_path=None, reset_num_timesteps=True): - timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, + n_eval_episodes, eval_log_path, reset_num_timesteps) actor_steps = 0 + continue_training = True + + callback.on_training_start(locals(), globals()) while self.num_timesteps < total_timesteps: self.fitnesses = [] self.es_params = self.es.ask(self.pop_size) - if callback is not None: - # Only stop training if return value is False, not when it is None. - if callback(locals(), globals()) is False: - break - if self.num_timesteps > 0: # self.train(episode_timesteps) # Gradient steps for half of the population @@ -157,46 +157,35 @@ class CEMRL(TD3): # Get the params back in the population self.es_params[i] = self.actor.parameters_to_vector() - # Evaluate agent - if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: - timesteps_since_eval %= eval_freq - - self.actor.load_from_vector(self.es.mu) - sync_envs_normalization(self.env, eval_env) - - mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes) - evaluations.append(mean_reward) - - 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))) - actor_steps = 0 # evaluate all actors for params in self.es_params: - self.actor.load_from_vector(params) rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout, n_steps=-1, action_noise=self.action_noise, - deterministic=False, callback=None, + deterministic=False, callback=callback, learning_starts=self.learning_starts, - num_timesteps=self.num_timesteps, replay_buffer=self.replay_buffer, obs=obs, episode_num=episode_num, log_interval=log_interval) # Unpack - episode_reward, episode_timesteps, n_episodes, obs = rollout + episode_reward, episode_timesteps, n_episodes, obs, continue_training = rollout + + if continue_training is False: + break episode_num += n_episodes - self.num_timesteps += episode_timesteps - timesteps_since_eval += episode_timesteps actor_steps += episode_timesteps self.fitnesses.append(episode_reward) + if continue_training is False: + break + self._update_current_progress(self.num_timesteps, total_timesteps) self.es.tell(self.es_params, self.fitnesses) - timesteps_since_eval += actor_steps + + callback.on_training_end() + return self diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 41743b4..9a3222f 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -2,8 +2,7 @@ import time import os import io import zipfile -import typing -from typing import Union, Type, Optional, Dict, Any, List, Tuple +from typing import Union, Type, Optional, Dict, Any, List, Tuple, Callable from abc import ABC, abstractmethod from collections import deque @@ -14,14 +13,12 @@ import numpy as np from torchy_baselines.common import logger 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.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize 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, recursive_getattr, recursive_setattr - -# TODO: define aliases, ex GymEnv = Union[gym.Env, VecEnv] -if typing.TYPE_CHECKING: - from torchy_baselines.common.noise import ActionNoise +from torchy_baselines.common.type_aliases import GymEnv, TensorDict, OptimizerStateDict +from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback, EvalCallback +from torchy_baselines.common.noise import ActionNoise class BaseRLModel(ABC): @@ -52,7 +49,7 @@ class BaseRLModel(ABC): def __init__(self, policy: Type[BasePolicy], - env: Union[gym.Env, VecEnv, str], + env: Union[GymEnv, str], policy_base: Type[BasePolicy], policy_kwargs: Dict[str, Any] = None, verbose: int = 0, @@ -76,7 +73,7 @@ class BaseRLModel(ABC): if verbose > 0: print(f"Using {self.device} device") - self.env = None # type: Union[gym.Env, VecEnv] + self.env = None # type: GymEnv # get VecNormalize object if needed self._vec_normalize_env = unwrap_vec_normalize(env) self.verbose = verbose @@ -134,16 +131,12 @@ class BaseRLModel(ABC): def _setup_model(self) -> None: """ Setup model so state_dict can be loaded - """ raise NotImplementedError() - def _get_eval_env(self, eval_env: Union[gym.Env, VecEnv, None]) -> Union[gym.Env, VecEnv, None]: + def _get_eval_env(self, eval_env: Optional[GymEnv]) -> Optional[GymEnv]: """ Return the environment that will be used for evaluation. - - :param eval_env: - :return: """ if eval_env is None: eval_env = self.eval_env @@ -159,8 +152,7 @@ class BaseRLModel(ABC): Rescale the action from [low, high] to [-1, 1] (no need for symmetric action space) - :param action: - :return: + :param action: Action to scale """ low, high = self.action_space.low, self.action_space.high return 2.0 * ((action - low) / (high - low)) - 1.0 @@ -170,8 +162,7 @@ class BaseRLModel(ABC): Rescale the action from [-1, 1] to [low, high] (no need for symmetric action space) - :param scaled_action: - :return: + :param scaled_action: Action to un-scale """ low, high = self.action_space.low, self.action_space.high return low + (0.5 * (scaled_action + 1.0) * (high - low)) @@ -215,7 +206,7 @@ class BaseRLModel(ABC): """ return np.nan if len(arr) == 0 else np.mean(arr) - def get_env(self) -> Union[VecEnv, None]: + def get_env(self) -> Optional[VecEnv]: """ Returns the current environment (can be None if not defined). @@ -230,7 +221,10 @@ class BaseRLModel(ABC): Checked parameters: - observation_space - action_space - :return: True if environment seems to be coherent + + :param observation_space: (gym.spaces.Space) + :param action_space: (gym.spaces.Space) + :return: (bool) True if environment seems to be coherent """ if observation_space != env.observation_space: return False @@ -239,7 +233,7 @@ class BaseRLModel(ABC): # return true if no check failed return True - def set_env(self, env: Union[gym.Env, VecEnv]) -> None: + def set_env(self, env: GymEnv) -> 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 @@ -276,11 +270,13 @@ class BaseRLModel(ABC): @abstractmethod def learn(self, total_timesteps: int, - callback=None, log_interval: int = 100, + callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None, + log_interval: int = 100, tb_log_name: str = "run", - eval_env: Union[gym.Env, VecEnv, None] = None, + eval_env: Optional[GymEnv] = None, eval_freq: int = -1, n_eval_episodes: int = 5, + eval_log_path: Optional[str] = None, reset_num_timesteps: bool = True): """ Return a trained model. @@ -294,6 +290,8 @@ class BaseRLModel(ABC): :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 + :param eval_log_path: (Optional[str]) Path to a folder where the evaluations will be saved + :param reset_num_timesteps: (bool) :return: (BaseRLModel) the trained model """ raise NotImplementedError() @@ -315,7 +313,7 @@ class BaseRLModel(ABC): raise NotImplementedError() @classmethod - def load(cls, load_path: str, env: Union[gym.Env, VecEnv, None] = None, **kwargs): + def load(cls, load_path: str, env: Optional[GymEnv] = None, **kwargs): """ Load the model from a zip-file @@ -364,7 +362,9 @@ class BaseRLModel(ABC): return model @staticmethod - def _load_from_file(load_path: str, load_data: bool = True): + def _load_from_file(load_path: str, load_data: bool = True) -> (Tuple[Optional[Dict[str, Any]], + Optional[TensorDict], + Optional[TensorDict]]): """ Load model data from a .zip archive :param load_path: Where to load the model from @@ -397,7 +397,7 @@ class BaseRLModel(ABC): data = None tensors = None params = {} - + if "data" in namelist and load_data: # Load class parameters and convert to string json_data = archive.read("data").decode() @@ -467,12 +467,52 @@ class BaseRLModel(ABC): if self.eval_env is not None: self.eval_env.seed(seed) - def _setup_learn(self, eval_env): + def _init_callback(self, + callback: Union[None, Callable, List[BaseCallback], BaseCallback], + eval_env: Optional[VecEnv] = None, + eval_freq: int = 10000, + n_eval_episodes: int = 5, + log_path: Optional[str] = None) -> BaseCallback: + """ + :param callback: (Union[callable, [BaseCallback], BaseCallback, None]) + :return: (BaseCallback) + """ + # Convert a list of callbacks into a callback + if isinstance(callback, list): + callback = CallbackList(callback) + + # Convert functional callback to object + if not isinstance(callback, BaseCallback): + callback = ConvertCallback(callback) + + # Create eval callback in charge of the evaluation + if eval_env is not None: + eval_callback = EvalCallback(eval_env, + best_model_save_path=log_path, + log_path=log_path, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes) + callback = CallbackList([callback, eval_callback]) + + callback.init_callback(self) + return callback + + def _setup_learn(self, + eval_env: Optional[GymEnv], + callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None, + eval_freq: int = 10000, + n_eval_episodes: int = 5, + log_path: Optional[str] = None, + reset_num_timesteps: bool = True, + ) -> Tuple[int, np.ndarray, BaseCallback]: """ Initialize different variables needed for training. - :param eval_env: (gym.Env or VecEnv) - :return: (int, int, [float], np.ndarray, VecEnv) + :param eval_env: (Optional[GymEnv]) + :param callback: (Union[None, BaseCallback, List[BaseCallback, Callable]]) + :param eval_freq: (int) + :param n_eval_episodes: (int) + :param log_path (Optional[str]): Path to a log folder + :param reset_num_timesteps: (bool) Whether to reset or not the `num_timesteps` attribute + :return: (Tuple[int, np.ndarray, BaseCallback]) """ self.start_time = time.time() self.ep_info_buffer = deque(maxlen=100) @@ -481,16 +521,22 @@ class BaseRLModel(ABC): self.action_noise.reset() timesteps_since_eval, episode_num = 0, 0 - evaluations = [] + + if reset_num_timesteps: + self.num_timesteps = 0 if eval_env is not None and self.seed is not None: eval_env.seed(self.seed) eval_env = self._get_eval_env(eval_env) - obs = self.env.reset() # type: Union[gym.Env, VecEnv] - return timesteps_since_eval, episode_num, evaluations, obs, eval_env + obs = self.env.reset() - def _update_info_buffer(self, infos): + # Create eval callback if needed + callback = self._init_callback(callback, eval_env, eval_freq, n_eval_episodes, log_path) + + return episode_num, obs, callback + + def _update_info_buffer(self, infos: List[Dict[str, Any]]) -> None: """ Retrieve reward and episode length and update the buffer if using Monitor wrapper. @@ -502,11 +548,18 @@ class BaseRLModel(ABC): if maybe_ep_info is not None: self.ep_info_buffer.extend([maybe_ep_info]) - def collect_rollouts(self, env, n_episodes=1, n_steps=-1, action_noise=None, - deterministic=False, callback=None, - learning_starts=0, num_timesteps=0, - replay_buffer=None, obs=None, - episode_num=0, log_interval=None): + def collect_rollouts(self, + env: VecEnv, + callback: 'BaseCallback', # Type hint as string to avoid circular import + n_episodes: int = 1, + n_steps: int = -1, + action_noise: Optional[ActionNoise] = None, + deterministic: bool = False, + learning_starts: int = 0, + replay_buffer=None, + obs: Optional[np.ndarray] = None, + episode_num: int = 0, + log_interval: Optional[int] = None) -> Tuple[float, int, int, Optional[np.ndarray], bool]: """ Collect rollout using the current policy (and possibly fill the replay buffer) TODO: move this method to off-policy base class. @@ -516,9 +569,8 @@ class BaseRLModel(ABC): :param n_steps: (int) :param action_noise: (ActionNoise) :param deterministic: (bool) - :param callback: (callable) + :param callback: (BaseCallback) :param learning_starts: (int) - :param num_timesteps: (int) :param replay_buffer: (ReplayBuffer) :param obs: (np.ndarray) :param episode_num: (int) @@ -541,6 +593,9 @@ class BaseRLModel(ABC): if self.on_policy_exploration: self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones', 'values']} + callback.on_rollout_start() + continue_training = True + while total_steps < n_steps or total_episodes < n_episodes: done = False # Reset environment: not needed for VecEnv @@ -548,6 +603,12 @@ class BaseRLModel(ABC): episode_reward, episode_timesteps = 0.0, 0 while not done: + + # Only stop training if return value is False, not when it is None. + if callback() is False: + continue_training = False + return 0.0, total_steps, total_episodes, None, continue_training + if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0: # Sample a new noise matrix self.actor.reset_noise() @@ -555,7 +616,7 @@ class BaseRLModel(ABC): # Select action randomly or according to policy # TODO: use action from policy when using SDE during the warmup phase? # if num_timesteps < learning_starts and not self.use_sde: - if num_timesteps < learning_starts: + if self.num_timesteps < learning_starts: # Warmup phase unscaled_action = np.array([self.action_space.sample()]) else: @@ -614,7 +675,7 @@ class BaseRLModel(ABC): if self._vec_normalize_env is not None: obs_ = new_obs_ - num_timesteps += 1 + self.num_timesteps += 1 episode_timesteps += 1 total_steps += 1 if 0 < n_steps <= total_steps: @@ -630,8 +691,8 @@ class BaseRLModel(ABC): # Display training infos if self.verbose >= 1 and log_interval is not None and ( - episode_num + total_episodes) % log_interval == 0: - fps = int(num_timesteps / (time.time() - self.start_time)) + episode_num + total_episodes) % log_interval == 0: + fps = int(self.num_timesteps / (time.time() - self.start_time)) logger.logkv("episodes", episode_num + total_episodes) if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0: logger.logkv('ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer])) @@ -639,7 +700,7 @@ class BaseRLModel(ABC): # logger.logkv("n_updates", n_updates) logger.logkv("fps", fps) logger.logkv('time_elapsed', int(time.time() - self.start_time)) - logger.logkv("total timesteps", num_timesteps) + logger.logkv("total timesteps", self.num_timesteps) if self.use_sde: logger.logkv("std", (self.actor.get_std()).mean().item()) logger.dumpkvs() @@ -667,7 +728,9 @@ class BaseRLModel(ABC): self.rollout_data['returns'][step] = last_return self.rollout_data['advantage'] = self.rollout_data['returns'] - self.rollout_data['values'] - return mean_reward, total_steps, total_episodes, obs + callback.on_rollout_end() + + return mean_reward, total_steps, total_episodes, obs, continue_training @staticmethod def _save_to_file_zip(save_path: str, data: Dict[str, Any] = None, @@ -767,27 +830,3 @@ class BaseRLModel(ABC): params_to_save[name] = attr.state_dict() self._save_to_file_zip(path, data=data, params=params_to_save, tensors=tensors) - - 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_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, - render=render, deterministic=deterministic) - if self.verbose > 0: - 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 diff --git a/torchy_baselines/common/callbacks.py b/torchy_baselines/common/callbacks.py new file mode 100644 index 0000000..2481299 --- /dev/null +++ b/torchy_baselines/common/callbacks.py @@ -0,0 +1,345 @@ +import os +from abc import ABC, abstractmethod +import typing +from typing import Union, List, Dict, Any, Optional + +import gym +import numpy as np + +from torchy_baselines.common.vec_env import VecEnv, sync_envs_normalization +from torchy_baselines.common.evaluation import evaluate_policy +from torchy_baselines.common.logger import Logger + +if typing.TYPE_CHECKING: + from torchy_baselines.common.base_class import BaseRLModel # pytype: disable=pyi-error + + +class BaseCallback(ABC): + """ + Base class for callback. + + :param verbose: (int) + """ + def __init__(self, verbose: int = 0): + super(BaseCallback, self).__init__() + self.model = None # type: BaseRLModel + self.training_env = None # type: Union[gym.Env, VecEnv, None] + self.n_calls = 0 # type: int + self.num_timesteps = 0 # type: int + self.verbose = verbose + self.locals = None # type: Dict[str, Any] + self.globals = None # type: Dict[str, Any] + self.logger = None # type: Logger + # Sometimes, for event callback, it is useful + # to have access to the parent object + self.parent = None # type: Optional[BaseCallback] + + # Type hint as string to avoid circular import + def init_callback(self, model: 'BaseRLModel') -> None: + """ + Initialize the callback by saving references to the + RL model and the training environment for convenience. + """ + self.model = model + self.training_env = model.get_env() + self.logger = Logger.CURRENT + self._init_callback() + + def _init_callback(self) -> None: + pass + + def on_training_start(self, locals_: Dict[str, Any], globals_: Dict[str, Any]) -> None: + # Those are reference and will be updated automatically + self.locals = locals_ + self.globals = globals_ + self._on_training_start() + + def _on_training_start(self) -> None: + pass + + def on_rollout_start(self) -> None: + self._on_rollout_start() + + def _on_rollout_start(self) -> None: + pass + + @abstractmethod + def _on_step(self) -> bool: + """ + :return: (bool) If the callback returns False, training is aborted early. + """ + return True + + def __call__(self) -> bool: + """ + This method will be called by the model. This is the equivalent to the callback function. + :return: (bool) If the callback returns False, training is aborted early. + """ + self.n_calls += 1 + # timesteps start at zero + self.num_timesteps = self.model.num_timesteps + 1 + + return self._on_step() + + def on_training_end(self) -> None: + self._on_training_end() + + def _on_training_end(self) -> None: + pass + + def on_rollout_end(self) -> None: + self._on_rollout_end() + + def _on_rollout_end(self) -> None: + pass + + +class EventCallback(BaseCallback): + """ + Base class for triggering callback on event. + + :param callback: (Optional[BaseCallback]) Callback that will be called + when an event is triggered. + :param verbose: (int) + """ + def __init__(self, callback: Optional[BaseCallback] = None, verbose: int = 0): + super(EventCallback, self).__init__(verbose=verbose) + self.callback = callback + # Give access to the parent + if callback is not None: + self.callback.parent = self + + def init_callback(self, model: 'BaseRLModel') -> None: + super(EventCallback, self).init_callback(model) + if self.callback is not None: + self.callback.init_callback(self.model) + + def _on_training_start(self) -> None: + if self.callback is not None: + self.callback.on_training_start(self.locals, self.globals) + + def _on_event(self) -> bool: + if self.callback is not None: + return self.callback() + return True + + def _on_step(self) -> bool: + return True + + +class CallbackList(BaseCallback): + def __init__(self, callbacks: List[BaseCallback]): + super(CallbackList, self).__init__() + assert isinstance(callbacks, list) + self.callbacks = callbacks + + def _init_callback(self) -> None: + for callback in self.callbacks: + callback.init_callback(self.model) + + def _on_training_start(self) -> None: + for callback in self.callbacks: + callback.on_training_start(self.locals, self.globals) + + def _on_step(self) -> bool: + continue_training = True + for callback in self.callbacks: + # # Update variables + # callback.num_timesteps = self.num_timesteps + # callback.n_calls = self.n_calls + # Return False (stop training) if at least one callback returns False + continue_training = callback() and continue_training + return continue_training + + def _on_training_end(self) -> None: + for callback in self.callbacks: + callback.on_training_end() + + +class CheckpointCallback(BaseCallback): + """ + Callback for saving a model every `save_freq` steps + + :param save_freq: (int) + :param save_path: (str) Path to the folder where the model will be saved. + :param name_prefix: (str) Common prefix to the saved models + """ + def __init__(self, save_freq: int, save_path: str, name_prefix='rl_model', verbose=0): + super(CheckpointCallback, self).__init__(verbose) + self.save_freq = save_freq + self.save_path = save_path + self.name_prefix = name_prefix + + def _init_callback(self) -> None: + # Create folder if needed + if self.save_path is not None: + os.makedirs(self.save_path, exist_ok=True) + + def _on_step(self) -> bool: + if self.n_calls % self.save_freq == 0: + path = os.path.join(self.save_path, f'{self.name_prefix}_{self.num_timesteps}_steps') + self.model.save(path) + if self.verbose > 1: + print(f"Saving model checkpoint to {path}") + return True + + +class ConvertCallback(BaseCallback): + """ + Convert functional callback (old-style) to object. + + :param callback: (callable) + :param verbose: (int) + """ + def __init__(self, callback, verbose=0): + super(ConvertCallback, self).__init__(verbose) + self.callback = callback + + def _on_step(self) -> bool: + if self.callback is not None: + return self.callback(self.locals, self.globals) + return True + + +class EvalCallback(EventCallback): + """ + Callback for evaluating an agent. + + :param eval_env: (Union[gym.Env, VecEnv]) The environment used for initialization + :param callback_on_new_best: (Optional[BaseCallback]) Callback to trigger + when there is a new best model according to the `mean_reward` + :param n_eval_episodes: (int) The number of episodes to test the agent + :param eval_freq: (int) Evaluate the agent every eval_freq call of the callback. + :param log_path: (str) Path to a folder where the evaluations (`evaluations.npz`) + will be saved. It will be updated at each evaluation. + :param best_model_save_path: (str) Path to a folder where the best model + according to performance on the eval env will be saved. + :param deterministic: (bool) Whether the evaluation should + use a stochastic or deterministic actions. + :param deterministic: (bool) Whether to render or not the environment during evaluation + :param verbose: (int) + """ + def __init__(self, eval_env: Union[gym.Env, VecEnv], + callback_on_new_best: Optional[BaseCallback] = None, + n_eval_episodes: int = 5, + eval_freq: int = 10000, + log_path: str = None, + best_model_save_path: str = None, + deterministic: bool = True, + render: bool = False, + verbose: int = 1): + super(EvalCallback, self).__init__(callback_on_new_best, verbose=verbose) + self.n_eval_episodes = n_eval_episodes + self.eval_freq = eval_freq + self.best_mean_reward = -np.inf + self.deterministic = deterministic + self.render = render + + if isinstance(eval_env, VecEnv): + assert eval_env.num_envs == 1, "You must pass only one environment for evaluation" + + self.eval_env = eval_env + self.best_model_save_path = best_model_save_path + # Logs will be written in `evaluations.npz` + if log_path is not None: + os.path.join(log_path, 'evaluations') + self.log_path = log_path + self.evaluations_results = [] + self.evaluations_timesteps = [] + self.evaluations_length = [] + + def _init_callback(self): + # Does not work when eval_env is a gym.Env and training_env is a VecEnv + # assert type(self.training_env) is type(self.eval_env), ("training and eval env are not of the same type", + # "{} != {}".format(self.training_env, self.eval_env)) + + # Create folders if needed + if self.best_model_save_path is not None: + os.makedirs(self.best_model_save_path, exist_ok=True) + if self.log_path is not None: + os.makedirs(os.path.dirname(self.log_path), exist_ok=True) + + def _on_step(self) -> bool: + + if self.eval_freq > 0 and self.n_calls % self.eval_freq == 0: + # Sync training and eval env if there is VecNormalize + sync_envs_normalization(self.training_env, self.eval_env) + + episode_rewards, episode_lengths = evaluate_policy(self.model, self.eval_env, + n_eval_episodes=self.n_eval_episodes, + render=self.render, + deterministic=self.deterministic, + return_episode_rewards=True) + + if self.log_path is not None: + self.evaluations_timesteps.append(self.num_timesteps) + self.evaluations_results.append(episode_rewards) + self.evaluations_length.append(episode_lengths) + np.savez(self.log_path, timesteps=self.evaluations_timesteps, + results=self.evaluations_results, ep_lengths=self.evaluations_length) + + mean_reward, std_reward = np.mean(episode_rewards), np.std(episode_rewards) + mean_ep_length, std_ep_length = np.mean(episode_lengths), np.std(episode_lengths) + + if self.verbose > 0: + print(f"Eval num_timesteps={self.num_timesteps}, " + f"episode_reward={mean_reward:.2f} +/- {std_reward:.2f}") + print(f"Episode length: {mean_ep_length:.2f} +/- {std_ep_length:.2f}") + + if mean_reward > self.best_mean_reward: + if self.verbose > 0: + print("New best mean reward!") + if self.best_model_save_path is not None: + self.model.save(os.path.join(self.best_model_save_path, 'best_model')) + self.best_mean_reward = mean_reward + # Trigger callback if needed + if self.callback is not None: + return self._on_event() + + return True + + +class StopTrainingOnRewardThreshold(BaseCallback): + """ + Stop the training once a threshold in episodic reward + has been reached (i.e. when the model is good enough). + + It must be used with the `EvalCallback`. + + :param reward_threshold: (float) Minimum expected reward per episode + to stop training. + :param verbose: (int) + """ + def __init__(self, reward_threshold: float, verbose: int = 0): + super(StopTrainingOnRewardThreshold, self).__init__(verbose=verbose) + self.reward_threshold = reward_threshold + + def _on_step(self) -> bool: + assert self.parent is not None, ("`StopTrainingOnMinimumReward` callback must be used " + "with an `EvalCallback`") + # Convert np.bool to bool, otherwise callback() is False won't work + continue_training = bool(self.parent.best_mean_reward < self.reward_threshold) + if self.verbose > 0 and not continue_training: + print(f"Stopping training because the mean reward {self.parent.best_mean_reward:.2f} " + f" is above the threshold {self.reward_threshold}") + return continue_training + + +class EveryNTimesteps(EventCallback): + """ + Trigger a callback every `n_steps` timesteps + + :param n_steps: (int) Number of timesteps between two trigger. + :param callback: (BaseCallback) Callback that will be called + when the event is triggered. + """ + def __init__(self, n_steps: int, callback: BaseCallback): + super(EveryNTimesteps, self).__init__(callback) + self.n_steps = n_steps + self.last_time_trigger = 0 + + def _on_step(self) -> bool: + if (self.num_timesteps - self.last_time_trigger) >= self.n_steps: + self.last_time_trigger = self.num_timesteps + return self._on_event() + return True diff --git a/torchy_baselines/common/evaluation.py b/torchy_baselines/common/evaluation.py index 0133621..ff9fa13 100644 --- a/torchy_baselines/common/evaluation.py +++ b/torchy_baselines/common/evaluation.py @@ -24,31 +24,33 @@ def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True, :param return_episode_rewards: (bool) If True, a list of reward per episode will be returned instead of the mean. :return: (float, float) Mean reward per episode, std of reward per episode - returns ([float], int) when `return_episode_rewards` is True + returns ([float], [int]) when `return_episode_rewards` is True """ if isinstance(env, VecEnv): assert env.num_envs == 1, "You must pass only one environment when using this function" - episode_rewards, n_steps = [], 0 + episode_rewards, episode_lengths = [], [] for _ in range(n_eval_episodes): obs = env.reset() done = False episode_reward = 0.0 + episode_length = 0 while not done: action = model.predict(obs, deterministic=deterministic) obs, reward, done, _info = env.step(action) episode_reward += reward if callback is not None: callback(locals(), globals()) - n_steps += 1 + episode_length += 1 if render: env.render() episode_rewards.append(episode_reward) + episode_lengths.append(episode_length) mean_reward = np.mean(episode_rewards) std_reward = np.std(episode_rewards) if reward_threshold is not None: 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 episode_rewards, episode_lengths return mean_reward, std_reward diff --git a/torchy_baselines/common/noise.py b/torchy_baselines/common/noise.py index dd7d8a3..fa25f42 100644 --- a/torchy_baselines/common/noise.py +++ b/torchy_baselines/common/noise.py @@ -1,19 +1,26 @@ """ Taken from stable-baselines """ +from abc import ABC, abstractmethod import numpy as np -class ActionNoise(object): +class ActionNoise(ABC): """ The action noise base class """ + def __init__(self): + super(ActionNoise, self).__init__() + def reset(self): """ call end of episode reset for the noise """ pass + @abstractmethod + def __call__(self): + pass class NormalActionNoise(ActionNoise): """ @@ -25,6 +32,7 @@ class NormalActionNoise(ActionNoise): def __init__(self, mean, sigma): self._mu = mean self._sigma = sigma + super(NormalActionNoise, self).__init__() def __call__(self): return np.random.normal(self._mu, self._sigma) @@ -54,6 +62,7 @@ class OrnsteinUhlenbeckActionNoise(ActionNoise): self.initial_noise = initial_noise self.noise_prev = None self.reset() + super(OrnsteinUhlenbeckActionNoise, self).__init__() def __call__(self): noise = self.noise_prev + self._theta * (self._mu - self.noise_prev) * self._dt + \ diff --git a/torchy_baselines/common/type_aliases.py b/torchy_baselines/common/type_aliases.py new file mode 100644 index 0000000..8378647 --- /dev/null +++ b/torchy_baselines/common/type_aliases.py @@ -0,0 +1,14 @@ +""" +Common aliases for type hing +""" +from typing import Union, Type, Optional, Dict, Any, List, Tuple + +import torch +import gym + +from torchy_baselines.common.vec_env import VecEnv + + +GymEnv = Union[gym.Env, VecEnv] +TensorDict = Dict[str, torch.Tensor] +OptimizerStateDict = Dict[str, Any] diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 9e7c48c..c5916ea 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -1,6 +1,6 @@ import os import time -from typing import List, Tuple +from typing import Optional, Tuple, List import gym from gym import spaces @@ -17,6 +17,8 @@ import numpy as np from torchy_baselines.common.base_class import BaseRLModel from torchy_baselines.common.buffers import RolloutBuffer from torchy_baselines.common.utils import explained_variance, get_schedule_fn +from torchy_baselines.common.vec_env import VecEnv +from torchy_baselines.common.callbacks import BaseCallback from torchy_baselines.common import logger from torchy_baselines.ppo.policies import PPOPolicy @@ -151,17 +153,30 @@ class PPO(BaseRLModel): clipped_actions = np.clip(clipped_actions, self.action_space.low, self.action_space.high) return clipped_actions - def collect_rollouts(self, env, rollout_buffer, n_rollout_steps=256, callback=None, - obs=None): + def collect_rollouts(self, + env: VecEnv, + callback: BaseCallback, + rollout_buffer: RolloutBuffer, + n_rollout_steps: int = 256, + obs: Optional[np.ndarray] = None) -> Tuple[Optional[np.ndarray], bool]: n_steps = 0 + continue_training = True rollout_buffer.reset() # Sample new weights for the state dependent exploration # TODO: ensure episodic setting? if self.use_sde: self.policy.reset_noise(env.num_envs) + callback.on_rollout_start() + while n_steps < n_rollout_steps: + + if callback() is False: + continue_training = False + return None, continue_training + + if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0: # Sample a new noise matrix self.policy.reset_noise(env.num_envs) @@ -179,6 +194,8 @@ class PPO(BaseRLModel): self._update_info_buffer(infos) n_steps += 1 + self.num_timesteps += env.num_envs + if isinstance(self.action_space, gym.spaces.Discrete): # Reshape in case of discrete action actions = actions.reshape(-1, 1) @@ -187,7 +204,9 @@ class PPO(BaseRLModel): rollout_buffer.compute_returns_and_advantage(values, dones=dones) - return obs + callback.on_rollout_end() + + return obs, continue_training def train(self, gradient_steps, batch_size=64): # Update optimizer learning rate @@ -268,25 +287,29 @@ class PPO(BaseRLModel): logger.logkv("std", th.exp(self.policy.log_std).mean().item()) def learn(self, total_timesteps, callback=None, log_interval=1, - eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="PPO", reset_num_timesteps=True): + eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="PPO", + eval_log_path=None, reset_num_timesteps=True): - timesteps_since_eval, iteration, evaluations, obs, eval_env = self._setup_learn(eval_env) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, + n_eval_episodes, eval_log_path, reset_num_timesteps) + iteration = 0 if self.tensorboard_log is not None and SummaryWriter is not None: self.tb_writer = SummaryWriter(log_dir=os.path.join(self.tensorboard_log, tb_log_name)) + callback.on_training_start(locals(), globals()) + while self.num_timesteps < total_timesteps: - if callback is not None: - # Only stop training if return value is False, not when it is None. - if callback(locals(), globals()) is False: - break + obs, continue_training = self.collect_rollouts(self.env, callback, + self.rollout_buffer, + n_rollout_steps=self.n_steps, + obs=obs) + + if continue_training is False: + break - obs = self.collect_rollouts(self.env, self.rollout_buffer, n_rollout_steps=self.n_steps, - obs=obs) iteration += 1 - self.num_timesteps += self.n_steps * self.n_envs - timesteps_since_eval += self.n_steps * self.n_envs self._update_current_progress(self.num_timesteps, total_timesteps) # Display training infos @@ -303,13 +326,12 @@ class PPO(BaseRLModel): self.train(self.n_epochs, batch_size=self.batch_size) - # Evaluate the agent - timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes, - timesteps_since_eval, deterministic=True) # For tensorboard integration # if self.tb_writer is not None: # self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps) + callback.on_training_end() + return self def get_torch_variables(self) -> Tuple[List[str], List[str]]: diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 7cd92a4..907e921 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -259,31 +259,27 @@ class SAC(BaseRLModel): def learn(self, total_timesteps, callback=None, log_interval=4, eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="SAC", - reset_num_timesteps=True): + eval_log_path=None, reset_num_timesteps=True): - timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, + n_eval_episodes, eval_log_path, reset_num_timesteps) + callback.on_training_start(locals(), globals()) while self.num_timesteps < total_timesteps: - - if callback is not None: - # Only stop training if return value is False, not when it is None. - if callback(locals(), globals()) is False: - break - rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout, n_steps=self.train_freq, action_noise=self.action_noise, - deterministic=False, callback=None, + deterministic=False, callback=callback, learning_starts=self.learning_starts, - num_timesteps=self.num_timesteps, replay_buffer=self.replay_buffer, obs=obs, episode_num=episode_num, log_interval=log_interval) # Unpack - episode_reward, episode_timesteps, n_episodes, obs = rollout + episode_reward, episode_timesteps, n_episodes, obs, continue_training = rollout + + if continue_training is False: + break - self.num_timesteps += episode_timesteps episode_num += n_episodes - timesteps_since_eval += episode_timesteps self._update_current_progress(self.num_timesteps, total_timesteps) if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: @@ -291,9 +287,7 @@ class SAC(BaseRLModel): self.train(gradient_steps, batch_size=self.batch_size) - timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes, - timesteps_since_eval, deterministic=True) - + callback.on_training_end() return self def excluded_save_params(self) -> List[str]: diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index f016dab..8bc5f60 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -1,3 +1,6 @@ +from typing import List, Tuple, Callable, Optional + +import torch import torch as th import torch.nn as nn @@ -27,9 +30,18 @@ class Actor(BaseNetwork): a positive standard deviation (cf paper). It allows to keep variance above zero and prevent it from growing too fast. In practice, `exp()` is usually enough. """ - def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU, - use_sde=False, log_std_init=-3, clip_noise=None, - lr_sde=3e-4, full_std=False, sde_net_arch=None, use_expln=False): + def __init__(self, + obs_dim: int, + action_dim: int, + net_arch: List[int], + activation_fn: nn.Module = nn.ReLU, + use_sde: bool = False, + log_std_init: float = -3, + clip_noise: Optional[float] = None, + lr_sde: float = 3e-4, + full_std: bool = False, + sde_net_arch: Optional[List[int]] = None, + use_expln: bool = False): super(Actor, self).__init__() self.latent_pi, self.log_std = None, None @@ -65,7 +77,7 @@ class Actor(BaseNetwork): actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True) self.mu = nn.Sequential(*actor_net) - def get_std(self): + def get_std(self) -> torch.Tensor: """ Retrieve the standard deviation of the action distribution. Only useful when using SDE. @@ -81,7 +93,7 @@ class Actor(BaseNetwork): mean_actions = self.mu(latent_pi) return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde) - def _get_latent(self, obs): + def _get_latent(self, obs) -> Tuple[torch.Tensor, torch.Tensor]: latent_pi = self.latent_pi(obs) if self.sde_feature_extractor is not None: @@ -90,7 +102,7 @@ class Actor(BaseNetwork): latent_sde = latent_pi return latent_pi, latent_sde - def evaluate_actions(self, obs, action): + def evaluate_actions(self, obs: torch.Tensor, action: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ Evaluate actions according to the current policy, given the observations. Only useful when using SDE. @@ -106,13 +118,13 @@ class Actor(BaseNetwork): # value = self.value_net(latent_vf) return log_prob, distribution.entropy() - def reset_noise(self): + def reset_noise(self) -> None: """ Sample new weights for the exploration matrix, when using SDE. """ self.action_dist.sample_weights(self.log_std) - def forward(self, obs, deterministic=True): + def forward(self, obs: torch.Tensor, deterministic: bool = True) -> torch.Tensor: if self.use_sde: latent_pi, latent_sde = self._get_latent(obs) if deterministic: @@ -141,8 +153,8 @@ class Critic(BaseNetwork): :param net_arch: ([int]) Network architecture :param activation_fn: (nn.Module) Activation function """ - def __init__(self, obs_dim, action_dim, - net_arch, activation_fn=nn.ReLU): + def __init__(self, obs_dim: int, action_dim: int, + net_arch: List[int], activation_fn: nn.Module = nn.ReLU): super(Critic, self).__init__() q1_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn) @@ -151,14 +163,12 @@ class Critic(BaseNetwork): q2_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn) self.q2_net = nn.Sequential(*q2_net) - self.q_networks = [self.q1_net, self.q2_net] - - def forward(self, obs, action): + def forward(self, obs: torch.Tensor, action: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: qvalue_input = th.cat([obs, action], dim=1) - return [q_net(qvalue_input) for q_net in self.q_networks] + return self.q1_net(qvalue_input), self.q2_net(qvalue_input) - def q1_forward(self, obs, action): - return self.q_networks[0](th.cat([obs, action], dim=1)) + def q1_forward(self, obs: torch.Tensor, action: torch.Tensor) -> torch.Tensor: + return self.q1_net(th.cat([obs, action], dim=1)) class ValueFunction(BaseNetwork): diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 5729ebf..1c2e417 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -251,31 +251,30 @@ class TD3(BaseRLModel): del self.rollout_data def learn(self, total_timesteps, callback=None, log_interval=4, - eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="TD3", reset_num_timesteps=True): + eval_env=None, eval_freq=-1, n_eval_episodes=5, + tb_log_name="TD3", eval_log_path=None, reset_num_timesteps=True): - timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, + n_eval_episodes, eval_log_path, reset_num_timesteps) + + callback.on_training_start(locals(), globals()) while self.num_timesteps < total_timesteps: - if callback is not None: - # Only stop training if return value is False, not when it is None. - if callback(locals(), globals()) is False: - break - rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout, n_steps=self.train_freq, action_noise=self.action_noise, - deterministic=False, callback=None, + deterministic=False, callback=callback, learning_starts=self.learning_starts, - num_timesteps=self.num_timesteps, replay_buffer=self.replay_buffer, obs=obs, episode_num=episode_num, log_interval=log_interval) # Unpack - episode_reward, episode_timesteps, n_episodes, obs = rollout + episode_reward, episode_timesteps, n_episodes, obs, continue_training = rollout + + if continue_training is False: + break episode_num += n_episodes - self.num_timesteps += episode_timesteps - timesteps_since_eval += episode_timesteps self._update_current_progress(self.num_timesteps, total_timesteps) if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: @@ -292,9 +291,7 @@ class TD3(BaseRLModel): gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay) - # Evaluate the agent - timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes, - timesteps_since_eval, deterministic=True) + callback.on_training_end() return self