diff --git a/torchy_baselines/common/buffers.py b/torchy_baselines/common/buffers.py index 0932ce6..fd78d3e 100644 --- a/torchy_baselines/common/buffers.py +++ b/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) diff --git a/torchy_baselines/common/preprocessing.py b/torchy_baselines/common/preprocessing.py index 024deca..2b06e6e 100644 --- a/torchy_baselines/common/preprocessing.py +++ b/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): diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 163e890..ef366ed 100644 --- a/torchy_baselines/ppo/policies.py +++ b/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 diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index b48e98a..9b6f5d7 100644 --- a/torchy_baselines/td3/policies.py +++ b/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)