AddPixelsWrapper

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

Bases: Wrapper

Adds rendered environment pixels to info dict with optional resizing and transforms.

Supports single images, dictionaries of images (multiview), or lists of images. Uses PIL for resizing and optional torchvision transforms.

Parameters:

  • env – The Gymnasium environment to wrap.

  • pixels_shape – Target (height, width) for resized images. Defaults to (84, 84).

  • torchvision_transform – Optional transform to apply to PIL images.

Info Keys Added:

  • pixels: Rendered image (single view).

  • pixels.{key}: Individual images (multiview dict).

  • pixels.{idx}: Individual images (multiview list).

  • render_time: Time taken to render in seconds.

reset(*args, **kwargs)[source]

Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]

Uses the step() of the env that can be overwritten to change the returned data.