class_docstr: Abstract base class for Atari wrappers.
  
  This assumes that the input environment is a dm_env.Environment instance in
  which observations are tuples whose first element is an RGB observation and
  the second element is the lives count.
  
  The wrapper itself performs the following modifications:
  
    1. Soft-termination (setting discount to zero) on loss of life.
    2. Action repeats.
    3. Frame pooling for action repeats.
    4. Conversion to grayscale and downscaling.
    5. Reward clipping.
    6. Observation stacking.
  
  The details of grayscale conversion, downscaling, and frame pooling are
  delegated to the concrete subclasses.
  
  This wrapper will raise an error if the underlying Atari environment does not:
  
  - Exposes RGB observations in interleaved format (shape `(H, W, C)`).
  - Expose zero-indexed actions.
  
  Note that this class does not expose a configurable rescale method (defaults
  to bilinear internally).
  
  This class also exposes an additional option `to_float` that doesn't feature
  in other wrappers, which rescales pixel values to floats in the range [0, 1].
docstr: Initializes a new AtariWrapper.
  
  Args:
    environment: An Atari environment.
    max_abs_reward: Maximum absolute reward value before clipping is applied.
      If set to `None` (default), no clipping is applied.
    scale_dims: Image size for the rescaling step after grayscaling, given as
      `(height, width)`. Set to `None` to disable resizing.
    action_repeats: Number of times to step wrapped environment for each given
      action.
    pooled_frames: Number of observations to pool over. Set to 1 to disable
      frame pooling.
    zero_discount_on_life_loss: If `True`, sets the discount to zero when the
      number of lives decreases in in Atari environment.
    expose_lives_observation: If `False`, the `lives` part of the observation
      is discarded, otherwise it is kept as part of an observation tuple. This
      does not affect the `zero_discount_on_life_loss` feature. When enabled,
      the observation consists of a single pixel array, otherwise it is a
      tuple (pixel_array, lives).
    num_stacked_frames: Number of recent (pooled) observations to stack into
      the returned observation.
    flatten_frame_stack: Whether to flatten the stack of frames such that
      the channel (RGB) and stacking dimensions are merged.
    max_episode_len: Number of frames before truncating episode. By default,
      there is no maximum length.
    to_float: If `True`, rescales RGB observations to floats in [0, 1].
    grayscaling: If `True` returns a grayscale version of the observations. In
      this case, the observation is 3D (H, W, num_stacked_frames). If `False`
      the observations are RGB and have shape (H, W, C, num_stacked_frames).
  
  Raises:
    ValueError: For various invalid inputs.
