# Copyright 2018 The trfl Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""TensorFlow ops for implementing Pixel Control.

Pixel Control is an auxiliary task introduced in the UNREAL agent.
In Pixel Control an additional agent head is trained off-policy to predict
action-value functions for a host of pseudo rewards derived from the stream of
observations. This leads to better state representations and therefore improved
performance, both in terms of data efficiency and final performance.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

# Dependency imports
import tensorflow.compat.v1 as tf
from trfl import action_value_ops
from trfl import base_ops


PixelControlExtra = collections.namedtuple(
    "pixel_control_extra", ["spatial_loss", "pseudo_rewards"])


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


def pixel_control_loss(
    observations, actions, action_values, cell_size, discount_factor,
    scale, crop_height_dim=(None, None), crop_width_dim=(None, None)):
  """Calculate n-step Q-learning loss for pixel control auxiliary task.

  For each pixel-based pseudo reward signal, the corresponding action-value
  function is trained off-policy, using Q(lambda). A discount of 0.9 is
  commonly used for learning the value functions.

  Note that, since pseudo rewards have a spatial structure, with neighbouring
  cells exhibiting strong correlations, it is convenient to predict the action
  values for all the cells through a deconvolutional head.

  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, ...]`; `...` is the observation
      shape, `T` the sequence length, and `B` the batch size. `T` and `B` can
      be statically unknown for `observations`, `actions` and `action_values`.
    actions: A tensor, shape `[T,B]`, of the actions across each sequence.
    action_values: A tensor, shape `[T+1,B,H,W,N]` of pixel control action
      values, where `H`, `W` are the number of pixel control cells/tasks, and
      `N` is the number of actions.
    cell_size: size of the cells used to derive the pixel based pseudo-rewards.
    discount_factor: discount used for learning the value function associated
      to the pseudo rewards; must be a scalar or a Tensor of shape [T,B].
    scale: scale factor for pixels in `observations`.
    crop_height_dim: tuple (min_height, max_height) specifying how
      to crop the input observations before computing the pseudo-rewards.
    crop_width_dim: tuple (min_width, max_width) specifying how
      to crop the input observations before computing the pseudo-rewards.

  Returns:
    A namedtuple with fields:

    * `loss`: a tensor containing the batch of losses, shape [B].
    * `extra`: a namedtuple with fields:
        * `target`: batch of target values for `q_tm1[a_tm1]`, shape [B].
        * `td_error`: batch of temporal difference errors, shape [B].

  Raises:
    ValueError: if the shape of `action_values` is not compatible with that of
      the pseudo-rewards derived from the observations.
  """
  # Useful shapes.
  sequence_length, batch_size = base_ops.best_effort_shape(actions)
  num_actions = action_values.get_shape().as_list()[-1]
  height_width_q = action_values.get_shape().as_list()[2:-1]
  # Calculate rewards using the observations. Crop observations if appropriate.
  if crop_height_dim[0] is not None:
    h_low, h_high = crop_height_dim
    observations = observations[:, :, h_low:h_high, :]
  if crop_width_dim[0] is not None:
    w_low, w_high = crop_width_dim
    observations = observations[:, :, :, w_low:w_high]
  # Rescale observations by a constant factor.
  observations *= tf.constant(scale)
  # Compute pseudo-rewards and get their shape.
  pseudo_rewards = pixel_control_rewards(observations, cell_size)
  height_width = pseudo_rewards.get_shape().as_list()[2:]
  # Check that pseudo-rewards and Q-values are compatible in shape.
  if height_width != height_width_q:
    raise ValueError(
        "Pixel Control values are not compatible with the shape of the"
        "pseudo-rewards derived from the observation. Pseudo-rewards have shape"
        "{}, while Pixel Control values have shape {}".format(
            height_width, height_width_q))
  # We now have Q(s,a) and rewards, so can calculate the n-step loss. The
  # QLambda loss op expects inputs of shape [T,B,N] and [T,B], but our tensors
  # are in a variety of incompatible shapes. The state-action values have
  # shape [T,B,H,W,N] and rewards have shape [T,B,H,W]. We can think of the
  # [H,W] dimensions as extra batch dimensions for the purposes of the loss
  # calculation, so we first collapse [B,H,W] into a single dimension.
  q_tm1 = tf.reshape(
      action_values[:-1],  # [T,B,H,W,N].
      [sequence_length, -1, num_actions],
      name="q_tm1")  # [T,BHW,N].
  r_t = tf.reshape(
      pseudo_rewards,  # [T,B,H,W].
      [sequence_length, -1],
      name="r_t")  # [T,BHW].
  q_t = tf.reshape(
      action_values[1:],  # [T,B,H,W,N].
      [sequence_length, -1, num_actions],
      name="q_t")  # [T,BHW,N].
  # The actions tensor is of shape [T,B], and is the same for each H and W.
  # We thus expand it to be same shape as the reward tensor, [T,BHW].
  expanded_actions = tf.expand_dims(tf.expand_dims(actions, -1), -1)
  a_tm1 = tf.tile(
      expanded_actions, multiples=[1, 1] + height_width)  # [T,B,H,W].
  a_tm1 = tf.reshape(a_tm1, [sequence_length, -1])  # [T,BHW].
  # We similarly expand-and-tile the discount to [T,BHW].
  discount_factor = tf.convert_to_tensor(discount_factor)
  if discount_factor.shape.ndims == 0:
    pcont_t = tf.reshape(discount_factor, [1, 1])  # [1,1].
    pcont_t = tf.tile(pcont_t, tf.shape(a_tm1))  # [T,BHW].
  elif discount_factor.shape.ndims == 2:
    tiled_pcont = tf.tile(
        tf.expand_dims(tf.expand_dims(discount_factor, -1), -1),
        [1, 1] + height_width)
    pcont_t = tf.reshape(tiled_pcont, [sequence_length, -1])
  else:
    raise ValueError(
        "The discount_factor must be a scalar or a tensor of rank 2."
        "instead is a tensor of shape {}".format(
            discount_factor.shape.as_list()))
  # Compute a QLambda loss of shape [T,BHW]
  loss, _ = action_value_ops.qlambda(q_tm1, a_tm1, r_t, pcont_t, q_t, lambda_=1)
  # Take sum over sequence, sum over cells.
  expanded_shape = [sequence_length, batch_size] + height_width
  spatial_loss = tf.reshape(loss, expanded_shape)  # [T,B,H,W].
  # Return.
  extra = PixelControlExtra(
      spatial_loss=spatial_loss, pseudo_rewards=pseudo_rewards)
  return base_ops.LossOutput(
      tf.reduce_sum(spatial_loss, axis=[0, 2, 3]), extra)  # [B]
