code: def pixel_control_rewards(
      observations: Array,
      cell_size: int,
  ) -> Array:
      """Calculates cumulants for pixel control tasks from an observation sequence.
  
      The observations are first split in a grid of KxK cells. For each cell a
      distinct pseudo reward is computed as the average absolute change in pixel
      intensity across all pixels in the cell. The change in intensity is averaged
      across both pixels and channels (e.g. RGB).
  
      The `observations` provided to this function should be cropped suitably, to
      ensure that the observations' height and width are a multiple of `cell_size`.
      The values of the `observations` tensor should be rescaled to [0, 1].
  
      See "Reinforcement Learning with Unsupervised Auxiliary Tasks" by Jaderberg,
      Mnih, Czarnecki et al. (https://arxiv.org/abs/1611.05397).
  
      Args:
        observations: A tensor of shape `[T+1,H,W,C]`, where
          * `T` is the sequence length,
          * `H` is height,
          * `W` is width,
          * `C` is a channel dimension.
        cell_size: The size of each cell.
  
      Returns:
        A tensor of pixel control rewards calculated from the observation. The
        shape is `[T,H',W']`, where `H'=H/cell_size` and `W'=W/cell_size`.
      """
      chex.assert_rank(observations, 4)
      chex.assert_type(observations, float)
  
      # Shape info.
      h = observations.shape[1] // cell_size  # new height.
      w = observations.shape[2] // cell_size  # new width.
      # Calculate the absolute differences across the sequence.
      abs_diff = jnp.abs(observations[1:] - observations[:-1])
      # Average within cells to get the cumulants.
      abs_diff = abs_diff.reshape(
          (-1, h, cell_size, w, cell_size, observations.shape[3])
      )
      return abs_diff.mean(axis=(2, 4, 5))
  
label_desc: Use image node_id 1 for calling a global function with example usage: pixel_control_rewards(observations, cell_size) and returns: abs_diff