code: def pixel_control_rewards(observations, cell_size):
      """Calculates pixel control task rewards from 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 for 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]. In the
      UNREAL agent observations are cropped to 80x80, and each cell is 4x4 in size.
  
      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,B,H,W,C...]`, where
          * `T` is the sequence length, `B` is the batch size.
          * `H` is height, `W` is width.
          * `C...` is at least one channel dimension (e.g., colour, stack).
          * `T` and `B` can be statically unknown.
        cell_size: The size of each cell.
  
      Returns:
        A tensor of pixel control rewards calculated from the observation. The
        shape is `[T,B,H',W']`, where `H'` and `W'` are determined by the
        `cell_size`. If evenly-divisible, `H' = H/cell_size`, and similar for `W`.
      """
      # Calculate the absolute differences across the sequence.
      abs_diff = tf.abs(observations[1:] - observations[:-1])
      # Average over cells. `abs_diff` has shape [T,B,H,W,C...], e.g.,
      # [T,B,H,W,C] if we have a colour channel. We want to use the TF avg_pool3d
      # op, but it expects 5D inputs so we collapse all channel dimensions.
      # Merge remaining dimensions after W: [T,B,H,W,C'].
      full_shape = tf.shape(abs_diff)
      preserved_shape = full_shape[:4]
      trailing_shape = (tf.reduce_prod(full_shape[4:]),)
      shape = tf.concat([preserved_shape, trailing_shape], 0)
      abs_diff = tf.reshape(abs_diff, shape)
      # Apply the averaging using average pooling and reducing over channel.
      avg_abs_diff = tf.nn.avg_pool3d(
          abs_diff,
          ksize=[1, 1, cell_size, cell_size, 1],
          strides=[1, 1, cell_size, cell_size, 1],
          padding="VALID",
      )  # [T,B,H',W',C'].
      pseudo_rewards = tf.reduce_mean(
          avg_abs_diff, axis=[4], name="pseudo_rewards"
      )  # [T,B,H',W'].
      sequence_batch = abs_diff.get_shape()[:2]
      new_height_width = avg_abs_diff.get_shape()[2:4]
      pseudo_rewards.set_shape(
          sequence_batch.concatenate(new_height_width)
      )
      return pseudo_rewards
  
file_path: trfl/trfl/pixel_control_ops.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import tensorflow.compat.v1
  from trfl import action_value_ops
  from trfl import base_ops
label_desc: Use image node_id 1 for calling a global function with example usage: pixel_control_rewards(observations, cell_size) and returns: pseudo_rewards