Observation Wrappers#

Base Class#

class gymnasium.ObservationWrapper(env: Env[ObsType, ActType])[source]#

Superclass of wrappers that can modify observations using observation() for reset() and step().

If you would like to apply a function to only the observation before passing it to the learning code, you can simply inherit from ObservationWrapper and overwrite the method observation() to implement that transformation. The transformation defined in that method must be reflected by the env observation space. Otherwise, you need to specify the new observation space of the wrapper by setting self.observation_space in the __init__() method of your wrapper.

Among others, Gymnasium provides the observation wrapper TimeAwareObservation, which adds information about the index of the timestep to the observation.

Constructor for the observation wrapper.

observation(observation: ObsType) → WrapperObsType[source]#

Returns a modified observation.

Parameters:

observation – The env observation

Returns:

The modified observation

Available Observation Wrappers#

class gymnasium.wrappers.TransformObservation(env: Env, f: Callable[[Any], Any])[source]#

Transform the observation via an arbitrary function f.

The function f should be defined on the observation space of the base environment, env, and should, ideally, return values in the same space.

If the transformation you wish to apply to observations returns values in a different space, you should subclass ObservationWrapper, implement the transformation, and set the new observation space accordingly. If you were to use this wrapper instead, the observation space would be set incorrectly.

Example

>>> import gymnasium as gym
>>> from gymnasium.wrappers import TransformObservation
>>> import numpy as np
>>> np.random.seed(0)
>>> env = gym.make("CartPole-v1")
>>> env = TransformObservation(env, lambda obs: obs + 0.1 * np.random.randn(*obs.shape))
>>> env.reset(seed=42)
(array([0.20380084, 0.03390356, 0.13373359, 0.24382612]), {})
Parameters:

  • env – The environment to apply the wrapper

  • f – A function that transforms the observation

class gymnasium.wrappers.FilterObservation(env: Env, filter_keys: Sequence[str] | None = None)[source]#

Filter Dict observation space by the keys.

Example

>>> import gymnasium as gym
>>> from gymnasium.wrappers import TransformObservation
>>> env = gym.make("CartPole-v1")
>>> env = TransformObservation(env, lambda obs: {'obs': obs, 'time': 0})
>>> env.observation_space = gym.spaces.Dict(obs=env.observation_space, time=gym.spaces.Discrete(1))
>>> env.reset(seed=42)
({'obs': array([ 0.0273956 , -0.00611216,  0.03585979,  0.0197368 ], dtype=float32), 'time': 0}, {})
>>> env = FilterObservation(env, filter_keys=['obs'])
>>> env.reset(seed=42)
({'obs': array([ 0.0273956 , -0.00611216,  0.03585979,  0.0197368 ], dtype=float32)}, {})
>>> env.step(0)
({'obs': array([ 0.02727336, -0.20172954,  0.03625453,  0.32351476], dtype=float32)}, 1.0, False, False, {})
Parameters:

  • env – The environment to apply the wrapper

  • filter_keys – List of keys to be included in the observations. If None, observations will not be filtered and this wrapper has no effect

Raises:

  • ValueError – If the environment’s observation space is not spaces.Dict

  • ValueError – If any of the filter_keys are not included in the original env’s observation space

class gymnasium.wrappers.FlattenObservation(env: Env)[source]#

Observation wrapper that flattens the observation.

Example

>>> import gymnasium as gym
>>> from gymnasium.wrappers import FlattenObservation
>>> env = gym.make("CarRacing-v2")
>>> env.observation_space.shape
(96, 96, 3)
>>> env = FlattenObservation(env)
>>> env.observation_space.shape
(27648,)
>>> obs, _ = env.reset()
>>> obs.shape
(27648,)
Parameters:

env – The environment to apply the wrapper

class gymnasium.wrappers.FrameStack(env: Env, num_stack: int, lz4_compress: bool = False)[source]#

Observation wrapper that stacks the observations in a rolling manner.

For example, if the number of stacks is 4, then the returned observation contains the most recent 4 observations. For environment ‘Pendulum-v1’, the original observation is an array with shape [3], so if we stack 4 observations, the processed observation has shape [4, 3].

Note

  • To be memory efficient, the stacked observations are wrapped by LazyFrame.

  • The observation space must be Box type. If one uses Dict as observation space, it should apply FlattenObservation wrapper first.

  • After reset() is called, the frame buffer will be filled with the initial observation. I.e. the observation returned by reset() will consist of num_stack many identical frames.

Example

>>> import gymnasium as gym
>>> from gymnasium.wrappers import FrameStack
>>> env = gym.make("CarRacing-v2")
>>> env = FrameStack(env, 4)
>>> env.observation_space
Box(0, 255, (4, 96, 96, 3), uint8)
>>> obs, _ = env.reset()
>>> obs.shape
(4, 96, 96, 3)
Parameters:

  • env (Env) – The environment to apply the wrapper

  • num_stack (int) – The number of frames to stack

  • lz4_compress (bool) – Use lz4 to compress the frames internally

class gymnasium.wrappers.GrayScaleObservation(env: Env, keep_dim: bool = False)[source]#

Convert the image observation from RGB to gray scale.

Example

