diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index cb5d7aa..5930131 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -23,6 +23,7 @@ Deprecations: Others: ^^^^^^^ - Refactor handling of observation and action spaces +- Refactored features extraction to have proper preprocessing Documentation: ^^^^^^^^^^^^^^ diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 54024a3..cda1866 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -78,7 +78,7 @@ class BaseRLModel(ABC): if verbose > 0: print(f"Using {self.device} device") - self.env = None # type: GymEnv + self.env = None # type: Optional[GymEnv] # get VecNormalize object if needed self._vec_normalize_env = unwrap_vec_normalize(env) self.verbose = verbose @@ -380,8 +380,10 @@ class BaseRLModel(ABC): :param state: (Optional[np.ndarray]) The last states (can be None, used in recurrent policies) :param mask: (Optional[np.ndarray]) The last masks (can be None, used in recurrent policies) :param deterministic: (bool) Whether or not to return deterministic actions. - :return: (Tuple[np.ndarray, Optional[np.ndarray]]) the model's action and the next state (used in recurrent policies) + :return: (Tuple[np.ndarray, Optional[np.ndarray]]) the model's action and the next state + (used in recurrent policies) """ + # TODO: move this block to BasePolicy # if state is None: # state = self.initial_state # if mask is None: @@ -390,9 +392,7 @@ class BaseRLModel(ABC): vectorized_env = self._is_vectorized_observation(observation, self.observation_space) observation = observation.reshape((-1,) + self.observation_space.shape) - # Convert to float pytorch - # TODO: replace with preprocessing - observation = th.as_tensor(observation).float().to(self.device) + observation = th.as_tensor(observation).to(self.device) with th.no_grad(): actions = self.policy.predict(observation, deterministic=deterministic) # Convert to numpy @@ -812,7 +812,7 @@ class OffPolicyRLModel(BaseRLModel): policy_kwargs, verbose, device, support_multi_env, create_eval_env, monitor_wrapper, seed, use_sde, sde_sample_freq) - self.buffer_size = buffer_size + self.buffer_size = buffer_size self.batch_size = batch_size self.learning_starts = learning_starts self.actor = None diff --git a/torchy_baselines/common/identity_env.py b/torchy_baselines/common/identity_env.py index d815220..aa724e1 100644 --- a/torchy_baselines/common/identity_env.py +++ b/torchy_baselines/common/identity_env.py @@ -46,7 +46,6 @@ class IdentityEnvBox(IdentityEnv): """ Identity environment for testing purposes - :param dim: (int) the size of the dimensions you want to learn :param low: (float) the lower bound of the box dim :param high: (float) the upper bound of the box dim :param eps: (float) the epsilon bound for correct value diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 751f49a..8711cfc 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 +from typing import Union, Type, Dict, List, Tuple, Optional from itertools import zip_longest @@ -7,6 +7,8 @@ import torch as th import torch.nn as nn import numpy as np +from torchy_baselines.common.preprocessing import preprocess_obs + class BasePolicy(nn.Module): """ @@ -17,17 +19,36 @@ class BasePolicy(nn.Module): :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. + :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, + dividing by 255.0 (True by default) """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, device: Union[th.device, str] = 'cpu', - squash_output: bool = False): + 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.features_extractor = features_extractor + self.normalize_images = normalize_images self._squash_output = squash_output + def extract_features(self, obs: th.Tensor) -> th.Tensor: + """ + Preprocess the observation if needed and extract features. + + :param obs: (th.Tensor) + :return: (th.Tensor) + """ + assert self.features_extractor is not None, 'No feature extractor was set' + preprocessed_obs = preprocess_obs(obs, self.observation_space, normalize_images=self.normalize_images) + return self.features_extractor(preprocessed_obs) + @property def squash_output(self) -> bool: """ (bool) Getter for squash_output.""" @@ -45,6 +66,10 @@ class BasePolicy(nn.Module): def predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor: """ Get the action according to the policy for a given observation. + + :param observation: (th.Tensor) + :param deterministic: (bool) Whether to use stochastic or deterministic actions + :return: (th.Tensor) Taken action according to the policy """ raise NotImplementedError() @@ -118,12 +143,12 @@ def create_mlp(input_dim: int, return modules -def create_sde_feature_extractor(features_dim: int, - sde_net_arch: List[int], - activation_fn: nn.Module) -> Tuple[nn.Sequential, int]: +def create_sde_features_extractor(features_dim: int, + sde_net_arch: List[int], + activation_fn: nn.Module) -> Tuple[nn.Sequential, int]: """ Create the neural network that will be used to extract features - for the SDE. + for the SDE exploration function. :param features_dim: (int) :param sde_net_arch: ([int]) @@ -135,8 +160,8 @@ def create_sde_feature_extractor(features_dim: int, sde_activation = activation_fn if len(sde_net_arch) > 0 else None latent_sde_net = create_mlp(features_dim, -1, sde_net_arch, activation_fn=sde_activation, squash_output=False) latent_sde_dim = sde_net_arch[-1] if len(sde_net_arch) > 0 else features_dim - sde_feature_extractor = nn.Sequential(*latent_sde_net) - return sde_feature_extractor, latent_sde_dim + sde_features_extractor = nn.Sequential(*latent_sde_net) + return sde_features_extractor, latent_sde_dim _policy_registry = dict() # type: Dict[Type[BasePolicy], Dict[str, Type[BasePolicy]]] diff --git a/torchy_baselines/common/preprocessing.py b/torchy_baselines/common/preprocessing.py index 2b06e6e..e117287 100644 --- a/torchy_baselines/common/preprocessing.py +++ b/torchy_baselines/common/preprocessing.py @@ -34,7 +34,7 @@ def is_image_space(observation_space: spaces.Space) -> bool: def preprocess_obs(obs: th.Tensor, observation_space: spaces.Space, - normalize_image: bool = True) -> th.Tensor: + normalize_images: bool = True) -> th.Tensor: """ Preprocess observation to be to a neural network. For images, it normalizes the values by dividing them by 255 (to have values in [0, 1]) @@ -42,14 +42,14 @@ def preprocess_obs(obs: th.Tensor, observation_space: spaces.Space, :param obs: (th.Tensor) Observation :param observation_space: (spaces.Space) - :param normalize_image: (bool) Whether to normalize images or not + :param normalize_images: (bool) Whether to normalize images or not (True by default) :return: (th.Tensor) """ if isinstance(observation_space, spaces.Box): - if is_image_space(observation_space) and normalize_image: - return obs / 255.0 - return obs + if is_image_space(observation_space) and normalize_images: + return obs.float() / 255.0 + return obs.float() elif isinstance(observation_space, spaces.Discrete): # One hot encoding and convert to float to avoid errors return F.one_hot(obs, num_classes=observation_space.n).float() @@ -69,7 +69,7 @@ def get_obs_shape(observation_space: spaces.Space) -> Tuple[int, ...]: return observation_space.shape elif isinstance(observation_space, spaces.Discrete): # Observation is an int - return (1,) + return 1, else: # TODO: Multidiscrete, Binary, MultiBinary, Tuple, Dict raise NotImplementedError() diff --git a/torchy_baselines/common/save_util.py b/torchy_baselines/common/save_util.py index 85fb6fd..a3e7e22 100644 --- a/torchy_baselines/common/save_util.py +++ b/torchy_baselines/common/save_util.py @@ -5,9 +5,8 @@ used to serialize data (class parameters) of model classes import json import base64 import functools -from typing import Dict, Any, Optional, Union +from typing import Dict, Any, Optional -import torch as th import cloudpickle import warnings diff --git a/torchy_baselines/common/type_aliases.py b/torchy_baselines/common/type_aliases.py index 53c152c..d116ac6 100644 --- a/torchy_baselines/common/type_aliases.py +++ b/torchy_baselines/common/type_aliases.py @@ -1,7 +1,6 @@ """ Common aliases for type hint """ -import typing from typing import Union, Dict, Any, NamedTuple, Optional, List, Callable import numpy as np diff --git a/torchy_baselines/common/vec_env/__init__.py b/torchy_baselines/common/vec_env/__init__.py index 10f9288..8748923 100644 --- a/torchy_baselines/common/vec_env/__init__.py +++ b/torchy_baselines/common/vec_env/__init__.py @@ -1,6 +1,6 @@ # flake8: noqa F401 import typing -from typing import Optional +from typing import Optional, Union from copy import deepcopy from torchy_baselines.common.vec_env.base_vec_env import (AlreadySteppingError, NotSteppingError, @@ -15,7 +15,7 @@ if typing.TYPE_CHECKING: from torchy_baselines.common.type_aliases import GymEnv -def unwrap_vec_normalize(env: 'GymEnv') -> Optional[VecNormalize]: +def unwrap_vec_normalize(env: Union['GymEnv', VecEnv]) -> Optional[VecNormalize]: """ :param env: (gym.Env) :return: (VecNormalize) diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index ef366ed..750cf21 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -8,7 +8,7 @@ import numpy as np from torchy_baselines.common.preprocessing import get_obs_dim from torchy_baselines.common.policies import (BasePolicy, register_policy, MlpExtractor, - create_sde_feature_extractor) + create_sde_features_extractor) from torchy_baselines.common.distributions import (make_proba_distribution, Distribution, DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution) @@ -38,6 +38,8 @@ class PPOPolicy(BasePolicy): above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough. :param squash_output: (bool) Whether to squash the output using a tanh function, 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) """ def __init__(self, observation_space: gym.spaces.Space, @@ -53,9 +55,9 @@ class PPOPolicy(BasePolicy): full_std: bool = True, sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, - squash_output: bool = False): + squash_output: bool = False, + normalize_images: bool = True): super(PPOPolicy, self).__init__(observation_space, action_space, device, squash_output=squash_output) - self.obs_dim = get_obs_dim(self.observation_space) # Default network architecture, from stable-baselines if net_arch is None: @@ -67,7 +69,8 @@ class PPOPolicy(BasePolicy): self.ortho_init = ortho_init # In the future, feature_extractor will be replaced with a CNN self.features_extractor = nn.Flatten() - self.features_dim = self.obs_dim + self.features_dim = get_obs_dim(self.observation_space) + self.normalize_images = normalize_images self.log_std_init = log_std_init dist_kwargs = None # Keyword arguments for SDE distribution @@ -79,7 +82,7 @@ class PPOPolicy(BasePolicy): 'learn_features': sde_net_arch is not None } - self.sde_feature_extractor = None + self.sde_features_extractor = None self.sde_net_arch = sde_net_arch self.use_sde = use_sde @@ -98,6 +101,12 @@ class PPOPolicy(BasePolicy): self.action_dist.sample_weights(self.log_std, batch_size=n_envs) def _build(self, lr_schedule: Callable) -> None: + """ + Create the networks and the optimizer. + + :param lr_schedule: (Callable) Learning rate schedule + lr_schedule(1) is the initial learning rate + """ self.mlp_extractor = MlpExtractor(self.features_dim, net_arch=self.net_arch, activation_fn=self.activation_fn, device=self.device) @@ -105,9 +114,9 @@ class PPOPolicy(BasePolicy): # Separate feature extractor for SDE if self.sde_net_arch is not None: - self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(self.features_dim, - self.sde_net_arch, - self.activation_fn) + self.sde_features_extractor, latent_sde_dim = create_sde_features_extractor(self.features_dim, + self.sde_net_arch, + self.activation_fn) if isinstance(self.action_dist, DiagGaussianDistribution): self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi, @@ -132,12 +141,20 @@ class PPOPolicy(BasePolicy): self.value_net: 1 }[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) - def forward(self, obs: th.Tensor, deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, th.Tensor]: - if not isinstance(obs, th.Tensor): - obs = th.FloatTensor(obs).to(self.device) + def forward(self, obs: th.Tensor, + deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, th.Tensor]: + """ + Forward pass in all the networks (actor and critic) + + :param obs: (th.Tensor) Observation + :param deterministic: (bool) Whether to sample or use deterministic actions + :return: (Tuple[th.Tensor, th.Tensor, th.Tensor]) action, value and log probability of the action + """ latent_pi, latent_vf, latent_sde = self._get_latent(obs) + # Evaluate the values for the given observations value = self.value_net(latent_vf) action, action_distribution = self._get_action_dist_from_latent(latent_pi, latent_sde=latent_sde, deterministic=deterministic) @@ -145,17 +162,35 @@ class PPOPolicy(BasePolicy): return action, value, log_prob def _get_latent(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]: - features = self.features_extractor(obs) + """ + Get the latent code (i.e., activations of the last layer of each network) + for the different networks. + + :param obs: (th.Tensor) Observation + :return: (Tuple[th.Tensor, th.Tensor, th.Tensor]) Latent codes + for the actor, the value function and for SDE function + """ + # Preprocess the observation if needed + features = self.extract_features(obs) latent_pi, latent_vf = self.mlp_extractor(features) # Features for sde latent_sde = latent_pi - if self.sde_feature_extractor is not None: - latent_sde = self.sde_feature_extractor(features) + if self.sde_features_extractor is not None: + latent_sde = self.sde_features_extractor(features) return latent_pi, latent_vf, latent_sde def _get_action_dist_from_latent(self, latent_pi: th.Tensor, latent_sde: Optional[th.Tensor] = None, deterministic: bool = False) -> Tuple[th.Tensor, Distribution]: + """ + Retrieve action and associated action distribution + given the latent codes. + + :param latent_pi: (th.Tensor) Latent code for the actor + :param latent_sde: (Optional[th.Tensor]) Latent code for the SDE exploration function + :param deterministic: (bool) Whether to sample or use deterministic actions + :return: (Tuple[th.Tensor, Distribution]) Action and action distribution + """ mean_actions = self.action_net(latent_pi) if isinstance(self.action_dist, DiagGaussianDistribution): @@ -172,6 +207,13 @@ class PPOPolicy(BasePolicy): raise ValueError('Invalid action distribution') def predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor: + """ + Get the action according to the policy for a given observation. + + :param observation: (th.Tensor) + :param deterministic: (bool) Whether to use stochastic or deterministic actions + :return: (th.Tensor) Taken action according to the policy + """ latent_pi, _, latent_sde = self._get_latent(observation) action, _ = self._get_action_dist_from_latent(latent_pi, latent_sde, deterministic=deterministic) return action diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 50efaf5..0aea585 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -1,4 +1,3 @@ -import os import time from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any @@ -8,10 +7,11 @@ import torch as th import torch.nn.functional as F # Check if tensorboard is available for pytorch -try: - from torch.utils.tensorboard import SummaryWriter -except ImportError: - SummaryWriter = None +# TODO: finish tensorboard integration +# try: +# from torch.utils.tensorboard import SummaryWriter +# except ImportError: +# SummaryWriter = None import numpy as np from torchy_baselines.common import logger @@ -144,6 +144,7 @@ class PPO(BaseRLModel): n_rollout_steps: int = 256, obs: Optional[np.ndarray] = None) -> Tuple[Optional[np.ndarray], bool]: + assert obs is not None, "No previous observation was provided" n_steps = 0 continue_training = True rollout_buffer.reset() @@ -160,7 +161,10 @@ class PPO(BaseRLModel): self.policy.reset_noise(env.num_envs) with th.no_grad(): - actions, values, log_probs = self.policy.forward(obs) + # Convert to pytorch tensor + obs_tensor = obs.reshape((-1,) + self.observation_space.shape) + obs_tensor = th.as_tensor(obs_tensor).to(self.device) + actions, values, log_probs = self.policy.forward(obs_tensor) actions = actions.cpu().numpy() # Rescale and perform action @@ -308,8 +312,8 @@ class PPO(BaseRLModel): n_eval_episodes, eval_log_path, reset_num_timesteps) iteration = 0 - if self.tensorboard_log is not None and SummaryWriter is not None: - self.tb_writer = SummaryWriter(log_dir=os.path.join(self.tensorboard_log, tb_log_name)) + # 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()) diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index 3467d23..5a90e62 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -6,7 +6,7 @@ import torch.nn as nn from torchy_baselines.common.preprocessing import get_action_dim, get_obs_dim from torchy_baselines.common.policies import (BasePolicy, register_policy, create_mlp, - create_sde_feature_extractor) + create_sde_features_extractor) from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution # CAP the standard deviation of the actor @@ -21,6 +21,9 @@ class Actor(BasePolicy): :param observation_space: (gym.spaces.Space) Obervation space :param action_space: (gym.spaces.Space) Action space :param net_arch: ([int]) Network architecture + :param features_extractor: (nn.Module) Network to extract features + (a CNN when using images, a nn.Flatten() layer otherwise) + :param features_dim: (int) Number of features :param activation_fn: (nn.Module) Activation function :param use_sde: (bool) Whether to use State Dependent Exploration or not :param log_std_init: (float) Initial value for the log standard deviation @@ -33,33 +36,39 @@ class Actor(BasePolicy): 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. :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) """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, net_arch: List[int], + features_extractor: nn.Module, + features_dim: int, activation_fn: nn.Module = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, full_std: bool = True, sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, - clip_mean: float = 2.0): - super(Actor, self).__init__(observation_space, action_space) + clip_mean: float = 2.0, + normalize_images: bool = True): + super(Actor, self).__init__(observation_space, action_space, + features_extractor=features_extractor, + normalize_images=normalize_images) - obs_dim = get_obs_dim(self.observation_space) action_dim = get_action_dim(self.action_space) - latent_pi_net = create_mlp(obs_dim, -1, net_arch, activation_fn) + latent_pi_net = create_mlp(features_dim, -1, net_arch, activation_fn) self.latent_pi = nn.Sequential(*latent_pi_net) self.use_sde = use_sde - self.sde_feature_extractor = None + self.sde_features_extractor = None if self.use_sde: latent_sde_dim = net_arch[-1] # Separate feature extractor for SDE if sde_net_arch is not None: - self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch, - activation_fn) + self.sde_features_extractor, latent_sde_dim = create_sde_features_extractor(features_dim, sde_net_arch, + activation_fn) self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=use_expln, learn_features=True, squash_output=True) @@ -85,7 +94,8 @@ class Actor(BasePolicy): :return: (th.Tensor) """ - assert isinstance(self.action_dist, StateDependentNoiseDistribution), 'get_std() is only available when using SDE' + assert isinstance(self.action_dist, StateDependentNoiseDistribution), \ + 'get_std() is only available when using SDE' return self.action_dist.get_std(self.log_std) def reset_noise(self, batch_size: int = 1) -> None: @@ -94,14 +104,14 @@ class Actor(BasePolicy): :param batch_size: (int) """ - assert isinstance(self.action_dist, StateDependentNoiseDistribution), 'reset_noise() is only available when using SDE' + assert isinstance(self.action_dist, StateDependentNoiseDistribution), \ + 'reset_noise() is only available when using SDE' self.action_dist.sample_weights(self.log_std, batch_size=batch_size) def _get_latent(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: - latent_pi = self.latent_pi(obs) - - latent_sde = self.sde_feature_extractor(obs) if self.sde_feature_extractor is not None else latent_pi - + features = self.extract_features(obs) + latent_pi = self.latent_pi(features) + latent_sde = self.sde_features_extractor(features) if self.sde_features_extractor is not None else latent_pi return latent_pi, latent_sde def get_action_dist_params(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]: @@ -138,27 +148,37 @@ class Critic(BasePolicy): :param observation_space: (gym.spaces.Space) Obervation space :param action_space: (gym.spaces.Space) Action space :param net_arch: ([int]) Network architecture + :param features_extractor: (nn.Module) Network to extract features + (a CNN when using images, a nn.Flatten() layer otherwise) + :param features_dim: (int) Number of features :param activation_fn: (nn.Module) Activation function + :param normalize_images: (bool) Whether to normalize images or not, + dividing by 255.0 (True by default) """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, net_arch: List[int], - activation_fn: nn.Module = nn.ReLU): - super(Critic, self).__init__(observation_space, action_space) + features_extractor: nn.Module, + features_dim: int, + activation_fn: nn.Module = nn.ReLU, + normalize_images: bool = True): + super(Critic, self).__init__(observation_space, action_space, + features_extractor=features_extractor, + normalize_images=normalize_images) - obs_dim = get_obs_dim(self.observation_space) action_dim = get_action_dim(self.action_space) - q1_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn) + q1_net = create_mlp(features_dim + action_dim, 1, net_arch, activation_fn) self.q1_net = nn.Sequential(*q1_net) - q2_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn) + q2_net = create_mlp(features_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: th.Tensor, action: th.Tensor) -> List[th.Tensor]: - qvalue_input = th.cat([obs, action], dim=1) + features = self.extract_features(obs) + qvalue_input = th.cat([features, action], dim=1) return [q_net(qvalue_input) for q_net in self.q_networks] @@ -181,6 +201,8 @@ class SACPolicy(BasePolicy): 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. :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) """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, @@ -192,19 +214,27 @@ class SACPolicy(BasePolicy): log_std_init: float = -3, sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, - clip_mean: float = 2.0): + clip_mean: float = 2.0, + normalize_images: bool = True): super(SACPolicy, self).__init__(observation_space, action_space, device, squash_output=True) if net_arch is None: net_arch = [256, 256] + # 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) + self.net_arch = net_arch self.activation_fn = activation_fn self.net_args = { 'observation_space': self.observation_space, 'action_space': self.action_space, + 'features_extractor': self.features_extractor, + 'features_dim': self.features_dim, 'net_arch': self.net_arch, - 'activation_fn': self.activation_fn + 'activation_fn': self.activation_fn, + 'normalize_images': normalize_images } self.actor_kwargs = self.net_args.copy() sde_kwargs = { diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index a556f2b..3c9698c 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -6,7 +6,7 @@ import torch.nn as nn from torchy_baselines.common.preprocessing import get_action_dim, get_obs_dim from torchy_baselines.common.policies import (BasePolicy, register_policy, create_mlp, - create_sde_feature_extractor) + create_sde_features_extractor) from torchy_baselines.common.distributions import StateDependentNoiseDistribution, Distribution @@ -17,6 +17,9 @@ class Actor(BasePolicy): :param observation_space: (gym.spaces.Space) Obervation space :param action_space: (gym.spaces.Space) Action space :param net_arch: ([int]) Network architecture + :param features_extractor: (nn.Module) Network to extract features + (a CNN when using images, a nn.Flatten() layer otherwise) + :param features_dim: (int) Number of features :param activation_fn: (nn.Module) Activation function :param use_sde: (bool) Whether to use State Dependent Exploration or not :param log_std_init: (float) Initial value for the log standard deviation @@ -30,11 +33,15 @@ class Actor(BasePolicy): :param use_expln: (bool) Use ``expln()`` function instead of ``exp()`` when using SDE to ensure 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. + :param normalize_images: (bool) Whether to normalize images or not, + dividing by 255.0 (True by default) """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, net_arch: List[int], + features_extractor: nn.Module, + features_dim: int, activation_fn: nn.Module = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, @@ -42,32 +49,37 @@ class Actor(BasePolicy): lr_sde: float = 3e-4, full_std: bool = False, sde_net_arch: Optional[List[int]] = None, - use_expln: bool = False): - super(Actor, self).__init__(observation_space, action_space) + use_expln: bool = False, + normalize_images: bool = True): + super(Actor, self).__init__(observation_space, action_space, + features_extractor=features_extractor, + normalize_images=normalize_images) self.latent_pi, self.log_std = None, None self.weights_dist, self.exploration_mat = None, None self.use_sde, self.sde_optimizer = use_sde, None self.full_std = full_std - self.sde_feature_extractor = None + self.sde_features_extractor = None + self.features_extractor = features_extractor + self.normalize_images = normalize_images - obs_dim = get_obs_dim(self.observation_space) action_dim = get_action_dim(self.action_space) if use_sde: - latent_pi_net = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_output=False) + latent_pi_net = create_mlp(features_dim, -1, net_arch, activation_fn, squash_output=False) self.latent_pi = nn.Sequential(*latent_pi_net) latent_sde_dim = net_arch[-1] learn_features = sde_net_arch is not None # Separate feature extractor for SDE if sde_net_arch is not None: - self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch, - activation_fn) + self.sde_features_extractor, latent_sde_dim = create_sde_features_extractor(features_dim, sde_net_arch, + activation_fn) # Create state dependent noise matrix (SDE) self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=use_expln, squash_output=False, learn_features=learn_features) + action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1], latent_sde_dim=latent_sde_dim, log_std_init=log_std_init) @@ -77,7 +89,7 @@ class Actor(BasePolicy): self.sde_optimizer = th.optim.Adam([self.log_std], lr=lr_sde) self.reset_noise() else: - actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_output=True) + actor_net = create_mlp(features_dim, action_dim, net_arch, activation_fn, squash_output=True) self.mu = nn.Sequential(*actor_net) def get_std(self) -> th.Tensor: @@ -98,12 +110,9 @@ class Actor(BasePolicy): return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde) def _get_latent(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: - latent_pi = self.latent_pi(obs) - - if self.sde_feature_extractor is not None: - latent_sde = self.sde_feature_extractor(obs) - else: - latent_sde = latent_pi + features = self.extract_features(obs) + latent_pi = self.latent_pi(features) + latent_sde = self.sde_features_extractor(features) if self.sde_features_extractor is not None else latent_pi return latent_pi, latent_sde def evaluate_actions(self, obs: th.Tensor, action: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: @@ -144,7 +153,8 @@ class Actor(BasePolicy): # action, _ = self._get_action_dist_from_latent(latent_pi) # return action else: - return self.mu(obs) + features = self.extract_features(obs) + return self.mu(features) class Critic(BasePolicy): @@ -155,29 +165,40 @@ class Critic(BasePolicy): :param observation_space: (gym.spaces.Space) Obervation space :param action_space: (gym.spaces.Space) Action space :param net_arch: ([int]) Network architecture + :param features_extractor: (nn.Module) Network to extract features + (a CNN when using images, a nn.Flatten() layer otherwise) + :param features_dim: (int) Number of features :param activation_fn: (nn.Module) Activation function + :param normalize_images: (bool) Whether to normalize images or not, + dividing by 255.0 (True by default) """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, net_arch: List[int], - activation_fn: nn.Module = nn.ReLU): - super(Critic, self).__init__(observation_space, action_space) + features_extractor: nn.Module, + features_dim: int, + activation_fn: nn.Module = nn.ReLU, + normalize_images: bool = True): + super(Critic, self).__init__(observation_space, action_space, + features_extractor=features_extractor, + normalize_images=normalize_images) - obs_dim = get_obs_dim(self.observation_space) action_dim = get_action_dim(self.action_space) - q1_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn) + q1_net = create_mlp(features_dim + action_dim, 1, net_arch, activation_fn) self.q1_net = nn.Sequential(*q1_net) - q2_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn) + q2_net = create_mlp(features_dim + action_dim, 1, net_arch, activation_fn) self.q2_net = nn.Sequential(*q2_net) def forward(self, obs: th.Tensor, action: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: - qvalue_input = th.cat([obs, action], dim=1) + features = self.extract_features(obs) + qvalue_input = th.cat([features, action], dim=1) return self.q1_net(qvalue_input), self.q2_net(qvalue_input) def q1_forward(self, obs: th.Tensor, action: th.Tensor) -> th.Tensor: - return self.q1_net(th.cat([obs, action], dim=1)) + features = self.extract_features(obs) + return self.q1_net(th.cat([features, action], dim=1)) class ValueFunction(BasePolicy): @@ -186,24 +207,34 @@ class ValueFunction(BasePolicy): :param observation_space: (gym.spaces.Space) Obervation space :param action_space: (gym.spaces.Space) Action space + :param features_extractor: (nn.Module) Network to extract features + (a CNN when using images, a nn.Flatten() layer otherwise) + :param features_dim: (int) Number of features :param net_arch: (Optional[List[int]]) Network architecture :param activation_fn: (nn.Module) Activation function + :param normalize_images: (bool) Whether to normalize images or not, + dividing by 255.0 (True by default) """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, + features_extractor: nn.Module, + features_dim: int, net_arch: Optional[List[int]] = None, - activation_fn: nn.Module = nn.Tanh): - super(ValueFunction, self).__init__(observation_space, action_space) + activation_fn: nn.Module = nn.Tanh, + normalize_images: bool = True): + super(ValueFunction, self).__init__(observation_space, action_space, + features_extractor=features_extractor, + normalize_images=normalize_images) - obs_dim = get_obs_dim(self.observation_space) if net_arch is None: net_arch = [64, 64] - vf_net = create_mlp(obs_dim, 1, net_arch, activation_fn) + vf_net = create_mlp(features_dim, 1, net_arch, activation_fn) self.vf_net = nn.Sequential(*vf_net) def forward(self, obs: th.Tensor) -> th.Tensor: - return self.vf_net(obs) + features = self.extract_features(obs) + return self.vf_net(features) class TD3Policy(BasePolicy): @@ -224,6 +255,8 @@ class TD3Policy(BasePolicy): :param use_expln: (bool) Use ``expln()`` function instead of ``exp()`` when using SDE to ensure 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. + :param normalize_images: (bool) Whether to normalize images or not, + dividing by 255.0 (True by default) """ def __init__(self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, @@ -236,20 +269,28 @@ class TD3Policy(BasePolicy): clip_noise: Optional[float] = None, lr_sde: float = 3e-4, sde_net_arch: Optional[List[int]] = None, - use_expln: bool = False): + use_expln: bool = False, + normalize_images: bool = True): 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] + # 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) + self.net_arch = net_arch self.activation_fn = activation_fn self.net_args = { 'observation_space': self.observation_space, 'action_space': self.action_space, + 'features_extractor': self.features_extractor, + 'features_dim': self.features_dim, 'net_arch': self.net_arch, - 'activation_fn': self.activation_fn + 'activation_fn': self.activation_fn, + 'normalize_images': normalize_images } self.actor_kwargs = self.net_args.copy() sde_kwargs = { @@ -282,7 +323,9 @@ class TD3Policy(BasePolicy): self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=lr_schedule(1)) if self.use_sde: - self.vf_net = ValueFunction(self.observation_space, self.action_space) + 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 reset_noise(self) -> None: diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 9158a0d..d39b9f5 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -4,9 +4,8 @@ from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any from torchy_baselines.common import logger from torchy_baselines.common.base_class import OffPolicyRLModel -from torchy_baselines.common.buffers import ReplayBuffer from torchy_baselines.common.noise import ActionNoise -from torchy_baselines.common.type_aliases import ReplayBufferSamples, GymEnv, MaybeCallback +from torchy_baselines.common.type_aliases import GymEnv, MaybeCallback from torchy_baselines.td3.policies import TD3Policy @@ -161,7 +160,8 @@ class TD3(OffPolicyRLModel): # Delayed policy updates if gradient_step % policy_delay == 0: # Compute actor loss - actor_loss = -self.critic.q1_forward(replay_data.observations, self.actor(replay_data.observations)).mean() + actor_loss = -self.critic.q1_forward(replay_data.observations, + self.actor(replay_data.observations)).mean() # Optimize the actor self.actor.optimizer.zero_grad() @@ -178,7 +178,6 @@ class TD3(OffPolicyRLModel): self._n_updates += gradient_steps logger.logkv("n_updates", self._n_updates) - def train_sde(self) -> None: # Update optimizer learning rate # self._update_learning_rate(self.policy.optimizer) diff --git a/torchy_baselines/version.txt b/torchy_baselines/version.txt index 6495aed..db79cdc 100644 --- a/torchy_baselines/version.txt +++ b/torchy_baselines/version.txt @@ -1 +1 @@ -0.4.0a1 +0.4.0a2