From b66003cfb37a4af9062289c964ff5c6f43927615 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Mon, 27 Jan 2020 14:32:31 +0100 Subject: [PATCH 1/9] Add callback support --- .coveragerc | 1 + docs/misc/changelog.rst | 2 + tests/test_callbacks.py | 37 +++ tests/test_run.py | 5 +- torchy_baselines/cem_rl/cem_rl.py | 23 +- torchy_baselines/common/base_class.py | 107 +++++--- torchy_baselines/common/callbacks.py | 320 ++++++++++++++++++++++++ torchy_baselines/common/noise.py | 11 +- torchy_baselines/common/type_aliases.py | 14 ++ torchy_baselines/ppo/ppo.py | 41 ++- torchy_baselines/sac/sac.py | 18 +- torchy_baselines/td3/policies.py | 42 ++-- torchy_baselines/td3/td3.py | 18 +- 13 files changed, 555 insertions(+), 84 deletions(-) create mode 100644 tests/test_callbacks.py create mode 100644 torchy_baselines/common/callbacks.py create mode 100644 torchy_baselines/common/type_aliases.py 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/docs/misc/changelog.rst b/docs/misc/changelog.rst index 4f95650..37e0f27 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -13,6 +13,7 @@ Breaking Changes: New Features: ^^^^^^^^^^^^^ - Add `seed()` method to `VecEnv` class +- Add support for Callback (cf https://github.com/hill-a/stable-baselines/pull/644) Bug Fixes: ^^^^^^^^^^ @@ -24,6 +25,7 @@ Others: ^^^^^^^ - Add type check - Converted all format string to f-strings +- Add test for `OrnsteinUhlenbeckActionNoise` Documentation: ^^^^^^^^^^^^^^ diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py new file mode 100644 index 0000000..533c548 --- /dev/null +++ b/tests/test_callbacks.py @@ -0,0 +1,37 @@ +import pytest + +from torchy_baselines import SAC +from torchy_baselines.common.callbacks import (CallbackList, CheckpointCallback, EvalCallback, + EveryNTimesteps, StopTrainingOnRewardThreshold) + + +@pytest.mark.parametrize("model_class", [SAC]) +def test_callbacks(model_class): + # Create RL model + model = model_class('MlpPolicy', 'Pendulum-v0') + + checkpoint_callback = CheckpointCallback(save_freq=1000, save_path='./logs/') + + # For testing: use the same training env + eval_env = model.get_env() + # 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='./logs/best_model', + log_path='./logs/results', eval_freq=100) + + # Equivalent to the `checkpoint_callback` + # but here in an event-driven manner + checkpoint_on_event = CheckpointCallback(save_freq=1, save_path='./logs/', + name_prefix='event') + event_callback = EveryNTimesteps(n_steps=1000, callback=checkpoint_on_event) + + callback = CallbackList([checkpoint_callback, eval_callback, event_callback]) + + model.learn(1000, 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) 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/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 8ae7bb4..1e26d15 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -102,19 +102,17 @@ class CEMRL(TD3): 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): - timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env) + timesteps_since_eval, episode_num, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback) 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 @@ -180,7 +178,7 @@ class CEMRL(TD3): 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, @@ -188,7 +186,10 @@ class CEMRL(TD3): 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 @@ -196,7 +197,13 @@ class CEMRL(TD3): 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 563e419..5bc2abd 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -18,10 +18,11 @@ from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_norm 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 +from torchy_baselines.common.type_aliases import GymEnv, TensorDict, OptimizerStateDict +from torchy_baselines.common.noise import ActionNoise -# TODO: define aliases, ex GymEnv = Union[gym.Env, VecEnv] if typing.TYPE_CHECKING: - from torchy_baselines.common.noise import ActionNoise + from torchy_baselines.common.callbacks import BaseCallback class BaseRLModel(ABC): @@ -51,7 +52,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, @@ -75,7 +76,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 @@ -129,7 +130,7 @@ 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: 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. @@ -145,6 +146,27 @@ class BaseRLModel(ABC): assert eval_env.num_envs == 1 return eval_env + # Type hint as string to avoid circular import + def _init_callback(self, callback) -> 'BaseCallback': + """ + Note: we cannot use type hint here because of circular import. + + :param callback: (Union[callable, [BaseCallback], BaseCallback, None]) + :return: (BaseCallback) + """ + # Avoid circular import + from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback + + # 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) + + callback.init_callback(self) + return callback + def scale_action(self, action: np.ndarray) -> np.ndarray: """ Rescale the action from [low, high] to [-1, 1] @@ -206,7 +228,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 +252,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 @@ -252,7 +274,7 @@ class BaseRLModel(ABC): self.n_envs = env.num_envs self.env = env - def get_parameters(self) -> Tuple[Dict[str, Any], Dict[str, Any]]: + def get_parameters(self) -> Tuple[TensorDict, OptimizerStateDict]: """ Returns policy and optimizer parameters as a tuple @@ -260,7 +282,7 @@ class BaseRLModel(ABC): """ return self.get_policy_parameters(), self.get_opt_parameters() - def get_policy_parameters(self) -> Dict[str, Any]: + def get_policy_parameters(self) -> TensorDict: """ Get current model policy parameters as dictionary of variable name -> tensors. @@ -269,7 +291,7 @@ class BaseRLModel(ABC): return self.policy.state_dict() @abstractmethod - def get_opt_parameters(self)-> Dict[str, Any]: + def get_opt_parameters(self)-> OptimizerStateDict: """ Get current model optimizer parameters as dictionary of variable names -> tensors :return: (dict) Dictionary of variable name -> tensor of model's optimizer parameters @@ -280,7 +302,7 @@ class BaseRLModel(ABC): 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_env: Optional[GymEnv] = None, eval_freq: int = -1, n_eval_episodes: int = 5, reset_num_timesteps: bool = True): @@ -316,7 +338,7 @@ class BaseRLModel(ABC): """ raise NotImplementedError() - def load_parameters(self, load_dict: Dict[str, Any], opt_params: Dict[str, Any]) -> None: + def load_parameters(self, load_dict: TensorDict, opt_params: OptimizerStateDict) -> None: """ Load model parameters from a dictionary load_dict should contain all keys from torch.model.state_dict() @@ -325,14 +347,14 @@ class BaseRLModel(ABC): :param load_dict: dict of parameters from model.state_dict() - :param opt_params: dict of optimizer state_dicts should be handled in child_class + :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: 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 @@ -368,7 +390,8 @@ 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[OptimizerStateDict]]): """ Load model data from a .zip archive :param load_path: Where to load the model from @@ -450,12 +473,14 @@ class BaseRLModel(ABC): if self.eval_env is not None: self.eval_env.seed(seed) - def _setup_learn(self, eval_env): + def _setup_learn(self, eval_env: Optional[GymEnv], callback=None) -> (Tuple[int, int, + List[Any], np.ndarray, Optional[VecEnv], Any]): """ 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]]) + :return: (int, int, [float], np.ndarray, VecEnv, BaseCallback) """ self.start_time = time.time() self.ep_info_buffer = deque(maxlen=100) @@ -463,6 +488,8 @@ class BaseRLModel(ABC): if self.action_noise is not None: self.action_noise.reset() + callback = self._init_callback(callback) + timesteps_since_eval, episode_num = 0, 0 evaluations = [] @@ -470,10 +497,11 @@ class BaseRLModel(ABC): 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() # type: GymEnv - def _update_info_buffer(self, infos): + return timesteps_since_eval, episode_num, evaluations, obs, eval_env, 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. @@ -485,11 +513,19 @@ 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, + num_timesteps: 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. @@ -499,7 +535,7 @@ 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) @@ -524,6 +560,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 @@ -531,6 +570,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() @@ -650,11 +695,13 @@ 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, - params: Dict[str, Any] = None, opt_params: Dict[str, Any] = None) -> None: + params: TensorDict = None, opt_params: OptimizerStateDict = None) -> None: """ Save model to a zip archive. @@ -730,7 +777,7 @@ 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: int, eval_env: int, n_eval_episodes: int, + def _eval_policy(self, eval_freq: int, eval_env: GymEnv, n_eval_episodes: int, timesteps_since_eval: int, render: bool = False, deterministic: bool = True) -> int: """ Evaluate the current policy on a test environment. diff --git a/torchy_baselines/common/callbacks.py b/torchy_baselines/common/callbacks.py new file mode 100644 index 0000000..061d6e5 --- /dev/null +++ b/torchy_baselines/common/callbacks.py @@ -0,0 +1,320 @@ +import os +from abc import ABC, abstractmethod +from typing import Union, List, Dict, Any, Optional + +import gym +import numpy as np + +from torchy_baselines.common.base_class import BaseRLModel # pytype: disable=pyi-error +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 + + +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] + + 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) + self.callback.init_callback(self.model) + + def _on_training_start(self) -> 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 + + +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 on_step: (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 log file (.npz) where the evaluations + 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 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, + 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 + 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 + self.log_path = log_path + self.evaluations_results = [] + self.evaluations_timesteps = [] + + 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.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, _ = evaluate_policy(self.model, self.eval_env, n_eval_episodes=self.n_eval_episodes, + 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) + np.savez(self.log_path, timesteps=self.evaluations_timesteps, results=self.evaluations_results) + + mean_reward, std_reward = np.mean(episode_rewards), np.std(episode_rewards) + if self.verbose > 0: + print(f"Eval num_timesteps={self.num_timesteps}, " + f"episode_reward={mean_reward:.2f} +/- {std_reward:.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/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 77a47cf..77c19b4 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -1,5 +1,6 @@ import os import time +from typing import Optional, Tuple import gym from gym import spaces @@ -16,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 @@ -149,17 +152,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) @@ -185,7 +201,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,20 +286,21 @@ class PPO(BaseRLModel): 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): - timesteps_since_eval, iteration, evaluations, obs, eval_env = self._setup_learn(eval_env) + timesteps_since_eval, iteration, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback) 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 @@ -308,6 +327,8 @@ class PPO(BaseRLModel): # 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_opt_parameters(self): diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 1e682cf..3e2b96f 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -259,25 +259,24 @@ class SAC(BaseRLModel): eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="SAC", reset_num_timesteps=True): - timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env) + timesteps_since_eval, episode_num, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback) + + 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 @@ -292,6 +291,7 @@ class SAC(BaseRLModel): 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 get_opt_parameters(self): 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 32cc1b2..4c0075d 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -251,25 +251,25 @@ class TD3(BaseRLModel): 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): - timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env) + timesteps_since_eval, episode_num, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback) + + 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 @@ -294,6 +294,8 @@ class TD3(BaseRLModel): 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 get_opt_parameters(self): From d514cd9126842b4f83f855be0fc8d1f73e540819 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Mon, 27 Jan 2020 14:36:11 +0100 Subject: [PATCH 2/9] Add templates --- .github/ISSUE_TEMPLATE/issue-template.md | 52 ++++++++++++++++++++++++ .github/PULL_REQUEST_TEMPLATE.md | 29 +++++++++++++ 2 files changed, 81 insertions(+) create mode 100644 .github/ISSUE_TEMPLATE/issue-template.md create mode 100644 .github/PULL_REQUEST_TEMPLATE.md 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. + + From a628354721e9b6560c5468eb82b28f154443d9b4 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Mon, 27 Jan 2020 15:53:27 +0100 Subject: [PATCH 3/9] Refactor evaluation --- torchy_baselines/a2c/a2c.py | 6 +- torchy_baselines/cem_rl/cem_rl.py | 24 +---- torchy_baselines/common/base_class.py | 147 ++++++++++++-------------- torchy_baselines/common/callbacks.py | 42 ++++++-- torchy_baselines/common/evaluation.py | 10 +- torchy_baselines/ppo/ppo.py | 17 +-- torchy_baselines/sac/sac.py | 10 +- torchy_baselines/td3/td3.py | 12 +-- 8 files changed, 125 insertions(+), 143 deletions(-) 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 1e26d15..2e3be4d 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -100,9 +100,10 @@ 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, callback = self._setup_learn(eval_env, callback) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) actor_steps = 0 continue_training = True @@ -155,21 +156,6 @@ 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: @@ -180,7 +166,6 @@ class CEMRL(TD3): n_steps=-1, action_noise=self.action_noise, 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) @@ -192,8 +177,6 @@ class CEMRL(TD3): break episode_num += n_episodes - self.num_timesteps += episode_timesteps - timesteps_since_eval += episode_timesteps actor_steps += episode_timesteps self.fitnesses.append(episode_reward) @@ -202,7 +185,6 @@ class CEMRL(TD3): 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() diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 5bc2abd..6cfdeb2 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,16 +13,13 @@ 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 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 -if typing.TYPE_CHECKING: - from torchy_baselines.common.callbacks import BaseCallback - class BaseRLModel(ABC): """ @@ -50,11 +46,12 @@ class BaseRLModel(ABC): :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[GymEnv, str], policy_base: Type[BasePolicy], - policy_kwargs : Dict[str, Any] = None, + policy_kwargs: Dict[str, Any] = None, verbose: int = 0, device: Union[th.device, str] = 'auto', support_multi_env: bool = False, @@ -133,9 +130,6 @@ class BaseRLModel(ABC): 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 @@ -146,34 +140,10 @@ class BaseRLModel(ABC): assert eval_env.num_envs == 1 return eval_env - # Type hint as string to avoid circular import - def _init_callback(self, callback) -> 'BaseCallback': - """ - Note: we cannot use type hint here because of circular import. - - :param callback: (Union[callable, [BaseCallback], BaseCallback, None]) - :return: (BaseCallback) - """ - # Avoid circular import - from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback - - # 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) - - callback.init_callback(self) - return callback - 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: - :return: """ low, high = self.action_space.low, self.action_space.high return 2.0 * ((action - low) / (high - low)) - 1.0 @@ -182,9 +152,6 @@ class BaseRLModel(ABC): """ Rescale the action from [-1, 1] to [low, high] (no need for symmetric action space) - - :param scaled_action: - :return: """ low, high = self.action_space.low, self.action_space.high return low + (0.5 * (scaled_action + 1.0) * (high - low)) @@ -291,7 +258,7 @@ class BaseRLModel(ABC): return self.policy.state_dict() @abstractmethod - def get_opt_parameters(self)-> OptimizerStateDict: + def get_opt_parameters(self) -> OptimizerStateDict: """ Get current model optimizer parameters as dictionary of variable names -> tensors :return: (dict) Dictionary of variable name -> tensor of model's optimizer parameters @@ -300,11 +267,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: 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. @@ -318,6 +287,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() @@ -391,7 +362,8 @@ class BaseRLModel(ABC): @staticmethod def _load_from_file(load_path: str, load_data: bool = True) -> (Tuple[Optional[Dict[str, Any]], - Optional[TensorDict], Optional[OptimizerStateDict]]): + Optional[TensorDict], + Optional[OptimizerStateDict]]): """ Load model data from a .zip archive :param load_path: Where to load the model from @@ -473,14 +445,53 @@ class BaseRLModel(ABC): if self.eval_env is not None: self.eval_env.seed(seed) - def _setup_learn(self, eval_env: Optional[GymEnv], callback=None) -> (Tuple[int, int, - List[Any], np.ndarray, Optional[VecEnv], Any]): + 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: + # Same folder as the rest + best_model_save_path = os.path.dirname(log_path) if log_path is not None else None + + eval_callback = EvalCallback(eval_env, + best_model_save_path=best_model_save_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 + ) -> Tuple[int, np.ndarray, BaseCallback]: """ Initialize different variables needed for training. :param eval_env: (Optional[GymEnv]) :param callback: (Union[None, BaseCallback, List[BaseCallback, Callable]]) - :return: (int, int, [float], np.ndarray, VecEnv, BaseCallback) + :param eval_freq: (int) + :param n_eval_episodes: (int) + :param log_path (Optional[str]): + :return: (Tuple[int, np.ndarray, BaseCallback]) """ self.start_time = time.time() self.ep_info_buffer = deque(maxlen=100) @@ -488,18 +499,18 @@ class BaseRLModel(ABC): if self.action_noise is not None: self.action_noise.reset() - callback = self._init_callback(callback) - timesteps_since_eval, episode_num = 0, 0 - evaluations = [] 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: GymEnv + obs = self.env.reset() - return timesteps_since_eval, episode_num, evaluations, obs, eval_env, callback + # 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: """ @@ -521,7 +532,6 @@ class BaseRLModel(ABC): action_noise: Optional[ActionNoise] = None, deterministic: bool = False, learning_starts: int = 0, - num_timesteps: int = 0, replay_buffer=None, obs: Optional[np.ndarray] = None, episode_num: int = 0, @@ -537,7 +547,6 @@ class BaseRLModel(ABC): :param deterministic: (bool) :param callback: (BaseCallback) :param learning_starts: (int) - :param num_timesteps: (int) :param replay_buffer: (ReplayBuffer) :param obs: (np.ndarray) :param episode_num: (int) @@ -583,7 +592,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: @@ -642,7 +651,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: @@ -658,8 +667,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])) @@ -667,7 +676,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() @@ -701,7 +710,7 @@ class BaseRLModel(ABC): @staticmethod def _save_to_file_zip(save_path: str, data: Dict[str, Any] = None, - params: TensorDict = None, opt_params: OptimizerStateDict = None) -> None: + params: TensorDict = None, opt_params: OptimizerStateDict = None) -> None: """ Save model to a zip archive. @@ -776,27 +785,3 @@ class BaseRLModel(ABC): params_to_save = self.get_policy_parameters() 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: int, eval_env: GymEnv, 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 index 061d6e5..89b1e22 100644 --- a/torchy_baselines/common/callbacks.py +++ b/torchy_baselines/common/callbacks.py @@ -1,15 +1,18 @@ 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.base_class import BaseRLModel # pytype: disable=pyi-error 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): """ @@ -31,7 +34,8 @@ class BaseCallback(ABC): # to have access to the parent object self.parent = None # type: Optional[BaseCallback] - def init_callback(self, model: BaseRLModel) -> None: + # 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. @@ -105,18 +109,23 @@ class EventCallback(BaseCallback): if callback is not None: self.callback.parent = self - def init_callback(self, model: BaseRLModel) -> None: + def init_callback(self, model: 'BaseRLModel') -> None: super(EventCallback, self).init_callback(model) - self.callback.init_callback(self.model) + if self.callback is not None: + self.callback.init_callback(self.model) def _on_training_start(self) -> None: - self.callback.on_training_start(self.locals, self.globals) + 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]): @@ -179,7 +188,7 @@ class ConvertCallback(BaseCallback): """ Convert functional callback (old-style) to object. - :param on_step: (callable) + :param callback: (callable) :param verbose: (int) """ def __init__(self, callback, verbose=0): @@ -207,6 +216,7 @@ class EvalCallback(EventCallback): 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], @@ -216,12 +226,15 @@ class EvalCallback(EventCallback): 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" @@ -230,6 +243,7 @@ class EvalCallback(EventCallback): 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 @@ -244,22 +258,30 @@ class EvalCallback(EventCallback): def _on_step(self) -> bool: - if self.n_calls % self.eval_freq == 0: + 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, _ = evaluate_policy(self.model, self.eval_env, n_eval_episodes=self.n_eval_episodes, - deterministic=self.deterministic, return_episode_rewards=True) + 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) - np.savez(self.log_path, timesteps=self.evaluations_timesteps, results=self.evaluations_results) + 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: 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/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 77c19b4..39d165e 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -193,6 +193,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) @@ -284,9 +286,11 @@ 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, callback = self._setup_learn(eval_env, callback) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) + 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)) @@ -295,15 +299,15 @@ class PPO(BaseRLModel): while self.num_timesteps < total_timesteps: - obs, continue_training = self.collect_rollouts(self.env, callback, self.rollout_buffer, n_rollout_steps=self.n_steps, + 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 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 @@ -320,9 +324,6 @@ 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) diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 3e2b96f..ec24d86 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -257,9 +257,9 @@ 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, callback = self._setup_learn(eval_env, callback) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) callback.on_training_start(locals(), globals()) @@ -268,7 +268,6 @@ class SAC(BaseRLModel): n_steps=self.train_freq, action_noise=self.action_noise, 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) @@ -278,9 +277,7 @@ class SAC(BaseRLModel): 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: @@ -288,9 +285,6 @@ 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 diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 4c0075d..9465d96 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -249,9 +249,10 @@ 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, callback = self._setup_learn(eval_env, callback) + episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) callback.on_training_start(locals(), globals()) @@ -261,7 +262,6 @@ class TD3(BaseRLModel): n_steps=self.train_freq, action_noise=self.action_noise, 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) @@ -272,8 +272,6 @@ class TD3(BaseRLModel): 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: @@ -290,10 +288,6 @@ 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 From 98037352f5c1a3cf8776ddc51f10347e246f1caa Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Mon, 27 Jan 2020 15:57:34 +0100 Subject: [PATCH 4/9] Update changelog --- docs/misc/changelog.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 37e0f27..d38f7b2 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -9,6 +9,7 @@ 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: ^^^^^^^^^^^^^ @@ -26,6 +27,7 @@ Others: - Add type check - Converted all format string to f-strings - Add test for `OrnsteinUhlenbeckActionNoise` +- Add type aliases in `common.type_aliases` Documentation: ^^^^^^^^^^^^^^ From 5c94a225ef15583218a35cb25b1673eb1d46542c Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Tue, 28 Jan 2020 10:24:02 +0100 Subject: [PATCH 5/9] Minor edit to the doc --- Makefile | 5 ++++- docs/conf.py | 8 ++++++++ torchy_baselines/common/base_class.py | 5 ++++- 3 files changed, 16 insertions(+), 2 deletions(-) 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/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 6cfdeb2..bbcb951 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -210,7 +210,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 From 6ae842161508e4bfefa8f7b287346dac9eaeb896 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Tue, 28 Jan 2020 10:28:44 +0100 Subject: [PATCH 6/9] Update docstring --- torchy_baselines/common/base_class.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index bbcb951..b69e0aa 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -144,6 +144,8 @@ class BaseRLModel(ABC): """ Rescale the action from [low, high] to [-1, 1] (no need for symmetric action space) + + :param action: Action to scale """ low, high = self.action_space.low, self.action_space.high return 2.0 * ((action - low) / (high - low)) - 1.0 @@ -152,6 +154,8 @@ class BaseRLModel(ABC): """ Rescale the action from [-1, 1] to [low, high] (no need for symmetric action space) + + :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)) From 5d4e73544ca8af22310bb86129c325021065a579 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 31 Jan 2020 13:16:28 +0100 Subject: [PATCH 7/9] Fix `reset_num_timesteps` --- docs/misc/changelog.rst | 1 + torchy_baselines/cem_rl/cem_rl.py | 3 ++- torchy_baselines/common/base_class.py | 33 ++++++--------------------- torchy_baselines/ppo/ppo.py | 3 ++- torchy_baselines/sac/sac.py | 4 ++-- torchy_baselines/td3/td3.py | 3 ++- 6 files changed, 16 insertions(+), 31 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 917199c..d0d69c1 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -19,6 +19,7 @@ New Features: Bug Fixes: ^^^^^^^^^^ - Fix loading model on CPU that were trained on GPU +- Fix `reset_num_timesteps` that was not used Deprecations: ^^^^^^^^^^^^^ diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 2e3be4d..35c6e3c 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -103,7 +103,8 @@ class CEMRL(TD3): eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="CEMRL", eval_log_path=None, reset_num_timesteps=True): - episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) + 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 diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index e3575b0..1f711be 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -364,7 +364,7 @@ class BaseRLModel(ABC): @staticmethod def _load_from_file(load_path: str, load_data: bool = True) -> (Tuple[Optional[Dict[str, Any]], Optional[TensorDict], - Optional[OptimizerStateDict]]): + Optional[TensorDict]]): """ Load model data from a .zip archive :param load_path: Where to load the model from @@ -503,7 +503,8 @@ class BaseRLModel(ABC): callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None, eval_freq: int = 10000, n_eval_episodes: int = 5, - log_path: Optional[str] = None + log_path: Optional[str] = None, + reset_num_timesteps: bool = True, ) -> Tuple[int, np.ndarray, BaseCallback]: """ Initialize different variables needed for training. @@ -513,6 +514,7 @@ class BaseRLModel(ABC): :param eval_freq: (int) :param n_eval_episodes: (int) :param log_path (Optional[str]): + :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() @@ -523,6 +525,9 @@ class BaseRLModel(ABC): timesteps_since_eval, episode_num = 0, 0 + 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) @@ -828,27 +833,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/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 188dfca..c5916ea 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -290,7 +290,8 @@ class PPO(BaseRLModel): eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="PPO", eval_log_path=None, reset_num_timesteps=True): - episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) + 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: diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 03caf7e..907e921 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -261,8 +261,8 @@ class SAC(BaseRLModel): eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="SAC", eval_log_path=None, reset_num_timesteps=True): - episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) - + 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: diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index c39fdb7..1c2e417 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -254,7 +254,8 @@ class TD3(BaseRLModel): eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="TD3", eval_log_path=None, reset_num_timesteps=True): - episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, n_eval_episodes, eval_log_path) + 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()) From ec657cc34e31a9c7e4d91302a59924021e1b9ea2 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 31 Jan 2020 13:42:04 +0100 Subject: [PATCH 8/9] Fix tests and change `log_path` behavior for `EvalCallback` --- tests/test_callbacks.py | 36 ++++++++++++++++++---------- tests/test_save_load.py | 2 +- torchy_baselines/cem_rl/cem_rl.py | 1 - torchy_baselines/common/callbacks.py | 4 ++-- 4 files changed, 27 insertions(+), 16 deletions(-) diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index 533c548..96db45a 100644 --- a/tests/test_callbacks.py +++ b/tests/test_callbacks.py @@ -1,37 +1,49 @@ -import pytest +import os +import shutil -from torchy_baselines import SAC +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", [SAC]) +@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 - model = model_class('MlpPolicy', 'Pendulum-v0') + # 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='./logs/') + checkpoint_callback = CheckpointCallback(save_freq=1000, save_path=log_folder) - # For testing: use the same training env - eval_env = model.get_env() + 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='./logs/best_model', - log_path='./logs/results', eval_freq=100) + 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='./logs/', + checkpoint_on_event = CheckpointCallback(save_freq=1, save_path=log_folder, name_prefix='event') - event_callback = EveryNTimesteps(n_steps=1000, callback=checkpoint_on_event) + event_callback = EveryNTimesteps(n_steps=500, callback=checkpoint_on_event) callback = CallbackList([checkpoint_callback, eval_callback, event_callback]) - model.learn(1000, callback=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_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/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 35c6e3c..653b3e2 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -160,7 +160,6 @@ class CEMRL(TD3): 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, diff --git a/torchy_baselines/common/callbacks.py b/torchy_baselines/common/callbacks.py index 89b1e22..346f115 100644 --- a/torchy_baselines/common/callbacks.py +++ b/torchy_baselines/common/callbacks.py @@ -210,7 +210,7 @@ class EvalCallback(EventCallback): 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 log file (.npz) where the evaluations + :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. @@ -240,7 +240,7 @@ class EvalCallback(EventCallback): self.eval_env = eval_env self.best_model_save_path = best_model_save_path - self.log_path = log_path + self.log_path = os.path.join(log_path, 'evaluations') self.evaluations_results = [] self.evaluations_timesteps = [] self.evaluations_length = [] From 6710f1576cb3253279db5085cfb0c7021b7a5603 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 31 Jan 2020 13:48:25 +0100 Subject: [PATCH 9/9] Fix eval log path --- torchy_baselines/common/base_class.py | 7 ++----- torchy_baselines/common/callbacks.py | 5 ++++- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 1f711be..9a3222f 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -487,11 +487,8 @@ class BaseRLModel(ABC): # Create eval callback in charge of the evaluation if eval_env is not None: - # Same folder as the rest - best_model_save_path = os.path.dirname(log_path) if log_path is not None else None - eval_callback = EvalCallback(eval_env, - best_model_save_path=best_model_save_path, + 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]) @@ -513,7 +510,7 @@ class BaseRLModel(ABC): :param callback: (Union[None, BaseCallback, List[BaseCallback, Callable]]) :param eval_freq: (int) :param n_eval_episodes: (int) - :param log_path (Optional[str]): + :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]) """ diff --git a/torchy_baselines/common/callbacks.py b/torchy_baselines/common/callbacks.py index 346f115..2481299 100644 --- a/torchy_baselines/common/callbacks.py +++ b/torchy_baselines/common/callbacks.py @@ -240,7 +240,10 @@ class EvalCallback(EventCallback): self.eval_env = eval_env self.best_model_save_path = best_model_save_path - self.log_path = os.path.join(log_path, 'evaluations') + # 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 = []