>>> import gymnasium as gym
>>> from gymnasium.wrappers import GrayScaleObservation
>>> env = gym.make("CarRacing-v2")
>>> env.observation_space
Box(0, 255, (96, 96, 3), uint8)
>>> env = GrayScaleObservation(gym.make("CarRacing-v2"))
>>> env.observation_space
Box(0, 255, (96, 96), uint8)
>>> env = GrayScaleObservation(gym.make("CarRacing-v2"), keep_dim=True)
>>> env.observation_space
Box(0, 255, (96, 96, 1), uint8)
Parameters:

  • env (Env) – The environment to apply the wrapper

  • keep_dim (bool) – If True, a singleton dimension will be added, i.e. observations are of the shape AxBx1. Otherwise, they are of shape AxB.

class gymnasium.wrappers.NormalizeObservation(env: Env, epsilon: float = 1e-8)[source]#

This wrapper will normalize observations s.t. each coordinate is centered with unit variance.

Note

The normalization depends on past trajectories and observations will not be normalized correctly if the wrapper was newly instantiated or the policy was changed recently.

Parameters:

  • env (Env) – The environment to apply the wrapper

  • epsilon – A stability parameter that is used when scaling the observations.

class gymnasium.wrappers.PixelObservationWrapper(env: Env, pixels_only: bool = True, render_kwargs: Dict[str, Dict[str, Any]] | None = None, pixel_keys: Tuple[str, ...] = ('pixels',))[source]#

Augment observations by pixel values.

Observations of this wrapper will be dictionaries of images. You can also choose to add the observation of the base environment to this dictionary. In that case, if the base environment has an observation space of type Dict, the dictionary of rendered images will be updated with the base environment’s observation. If, however, the observation space is of type Box, the base environment’s observation (which will be an element of the Box space) will be added to the dictionary under the key “state”.

Example

>>> import gymnasium as gym
>>> from gymnasium.wrappers import PixelObservationWrapper
>>> env = PixelObservationWrapper(gym.make("CarRacing-v2", render_mode="rgb_array"))
>>> obs, _ = env.reset()
>>> obs.keys()
odict_keys(['pixels'])
>>> obs['pixels'].shape
(400, 600, 3)
>>> env = PixelObservationWrapper(gym.make("CarRacing-v2", render_mode="rgb_array"), pixels_only=False)
>>> obs, _ = env.reset()
>>> obs.keys()
odict_keys(['state', 'pixels'])
>>> obs['state'].shape
(96, 96, 3)
>>> obs['pixels'].shape
(400, 600, 3)
>>> env = PixelObservationWrapper(gym.make("CarRacing-v2", render_mode="rgb_array"), pixel_keys=('obs',))
>>> obs, _ = env.reset()
>>> obs.keys()
odict_keys(['obs'])
>>> obs['obs'].shape
(400, 600, 3)
Parameters:

  • env – The environment to wrap.

  • pixels_only (bool) – If True (default), the original observation returned by the wrapped environment will be discarded, and a dictionary observation will only include pixels. If False, the observation dictionary will contain both the original observations and the pixel observations.

  • render_kwargs (dict) – Optional dictionary containing that maps elements of pixel_keys to keyword arguments passed to the self.render() method.

  • pixel_keys – Optional custom string specifying the pixel observation’s key in the OrderedDict of observations. Defaults to (pixels,).

Raises:

  • AssertionError – If any of the keys in render_kwargs``do not show up in ``pixel_keys.

  • ValueError – If env’s observation space is not compatible with the wrapper. Supported formats are a single array, or a dict of arrays.

  • ValueError – If env’s observation already contains any of the specified pixel_keys.

  • TypeError – When an unexpected pixel type is used

class gymnasium.wrappers.ResizeObservation(env: gym.Env, shape: tuple[int, int] | int)[source]#

Resize the image observation.

This wrapper works on environments with image observations. More generally, the input can either be two-dimensional (AxB, e.g. grayscale images) or three-dimensional (AxBxC, e.g. color images). This resizes the observation to the shape given by the 2-tuple shape. The argument shape may also be an integer, in which case, the observation is scaled to a square of side-length shape.

Example

>>> import gymnasium as gym
>>> from gymnasium.wrappers import ResizeObservation
>>> env = gym.make("CarRacing-v2")
>>> env.observation_space.shape
(96, 96, 3)
>>> env = ResizeObservation(env, 64)
>>> env.observation_space.shape
(64, 64, 3)
Parameters:

  • env – The environment to apply the wrapper

  • shape – The shape of the resized observations

class gymnasium.wrappers.TimeAwareObservation(env: Env)[source]#

Augment the observation with the current time step in the episode.

The observation space of the wrapped environment is assumed to be a flat Box. In particular, pixel observations are not supported. This wrapper will append the current timestep within the current episode to the observation.

Example

>>> import gymnasium as gym
>>> from gymnasium.wrappers import TimeAwareObservation
>>> env = gym.make("CartPole-v1")
>>> env = TimeAwareObservation(env)
>>> env.reset(seed=42)
(array([ 0.0273956 , -0.00611216,  0.03585979,  0.0197368 ,  0.        ]), {})
>>> _ = env.action_space.seed(42)
>>> env.step(env.action_space.sample())[0]
array([ 0.02727336, -0.20172954,  0.03625453,  0.32351476,  1.        ])
Parameters:

env – The environment to apply the wrapper