EnsureImageShape

class EnsureImageShape(env, image_key, image_shape)[source]

Bases: Wrapper

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

Parameters:

  • env – The Gymnasium environment to wrap.

  • image_key – Key in info dict containing the image to validate.

  • image_shape – Expected (height, width) tuple for the image.

Raises:

RuntimeError – If the image shape doesn’t match the expected dimensions.

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.