Bug fixes at loading and predict time

torchy_baselines/common/base_class.py

@@ -240,8 +240,8 @@ class BaseRLModel(ABC):
         """
         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)
+            and not (is_image_space(env.observation_space)
+                and observation_space == VecTransposeImage.transpose_space(env.observation_space))
             ):
             return False
         if action_space != env.action_space:
@@ -352,15 +352,14 @@ class BaseRLModel(ABC):
         if env is None and "env" in data:
             env = data["env"]
 
-        # first create model, but only setup if a env was given
         # noinspection PyArgumentList
-        model = cls(policy=data["policy_class"], env=env, device='auto', _init_setup_model=env is not None)
+        model = cls(policy=data["policy_class"], env=env, device='auto', _init_setup_model=False)
 
         # load parameters
         model.__dict__.update(data)
         model.__dict__.update(kwargs)
         if not hasattr(model, "_setup_model") and len(params) > 0:
-            raise NotImplementedError("loading was executed on a model that has no means to create the policies")
+            raise NotImplementedError(f"{cls} has no `_setup_model()` method")
         model._setup_model()
 
         # put state_dicts back in place

torchy_baselines/common/policies.py

@@ -9,6 +9,7 @@ import numpy as np
 
 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
+from torchy_baselines.common.vec_env import VecTransposeImage
 
 
 class BasePolicy(nn.Module):
@@ -102,9 +103,17 @@ class BasePolicy(nn.Module):
         # if mask is None:
         #     mask = [False for _ in range(self.n_envs)]
         observation = np.array(observation)
+
+        if is_image_space(self.observation_space, channels_last=False):
+            # TODO: handle the different cases
+            if (observation.shape != self.observation_space.shape
+                and observation.shape[1:] != self.observation_space.shape):
+                observation = VecTransposeImage.transpose_image(observation)
+
         vectorized_env = self._is_vectorized_observation(observation, self.observation_space)
 
         observation = observation.reshape((-1,) + self.observation_space.shape)
+
         observation = th.as_tensor(observation).to(self.device)
         with th.no_grad():
             actions = self._predict(observation, deterministic=deterministic)

torchy_baselines/common/preprocessing.py

@@ -6,7 +6,7 @@ import torch.nn.functional as F
 from gym import spaces
 
 
-def is_image_space(observation_space: spaces.Space) -> bool:
+def is_image_space(observation_space: spaces.Space, channels_last: bool = True) -> bool:
     """
     Check if a observation space has the shape, limits and dtype
     of a valid image.
@@ -16,6 +16,7 @@ def is_image_space(observation_space: spaces.Space) -> bool:
     Valid images: RGB, RGBD, GrayScale with values in [0, 255]
 
     :param observation_space: (spaces.Space)
+    :param channels_last: (bool)
     :return: (bool)
     """
     if isinstance(observation_space, spaces.Box) and len(observation_space.shape) == 3:
@@ -28,7 +29,10 @@ def is_image_space(observation_space: spaces.Space) -> bool:
             return False
 
         # Check the number of channels
-        n_channels = observation_space.shape[-1]
+        if channels_last:
+            n_channels = observation_space.shape[-1]
+        else:
+            n_channels = observation_space.shape[0]
         return n_channels in [1, 3, 4]
     return False
 

torchy_baselines/common/vec_env/vec_transpose.py

@@ -27,6 +27,8 @@ class VecTransposeImage(VecEnvWrapper):
 
     @staticmethod
     def transpose_image(image: np.ndarray) -> np.ndarray:
+        if len(image.shape) == 3:
+            return np.transpose(image, (2, 0, 1))
         return np.transpose(image, (0, 3, 1, 2))
 
     def step_wait(self):