From b289aca5fe76456417407cb6bdc98328d84f8778 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Tue, 21 Apr 2020 21:06:07 +0200 Subject: [PATCH] Bug fixes at loading and predict time --- torchy_baselines/common/base_class.py | 9 ++++----- torchy_baselines/common/policies.py | 9 +++++++++ torchy_baselines/common/preprocessing.py | 8 ++++++-- torchy_baselines/common/vec_env/vec_transpose.py | 2 ++ 4 files changed, 21 insertions(+), 7 deletions(-) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index e2d3152..ac49453 100644 --- a/torchy_baselines/common/base_class.py +++ b/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 diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 90ef99f..675e9be 100644 --- a/torchy_baselines/common/policies.py +++ b/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) diff --git a/torchy_baselines/common/preprocessing.py b/torchy_baselines/common/preprocessing.py index ebef581..a52226e 100644 --- a/torchy_baselines/common/preprocessing.py +++ b/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 diff --git a/torchy_baselines/common/vec_env/vec_transpose.py b/torchy_baselines/common/vec_env/vec_transpose.py index 34b4669..f4ef501 100644 --- a/torchy_baselines/common/vec_env/vec_transpose.py +++ b/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):