diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 41bf91e..7dc6d84 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,7 +3,7 @@ Changelog ========== -Pre-Release 0.5.0a1 (WIP) +Pre-Release 0.5.0a2 (WIP) ------------------------------ Breaking Changes: @@ -15,6 +15,7 @@ New Features: - Added ``optimizer`` and ``optimizer_kwargs`` to ``policy_kwargs`` in order to easily customizer optimizers - Complete independent save/load for policies +- Add ``CnnPolicies`` to support images as input (caveat: only support Atari resolution for now) Bug Fixes: diff --git a/tests/test_cnn.py b/tests/test_cnn.py new file mode 100644 index 0000000..fb40aba --- /dev/null +++ b/tests/test_cnn.py @@ -0,0 +1,16 @@ +import pytest + +from torchy_baselines import A2C, PPO, SAC, TD3 +from torchy_baselines.common.identity_env import FakeImageEnv + + +@pytest.mark.parametrize('model_class', [A2C, PPO, SAC]) +def test_cnn(model_class): + # Fake grayscale with frameskip + env = FakeImageEnv(screen_height=84, screen_width=84, n_channels=1, + discrete = model_class not in {SAC, TD3}) + if model_class in {A2C, PPO}: + kwargs = dict(n_steps=100) + else: + kwargs = dict(buffer_size=500) + _ = model_class('CnnPolicy', env, **kwargs).learn(500) diff --git a/tests/test_custom_policy.py b/tests/test_custom_policy.py index d2ffab0..94cbf17 100644 --- a/tests/test_custom_policy.py +++ b/tests/test_custom_policy.py @@ -7,6 +7,7 @@ from torchy_baselines import A2C, PPO, SAC, TD3 @pytest.mark.parametrize('net_arch', [ [12, dict(vf=[16], pi=[8])], [4], + [], [4, 4], [12, dict(vf=[8, 4], pi=[8])], [12, dict(vf=[8], pi=[8, 4])], diff --git a/torchy_baselines/common/identity_env.py b/torchy_baselines/common/identity_env.py index aa724e1..526f5d3 100644 --- a/torchy_baselines/common/identity_env.py +++ b/torchy_baselines/common/identity_env.py @@ -1,5 +1,6 @@ -import numpy as np +from typing import List +import numpy as np from gym import Env from gym.spaces import Discrete, MultiDiscrete, MultiBinary, Box @@ -102,3 +103,37 @@ class IdentityEnvMultiBinary(IdentityEnv): self.action_space = MultiBinary(dim) self.observation_space = self.action_space self.reset() + + + +class FakeImageEnv(Env): + def __init__(self, action_dim: int = 6, + screen_height: int = 210, + screen_width: int = 160, + n_channels: int = 3, + discrete: bool = True): + """ + Fake atari environment for testing purposes. + """ + self.observation_space = Box(low=0, high=255, shape=(screen_height, screen_width, n_channels), dtype=np.uint8) + if discrete: + self.action_space = Discrete(action_dim) + else: + self.action_space = Box(low=-1, high=1, shape=(5,), dtype=np.float32) + self.ep_length = 10 + + def reset(self): + self.current_step = 0 + return self.observation_space.sample() + + def step(self, action: int): + reward = 0.0 + self.current_step += 1 + done = self.current_step >= self.ep_length + return self.observation_space.sample(), reward, done, {} + + def render(self, mode='human'): + pass + + def get_action_meanings(self) -> List[str]: + return ['NOOP'] diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 9aff338..a546d4d 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -7,7 +7,7 @@ import torch as th import torch.nn as nn import numpy as np -from torchy_baselines.common.preprocessing import preprocess_obs +from torchy_baselines.common.preprocessing import preprocess_obs, get_obs_dim from torchy_baselines.common.utils import get_device, get_schedule_fn @@ -289,7 +289,8 @@ def create_mlp(input_dim: int, modules.append(activation_fn()) if output_dim > 0: - modules.append(nn.Linear(net_arch[-1], output_dim)) + last_layer_dim = net_arch[-1] if len(net_arch) > 0 else input_dim + modules.append(nn.Linear(last_layer_dim, output_dim)) if squash_output: modules.append(nn.Tanh()) return modules @@ -453,3 +454,44 @@ class MlpExtractor(nn.Module): """ shared_latent = self.shared_net(features) return self.policy_net(shared_latent), self.value_net(shared_latent) + + +class BaseFeaturesExtractor(nn.Module): + def __init__(self, observation_space: gym.Space, features_dim: int = 0): + super(BaseFeaturesExtractor, self).__init__() + assert features_dim > 0 + self._observation_space = observation_space + self._features_dim = features_dim + + @property + def features_dim(self) -> int: + return self._features_dim + + def forward(self, observations: th.Tensor) -> th.Tensor: + raise NotImplementedError() + + +class FlattenExtractor(BaseFeaturesExtractor): + def __init__(self, observation_space: gym.Space): + super(FlattenExtractor, self).__init__(observation_space, get_obs_dim(observation_space)) + self.flatten = nn.Flatten() + + def forward(self, observations: th.Tensor) -> th.Tensor: + return self.flatten(observations) + + +class NatureCNN(BaseFeaturesExtractor): + def __init__(self, observation_space: gym.Space, + features_dim: int = 512): + super(NatureCNN, self).__init__(observation_space, features_dim) + # TODO: custom init? + # we assume WxHxC images + # TODO: compute shape before flatten + n_input_channels = observation_space.shape[-1] + self.cnn = nn.Sequential(nn.Conv2d(n_input_channels, 32, 8, stride=4), nn.ReLU(), + nn.Conv2d(32, 64, 4, stride=2), nn.ReLU(), + nn.Conv2d(64, 32, 3, stride=1), nn.ReLU(), nn.Flatten(), + nn.Linear(32 * 7 * 7, features_dim), nn.ReLU()) + + def forward(self, observations: th.Tensor) -> th.Tensor: + return self.cnn(observations) diff --git a/torchy_baselines/common/preprocessing.py b/torchy_baselines/common/preprocessing.py index ebef581..70fc59e 100644 --- a/torchy_baselines/common/preprocessing.py +++ b/torchy_baselines/common/preprocessing.py @@ -47,8 +47,12 @@ def preprocess_obs(obs: th.Tensor, observation_space: spaces.Space, :return: (th.Tensor) """ if isinstance(observation_space, spaces.Box): - if is_image_space(observation_space) and normalize_images: - return obs.float() / 255.0 + if is_image_space(observation_space): + # Re-order from BxWxHxC to BxCxWxH + obs = obs.permute(0, 3, 1, 2) # .contiguous()? + if normalize_images: + return obs.float() / 255.0 + return obs.float() return obs.float() elif isinstance(observation_space, spaces.Discrete): # One hot encoding and convert to float to avoid errors diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 44754a6..7aa8b88 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -6,9 +6,9 @@ import torch as th import torch.nn as nn 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_features_extractor) + create_sde_features_extractor, NatureCNN, + BaseFeaturesExtractor, FlattenExtractor) from torchy_baselines.common.distributions import (make_proba_distribution, Distribution, DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution) @@ -37,6 +37,7 @@ 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 features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use. :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, @@ -58,6 +59,7 @@ class PPOPolicy(BasePolicy): sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, squash_output: bool = False, + features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor, normalize_images: bool = True, optimizer: Type[th.optim.Optimizer] = th.optim.Adam, optimizer_kwargs: Optional[Dict[str, Any]] = None): @@ -65,7 +67,10 @@ class PPOPolicy(BasePolicy): # Default network architecture, from stable-baselines if net_arch is None: - net_arch = [dict(pi=[64, 64], vf=[64, 64])] + if features_extractor_class == FlattenExtractor: + net_arch = [dict(pi=[64, 64], vf=[64, 64])] + else: + net_arch = [] self.net_arch = net_arch self.activation_fn = activation_fn @@ -79,9 +84,11 @@ class PPOPolicy(BasePolicy): 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() - self.features_dim = get_obs_dim(self.observation_space) + + self.features_extractor_class = features_extractor_class + self.features_extractor = features_extractor_class(self.observation_space) + self.features_dim = self.features_extractor.features_dim + self.normalize_images = normalize_images self.log_std_init = log_std_init dist_kwargs = None @@ -119,7 +126,8 @@ class PPOPolicy(BasePolicy): 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 + ortho_init=self.ortho_init, + features_extractor_class=self.features_extractor_class )) return data @@ -266,4 +274,75 @@ class PPOPolicy(BasePolicy): MlpPolicy = PPOPolicy + + +class CnnPolicy(PPOPolicy): + """ + CnnPolicy class (with both actor and critic) for A2C and derivates (PPO). + + :param observation_space: (gym.spaces.Space) Observation space + :param action_space: (gym.spaces.Space) Action space + :param lr_schedule: (Callable) Learning rate schedule (could be constant) + :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 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 + :param full_std: (bool) Whether to use (n_features x n_actions) parameters + for the std instead of only (n_features,) when using SDE + :param sde_net_arch: ([int]) Network architecture for extracting features + when using SDE. If None, the latent features from the policy will be used. + Pass an empty list to use the states as features. + :param use_expln: (bool) Use ``expln()`` function instead of ``exp()`` 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 squash_output: (bool) Whether to squash the output using a tanh function, + this allows to ensure boundaries when using SDE. + :param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use. + :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] = 'auto', + activation_fn: Type[nn.Module] = nn.Tanh, + ortho_init: bool = True, + use_sde: bool = False, + log_std_init: float = 0.0, + full_std: bool = True, + sde_net_arch: Optional[List[int]] = None, + use_expln: bool = False, + squash_output: bool = False, + features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN, + normalize_images: bool = True, + optimizer: Type[th.optim.Optimizer] = th.optim.Adam, + optimizer_kwargs: Optional[Dict[str, Any]] = None): + super(CnnPolicy, self).__init__(observation_space, + action_space, + lr_schedule, + net_arch, + device, + activation_fn, + ortho_init, + use_sde, + log_std_init, + full_std, + sde_net_arch, + use_expln, + squash_output, + features_extractor_class, + normalize_images, + optimizer, + optimizer_kwargs) + + register_policy("MlpPolicy", MlpPolicy) +register_policy("CnnPolicy", CnnPolicy) diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index c0fd3a6..4f4c2ee 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -4,9 +4,10 @@ import gym import torch as th import torch.nn as nn -from torchy_baselines.common.preprocessing import get_action_dim, get_obs_dim +from torchy_baselines.common.preprocessing import get_action_dim from torchy_baselines.common.policies import (BasePolicy, register_policy, create_mlp, - create_sde_features_extractor) + create_sde_features_extractor, NatureCNN, + BaseFeaturesExtractor, FlattenExtractor) from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution # CAP the standard deviation of the actor @@ -75,9 +76,11 @@ class Actor(BasePolicy): self.use_expln = use_expln self.full_std = full_std self.clip_mean = clip_mean + last_layer_dim = net_arch[-1] if len(net_arch) > 0 else features_dim + if self.use_sde: - latent_sde_dim = net_arch[-1] + latent_sde_dim = last_layer_dim # Separate feature extractor for SDE if sde_net_arch is not None: self.sde_features_extractor, latent_sde_dim = create_sde_features_extractor(features_dim, sde_net_arch, @@ -94,8 +97,8 @@ class Actor(BasePolicy): self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-clip_mean, max_val=clip_mean)) else: self.action_dist = SquashedDiagGaussianDistribution(action_dim) - self.mu = nn.Linear(net_arch[-1], action_dim) - self.log_std = nn.Linear(net_arch[-1], action_dim) + self.mu = nn.Linear(last_layer_dim, action_dim) + self.log_std = nn.Linear(last_layer_dim, action_dim) def _get_data(self) -> Dict[str, Any]: data = super()._get_data() @@ -239,6 +242,7 @@ 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 features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use. :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, @@ -257,13 +261,17 @@ class SACPolicy(BasePolicy): sde_net_arch: Optional[List[int]] = None, use_expln: bool = False, clip_mean: float = 2.0, + features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor, 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 features_extractor_class == FlattenExtractor: + net_arch = [256, 256] + else: + net_arch = [] if optimizer_kwargs is None: optimizer_kwargs = {} @@ -271,9 +279,9 @@ class SACPolicy(BasePolicy): 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) + self.features_extractor_class = features_extractor_class + self.features_extractor = features_extractor_class(self.observation_space) + self.features_dim = self.features_extractor.features_dim self.net_arch = net_arch self.activation_fn = activation_fn @@ -325,7 +333,8 @@ class SACPolicy(BasePolicy): 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 + optimizer_kwargs=self.optimizer_kwargs, + features_extractor_class=self.features_extractor_class )) return data @@ -344,4 +353,66 @@ class SACPolicy(BasePolicy): MlpPolicy = SACPolicy + +class CnnPolicy(SACPolicy): + """ + Policy class (with both actor and critic) for SAC. + + :param observation_space: (gym.spaces.Space) Observation space + :param action_space: (gym.spaces.Space) Action space + :param lr_schedule: (callable) Learning rate schedule (could be constant) + :param net_arch: (Optional[List[int]]) 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 use_sde: (bool) Whether to use State Dependent Exploration or not + :param log_std_init: (float) Initial value for the log standard deviation + :param sde_net_arch: ([int]) Network architecture for extracting features + when using SDE. If None, the latent features from the policy will be used. + Pass an empty list to use the states as features. + :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 clip_mean: (float) Clip the mean output when using SDE to avoid numerical instability. + :param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use. + :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] = '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, + features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN, + normalize_images: bool = True, + optimizer: Type[th.optim.Optimizer] = th.optim.Adam, + optimizer_kwargs: Optional[Dict[str, Any]] = None): + super(CnnPolicy, self).__init__(observation_space, + action_space, + lr_schedule, + net_arch, + device, + activation_fn, + use_sde, + log_std_init, + sde_net_arch, + use_expln, + clip_mean, + features_extractor_class, + normalize_images, + optimizer, + optimizer_kwargs) + + + register_policy("MlpPolicy", MlpPolicy) +register_policy("CnnPolicy", CnnPolicy) diff --git a/torchy_baselines/version.txt b/torchy_baselines/version.txt index d443359..8413ee4 100644 --- a/torchy_baselines/version.txt +++ b/torchy_baselines/version.txt @@ -1 +1 @@ -0.5.0a1 +0.5.0a2