Add check for image space

torchy_baselines/common/buffers.py

@@ -6,7 +6,7 @@ from gym import spaces
 
 from torchy_baselines.common.vec_env import VecNormalize
 from torchy_baselines.common.type_aliases import RolloutBufferSamples, ReplayBufferSamples
-from torchy_baselines.common.preprocessing import get_obs_dim, get_action_dim
+from torchy_baselines.common.preprocessing import get_action_dim, get_obs_shape
 
 
 class BaseBuffer(object):
@@ -30,7 +30,7 @@ class BaseBuffer(object):
         self.buffer_size = buffer_size
         self.observation_space = observation_space
         self.action_space = action_space
-        self.obs_dim = get_obs_dim(observation_space)
+        self.obs_shape = get_obs_shape(observation_space)
         self.action_dim = get_action_dim(action_space)
         self.pos = 0
         self.full = False
@@ -157,9 +157,9 @@ class ReplayBuffer(BaseBuffer):
 
         assert n_envs == 1, "Replay buffer only support single environment for now"
 
-        self.observations = np.zeros((self.buffer_size, self.n_envs, self.obs_dim), dtype=np.float32)
+        self.observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, dtype=np.float32)
         self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=np.float32)
-        self.next_observations = np.zeros((self.buffer_size, self.n_envs, self.obs_dim), dtype=np.float32)
+        self.next_observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, dtype=np.float32)
         self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
         self.dones = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
 
@@ -226,7 +226,7 @@ class RolloutBuffer(BaseBuffer):
         self.reset()
 
     def reset(self) -> None:
-        self.observations = np.zeros((self.buffer_size, self.n_envs, self.obs_dim), dtype=np.float32)
+        self.observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, dtype=np.float32)
         self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=np.float32)
         self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
         self.returns = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)

torchy_baselines/common/preprocessing.py

@@ -2,32 +2,92 @@ from typing import Tuple, Union
 
 import numpy as np
 import torch as th
+import torch.nn.functional as F
 from gym import spaces
 
 
-def is_image(observation_space):
+def is_image_space(observation_space: spaces.Space) -> bool:
+    """
+    Check if a observation space has the shape, limits and dtype
+    of a valid image.
+    The check is conservative, so that it returns False
+    if there is a doubt.
+
+    Valid images: RGB, RGBD, GrayScale with values in [0, 255]
+
+    :param observation_space: (spaces.Space)
+    :return: (bool)
+    """
+    if isinstance(observation_space, spaces.Box) and len(observation_space.shape) == 3:
+        # Check the type
+        if observation_space.dtype != np.uint8:
+            return False
+
+        # Check the value range
+        if np.any(observation_space.low != 0) or np.any(observation_space.high != 255):
+            return False
+
+        # Check the number of channels
+        n_channels = observation_space.shape[-1]
+        return n_channels in [1, 3, 4]
     return False
 
 
-def preprocess_obs(obs: th.Tensor, observation_space: spaces.Space) -> th.Tensor:
+def preprocess_obs(obs: th.Tensor, observation_space: spaces.Space,
+                   normalize_image: bool = True) -> th.Tensor:
+    """
+    Preprocess observation to be to a neural network.
+    For images, it normalizes the values by dividing them by 255 (to have values in [0, 1])
+    For discrete observations, it create a one hot vector.
+
+    :param obs: (th.Tensor) Observation
+    :param observation_space: (spaces.Space)
+    :param normalize_image: (bool) Whether to normalize images or not
+        (True by default)
+    :return: (th.Tensor)
+    """
     if isinstance(observation_space, spaces.Box):
-        if is_image(observation_space):
+        if is_image_space(observation_space) and normalize_image:
             return obs / 255.0
         return obs
     elif isinstance(observation_space, spaces.Discrete):
-        # TODO: one hot encoding
+        # One hot encoding and convert to float to avoid errors
+        return F.one_hot(obs, num_classes=observation_space.n).float()
+    else:
+        # TODO: Multidiscrete, Binary, MultiBinary, Tuple, Dict
         raise NotImplementedError()
+
+
+def get_obs_shape(observation_space: spaces.Space) -> Tuple[int, ...]:
+    """
+    Get the shape of the observation (useful for the buffers).
+
+    :param observation_space: (spaces.Space)
+    :return: (Tuple[int, ...])
+    """
+    if isinstance(observation_space, spaces.Box):
+        return observation_space.shape
+    elif isinstance(observation_space, spaces.Discrete):
+        # Observation is an int
+        return (1,)
     else:
         # TODO: Multidiscrete, Binary, MultiBinary, Tuple, Dict
         raise NotImplementedError()
 
 
 def get_obs_dim(observation_space: spaces.Space) -> Union[int, Tuple[int, ...]]:
+    """
+    Get the dimension of the observation space.
+
+    :param observation_space: (spaces.Space)
+    :return: (Union[int, Tuple[int, ...]])
+    """
     if isinstance(observation_space, spaces.Box):
-        if is_image(observation_space):
-            return observation_space.shape
+        # if is_image_space(observation_space):
+        #     raise NotImplementedError()
         return np.prod(observation_space.shape)
     elif isinstance(observation_space, spaces.Discrete):
+        # Observation is an int
         return 1
     else:
         # TODO: Multidiscrete, Binary, MultiBinary, Tuple, Dict
@@ -35,6 +95,12 @@ def get_obs_dim(observation_space: spaces.Space) -> Union[int, Tuple[int, ...]]:
 
 
 def get_action_dim(action_space: spaces.Space) -> int:
+    """
+    Get the dimension of the action space.
+
+    :param action_space: (spaces.Space)
+    :return: (int)
+    """
     if isinstance(action_space, spaces.Box):
         return int(np.prod(action_space.shape))
     elif isinstance(action_space, spaces.Discrete):

torchy_baselines/ppo/policies.py

@@ -65,14 +65,6 @@ class PPOPolicy(BasePolicy):
         self.activation_fn = activation_fn
         self.adam_epsilon = adam_epsilon
         self.ortho_init = ortho_init
-        self.net_args = {
-            'input_dim': self.obs_dim,
-            'output_dim': -1,
-            'net_arch': self.net_arch,
-            'activation_fn': self.activation_fn
-        }
-        self.shared_net = None
-        self.pi_net, self.vf_net = None, None
         # In the future, feature_extractor will be replaced with a CNN
         self.features_extractor = nn.Flatten()
         self.features_dim = self.obs_dim

torchy_baselines/td3/policies.py

@@ -7,7 +7,7 @@ import torch.nn as nn
 from torchy_baselines.common.preprocessing import get_action_dim, get_obs_dim
 from torchy_baselines.common.policies import (BasePolicy, register_policy, create_mlp, BaseNetwork,
                                               create_sde_feature_extractor)
-from torchy_baselines.common.distributions import StateDependentNoiseDistribution
+from torchy_baselines.common.distributions import StateDependentNoiseDistribution, Distribution
 
 
 class Actor(BaseNetwork):
@@ -90,7 +90,8 @@ class Actor(BaseNetwork):
         """
         return self.action_dist.get_std(self.log_std)
 
-    def _get_action_dist_from_latent(self, latent_pi, latent_sde):
+    def _get_action_dist_from_latent(self, latent_pi: th.Tensor,
+                                     latent_sde: th.Tensor) -> Tuple[th.Tensor, Distribution]:
         mean_actions = self.mu(latent_pi)
         return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde)