# coding=utf-8
# Copyright 2020 The TF-Agents Authors.
#
# 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
#
#     https://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.

"""Behavioral Cloning Agents.

Behavioral cloning was proposed in the following articles:

Pomerleau, D.A., 1991. Efficient training of artificial neural networks for
autonomous navigation. Neural Computation, 3(1), pp.88-97.

Russell, S., 1998, July. Learning agents for uncertain environments.
In Proceedings of the eleventh annual conference on Computational learning
theory (pp. 101-103). ACM.
"""

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

import collections
from typing import Callable, Optional, Text

import gin
import tensorflow as tf
import tensorflow_probability as tfp
from tf_agents.agents import data_converter
from tf_agents.agents import tf_agent
from tf_agents.distributions import utils as distribution_utils
from tf_agents.networks import network
from tf_agents.policies import actor_policy
from tf_agents.policies import epsilon_greedy_policy
from tf_agents.policies import greedy_policy
from tf_agents.policies import q_policy
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step as ts
from tf_agents.typing import types
from tf_agents.utils import common
import typing_extensions


Literal = typing_extensions.Literal


class BehavioralCloningLossInfo(
    collections.namedtuple('BehavioralCloningLossInfo', 'loss')
):
  """Stores a per-batch-entry loss value."""

  pass


@gin.configurable
class BehavioralCloningAgent(tf_agent.TFAgent):
  """A Behavioral Cloning agent.

  # pyformat: disable
  Implements a generic form of BehavioralCloning that can also be used to pipe
  supervised learning through TF-Agents. By default the agent defines two types
  of losses.

  For discrete actions the agent uses:

  ```python
  def discrete_loss(agent, experience, training=False):
    bc_logits = self._cloning_network(experience.observation, training=training)

    return tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=experience.action - action_spec.minimum, logits=bc_logits)
  ```

  This requires a Network that generates `num_action` Q-values. In the case of
  continuous actions a simple MSE loss is used by default:

  ```python
  def continuous_loss_fn(agent, experience, training=False):
    bc_output, _ = self._cloning_network(
        experience.observation,
        step_type=experience.step_type,
        training=training,
        network_state=network_state)

    if isinstance(bc_output, tfp.distributions.Distribution):
      bc_action = bc_output.sample()
    else:
      bc_action = bc_output

    return tf.losses.mse(experience.action, bc_action)
  ```
  # pyformat: enable

  The implementation of these loss functions is slightly more complex to
  support nested action_specs.
  """

  def __init__(
      self,
      time_step_spec: ts.TimeStep,
      action_spec: types.NestedTensorSpec,
      cloning_network: network.Network,
      optimizer: types.Optimizer,
      num_outer_dims: Literal[1, 2] = 1,  # pylint: disable=bad-whitespace
      epsilon_greedy: types.Float = 0.1,
      loss_fn: Optional[
          Callable[[types.NestedTensor, bool], types.Tensor]
      ] = None,
      gradient_clipping: Optional[types.Float] = None,
      # Params for debugging.
      debug_summaries: bool = False,
      summarize_grads_and_vars: bool = False,
      train_step_counter: Optional[tf.Variable] = None,
      name: Optional[Text] = None,
  ):
    """Creates an instance of a Behavioral Cloning agent.

    Args:
      time_step_spec: A `TimeStep` spec of the expected time_steps.
      action_spec: A nest of BoundedTensorSpec representing the actions.
      cloning_network: A `tf_agents.networks.Network` to be used by the agent.
        The network will be called as  ``` network(observation,
        step_type=step_type, network_state=initial_state) ``` and must return a
        2-tuple with elements `(output, next_network_state)`
      optimizer: The optimizer to use for training.
      num_outer_dims: The number of outer dimensions for the agent. Must be
        either 1 or 2. If 2, training will require both a batch_size and time
        dimension on every Tensor; if 1, training will require only a batch_size
        outer dimension.
      epsilon_greedy: probability of choosing a random action in the default
        epsilon-greedy collect policy (used only if actions are discrete)
      loss_fn: A function for computing the error between the output of the
        cloning network and the action that was taken. If None, the loss depends
        on the action dtype. The `loss_fn` is called with parameters:
        `(experience, training)`, and must return a loss value for each element
        of the batch.
      gradient_clipping: Norm length to clip gradients.
      debug_summaries: A bool to gather debug summaries.
      summarize_grads_and_vars: If True, gradient and network variable summaries
        will be written during training.
      train_step_counter: An optional counter to increment every time the train
        op is run.  Defaults to the global_step.
      name: The name of this agent. All variables in this module will fall under
        that name. Defaults to the class name.
    """
    tf.Module.__init__(self, name=name)
    self._cloning_network = cloning_network
    self._optimizer = optimizer
    self._gradient_clipping = gradient_clipping

    action_spec = tensor_spec.from_spec(action_spec)
    flat_action_spec = tf.nest.flatten(action_spec)
    continuous_specs = [tensor_spec.is_continuous(s) for s in flat_action_spec]

    if not flat_action_spec:
      raise ValueError('The `action_spec` must contain at least one action.')

    single_discrete_scalar_action = (
        len(flat_action_spec) == 1
        and flat_action_spec[0].shape.rank == 0
        and not tensor_spec.is_continuous(flat_action_spec[0])
    )
    single_continuous_action = len(
        flat_action_spec
    ) == 1 and tensor_spec.is_continuous(flat_action_spec[0])

    if (
        not loss_fn
        and not single_discrete_scalar_action
        and not single_continuous_action
    ):
      raise ValueError(
          'A `loss_fn` must be provided unless there is a single, scalar '
          'discrete action or a single (scalar or non-scalar) continuous '
          'action.'
      )

    self._network_output_spec = cloning_network.create_variables(
        time_step_spec.observation
    )

    # If there is a mix of continuous and discrete actions we want to use an
    # actor policy so we can use the `setup_as_continuous` method as long as the
    # user provided a custom loss_fn which we verified above.
    if any(continuous_specs):
      policy, collect_policy = self._setup_as_continuous(
          time_step_spec, action_spec, loss_fn
      )
    else:
      policy, collect_policy = self._setup_as_discrete(
          time_step_spec, action_spec, loss_fn, epsilon_greedy
      )

    super(BehavioralCloningAgent, self).__init__(
        time_step_spec,
        action_spec,
        policy,
        collect_policy,
        train_sequence_length=None,
        debug_summaries=debug_summaries,
        summarize_grads_and_vars=summarize_grads_and_vars,
        train_step_counter=train_step_counter,
    )

    self._as_trajectory = data_converter.AsTrajectory(
        self.data_context, sequence_length=None, num_outer_dims=num_outer_dims
    )

  @property
  def cloning_network(self):
    return self._cloning_network

  def _discrete_loss(self, experience, training=False):
    batch_size = (
        tf.compat.dimension_value(experience.step_type.shape[0])
        or tf.shape(experience.step_type)[0]
    )

    network_state = self._cloning_network.get_initial_state(batch_size)
    action, _ = self._cloning_network(
        experience.observation,
        step_type=experience.step_type,
        training=training,
        network_state=network_state,
    )

    # Get logits if the output of the cloning network is a distribution.
    if isinstance(action, tfp.distributions.Distribution):
      bc_logits = action.logits_parameter()
    else:
      bc_logits = action

    def loss(action, bc_logits, spec):
      # Subtract the minimum so that we get a proper cross entropy loss on
      # [0, maximum - minimum).
      return tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=action - spec.minimum, logits=bc_logits
      )

    losses = tf.nest.map_structure(
        loss, experience.action, bc_logits, self.action_spec
    )
    losses = tf.nest.flatten(losses)
    return tf.add_n(losses)

  def _continuous_loss_fn(self, experience, training: bool = False):
    batch_size = (
        tf.compat.dimension_value(experience.step_type.shape[0])
        or tf.shape(experience.step_type)[0]
    )
    network_state = self._cloning_network.get_initial_state(batch_size)
    bc_output, _ = self._cloning_network(
        experience.observation,
        step_type=experience.step_type,
        training=training,
        network_state=network_state,
    )

    if isinstance(bc_output, tfp.distributions.Distribution):
      bc_action = bc_output.sample()
    else:
      bc_action = bc_output

    losses = tf.nest.map_structure(tf.losses.mse, experience.action, bc_action)
    losses = tf.nest.flatten(losses)
    return tf.add_n(losses)

  def _setup_as_discrete(
      self, time_step_spec, action_spec, loss_fn, epsilon_greedy
  ):
    self._bc_loss_fn = loss_fn or self._discrete_loss

    if any(
        isinstance(d, distribution_utils.DistributionSpecV2)
        for d in tf.nest.flatten([self._network_output_spec])
    ):
      # If the output of the cloning network contains a distribution.
      base_policy = actor_policy.ActorPolicy(
          time_step_spec, action_spec, self._cloning_network
      )
    else:
      # If the output of the cloning network is logits.
      base_policy = q_policy.QPolicy(
          time_step_spec,
          action_spec,
          q_network=self._cloning_network,
          validate_action_spec_and_network=False,
      )
    policy = greedy_policy.GreedyPolicy(base_policy)
    collect_policy = epsilon_greedy_policy.EpsilonGreedyPolicy(
        base_policy, epsilon=epsilon_greedy
    )
    return policy, collect_policy

  def _setup_as_continuous(self, time_step_spec, action_spec, loss_fn):
    self._bc_loss_fn = loss_fn or self._continuous_loss_fn
    collect_policy = actor_policy.ActorPolicy(
        time_step_spec, action_spec, actor_network=self._cloning_network
    )
    policy = greedy_policy.GreedyPolicy(collect_policy)
    return policy, collect_policy

  def _loss(self, experience, weights=None, training: bool = False):
    experience = self._as_trajectory(experience)

    per_example_loss = self._bc_loss_fn(experience, training=training)  # pytype: disable=wrong-keyword-args  # dynamic-method-lookup
    aggregated_losses = common.aggregate_losses(
        per_example_loss=per_example_loss,
        sample_weight=weights,
        regularization_loss=self._cloning_network.losses,
    )

    return tf_agent.LossInfo(
        loss=aggregated_losses.total_loss,
        extra=BehavioralCloningLossInfo(per_example_loss),
    )

  def _train(self, experience, weights=None):
    experience = self._as_trajectory(experience)

    with tf.GradientTape() as tape:
      per_example_loss = self._bc_loss_fn(experience, training=True)  # pytype: disable=wrong-keyword-args  # dynamic-method-lookup

      aggregated_losses = common.aggregate_losses(
          per_example_loss=per_example_loss,
          sample_weight=weights,
          regularization_loss=self._cloning_network.losses,
      )

    self._apply_loss(
        aggregated_losses,
        variables_to_train=self._cloning_network.trainable_weights,
        tape=tape,
        optimizer=self._optimizer,
    )

    self.train_step_counter.assign_add(1)
    return tf_agent.LossInfo(
        aggregated_losses.total_loss,
        BehavioralCloningLossInfo(per_example_loss),
    )
