diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 98d1a87..41bf91e 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,6 +3,37 @@ Changelog ========== +Pre-Release 0.5.0a1 (WIP) +------------------------------ + +Breaking Changes: +^^^^^^^^^^^^^^^^^ +- Previous loading of policy weights is broken and replace by the new saving/loading for policy + +New Features: +^^^^^^^^^^^^^ +- Added ``optimizer`` and ``optimizer_kwargs`` to ``policy_kwargs`` in order to easily + customizer optimizers +- Complete independent save/load for policies + + +Bug Fixes: +^^^^^^^^^^ +- Fixed ``reset_num_timesteps`` behavior, so ``env.reset()`` is not called if ``reset_num_timesteps=True`` +- Fixed ``squashed_output`` that was not pass to policy constructor for ``SAC`` and ``TD3`` (would result in scaled actions for unscaled action spaces) + +Deprecations: +^^^^^^^^^^^^^ + +Others: +^^^^^^^ +- Cleanup rollout return +- Added ``get_device`` util to manage PyTorch devices + +Documentation: +^^^^^^^^^^^^^^ + + Pre-Release 0.4.0 (2020-02-14) ------------------------------ @@ -32,10 +63,6 @@ Others: - Refactored action distributions -Documentation: -^^^^^^^^^^^^^^ - - Pre-Release 0.3.0 (2020-02-14) ------------------------------ @@ -57,9 +84,6 @@ Bug Fixes: - Fixed colors in ``results_plotter`` - Fix entropy computation (now summed over action dim) -Deprecations: -^^^^^^^^^^^^^ - Others: ^^^^^^^ - SAC with SDE now sample only one matrix @@ -106,9 +130,6 @@ Bug Fixes: - Fix entropy computation for squashed Gaussian (approximate it now) - Fix seeding when using multiple environments (different seed per env) -Deprecations: -^^^^^^^^^^^^^ - Others: ^^^^^^^ - Add type check @@ -125,25 +146,11 @@ Pre-Release 0.1.0 (2020-01-20) ------------------------------ **First Release: base algorithms and state-dependent exploration** -Breaking Changes: -^^^^^^^^^^^^^^^^^ - New Features: ^^^^^^^^^^^^^ - Initial release of A2C, CEM-RL, PPO, SAC and TD3, working only with ``Box`` input space - State-Dependent Exploration (SDE) for A2C, PPO, SAC and TD3 -Bug Fixes: -^^^^^^^^^^ - -Deprecations: -^^^^^^^^^^^^^ - -Others: -^^^^^^^ - -Documentation: -^^^^^^^^^^^^^^ Maintainers diff --git a/tests/test_custom_policy.py b/tests/test_custom_policy.py index 2555016..d2ffab0 100644 --- a/tests/test_custom_policy.py +++ b/tests/test_custom_policy.py @@ -1,6 +1,7 @@ import pytest +import torch as th -from torchy_baselines import PPO +from torchy_baselines import A2C, PPO, SAC, TD3 @pytest.mark.parametrize('net_arch', [ @@ -11,5 +12,22 @@ from torchy_baselines import PPO [12, dict(vf=[8], pi=[8, 4])], [12, dict(pi=[8])], ]) -def test_flexible_mlp(net_arch): - _ = PPO('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000) +@pytest.mark.parametrize('model_class', [A2C, PPO]) +def test_flexible_mlp(model_class, net_arch): + _ = model_class('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000) + + +@pytest.mark.parametrize('net_arch', [ + [4], + [4, 4], +]) +@pytest.mark.parametrize('model_class', [SAC, TD3]) +def test_custom_offpolicy(model_class, net_arch): + _ = model_class('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=net_arch)).learn(1000) + + +@pytest.mark.parametrize('model_class', [A2C, PPO, SAC, TD3]) +@pytest.mark.parametrize('optimizer_kwargs', [None, dict(weight_decay=0.0)]) +def test_custom_optimizer(model_class, optimizer_kwargs): + policy_kwargs = dict(optimizer=th.optim.AdamW, optimizer_kwargs=optimizer_kwargs, net_arch=[32]) + _ = model_class('MlpPolicy', 'Pendulum-v0', policy_kwargs=policy_kwargs).learn(1000) diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 3bc7d23..1e547c0 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -180,9 +180,11 @@ def test_save_load_policy(model_class): observations = observations.reshape(10, -1) policy = model.policy - actor = None + policy_class = policy.__class__ + actor, actor_class = None, None if model_class in [SAC, TD3]: actor = policy.actor + actor_class = actor.__class__ # Get dictionary of current parameters params = deepcopy(policy.state_dict()) @@ -207,14 +209,17 @@ def test_save_load_policy(model_class): selected_actions_actor, _ = actor.predict(observations, deterministic=True) # Save and load policy - policy.save("./logs/policy_weights.pkl") + policy.save("./logs/policy.pkl") # Save and load actor if actor is not None: - actor.save("./logs/actor_weights.pkl") + actor.save("./logs/actor.pkl") - policy.load("./logs/policy_weights.pkl") - if actor is not None: - actor.load("./logs/actor_weights.pkl") + del policy, actor + + + policy = policy_class.load("./logs/policy.pkl") + if actor_class is not None: + actor = actor_class.load("./logs/actor.pkl") # check if params are still the same after load new_params = policy.state_dict() @@ -227,12 +232,12 @@ def test_save_load_policy(model_class): new_selected_actions, _ = policy.predict(observations, deterministic=True) assert np.allclose(selected_actions, new_selected_actions, 1e-4) - if actor is not None: + if actor_class is not None: new_selected_actions_actor, _ = actor.predict(observations, deterministic=True) assert np.allclose(selected_actions_actor, new_selected_actions_actor, 1e-4) assert np.allclose(selected_actions_actor, new_selected_actions, 1e-4) # clear file from os - os.remove("./logs/policy_weights.pkl") - if actor is not None: - os.remove("./logs/actor_weights.pkl") + os.remove("./logs/policy.pkl") + if actor_class is not None: + os.remove("./logs/actor.pkl") diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index 634a941..78e9038 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/torchy_baselines/a2c/a2c.py @@ -81,19 +81,15 @@ class A2C(PPO): seed=seed, _init_setup_model=False) self.normalize_advantage = normalize_advantage - self.rms_prop_eps = rms_prop_eps - self.use_rms_prop = use_rms_prop + # Override PPO optimizer to match original implementation + if use_rms_prop and 'optimizer' not in self.policy_kwargs: + self.policy_kwargs['optimizer'] = th.optim.RMSprop + self.policy_kwargs['optimizer_kwargs'] = dict(alpha=0.99, eps=rms_prop_eps, + weight_decay=0) if _init_setup_model: self._setup_model() - def _setup_model(self) -> None: - super(A2C, self)._setup_model() - if self.use_rms_prop: - self.policy.optimizer = th.optim.RMSprop(self.policy.parameters(), - lr=self.lr_schedule(1), alpha=0.99, - eps=self.rms_prop_eps, weight_decay=0) - def train(self, gradient_steps: int, batch_size: Optional[int] = None) -> None: # Update optimizer learning rate self._update_learning_rate(self.policy.optimizer) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 3c36d49..0efd7d5 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -13,7 +13,7 @@ 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.utils import set_random_seed, get_schedule_fn, update_learning_rate, get_device from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, VecNormalize from torchy_baselines.common.save_util import data_to_json, json_to_data, recursive_getattr, recursive_setattr from torchy_baselines.common.type_aliases import GymEnv, TensorDict, RolloutReturn, MaybeCallback @@ -71,10 +71,7 @@ class BaseRLModel(ABC): else: self.policy_class = policy - if device == 'auto': - device = 'cuda' if th.cuda.is_available() else 'cpu' - - self.device = th.device(device) + self.device = get_device(device) if verbose > 0: print(f"Using {self.device} device") @@ -93,7 +90,11 @@ class BaseRLModel(ABC): self.start_time = None self.policy = None self.learning_rate = learning_rate - self.lr_schedule = None # type: Optional[Callable] + self.lr_schedule = None # type: Optional[Callable] + self._last_obs = None # type: Optional[np.ndarray] + # When using VecNormalize: + self._last_original_obs = None # type: Optional[np.ndarray] + self._episode_num = 0 # Used for SDE only self.use_sde = use_sde self.sde_sample_freq = sde_sample_freq @@ -383,7 +384,7 @@ class BaseRLModel(ABC): raise ValueError(f"Error: the file {load_path} could not be found") # set device to cpu if cuda is not available - device = th.device('cuda') if th.cuda.is_available() else th.device('cpu') + device = get_device() # Open the zip archive and load data try: @@ -486,7 +487,7 @@ class BaseRLModel(ABC): n_eval_episodes: int = 5, log_path: Optional[str] = None, reset_num_timesteps: bool = True, - ) -> Tuple[int, np.ndarray, BaseCallback]: + ) -> 'BaseCallback': """ Initialize different variables needed for training. @@ -496,7 +497,7 @@ class BaseRLModel(ABC): :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]) + :return: (BaseCallback) """ self.start_time = time.time() self.ep_info_buffer = deque(maxlen=100) @@ -505,21 +506,26 @@ class BaseRLModel(ABC): if self.action_noise is not None: self.action_noise.reset() - timesteps_since_eval, episode_num = 0, 0 - if reset_num_timesteps: self.num_timesteps = 0 + self._episode_num = 0 + + # Avoid resetting the environment when calling `.learn()` consecutive times + if reset_num_timesteps or self._last_obs is None: + self._last_obs = self.env.reset() + # Retrieve unnormalized observation for saving into the buffer + if self._vec_normalize_env is not None: + self._last_original_obs = self._vec_normalize_env.get_original_obs() 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() # Create eval callback if needed callback = self._init_callback(callback, eval_env, eval_freq, n_eval_episodes, log_path) - return episode_num, obs, callback + return callback def _update_info_buffer(self, infos: List[Dict[str, Any]], dones: Optional[np.ndarray] = None) -> None: """ @@ -744,8 +750,6 @@ class OffPolicyRLModel(BaseRLModel): action_noise: Optional[ActionNoise] = None, learning_starts: int = 0, replay_buffer: Optional[ReplayBuffer] = None, - obs: Optional[np.ndarray] = None, - episode_num: int = 0, log_interval: Optional[int] = None) -> RolloutReturn: """ Collect rollout using the current policy (and possibly fill the replay buffer) @@ -762,8 +766,6 @@ class OffPolicyRLModel(BaseRLModel): (and at the beginning and end of the rollout) :param learning_starts: (int) Number of steps before learning for the warm-up phase. :param replay_buffer: (ReplayBuffer) - :param obs: (np.ndarray) Last observation from the environment - :param episode_num: (int) Episode index :param log_interval: (int) Log data every `log_interval` episodes :return: (RolloutReturn) """ @@ -773,10 +775,6 @@ class OffPolicyRLModel(BaseRLModel): assert isinstance(env, VecEnv), "You must pass a VecEnv" assert env.num_envs == 1, "OffPolicyRLModel only support single environment" - # Retrieve unnormalized observation for saving into the buffer - if self._vec_normalize_env is not None: - obs_ = self._vec_normalize_env.get_original_obs() - self.rollout_data = None if self.use_sde: self.actor.reset_noise() @@ -804,7 +802,7 @@ class OffPolicyRLModel(BaseRLModel): else: # Note: we assume that the policy uses tanh to scale the action # We use non-deterministic action in the case of SAC, for TD3, it does not matter - unscaled_action, _ = self.predict(obs, deterministic=False) + unscaled_action, _ = self.predict(self._last_obs, deterministic=False) # Rescale the action from [low, high] to [-1, 1] scaled_action = self.policy.scale_action(unscaled_action) @@ -827,7 +825,7 @@ class OffPolicyRLModel(BaseRLModel): # Only stop training if return value is False, not when it is None. if callback.on_step() is False: - return RolloutReturn(0.0, total_steps, total_episodes, None, continue_training=False) + return RolloutReturn(0.0, total_steps, total_episodes, continue_training=False) episode_reward += reward @@ -842,25 +840,23 @@ class OffPolicyRLModel(BaseRLModel): reward_ = self._vec_normalize_env.get_original_reward() else: # Avoid changing the original ones - obs_, new_obs_, reward_ = obs, new_obs, reward + self._last_original_obs, new_obs_, reward_ = self._last_obs, new_obs, reward - replay_buffer.add(obs_, new_obs_, clipped_action, reward_, done) + replay_buffer.add(self._last_original_obs, new_obs_, clipped_action, reward_, done) if self.rollout_data is not None: # Assume only one env - self.rollout_data['observations'].append(obs[0].copy()) + self.rollout_data['observations'].append(self._last_obs[0].copy()) self.rollout_data['actions'].append(scaled_action[0].copy()) self.rollout_data['rewards'].append(reward[0].copy()) self.rollout_data['dones'].append(done[0].copy()) - obs_tensor = th.FloatTensor(obs).to(self.device) + obs_tensor = th.FloatTensor(self._last_obs).to(self.device) self.rollout_data['values'].append(self.vf_net(obs_tensor)[0].cpu().detach().numpy()) - obs = new_obs - # Save the true unnormalized observation - # otherwise obs_ = self._vec_normalize_env.unnormalize_obs(obs) - # is a good approximation + self._last_obs = new_obs + # Save the unnormalized observation if self._vec_normalize_env is not None: - obs_ = new_obs_ + self._last_original_obs = new_obs_ self.num_timesteps += 1 episode_timesteps += 1 @@ -870,16 +866,16 @@ class OffPolicyRLModel(BaseRLModel): if done: total_episodes += 1 + self._episode_num += 1 episode_rewards.append(episode_reward) total_timesteps.append(episode_timesteps) if action_noise is not None: action_noise.reset() # Display training infos - if self.verbose >= 1 and log_interval is not None and ( - episode_num + total_episodes) % log_interval == 0: + if self.verbose >= 1 and log_interval is not None and (self._episode_num) % log_interval == 0: fps = int(self.num_timesteps / (time.time() - self.start_time)) - logger.logkv("episodes", episode_num + total_episodes) + logger.logkv("episodes", self._episode_num) 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])) logger.logkv('ep_len_mean', self.safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer])) @@ -909,7 +905,7 @@ class OffPolicyRLModel(BaseRLModel): for step in reversed(range(len(self.rollout_data['rewards']))): if step == len(self.rollout_data['rewards']) - 1: next_non_terminal = 1.0 - done[0] - next_value = self.vf_net(th.FloatTensor(obs).to(self.device))[0].detach() + next_value = self.vf_net(th.FloatTensor(self._last_obs).to(self.device))[0].detach() last_return = self.rollout_data['rewards'][step] + next_non_terminal * next_value else: next_non_terminal = 1.0 - self.rollout_data['dones'][step + 1] @@ -919,4 +915,4 @@ class OffPolicyRLModel(BaseRLModel): callback.on_rollout_end() - return RolloutReturn(mean_reward, total_steps, total_episodes, obs, continue_training) + return RolloutReturn(mean_reward, total_steps, total_episodes, continue_training) diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 58661f1..9aff338 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -1,4 +1,4 @@ -from typing import Union, Type, Dict, List, Tuple, Optional +from typing import Union, Type, Dict, List, Tuple, Optional, Any from itertools import zip_longest @@ -8,6 +8,7 @@ import torch.nn as nn import numpy as np from torchy_baselines.common.preprocessing import preprocess_obs +from torchy_baselines.common.utils import get_device, get_schedule_fn class BasePolicy(nn.Module): @@ -18,7 +19,7 @@ class BasePolicy(nn.Module): :param action_space: (gym.spaces.Space) The action space of the environment :param device: (Union[th.device, str]) Device on which the code should run. :param squash_output: (bool) For continuous actions, whether the output is squashed - or not using a `tanh()` function. + or not using a ``tanh()`` function. :param features_extractor: (nn.Module) Network to extract features (a CNN when using images, a nn.Flatten() layer otherwise) :param normalize_images: (bool) Whether to normalize images or not, @@ -26,17 +27,18 @@ class BasePolicy(nn.Module): """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, - device: Union[th.device, str] = 'cpu', + device: Union[th.device, str] = 'auto', squash_output: bool = False, features_extractor: Optional[nn.Module] = None, normalize_images: bool = True): super(BasePolicy, self).__init__() self.observation_space = observation_space self.action_space = action_space - self.device = device + self.device = get_device(device) self.features_extractor = features_extractor self.normalize_images = normalize_images self._squash_output = squash_output + self.optimizer = None # type: Optional[th.optim.Optimizer] def extract_features(self, obs: th.Tensor) -> th.Tensor: """ @@ -55,11 +57,19 @@ class BasePolicy(nn.Module): return self._squash_output @staticmethod - def init_weights(module: nn.Module, gain: float = 1): + def init_weights(module: nn.Module, gain: float = 1) -> None: + """ + Orthogonal initialization (used in PPO and A2C) + """ if isinstance(module, nn.Linear): nn.init.orthogonal_(module.weight, gain=gain) module.bias.data.fill_(0.0) + @staticmethod + def _dummy_schedule(progress: float) -> float: + """ (float) Useful for pickling policy.""" + return 0.0 + def forward(self, *_args, **kwargs): raise NotImplementedError() @@ -189,24 +199,47 @@ class BasePolicy(nn.Module): raise ValueError("Error: Cannot determine if the observation is vectorized with the space type {}." .format(observation_space)) + def _get_data(self) -> Dict[str, Any]: + """ + Get data that need to be saved in order to re-create the policy. + This corresponds to the arguments of the constructor. + + :return: (Dict[str, Any]) + """ + return dict( + observation_space=self.observation_space, + action_space=self.action_space, + # Passed to the constructor by child class + # squash_output=self.squash_output, + # features_extractor=self.features_extractor + normalize_images=self.normalize_images, + ) def save(self, path: str) -> None: """ - Save policy weights to a given location. - NOTE: we don't save policy parameters + Save policy to a given location. :param path: (str) """ - th.save(self.state_dict(), path) + th.save({'state_dict': self.state_dict(), 'data': self._get_data()}, path) - def load(self, path: str) -> None: + @classmethod + def load(cls, path: str, device: Union[th.device, str] = 'auto') -> 'BasePolicy': """ - Load policy weights from path. - NOTE: we don't load policy parameters + Load policy from path. :param path: (str) + :param device: ( Union[th.device, str]) Device on which the policy should be loaded. + :return: (BasePolicy) """ - self.load_state_dict(th.load(path)) + device = get_device(device) + saved_variables = th.load(path, map_location=device) + # Create policy object + model = cls(**saved_variables['data']) + # Load weights + model.load_state_dict(saved_variables['state_dict']) + model.to(device) + return model def load_from_vector(self, vector: np.ndarray): """ @@ -358,8 +391,9 @@ class MlpExtractor(nn.Module): def __init__(self, feature_dim: int, net_arch: List[Union[int, Dict[str, List[int]]]], activation_fn: Type[nn.Module], - device: Union[th.device, str] = 'cpu'): + device: Union[th.device, str] = 'auto'): super(MlpExtractor, self).__init__() + device = get_device(device) shared_net, policy_net, value_net = [], [], [] policy_only_layers = [] # Layer sizes of the network that only belongs to the policy network diff --git a/torchy_baselines/common/type_aliases.py b/torchy_baselines/common/type_aliases.py index d116ac6..db6b453 100644 --- a/torchy_baselines/common/type_aliases.py +++ b/torchy_baselines/common/type_aliases.py @@ -38,5 +38,4 @@ class RolloutReturn(NamedTuple): episode_reward: float episode_timesteps: int n_episodes: int - obs: Optional[np.ndarray] continue_training: bool diff --git a/torchy_baselines/common/utils.py b/torchy_baselines/common/utils.py index 1da28c2..b6fbb64 100644 --- a/torchy_baselines/common/utils.py +++ b/torchy_baselines/common/utils.py @@ -84,3 +84,26 @@ def constant_fn(val: float) -> Callable: return val return func + + +def get_device(device: Union[th.device, str] = 'auto') -> th.device: + """ + Retrieve PyTorch device. + It checks that the requested device is available first. + For now, it supports only cpu and cuda. + By default, it tries to use the gpu. + + :param device: (Union[str, th.device]) One for 'auto', 'cuda', 'cpu' + :return: (th.device) + """ + # Cuda by default + if device == 'auto': + device = 'cuda' + # Force conversion to th.device + device = th.device(device) + + # Cuda not available + if device == th.device('cuda') and not th.cuda.is_available(): + return th.device('cpu') + + return device diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 5ea14a7..44754a6 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -1,4 +1,4 @@ -from typing import Optional, List, Tuple, Callable, Union, Dict, Type +from typing import Optional, List, Tuple, Callable, Union, Dict, Type, Any from functools import partial import gym @@ -24,7 +24,6 @@ class PPOPolicy(BasePolicy): :param net_arch: ([int or dict]) The specification of the policy and value networks. :param device: (str or th.device) Device on which the code should run. :param activation_fn: (Type[nn.Module]) Activation function - :param adam_epsilon: (float) Small values to avoid NaN in ADAM optimizer :param ortho_init: (bool) Whether to use or not orthogonal initialization :param use_sde: (bool) Whether to use State Dependent Exploration or not :param log_std_init: (float) Initial value for the log standard deviation @@ -40,15 +39,18 @@ class PPOPolicy(BasePolicy): this allows to ensure boundaries when using SDE. :param normalize_images: (bool) Whether to normalize images or not, dividing by 255.0 (True by default) + :param optimizer: (Type[th.optim.Optimizer]) The optimizer to use, + ``th.optim.Adam`` by default + :param optimizer_kwargs: (Optional[Dict[str, Any]]) Additional keyword arguments, + excluding the learning rate, to pass to the optimizer """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[Union[int, Dict[str, List[int]]]]] = None, - device: Union[th.device, str] = 'cpu', + device: Union[th.device, str] = 'auto', activation_fn: Type[nn.Module] = nn.Tanh, - adam_epsilon: float = 1e-5, ortho_init: bool = True, use_sde: bool = False, log_std_init: float = 0.0, @@ -56,7 +58,9 @@ class PPOPolicy(BasePolicy): sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, squash_output: bool = False, - normalize_images: bool = True): + normalize_images: bool = True, + optimizer: Type[th.optim.Optimizer] = th.optim.Adam, + optimizer_kwargs: Optional[Dict[str, Any]] = None): super(PPOPolicy, self).__init__(observation_space, action_space, device, squash_output=squash_output) # Default network architecture, from stable-baselines @@ -65,7 +69,15 @@ class PPOPolicy(BasePolicy): self.net_arch = net_arch self.activation_fn = activation_fn - self.adam_epsilon = adam_epsilon + + if optimizer_kwargs is None: + optimizer_kwargs = {} + # Small values to avoid NaN in ADAM optimizer + if optimizer == th.optim.Adam: + optimizer_kwargs['eps'] = 1e-5 + + self.optimizer_class = optimizer + self.optimizer_kwargs = optimizer_kwargs self.ortho_init = ortho_init # In the future, feature_extractor will be replaced with a CNN self.features_extractor = nn.Flatten() @@ -85,12 +97,32 @@ class PPOPolicy(BasePolicy): self.sde_features_extractor = None self.sde_net_arch = sde_net_arch self.use_sde = use_sde + self.dist_kwargs = dist_kwargs # Action distribution self.action_dist = make_proba_distribution(action_space, use_sde=use_sde, dist_kwargs=dist_kwargs) self._build(lr_schedule) + def _get_data(self) -> Dict[str, Any]: + data = super()._get_data() + + data.update(dict( + net_arch=self.net_arch, + activation_fn=self.activation_fn, + use_sde=self.use_sde, + log_std_init=self.log_std_init, + squash_output=self.dist_kwargs['squash_output'] if self.dist_kwargs else None, + full_std=self.dist_kwargs['full_std'] if self.dist_kwargs else None, + sde_net_arch=self.dist_kwargs['sde_net_arch'] if self.dist_kwargs else None, + use_expln=self.dist_kwargs['use_expln'] if self.dist_kwargs else None, + lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone + optimizer=self.optimizer_class, + optimizer_kwargs=self.optimizer_kwargs, + ortho_init=self.ortho_init + )) + return data + def reset_noise(self, n_envs: int = 1) -> None: """ Sample new weights for the exploration matrix. @@ -142,7 +174,7 @@ class PPOPolicy(BasePolicy): }[module] module.apply(partial(self.init_weights, gain=gain)) # Setup optimizer with initial learning rate - self.optimizer = th.optim.Adam(self.parameters(), lr=lr_schedule(1), eps=self.adam_epsilon) + self.optimizer = self.optimizer_class(self.parameters(), lr=lr_schedule(1), **self.optimizer_kwargs) def forward(self, obs: th.Tensor, deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, th.Tensor]: diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 51473af..e9d8bdf 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -141,12 +141,10 @@ class PPO(BaseRLModel): env: VecEnv, callback: BaseCallback, rollout_buffer: RolloutBuffer, - n_rollout_steps: int = 256, - obs: Optional[np.ndarray] = None) -> Tuple[Optional[np.ndarray], bool]: + n_rollout_steps: int = 256) -> bool: - assert obs is not None, "No previous observation was provided" + assert self._last_obs is not None, "No previous observation was provided" n_steps = 0 - continue_training = True rollout_buffer.reset() # Sample new weights for the state dependent exploration if self.use_sde: @@ -162,7 +160,7 @@ class PPO(BaseRLModel): with th.no_grad(): # Convert to pytorch tensor - obs_tensor = th.as_tensor(obs).to(self.device) + obs_tensor = th.as_tensor(self._last_obs).to(self.device) actions, values, log_probs = self.policy.forward(obs_tensor) actions = actions.cpu().numpy() @@ -171,11 +169,11 @@ class PPO(BaseRLModel): # Clip the actions to avoid out of bound error if isinstance(self.action_space, gym.spaces.Box): clipped_actions = np.clip(actions, self.action_space.low, self.action_space.high) + new_obs, rewards, dones, infos = env.step(clipped_actions) if callback.on_step() is False: - continue_training = False - return None, continue_training + return False self._update_info_buffer(infos) n_steps += 1 @@ -184,14 +182,14 @@ class PPO(BaseRLModel): if isinstance(self.action_space, gym.spaces.Discrete): # Reshape in case of discrete action actions = actions.reshape(-1, 1) - rollout_buffer.add(obs, actions, rewards, dones, values, log_probs) - obs = new_obs + rollout_buffer.add(self._last_obs, actions, rewards, dones, values, log_probs) + self._last_obs = new_obs rollout_buffer.compute_returns_and_advantage(values, dones=dones) callback.on_rollout_end() - return obs, continue_training + return True def train(self, n_epochs: int, batch_size: int = 64) -> None: # Update optimizer learning rate @@ -307,9 +305,9 @@ class PPO(BaseRLModel): eval_log_path: Optional[str] = None, reset_num_timesteps: bool = True) -> 'PPO': - episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, - n_eval_episodes, eval_log_path, reset_num_timesteps) iteration = 0 + callback = self._setup_learn(eval_env, callback, eval_freq, + n_eval_episodes, eval_log_path, reset_num_timesteps) # 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)) @@ -318,10 +316,9 @@ 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=obs) + continue_training = self.collect_rollouts(self.env, callback, + self.rollout_buffer, + n_rollout_steps=self.n_steps) if continue_training is False: break diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index e83292f..c0fd3a6 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -1,4 +1,4 @@ -from typing import Optional, List, Tuple, Callable, Union, Type, Dict +from typing import Optional, List, Tuple, Callable, Union, Type, Dict, Any import gym import torch as th @@ -53,11 +53,12 @@ class Actor(BasePolicy): use_expln: bool = False, clip_mean: float = 2.0, normalize_images: bool = True, - device: Union[th.device, str] = 'cpu'): + device: Union[th.device, str] = 'auto'): super(Actor, self).__init__(observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, - device=device) + device=device, + squash_output=True) action_dim = get_action_dim(self.action_space) @@ -65,6 +66,15 @@ class Actor(BasePolicy): self.latent_pi = nn.Sequential(*latent_pi_net) self.use_sde = use_sde self.sde_features_extractor = None + self.sde_net_arch = sde_net_arch + self.net_arch = net_arch + self.features_dim = features_dim + self.activation_fn = activation_fn + self.log_std_init = log_std_init + self.sde_net_arch = sde_net_arch + self.use_expln = use_expln + self.full_std = full_std + self.clip_mean = clip_mean if self.use_sde: latent_sde_dim = net_arch[-1] @@ -87,6 +97,23 @@ class Actor(BasePolicy): self.mu = nn.Linear(net_arch[-1], action_dim) self.log_std = nn.Linear(net_arch[-1], action_dim) + def _get_data(self) -> Dict[str, Any]: + data = super()._get_data() + + data.update(dict( + net_arch=self.net_arch, + features_dim=self.features_dim, + activation_fn=self.activation_fn, + use_sde=self.use_sde, + log_std_init=self.log_std_init, + full_std=self.full_std, + sde_net_arch=self.sde_net_arch, + use_expln=self.use_expln, + features_extractor=self.features_extractor, + clip_mean=self.clip_mean + )) + return data + def get_std(self) -> th.Tensor: """ Retrieve the standard deviation of the action distribution. @@ -171,7 +198,7 @@ class Critic(BasePolicy): features_dim: int, activation_fn: Type[nn.Module] = nn.ReLU, normalize_images: bool = True, - device: Union[th.device, str] = 'cpu'): + device: Union[th.device, str] = 'auto'): super(Critic, self).__init__(observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, @@ -214,24 +241,36 @@ class SACPolicy(BasePolicy): :param clip_mean: (float) Clip the mean output when using SDE to avoid numerical instability. :param normalize_images: (bool) Whether to normalize images or not, dividing by 255.0 (True by default) + :param optimizer: (Type[th.optim.Optimizer]) The optimizer to use, + ``th.optim.Adam`` by default + :param optimizer_kwargs: (Optional[Dict[str, Any]]) Additional keyword arguments, + excluding the learning rate, to pass to the optimizer """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[int]] = None, - device: Union[th.device, str] = 'cpu', + device: Union[th.device, str] = 'auto', activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, clip_mean: float = 2.0, - normalize_images: bool = True): + normalize_images: bool = True, + optimizer: Type[th.optim.Optimizer] = th.optim.Adam, + optimizer_kwargs: Optional[Dict[str, Any]] = None): super(SACPolicy, self).__init__(observation_space, action_space, device, squash_output=True) if net_arch is None: net_arch = [256, 256] + if optimizer_kwargs is None: + optimizer_kwargs = {} + + self.optimizer_class = optimizer + self.optimizer_kwargs = optimizer_kwargs + # In the future, features_extractor will be replaced with a CNN self.features_extractor = nn.Flatten() self.features_dim = get_obs_dim(self.observation_space) @@ -264,12 +303,31 @@ class SACPolicy(BasePolicy): def _build(self, lr_schedule: Callable) -> None: self.actor = self.make_actor() - self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=lr_schedule(1)) + self.actor.optimizer = self.optimizer_class(self.actor.parameters(), lr=lr_schedule(1), + **self.optimizer_kwargs) self.critic = self.make_critic() self.critic_target = self.make_critic() self.critic_target.load_state_dict(self.critic.state_dict()) - self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=lr_schedule(1)) + self.critic.optimizer = self.optimizer_class(self.critic.parameters(), lr=lr_schedule(1), + **self.optimizer_kwargs) + + def _get_data(self) -> Dict[str, Any]: + data = super()._get_data() + + data.update(dict( + net_arch=self.net_args['net_arch'], + activation_fn=self.net_args['activation_fn'], + use_sde=self.actor_kwargs['use_sde'], + log_std_init=self.actor_kwargs['log_std_init'], + sde_net_arch=self.actor_kwargs['sde_net_arch'], + use_expln=self.actor_kwargs['use_expln'], + clip_mean=self.actor_kwargs['clip_mean'], + lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone + optimizer=self.optimizer_class, + optimizer_kwargs=self.optimizer_kwargs + )) + return data def make_actor(self) -> Actor: return Actor(**self.actor_kwargs).to(self.device) diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index d81bbbb..b37658e 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -256,8 +256,8 @@ class SAC(OffPolicyRLModel): eval_log_path: Optional[str] = None, reset_num_timesteps: bool = True) -> OffPolicyRLModel: - episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, - n_eval_episodes, eval_log_path, reset_num_timesteps) + 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: @@ -266,14 +266,11 @@ class SAC(OffPolicyRLModel): callback=callback, learning_starts=self.learning_starts, replay_buffer=self.replay_buffer, - obs=obs, episode_num=episode_num, log_interval=log_interval) if rollout.continue_training is False: break - obs = rollout.obs - episode_num += rollout.n_episodes self._update_current_progress(self.num_timesteps, total_timesteps) if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 994650c..c1ecbc8 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -1,4 +1,4 @@ -from typing import Optional, List, Tuple, Callable, Union, Type +from typing import Optional, List, Tuple, Callable, Union, Type, Any, Dict import gym import torch as th @@ -52,11 +52,12 @@ class Actor(BasePolicy): sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, normalize_images: bool = True, - device: Union[th.device, str] = 'cpu'): + device: Union[th.device, str] = 'auto'): super(Actor, self).__init__(observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, - device=device) + device=device, + squash_output=not use_sde) self.latent_pi, self.log_std = None, None self.weights_dist, self.exploration_mat = None, None @@ -65,6 +66,15 @@ class Actor(BasePolicy): self.sde_features_extractor = None self.features_extractor = features_extractor self.normalize_images = normalize_images + self.net_arch = net_arch + self.features_dim = features_dim + self.activation_fn = activation_fn + self.clip_noise = clip_noise + self.lr_sde = lr_sde + self.log_std_init = log_std_init + self.sde_net_arch = sde_net_arch + self.use_expln = use_expln + self.full_std = full_std action_dim = get_action_dim(self.action_space) @@ -88,13 +98,30 @@ class Actor(BasePolicy): log_std_init=log_std_init) # Squash output self.mu = nn.Sequential(action_net, nn.Tanh()) - self.clip_noise = clip_noise self.sde_optimizer = th.optim.Adam([self.log_std], lr=lr_sde) self.reset_noise() else: actor_net = create_mlp(features_dim, action_dim, net_arch, activation_fn, squash_output=True) self.mu = nn.Sequential(*actor_net) + def _get_data(self) -> Dict[str, Any]: + data = super()._get_data() + + data.update(dict( + net_arch=self.net_arch, + features_dim=self.features_dim, + activation_fn=self.activation_fn, + use_sde=self.use_sde, + log_std_init=self.log_std_init, + clip_noise=self.clip_noise, + lr_sde=self.lr_sde, + full_std=self.full_std, + sde_net_arch=self.sde_net_arch, + use_expln=self.use_expln, + features_extractor=self.features_extractor + )) + return data + def get_std(self) -> th.Tensor: """ Retrieve the standard deviation of the action distribution. @@ -179,7 +206,7 @@ class Critic(BasePolicy): features_dim: int, activation_fn: Type[nn.Module] = nn.ReLU, normalize_images: bool = True, - device: Union[th.device, str] = 'cpu'): + device: Union[th.device, str] = 'auto'): super(Critic, self).__init__(observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, @@ -259,12 +286,16 @@ class TD3Policy(BasePolicy): above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough. :param normalize_images: (bool) Whether to normalize images or not, dividing by 255.0 (True by default) + :param optimizer: (Type[th.optim.Optimizer]) The optimizer to use, + ``th.optim.Adam`` by default + :param optimizer_kwargs: (Optional[Dict[str, Any]]) Additional keyword arguments, + excluding the learning rate, to pass to the optimizer """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, lr_schedule: Callable, net_arch: Optional[List[int]] = None, - device: Union[th.device, str] = 'cpu', + device: Union[th.device, str] = 'auto', activation_fn: Type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, @@ -272,13 +303,21 @@ class TD3Policy(BasePolicy): lr_sde: float = 3e-4, sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, - normalize_images: bool = True): + normalize_images: bool = True, + optimizer: Type[th.optim.Optimizer] = th.optim.Adam, + optimizer_kwargs: Optional[Dict[str, Any]] = None): super(TD3Policy, self).__init__(observation_space, action_space, device, squash_output=True) # Default network architecture, from the original paper if net_arch is None: net_arch = [400, 300] + if optimizer_kwargs is None: + optimizer_kwargs = {} + + self.optimizer_class = optimizer + self.optimizer_kwargs = optimizer_kwargs + # In the future, features_extractor will be replaced with a CNN self.features_extractor = nn.Flatten() self.features_dim = get_obs_dim(self.observation_space) @@ -318,19 +357,37 @@ class TD3Policy(BasePolicy): self.actor = self.make_actor() self.actor_target = self.make_actor() self.actor_target.load_state_dict(self.actor.state_dict()) - self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=lr_schedule(1)) - + self.actor.optimizer = self.optimizer_class(self.actor.parameters(), lr=lr_schedule(1), + **self.optimizer_kwargs) self.critic = self.make_critic() self.critic_target = self.make_critic() self.critic_target.load_state_dict(self.critic.state_dict()) - self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=lr_schedule(1)) - + self.critic.optimizer = self.optimizer_class(self.critic.parameters(), lr=lr_schedule(1), + **self.optimizer_kwargs) if self.use_sde: self.vf_net = ValueFunction(self.observation_space, self.action_space, features_extractor=self.features_extractor, features_dim=self.features_dim) self.actor.sde_optimizer.add_param_group({'params': self.vf_net.parameters()}) # pytype: disable=attribute-error + def _get_data(self) -> Dict[str, Any]: + data = super()._get_data() + + data.update(dict( + net_arch=self.net_args['net_arch'], + activation_fn=self.net_args['activation_fn'], + use_sde=self.actor_kwargs['use_sde'], + log_std_init=self.actor_kwargs['log_std_init'], + clip_noise=self.actor_kwargs['clip_noise'], + lr_sde=self.actor_kwargs['lr_sde'], + sde_net_arch=self.actor_kwargs['sde_net_arch'], + use_expln=self.actor_kwargs['use_expln'], + lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone + optimizer=self.optimizer_class, + optimizer_kwargs=self.optimizer_kwargs + )) + return data + def reset_noise(self) -> None: return self.actor.reset_noise() diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index d39b9f5..09742df 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -235,8 +235,8 @@ class TD3(OffPolicyRLModel): eval_log_path: Optional[str] = None, reset_num_timesteps: bool = True) -> OffPolicyRLModel: - episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq, - n_eval_episodes, eval_log_path, reset_num_timesteps) + callback = self._setup_learn(eval_env, callback, eval_freq, + n_eval_episodes, eval_log_path, reset_num_timesteps) callback.on_training_start(locals(), globals()) @@ -247,14 +247,11 @@ class TD3(OffPolicyRLModel): callback=callback, learning_starts=self.learning_starts, replay_buffer=self.replay_buffer, - obs=obs, episode_num=episode_num, log_interval=log_interval) if rollout.continue_training is False: break - obs = rollout.obs - episode_num += rollout.n_episodes self._update_current_progress(self.num_timesteps, total_timesteps) if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: diff --git a/torchy_baselines/version.txt b/torchy_baselines/version.txt index 1d0ba9e..d443359 100644 --- a/torchy_baselines/version.txt +++ b/torchy_baselines/version.txt @@ -1 +1 @@ -0.4.0 +0.5.0a1