diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 1872c23..69fce7d 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -28,6 +28,7 @@ Others: ^^^^^^^ - Enabled Python 3.9 in GitHub CI - Fixed type annotations +- Refactored ``predict()`` by moving the preprocessing to ``obs_to_tensor()`` method Documentation: ^^^^^^^^^^^^^^ diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 0ed6be5..1550e23 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -204,6 +204,46 @@ class BaseModel(nn.Module, ABC): """ self.train(mode) + def obs_to_tensor(self, observation: Union[np.ndarray, Dict[str, np.ndarray]]) -> Tuple[th.Tensor, bool]: + """ + Convert an input observation to a PyTorch tensor that can be fed to a model. + Includes sugar-coating to handle different observations (e.g. normalizing images). + + :param observation: the input observation + :return: The observation as PyTorch tensor + and whether the observation is vectorized or not + """ + vectorized_env = False + if isinstance(observation, dict): + # need to copy the dict as the dict in VecFrameStack will become a torch tensor + observation = copy.deepcopy(observation) + for key, obs in observation.items(): + obs_space = self.observation_space.spaces[key] + if is_image_space(obs_space): + obs_ = maybe_transpose(obs, obs_space) + else: + obs_ = np.array(obs) + vectorized_env = vectorized_env or is_vectorized_observation(obs_, obs_space) + # Add batch dimension if needed + observation[key] = obs_.reshape((-1,) + self.observation_space[key].shape) + + elif is_image_space(self.observation_space): + # Handle the different cases for images + # as PyTorch use channel first format + observation = maybe_transpose(observation, self.observation_space) + + else: + observation = np.array(observation) + + if not isinstance(observation, dict): + # Dict obs need to be handled separately + vectorized_env = is_vectorized_observation(observation, self.observation_space) + # Add batch dimension if needed + observation = observation.reshape((-1,) + self.observation_space.shape) + + observation = obs_as_tensor(observation, self.device) + return observation, vectorized_env + class BasePolicy(BaseModel): """The base policy object. @@ -280,35 +320,7 @@ class BasePolicy(BaseModel): # Switch to eval mode (this affects batch norm / dropout) self.set_training_mode(False) - vectorized_env = False - if isinstance(observation, dict): - # need to copy the dict as the dict in VecFrameStack will become a torch tensor - observation = copy.deepcopy(observation) - for key, obs in observation.items(): - obs_space = self.observation_space.spaces[key] - if is_image_space(obs_space): - obs_ = maybe_transpose(obs, obs_space) - else: - obs_ = np.array(obs) - vectorized_env = vectorized_env or is_vectorized_observation(obs_, obs_space) - # Add batch dimension if needed - observation[key] = obs_.reshape((-1,) + self.observation_space[key].shape) - - elif is_image_space(self.observation_space): - # Handle the different cases for images - # as PyTorch use channel first format - observation = maybe_transpose(observation, self.observation_space) - - else: - observation = np.array(observation) - - if not isinstance(observation, dict): - # Dict obs need to be handled separately - vectorized_env = is_vectorized_observation(observation, self.observation_space) - # Add batch dimension if needed - observation = observation.reshape((-1,) + self.observation_space.shape) - - observation = obs_as_tensor(observation, self.device) + observation, vectorized_env = self.obs_to_tensor(observation) with th.no_grad(): actions = self._predict(observation, deterministic=deterministic) @@ -324,9 +336,8 @@ class BasePolicy(BaseModel): # out of bound error (e.g. if sampling from a Gaussian distribution) actions = np.clip(actions, self.action_space.low, self.action_space.high) + # Remove batch dimension if needed if not vectorized_env: - if state is not None: - raise ValueError("Error: The environment must be vectorized when using recurrent policies.") actions = actions[0] return actions, state