Add VecTransposeImage and fix for SAC

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:

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)

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
 

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))

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

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:

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()

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)