docstr_tok: ["Initializes","a","new","AtariWrapper",".","Args",":","environment",":","An","Atari","environment",".","max_abs_reward",":","Maximum","absolute","reward","value","before","clipping","is","applied",".","If","set","to","`","None","`","(","default",")",",","no","clipping","is","applied",".","scale_dims",":","Image","size","for","the","rescaling","step","after","grayscaling",",","given","as","`","(","height",",","width",")","`",".","Set","to","`","None","`","to","disable","resizing",".","action_repeats",":","Number","of","times","to","step","wrapped","environment","for","each","given","action",".","pooled_frames",":","Number","of","observations","to","pool","over",".","Set","to","1","to","disable","frame","pooling",".","zero_discount_on_life_loss",":","If","`","True","`",",","sets","the","discount","to","zero","when","the","number","of","lives","decreases","in","in","Atari","environment",".","expose_lives_observation",":","If","`","False","`",",","the","`","lives","`","part","of","the","observation","is","discarded",",","otherwise","it","is","kept","as","part","of","an","observation","tuple",".","This","does","not","affect","the","`","zero_discount_on_life_loss","`","feature",".","When","enabled",",","the","observation","consists","of","a","single","pixel","array",",","otherwise","it","is","a","tuple","(","pixel_array",",","lives",")",".","num_stacked_frames",":","Number","of","recent","(","pooled",")","observations","to","stack","into","the","returned","observation",".","flatten_frame_stack",":","Whether","to","flatten","the","stack","of","frames","such","that","the","channel","(","RGB",")","and","stacking","dimensions","are","merged",".","max_episode_len",":","Number","of","frames","before","truncating","episode",".","By","default",",","there","is","no","maximum","length",".","to_float",":","If","`","True","`",",","rescales","RGB","observations","to","floats","in","[","0",",","1","]",".","grayscaling",":","If","`","True","`","returns","a","grayscale","version","of","the","observations",".","In","this","case",",","the","observation","is","3D","(","H",",","W",",","num_stacked_frames",")",".","If","`","False","`","the","observations","are","RGB","and","have","shape","(","H",",","W",",","C",",","num_stacked_frames",")",".","Raises",":","ValueError",":","For","various","invalid","inputs","."]
code: def __init__(
      self,
      environment: dm_env.Environment,
      *,
      max_abs_reward: Optional[float] = None,
      scale_dims: Optional[Tuple[int, int]] = (84, 84),
      action_repeats: int = 4,
      pooled_frames: int = 2,
      zero_discount_on_life_loss: bool = False,
      expose_lives_observation: bool = False,
      num_stacked_frames: int = 4,
      flatten_frame_stack: bool = False,
      max_episode_len: Optional[int] = None,
      to_float: bool = False,
      grayscaling: bool = True,
  ):
      """Initializes a new AtariWrapper.
  
      Args:
        environment: An Atari environment.
        max_abs_reward: Maximum absolute reward value before clipping is applied.
          If set to `None` (default), no clipping is applied.
        scale_dims: Image size for the rescaling step after grayscaling, given as
          `(height, width)`. Set to `None` to disable resizing.
        action_repeats: Number of times to step wrapped environment for each given
          action.
        pooled_frames: Number of observations to pool over. Set to 1 to disable
          frame pooling.
        zero_discount_on_life_loss: If `True`, sets the discount to zero when the
          number of lives decreases in in Atari environment.
        expose_lives_observation: If `False`, the `lives` part of the observation
          is discarded, otherwise it is kept as part of an observation tuple. This
          does not affect the `zero_discount_on_life_loss` feature. When enabled,
          the observation consists of a single pixel array, otherwise it is a
          tuple (pixel_array, lives).
        num_stacked_frames: Number of recent (pooled) observations to stack into
          the returned observation.
        flatten_frame_stack: Whether to flatten the stack of frames such that
          the channel (RGB) and stacking dimensions are merged.
        max_episode_len: Number of frames before truncating episode. By default,
          there is no maximum length.
        to_float: If `True`, rescales RGB observations to floats in [0, 1].
        grayscaling: If `True` returns a grayscale version of the observations. In
          this case, the observation is 3D (H, W, num_stacked_frames). If `False`
          the observations are RGB and have shape (H, W, C, num_stacked_frames).
  
      Raises:
        ValueError: For various invalid inputs.
      """
      if not 1 <= pooled_frames <= action_repeats:
          raise ValueError(
              "pooled_frames ({}) must be between 1 and "
              "action_repeats ({}) inclusive".format(
                  pooled_frames, action_repeats
              )
          )
  
      if zero_discount_on_life_loss:
          super().__init__(_ZeroDiscountOnLifeLoss(environment))
      else:
          super().__init__(environment)
  
      if not max_episode_len:
          max_episode_len = np.inf
  
      self._frame_stacker = frame_stacking.FrameStacker(
          num_frames=num_stacked_frames, flatten=flatten_frame_stack
      )
      self._action_repeats = action_repeats
      self._pooled_frames = pooled_frames
      self._scale_dims = scale_dims
      self._max_abs_reward = max_abs_reward or np.inf
      self._to_float = to_float
      self._expose_lives_observation = expose_lives_observation
  
      if scale_dims:
          self._height, self._width = scale_dims
      else:
          spec = environment.observation_spec()
          self._height, self._width = spec[RGB_INDEX].shape[:2]
  
      self._episode_len = 0
      self._max_episode_len = max_episode_len
      self._reset_next_step = True
  
      self._grayscaling = grayscaling
  
      # Based on underlying observation spec, decide whether lives are to be
      # included in output observations.
      observation_spec = self._environment.observation_spec()
      spec_names = [spec.name for spec in observation_spec]
      if "lives" in spec_names and spec_names.index("lives") != 1:
          raise ValueError(
              "`lives` observation needs to have index 1 in Atari."
          )
  
      self._observation_spec = self._init_observation_spec()
  
      self._raw_observation = None
  
imports: import abc
  from typing import Tuple, List, Optional, Sequence, Union
  from acme.wrappers import base
  from acme.wrappers import frame_stacking
  import dm_env
  from dm_env import specs
  import numpy
  from PIL import Image
label_desc: Use image node_id 1 to create a new BaseAtariWrapper object from inherited base classes: abc, base with example: obj = BaseAtariWrapper(environment)