From 02b452a9f958cc857d9de85ba8eec385c9c91410 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Wed, 22 Apr 2020 18:28:00 +0200 Subject: [PATCH] Avoid double wrapping --- torchy_baselines/common/base_class.py | 2 +- torchy_baselines/common/preprocessing.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index ac49453..48936e1 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -141,7 +141,7 @@ class BaseRLModel(ABC): print("Wrapping the env in a DummyVecEnv.") env = DummyVecEnv([lambda: env]) - if is_image_space(env.observation_space): + if is_image_space(env.observation_space) and not isinstance(env, VecTransposeImage): if self.verbose >= 1: print("Wrapping the env in a VecTransposeImage.") env = VecTransposeImage(env) diff --git a/torchy_baselines/common/preprocessing.py b/torchy_baselines/common/preprocessing.py index af81684..332bea6 100644 --- a/torchy_baselines/common/preprocessing.py +++ b/torchy_baselines/common/preprocessing.py @@ -20,7 +20,7 @@ def is_image_space(observation_space: spaces.Space, :param observation_space: (spaces.Space) :param channels_last: (bool) :param check_channels: (bool) Whether to do or not the check for the number of channels. - Because of frame-skip, the observation space may have more channels than expected. + e.g., with frame-stacking, the observation space may have more channels than expected. :return: (bool) """ if isinstance(observation_space, spaces.Box) and len(observation_space.shape) == 3: