Wrapper

Previous

Custom environment wrappers

MegaWrapper 

MegaWrapper(
    env: Env,
    image_shape: tuple[int, int] = (84, 84),
    pixels_transform: Callable[[Any], Any] | None = None,
    goal_transform: Callable[[Any], Any] | None = None,
    required_keys: Iterable[str] | None = None,
    separate_goal: bool = True,
    history_size: int = 1,
    frame_skip: int = 1,
)

Bases: Wrapper

Combines multiple wrappers for comprehensive environment preprocessing.

Parameters:

  • env (Env) –

    The environment to wrap.

  • image_shape (tuple[int, int], default: (84, 84) ) –

    Target (height, width) for all image processing.

  • pixels_transform (Callable[[Any], Any] | None, default: None ) –

    Optional transform for rendered pixels.

  • goal_transform (Callable[[Any], Any] | None, default: None ) –

    Optional transform for goal images.

  • required_keys (Iterable[str] | None, default: None ) –

    Keys that must be present in info dict.

  • separate_goal (bool, default: True ) –

    Whether to handle goal separately.

  • history_size (int, default: 1 ) –

    Number of frames to stack.

  • frame_skip (int, default: 1 ) –

    Number of frames to skip for stacking.

EnsureInfoKeysWrapper 

EnsureInfoKeysWrapper(
    env: Env, required_keys: Iterable[str]
)

Bases: Wrapper

Validates that required keys are present in the info dict.

Parameters:

  • env (Env) –

    The environment to wrap.

  • required_keys (Iterable[str]) –

    Iterable of regex patterns that must match keys in info.

EnsureImageShape 

EnsureImageShape(
    env: Env, image_key: str, image_shape: tuple[int, int]
)

Bases: Wrapper

Validates that an image in the info dict has the expected spatial dimensions.

Parameters:

  • env (Env) –

    The environment to wrap.

  • image_key (str) –

    Key in info dict containing the image.

  • image_shape (tuple[int, int]) –

    Expected (height, width) of the image.

EnsureGoalInfoWrapper 

EnsureGoalInfoWrapper(
    env: Env, check_reset: bool, check_step: bool = False
)

Bases: Wrapper

Validates that 'goal' key is present in info dict.

Parameters:

  • env (Env) –

    The environment to wrap.

  • check_reset (bool) –

    Whether to check 'goal' presence on reset.

  • check_step (bool, default: False ) –

    Whether to check 'goal' presence on each step.

EverythingToInfoWrapper 

EverythingToInfoWrapper(env: Env)

Bases: Wrapper

Moves all transition information into the info dict.

Parameters:

  • env (Env) –

    The environment to wrap.

AddPixelsWrapper 

AddPixelsWrapper(
    env: Env,
    pixels_shape: tuple[int, int] = (84, 84),
    torchvision_transform: Callable[[Any], Any]
    | None = None,
)

Bases: Wrapper

Adds rendered environment pixels to info dict.

Parameters:

  • env (Env) –

    The environment to wrap.

  • pixels_shape (tuple[int, int], default: (84, 84) ) –

    Target (height, width) for rendered pixels.

  • torchvision_transform (Callable[[Any], Any] | None, default: None ) –

    Optional transform to apply to the pixels.

ResizeGoalWrapper 

ResizeGoalWrapper(
    env: Env,
    pixels_shape: tuple[int, int] = (84, 84),
    torchvision_transform: Callable[[Any], Any]
    | None = None,
)

Bases: Wrapper

Resizes goal images in info dict.

Parameters:

  • env (Env) –

    The environment to wrap.

  • pixels_shape (tuple[int, int], default: (84, 84) ) –

    Target (height, width) for resizing goal images.

  • torchvision_transform (Callable[[Any], Any] | None, default: None ) –

    Optional transform to apply to goal images.

StackedWrapper 

StackedWrapper(
    env: Env,
    key: str | list[str],
    history_size: int = 1,
    frameskip: int = 1,
)

Bases: Wrapper

Stacks specified key(s) in the info dict over the last k steps.

Parameters:

  • env (Env) –

    The environment to wrap.

  • key (str | list[str]) –

    Key(s) in info dict to stack.

  • history_size (int, default: 1 ) –

    Number of steps to stack.

  • frameskip (int, default: 1 ) –

    Number of steps to skip between stacked frames.

VariationWrapper 

VariationWrapper(
    env: VectorEnv, variation_mode: str | Space = 'same'
)

Bases: VectorWrapper

Manages variation spaces for vectorized environments.

Parameters:

  • env (VectorEnv) –

    The vectorized environment to wrap.

  • variation_mode (str | Space, default: 'same' ) –

    Either 'same', 'different', or a custom Gymnasium space.

Raises:

  • ValueError –

    If variation_mode is invalid or sub-environments are incompatible.

SyncWorld

Bases: SyncVectorEnv

Synchronous vectorized environment with per-environment options support.

Extends SyncVectorEnv to allow passing different options to each sub-environment during reset, enabling per-environment variations and configurations.

This is useful for scenarios where each environment in the vector needs different initialization parameters, such as different variation values or seeds.

Example

env_fns = [lambda: gym.make("MyEnv-v0") for _ in range(3)] vec_env = SyncWorld(env_fns)

Reset with per-env options

options = [{"variation": ["color"]}, {"variation": ["size"]}, None] obs, infos = vec_env.reset(options=options)

Dataset