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
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`.
  
trfl/pixel_control_ops.py