From 8aac9e819db6e4ea740f4f7704c2497dfac037b4 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Tue, 21 Apr 2020 20:41:58 +0200 Subject: [PATCH] Add `VecTransposeImage` and fix for SAC --- docs/misc/changelog.rst | 2 +- tests/test_cnn.py | 8 +++- torchy_baselines/common/base_class.py | 36 ++++++++++------ torchy_baselines/common/policies.py | 28 ++++++++---- torchy_baselines/common/preprocessing.py | 8 +--- torchy_baselines/common/vec_env/__init__.py | 1 + .../common/vec_env/vec_transpose.py | 43 +++++++++++++++++++ torchy_baselines/sac/policies.py | 22 ++++++++-- 8 files changed, 114 insertions(+), 34 deletions(-) create mode 100644 torchy_baselines/common/vec_env/vec_transpose.py diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 7dc6d84..57a0a33 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -15,7 +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) +- Add ``CnnPolicies`` to support images as input Bug Fixes: diff --git a/tests/test_cnn.py b/tests/test_cnn.py index fb40aba..d0b7ed7 100644 --- a/tests/test_cnn.py +++ b/tests/test_cnn.py @@ -7,10 +7,14 @@ 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, + # Atari after preprocessing: 84x84x1, here we are using lower resolution + # to check that the network handle it automatically + env = FakeImageEnv(screen_height=40, screen_width=40, 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) + # Avoid memory error when using replay buffer + # Reduce the size of the features + kwargs = dict(buffer_size=500, policy_kwargs=dict(features_extractor_kwargs=dict(features_dim=40))) _ = model_class('CnnPolicy', env, **kwargs).learn(500) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 0efd7d5..e2d3152 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -14,7 +14,8 @@ 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, get_device -from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, VecNormalize +from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, VecNormalize, VecTransposeImage +from torchy_baselines.common.preprocessing import is_image_space 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 from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback, EvalCallback @@ -123,12 +124,10 @@ class BaseRLModel(ABC): env = Monitor(env, filename=None) env = DummyVecEnv([lambda: env]) + env = self._wrap_env(env) + self.observation_space = env.observation_space self.action_space = env.action_space - if not isinstance(env, VecEnv): - if self.verbose >= 1: - print("Wrapping the env in a DummyVecEnv.") - env = DummyVecEnv([lambda: env]) self.n_envs = env.num_envs self.env = env @@ -136,6 +135,18 @@ class BaseRLModel(ABC): raise ValueError("Error: the model does not support multiple envs requires a single vectorized" " environment.") + def _wrap_env(self, env: GymEnv) -> VecEnv: + if not isinstance(env, VecEnv): + if self.verbose >= 1: + print("Wrapping the env in a DummyVecEnv.") + env = DummyVecEnv([lambda: env]) + + if is_image_space(env.observation_space): + if self.verbose >= 1: + print("Wrapping the env in a VecTransposeImage.") + env = VecTransposeImage(env) + return env + @abstractmethod def _setup_model(self) -> None: """ @@ -154,8 +165,7 @@ class BaseRLModel(ABC): eval_env = self.eval_env if eval_env is not None: - if not isinstance(eval_env, VecEnv): - eval_env = DummyVecEnv([lambda: eval_env]) + eval_env = self._wrap_env(eval_env) assert eval_env.num_envs == 1 return eval_env @@ -228,7 +238,11 @@ class BaseRLModel(ABC): :param action_space: (gym.spaces.Space) :return: (bool) True if environment seems to be coherent """ - if observation_space != env.observation_space: + if (observation_space != env.observation_space + # Special cases for images that need to be transposed + or (is_image_space(observation_space) + and VecTransposeImage.transpose_space(observation_space) != env.observation_space) + ): return False if action_space != env.action_space: return False @@ -250,10 +264,8 @@ class BaseRLModel(ABC): "observation and action spaces do not match") # it must be coherent now # if it is not a VecEnv, make it a VecEnv - if not isinstance(env, VecEnv): - if self.verbose >= 1: - print("Wrapping the env in a DummyVecEnv.") - env = DummyVecEnv([lambda: env]) + env = self._wrap_env(env) + self.n_envs = env.num_envs self.env = env diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index a546d4d..90ef99f 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, get_obs_dim +from torchy_baselines.common.preprocessing import preprocess_obs, get_obs_dim, is_image_space from torchy_baselines.common.utils import get_device, get_schedule_fn @@ -485,13 +485,23 @@ class NatureCNN(BaseFeaturesExtractor): 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()) + # TODO: check that the observation space is an image + # we assume CxWxH images + # assert is_image_space(observation_space), observation_space + n_input_channels = observation_space.shape[0] + self.cnn = nn.Sequential(nn.Conv2d(n_input_channels, 32, kernel_size=8, stride=4, padding=0), + nn.ReLU(), + nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=0), + nn.ReLU(), + nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=0), + nn.ReLU(), + nn.Flatten()) + + # Compute shape by doing one forward pass + with th.no_grad(): + n_flatten = self.cnn(th.as_tensor(observation_space.sample()[None]).float()).shape[1] + + self.linear = nn.Sequential(nn.Linear(n_flatten, features_dim), nn.ReLU()) def forward(self, observations: th.Tensor) -> th.Tensor: - return self.cnn(observations) + return self.linear(self.cnn(observations)) diff --git a/torchy_baselines/common/preprocessing.py b/torchy_baselines/common/preprocessing.py index 70fc59e..ebef581 100644 --- a/torchy_baselines/common/preprocessing.py +++ b/torchy_baselines/common/preprocessing.py @@ -47,12 +47,8 @@ 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): - # 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() + 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 diff --git a/torchy_baselines/common/vec_env/__init__.py b/torchy_baselines/common/vec_env/__init__.py index 8748923..b119a9f 100644 --- a/torchy_baselines/common/vec_env/__init__.py +++ b/torchy_baselines/common/vec_env/__init__.py @@ -9,6 +9,7 @@ from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv from torchy_baselines.common.vec_env.subproc_vec_env import SubprocVecEnv from torchy_baselines.common.vec_env.vec_frame_stack import VecFrameStack from torchy_baselines.common.vec_env.vec_normalize import VecNormalize +from torchy_baselines.common.vec_env.vec_transpose import VecTransposeImage # Avoid circular import if typing.TYPE_CHECKING: diff --git a/torchy_baselines/common/vec_env/vec_transpose.py b/torchy_baselines/common/vec_env/vec_transpose.py new file mode 100644 index 0000000..34b4669 --- /dev/null +++ b/torchy_baselines/common/vec_env/vec_transpose.py @@ -0,0 +1,43 @@ +import warnings + +import numpy as np +from gym import spaces + +from torchy_baselines.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper +from torchy_baselines.common.preprocessing import is_image_space + + +class VecTransposeImage(VecEnvWrapper): + """ + Re-order channels, from WxHxC to CxWxH. + :param venv: (VecEnv) + """ + def __init__(self, venv: VecEnv): + assert is_image_space(venv.observation_space), 'The observation space must be an image' + + observation_space = self.transpose_space(venv.observation_space) + super(VecTransposeImage, self).__init__(venv, observation_space=observation_space) + + @staticmethod + def transpose_space(observation_space: spaces.Box) -> spaces.Box: + assert is_image_space(observation_space), 'The observation space must be an image' + width, height, channels = observation_space.shape + new_shape = (channels, width, height) + return spaces.Box(low=0, high=255, shape=new_shape, dtype=observation_space.dtype) + + @staticmethod + def transpose_image(image: np.ndarray) -> np.ndarray: + return np.transpose(image, (0, 3, 1, 2)) + + def step_wait(self): + observations, rewards, dones, infos = self.venv.step_wait() + return self.transpose_image(observations), rewards, dones, infos + + def reset(self) -> np.ndarray: + """ + Reset all environments + """ + return self.transpose_image(self.venv.reset()) + + def close(self) -> None: + self.venv.close() diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index 4f4c2ee..f066da0 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -88,7 +88,7 @@ class Actor(BasePolicy): self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=use_expln, learn_features=True, squash_output=True) - self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1], + self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=last_layer_dim, latent_sde_dim=latent_sde_dim, log_std_init=log_std_init) # Avoid numerical issues by limiting the mean of the Gaussian @@ -243,6 +243,8 @@ class SACPolicy(BasePolicy): 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 features_extractor_kwargs: (Optional[Dict[str, Any]]) Keyword arguments + to pass to the feature extractor. :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, @@ -262,6 +264,7 @@ class SACPolicy(BasePolicy): use_expln: bool = False, clip_mean: float = 2.0, features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor, + features_extractor_kwargs: Optional[Dict[str, Any]] = None, normalize_images: bool = True, optimizer: Type[th.optim.Optimizer] = th.optim.Adam, optimizer_kwargs: Optional[Dict[str, Any]] = None): @@ -280,7 +283,10 @@ class SACPolicy(BasePolicy): self.optimizer_kwargs = optimizer_kwargs self.features_extractor_class = features_extractor_class - self.features_extractor = features_extractor_class(self.observation_space) + self.features_extractor_kwargs = features_extractor_kwargs + if features_extractor_kwargs is None: + features_extractor_kwargs = {} + self.features_extractor = features_extractor_class(self.observation_space, **features_extractor_kwargs) self.features_dim = self.features_extractor.features_dim self.net_arch = net_arch @@ -317,7 +323,12 @@ class SACPolicy(BasePolicy): self.critic = self.make_critic() self.critic_target = self.make_critic() self.critic_target.load_state_dict(self.critic.state_dict()) - self.critic.optimizer = self.optimizer_class(self.critic.parameters(), lr=lr_schedule(1), + # Do not optimize the shared feature extractor with the critic loss + # otherwise, there are gradient computation issues + # another solution: having duplicated features extractor but requires more memory and computation + # Note: check gradients, they are maybe computed but not zeroed by the critic + critic_parameters = [param for name, param in self.critic.named_parameters() if 'features_extractor' not in name] + self.critic.optimizer = self.optimizer_class(critic_parameters, lr=lr_schedule(1), **self.optimizer_kwargs) def _get_data(self) -> Dict[str, Any]: @@ -334,7 +345,8 @@ class SACPolicy(BasePolicy): lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone optimizer=self.optimizer_class, optimizer_kwargs=self.optimizer_kwargs, - features_extractor_class=self.features_extractor_class + features_extractor_class=self.features_extractor_class, + features_extractor_kwargs=self.features_extractor_kwargs )) return data @@ -393,6 +405,7 @@ class CnnPolicy(SACPolicy): use_expln: bool = False, clip_mean: float = 2.0, features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN, + features_extractor_kwargs: Optional[Dict[str, Any]] = None, normalize_images: bool = True, optimizer: Type[th.optim.Optimizer] = th.optim.Adam, optimizer_kwargs: Optional[Dict[str, Any]] = None): @@ -408,6 +421,7 @@ class CnnPolicy(SACPolicy): use_expln, clip_mean, features_extractor_class, + features_extractor_kwargs, normalize_images, optimizer, optimizer_kwargs)