def __init__()

in local_dm_control_suite/wrappers/pixels.py [0:0]

• 31 lines of code
• 8 McCabe index (conditional complexity)

  def __init__(self, env, pixels_only=True, render_kwargs=None,
               observation_key='pixels'):
    """Initializes a new pixel Wrapper.

    Args:
      env: The environment to wrap.
      pixels_only: If True (default), the original set of 'state' observations
        returned by the wrapped environment will be discarded, and the
        `OrderedDict` of observations will only contain pixels. If False, the
        `OrderedDict` will contain the original observations as well as the
        pixel observations.
      render_kwargs: Optional `dict` containing keyword arguments passed to the
        `mujoco.Physics.render` method.
      observation_key: Optional custom string specifying the pixel observation's
        key in the `OrderedDict` of observations. Defaults to 'pixels'.

    Raises:
      ValueError: If `env`'s observation spec is not compatible with the
        wrapper. Supported formats are a single array, or a dict of arrays.
      ValueError: If `env`'s observation already contains the specified
        `observation_key`.
    """
    if render_kwargs is None:
      render_kwargs = {}

    wrapped_observation_spec = env.observation_spec()

    if isinstance(wrapped_observation_spec, specs.Array):
      self._observation_is_dict = False
      invalid_keys = set([STATE_KEY])
    elif isinstance(wrapped_observation_spec, collections.MutableMapping):
      self._observation_is_dict = True
      invalid_keys = set(wrapped_observation_spec.keys())
    else:
      raise ValueError('Unsupported observation spec structure.')

    if not pixels_only and observation_key in invalid_keys:
      raise ValueError('Duplicate or reserved observation key {!r}.'
                       .format(observation_key))

    if pixels_only:
      self._observation_spec = collections.OrderedDict()
    elif self._observation_is_dict:
      self._observation_spec = wrapped_observation_spec.copy()
    else:
      self._observation_spec = collections.OrderedDict()
      self._observation_spec[STATE_KEY] = wrapped_observation_spec

    # Extend observation spec.
    pixels = env.physics.render(**render_kwargs)
    pixels_spec = specs.Array(
        shape=pixels.shape, dtype=pixels.dtype, name=observation_key)
    self._observation_spec[observation_key] = pixels_spec

    self._env = env
    self._pixels_only = pixels_only
    self._render_kwargs = render_kwargs
    self._observation_key = observation_key