lit_sb3.vec_env.vec_transpose
source module lit_sb3.vec_env.vec_transpose
Classes
-
VecTransposeImage — Re-order channels, from HxWxC to CxHxW.
source class VecTransposeImage(venv: VecEnv, skip: bool = False)
Bases : VecEnvWrapper
Re-order channels, from HxWxC to CxHxW.
It is required for PyTorch convolution layers.
Parameters
-
venv : VecEnv
-
skip : bool — Skip this wrapper if needed as we rely on heuristic to apply it or not, which may result in unwanted behavior, see GH issue #671.
Methods
-
transpose_space — Transpose an observation space (re-order channels).
-
transpose_image — Transpose an image or batch of images (re-order channels).
-
transpose_observations — Transpose (if needed) and return new observations.
-
reset — Reset all environments.
source staticmethod VecTransposeImage.transpose_space(observation_space: spaces.Box, key: str = '') → spaces.Box
Transpose an observation space (re-order channels).
Parameters
-
observation_space : spaces.Box
-
key : str — In case of dictionary space, the key of the observation space.
-
observation_space : spaces.Box — The observation space parameter.
source staticmethod VecTransposeImage.transpose_image(image: np.ndarray) → np.ndarray
Transpose an image or batch of images (re-order channels).
Parameters
-
image : np.ndarray
-
image : np.ndarray — The image parameter.
source method VecTransposeImage.transpose_observations(observations: np.ndarray | dict) → np.ndarray | dict
Transpose (if needed) and return new observations.
Parameters
-
observations : np.ndarray | dict
-
observations : np.ndarray | dict — The observations parameter.
Returns
-
np.ndarray | dict — Transposed observations
source method VecTransposeImage.step_wait() → VecEnvStepReturn
source method VecTransposeImage.reset() → np.ndarray | dict
Reset all environments.
source method VecTransposeImage.close() → None