# 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.

"""TensorFlow Policies API."""

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

import abc
from typing import Optional, Sequence, Text

import six
import tensorflow as tf
import tensorflow_probability as tfp
from tf_agents.distributions import reparameterized_sampling
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import policy_step
from tf_agents.trajectories import time_step as ts
from tf_agents.trajectories import trajectory
from tf_agents.typing import types
from tf_agents.utils import common
from tf_agents.utils import nest_utils


tfd = tfp.distributions


@six.add_metaclass(abc.ABCMeta)
class TFPolicy(tf.Module):
  """Abstract base class for TF Policies.

  The Policy represents a mapping from `time_steps` recieved from the
  environment to `actions` that can be applied to the environment.

  Agents expose two policies. A `policy` meant for deployment and evaluation,
  and a `collect_policy` for collecting data from the environment. The
  `collect_policy` is usually stochastic for exploring the environment better
  and may log auxilliary information such as log probabilities required for
  training as well. `Policy` objects can also be created directly by the users
  without using an `Agent`.

  The main methods of TFPolicy are:

  * `action`: Maps a `time_step` from the environment to an action.
  * `distribution`: Maps a `time_step` to a distribution over actions.
  * `get_initial_state`: Generates the initial state for stateful policies, e.g.
      RNN/LSTM policies.

  Example usage:

  ```
  env = SomeTFEnvironment()
  policy = TFRandomPolicy(env.time_step_spec(), env.action_spec())
  # Or policy = agent.policy or agent.collect_policy

  policy_state = policy.get_initial_state(env.batch_size)
  time_step = env.reset()

  while not time_step.is_last():
    policy_step = policy.action(time_step, policy_state)
    time_step = env.step(policy_step.action)

    policy_state = policy_step.state
    # policy_step.info may contain side info for logging, such as action log
    # probabilities.
  ```

  Policies can be saved to disk as SavedModels (see policy_saver.py and
  policy_loader.py) or as TF Checkpoints.

  A `PyTFEagerPolicy` can be used to wrap a `TFPolicy` so that it works with
  `PyEnvironment`s.


  **NOTE**: For API consistency, subclasses are not allowed to override public
  methods of `TFPolicy` class. Instead, they may implement the protected methods
  including `_get_initial_state`, `_action`, and `_distribution`. This
  public-calls-private convention allowed this base class to do things like
  properly add `spec` and shape checks, which provide users an easier experience
  when debugging their environments and networks.

  For researchers, and those developing new Policies, the `TFPolicy` base class
  constructor also accept a `validate_args` parameter.  If `False`, this
  disables all spec structure, dtype, and shape checks in the public methods of
  these classes.  It allows algorithm developers to iterate and try different
  input and output structures without worrying about overly restrictive
  requirements, or input and output states being in a certain format.  However,
  *disabling argument validation* can make it very hard to identify structural
  input or algorithmic errors; and should not be done for final, or
  production-ready, Policies.  In addition to having implementations that may
  disagree with specs, this mean that the resulting Policy may no longer
  interact well with other parts of TF-Agents.  Examples include impedance
  mismatches with Actor/Learner APIs, replay buffers, and the model export
  functionality in `PolicySaver.
  """

  # TODO(b/127327645) Remove this attribute.
  # This attribute allows subclasses to back out of automatic tf.function
  # attribute inside TF1 (for autodeps).
  _enable_functions = True

  def __init__(
      self,
      time_step_spec: ts.TimeStep,
      action_spec: types.NestedTensorSpec,
      policy_state_spec: types.NestedTensorSpec = (),
      info_spec: types.NestedTensorSpec = (),
      clip: bool = True,
      emit_log_probability: bool = False,
      automatic_state_reset: bool = True,
      observation_and_action_constraint_splitter: Optional[
          types.Splitter
      ] = None,
      validate_args: bool = True,
      name: Optional[Text] = None,
  ):
    """Initialization of TFPolicy class.

    Args:
      time_step_spec: A `TimeStep` spec of the expected time_steps. Usually
        provided by the user to the subclass.
      action_spec: A nest of BoundedTensorSpec representing the actions. Usually
        provided by the user to the subclass.
      policy_state_spec: A nest of TensorSpec representing the policy_state.
        Provided by the subclass, not directly by the user.
      info_spec: A nest of TensorSpec representing the policy info. Provided by
        the subclass, not directly by the user.
      clip: Whether to clip actions to spec before returning them.  Default
        True. Most policy-based algorithms (PCL, PPO, REINFORCE) use unclipped
        continuous actions for training.
      emit_log_probability: Emit log-probabilities of actions, if supported. If
        True, policy_step.info will have CommonFields.LOG_PROBABILITY set.
        Please consult utility methods provided in policy_step for setting and
        retrieving these. When working with custom policies, either provide a
        dictionary info_spec or a namedtuple with the field 'log_probability'.
      automatic_state_reset:  If `True`, then `get_initial_policy_state` is used
        to clear state in `action()` and `distribution()` for for time steps
        where `time_step.is_first()`.
      observation_and_action_constraint_splitter: A function used to process
        observations with action constraints. These constraints can indicate,
        for example, a mask of valid/invalid actions for a given state of the
        environment. The function takes in a full observation and returns a
        tuple consisting of 1) the part of the observation intended as input to
        the network and 2) the constraint. An example
        `observation_and_action_constraint_splitter` could be as simple as: ```
        def observation_and_action_constraint_splitter(observation): return
        observation['network_input'], observation['constraint'] ``` *Note*: when
        using `observation_and_action_constraint_splitter`, make sure the
        provided `q_network` is compatible with the network-specific half of the
        output of the `observation_and_action_constraint_splitter`. In
        particular, `observation_and_action_constraint_splitter` will be called
        on the observation before passing to the network. If
        `observation_and_action_constraint_splitter` is None, action constraints
        are not applied.
      validate_args: Python bool.  Whether to verify inputs to, and outputs of,
        functions like `action` and `distribution` against spec structures,
        dtypes, and shapes.  Research code may prefer to set this value to
        `False` to allow iterating on input and output structures without being
        hamstrung by overly rigid checking (at the cost of harder-to-debug
        errors).  See also `TFAgent.validate_args`.
      name: A name for this module. Defaults to the class name.
    """
    super(TFPolicy, self).__init__(name=name)
    common.check_tf1_allowed()
    common.tf_agents_gauge.get_cell('TFAPolicy').set(True)
    common.assert_members_are_not_overridden(base_cls=TFPolicy, instance=self)
    if not isinstance(time_step_spec, ts.TimeStep):
      raise ValueError(
          'The `time_step_spec` must be an instance of `TimeStep`, but is `{}`.'
          .format(type(time_step_spec))
      )

    self._time_step_spec = tensor_spec.from_spec(time_step_spec)
    self._action_spec = tensor_spec.from_spec(action_spec)
    self._policy_state_spec = tensor_spec.from_spec(policy_state_spec)
    self._emit_log_probability = emit_log_probability
    self._validate_args = validate_args

    if emit_log_probability:
      log_probability_spec = tensor_spec.BoundedTensorSpec(
          shape=(),
          dtype=tf.float32,
          maximum=0,
          minimum=-float('inf'),
          name='log_probability',
      )
      log_probability_spec = tf.nest.map_structure(
          lambda _: log_probability_spec, action_spec
      )
      info_spec = policy_step.set_log_probability(
          info_spec, log_probability_spec
      )  # pytype: disable=wrong-arg-types

    self._info_spec = tensor_spec.from_spec(info_spec)
    self._setup_specs()
    self._clip = clip
    self._action_fn = common.function_in_tf1(reduce_retracing=False)(
        self._action
    )
    self._automatic_state_reset = automatic_state_reset
    self._observation_and_action_constraint_splitter = (
        observation_and_action_constraint_splitter
    )

  def _setup_specs(self):
    self._policy_step_spec = policy_step.PolicyStep(
        action=self._action_spec,
        state=self._policy_state_spec,
        info=self._info_spec,
    )
    self._trajectory_spec = trajectory.from_transition(
        self._time_step_spec, self._policy_step_spec, self._time_step_spec
    )

  def variables(self) -> Sequence[tf.Variable]:
    """Returns the list of Variables that belong to the policy."""
    # Ignore self._variables() in favor of using tf.Module's tracking.
    return super(TFPolicy, self).variables

  @property
  def observation_and_action_constraint_splitter(self) -> types.Splitter:
    return self._observation_and_action_constraint_splitter

  @property
  def validate_args(self) -> bool:
    """Whether `action` & `distribution` validate input and output args."""
    return self._validate_args

  def get_initial_state(
      self, batch_size: Optional[types.Int]
  ) -> types.NestedTensor:
    """Returns an initial state usable by the policy.

    Args:
      batch_size: Tensor or constant: size of the batch dimension. Can be None
        in which case no dimensions gets added.

    Returns:
      A nested object of type `policy_state` containing properly
      initialized Tensors.
    """
    return self._get_initial_state(batch_size)

  def _maybe_reset_state(self, time_step, policy_state):
    if policy_state is ():  # pylint: disable=literal-comparison
      return policy_state

    batch_size = tf.compat.dimension_value(time_step.discount.shape[0])
    if batch_size is None:
      batch_size = tf.shape(time_step.discount)[0]

    # Make sure we call this with a kwarg as it may be wrapped in tf.function
    # which would expect a tensor if it was not a kwarg.
    zero_state = self.get_initial_state(batch_size=batch_size)
    condition = time_step.is_first()
    # When experience is a sequence we only reset automatically for the first
    # time_step in the sequence as we can't easily generalize how the policy is
    # unrolled over the sequence.
    if nest_utils.get_outer_rank(time_step, self._time_step_spec) > 1:
      condition = time_step.is_first()[:, 0, ...]
    return nest_utils.where(condition, zero_state, policy_state)

  def action(
      self,
      time_step: ts.TimeStep,
      policy_state: types.NestedTensor = (),
      seed: Optional[types.Seed] = None,
  ) -> policy_step.PolicyStep:
    """Generates next action given the time_step and policy_state.

    Args:
      time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
      policy_state: A Tensor, or a nested dict, list or tuple of Tensors
        representing the previous policy_state.
      seed: Seed to use if action performs sampling (optional).

    Returns:
      A `PolicyStep` named tuple containing:
        `action`: An action Tensor matching the `action_spec`.
        `state`: A policy state tensor to be fed into the next call to action.
        `info`: Optional side information such as action log probabilities.

    Raises:
      RuntimeError: If subclass __init__ didn't call super().__init__.
      ValueError or TypeError: If `validate_args is True` and inputs or
        outputs do not match `time_step_spec`, `policy_state_spec`,
        or `policy_step_spec`.
    """
    if self._enable_functions and getattr(self, '_action_fn', None) is None:
      raise RuntimeError(
          'Cannot find _action_fn.  Did %s.__init__ call super?'
          % type(self).__name__
      )
    if self._enable_functions:
      action_fn = self._action_fn
    else:
      action_fn = self._action

    if self._validate_args:
      time_step = nest_utils.prune_extra_keys(self._time_step_spec, time_step)
      policy_state = nest_utils.prune_extra_keys(
          self._policy_state_spec, policy_state
      )
      nest_utils.assert_same_structure(
          time_step,
          self._time_step_spec,
          message='time_step and time_step_spec structures do not match',
      )
      if common.safe_has_state(policy_state):
        nest_utils.assert_same_structure(
            policy_state,
            self._policy_state_spec,
            message=(
                'policy_state and policy_state_spec structures do not match'
            ),
        )

    if self._automatic_state_reset:
      policy_state = self._maybe_reset_state(time_step, policy_state)
    step = action_fn(time_step=time_step, policy_state=policy_state, seed=seed)

    def clip_action(action, action_spec):
      if isinstance(action_spec, tensor_spec.BoundedTensorSpec):
        return common.clip_to_spec(action, action_spec)
      return action

    if self._validate_args:
      nest_utils.assert_same_structure(
          step.action,
          self._action_spec,
          message='action and action_spec structures do not match',
      )

    if self._clip:
      clipped_actions = tf.nest.map_structure(
          clip_action, step.action, self._action_spec
      )
      step = step._replace(action=clipped_actions)

    if self._validate_args:
      nest_utils.assert_same_structure(
          step,
          self._policy_step_spec,
          message='action output and policy_step_spec structures do not match',
      )

      def compare_to_spec(value, spec):
        return value.dtype.is_compatible_with(spec.dtype)

      compatibility = [
          compare_to_spec(v, s)
          for (v, s) in zip(
              tf.nest.flatten(step.action), tf.nest.flatten(self.action_spec)
          )
      ]

      if not all(compatibility):
        get_dtype = lambda x: x.dtype
        action_dtypes = tf.nest.map_structure(get_dtype, step.action)
        spec_dtypes = tf.nest.map_structure(get_dtype, self.action_spec)

        raise TypeError(
            "Policy produced an action with a dtype that doesn't "
            'match its action_spec. Got action:\n  %s\n with '
            'action_spec:\n  %s' % (action_dtypes, spec_dtypes)
        )

    return step

  def distribution(
      self, time_step: ts.TimeStep, policy_state: types.NestedTensor = ()
  ) -> policy_step.PolicyStep:
    """Generates the distribution over next actions given the time_step.

    Args:
      time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
      policy_state: A Tensor, or a nested dict, list or tuple of Tensors
        representing the previous policy_state.

    Returns:
      A `PolicyStep` named tuple containing:

        `action`: A tf.distribution capturing the distribution of next actions.
        `state`: A policy state tensor for the next call to distribution.
        `info`: Optional side information such as action log probabilities.

    Raises:
      ValueError or TypeError: If `validate_args is True` and inputs or
        outputs do not match `time_step_spec`, `policy_state_spec`,
        or `policy_step_spec`.
    """
    if self._validate_args:
      time_step = nest_utils.prune_extra_keys(self._time_step_spec, time_step)
      policy_state = nest_utils.prune_extra_keys(
          self._policy_state_spec, policy_state
      )
      nest_utils.assert_same_structure(
          time_step,
          self._time_step_spec,
          message='time_step and time_step_spec structures do not match',
      )
      nest_utils.assert_same_structure(
          policy_state,
          self._policy_state_spec,
          message='policy_state and policy_state_spec structures do not match',
      )
    if self._automatic_state_reset:
      policy_state = self._maybe_reset_state(time_step, policy_state)
    step = self._distribution(time_step=time_step, policy_state=policy_state)
    if self.emit_log_probability:
      # This here is set only for compatibility with info_spec in constructor.
      info = policy_step.set_log_probability(
          step.info,
          tf.nest.map_structure(
              lambda _: tf.constant(0.0, dtype=tf.float32),
              policy_step.get_log_probability(self._info_spec),
          ),
      )
      step = step._replace(info=info)
    if self._validate_args:
      nest_utils.assert_same_structure(
          step,
          self._policy_step_spec,
          message=(
              'distribution output and policy_step_spec structures do not match'
          ),
      )
    return step

  def update(
      self,
      policy,
      tau: float = 1.0,
      tau_non_trainable: Optional[float] = None,
      sort_variables_by_name: bool = False,
  ) -> tf.Operation:
    """Update the current policy with another policy.

    This would include copying the variables from the other policy.

    Args:
      policy: Another policy it can update from.
      tau: A float scalar in [0, 1]. When tau is 1.0 (the default), we do a hard
        update. This is used for trainable variables.
      tau_non_trainable: A float scalar in [0, 1] for non_trainable variables.
        If None, will copy from tau.
      sort_variables_by_name: A bool, when True would sort the variables by name
        before doing the update.

    Returns:
      An TF op to do the update.
    """
    if self.variables():
      return common.soft_variables_update(
          policy.variables(),
          self.variables(),
          tau=tau,
          tau_non_trainable=tau_non_trainable,
          sort_variables_by_name=sort_variables_by_name,
      )
    else:
      return tf.no_op()

  @property
  def emit_log_probability(self) -> bool:
    """Whether this policy instance emits log probabilities or not."""
    return self._emit_log_probability

  @property
  def time_step_spec(self) -> ts.TimeStep:
    """Describes the `TimeStep` tensors returned by `step()`.

    Returns:
      A `TimeStep` namedtuple with `TensorSpec` objects instead of Tensors,
      which describe the shape, dtype and name of each tensor returned by
      `step()`.
    """
    return self._time_step_spec

  @property
  def action_spec(self) -> types.NestedTensorSpec:
    """Describes the TensorSpecs of the Tensors expected by `step(action)`.

    `action` can be a single Tensor, or a nested dict, list or tuple of
    Tensors.

    Returns:
      An single BoundedTensorSpec, or a nested dict, list or tuple of
      `BoundedTensorSpec` objects, which describe the shape and
      dtype of each Tensor expected by `step()`.
    """
    return self._action_spec

  @property
  def policy_state_spec(self) -> types.NestedTensorSpec:
    """Describes the Tensors expected by `step(_, policy_state)`.

    `policy_state` can be an empty tuple, a single Tensor, or a nested dict,
    list or tuple of Tensors.

    Returns:
      An single TensorSpec, or a nested dict, list or tuple of
      `TensorSpec` objects, which describe the shape and
      dtype of each Tensor expected by `step(_, policy_state)`.
    """
    return self._policy_state_spec

  @property
  def info_spec(self) -> types.NestedTensorSpec:
    """Describes the Tensors emitted as info by `action` and `distribution`.

    `info` can be an empty tuple, a single Tensor, or a nested dict,
    list or tuple of Tensors.

    Returns:
      An single TensorSpec, or a nested dict, list or tuple of
      `TensorSpec` objects, which describe the shape and
      dtype of each Tensor expected by `step(_, policy_state)`.
    """
    return self._info_spec

  @property
  def policy_step_spec(self) -> policy_step.PolicyStep:
    """Describes the output of `action()`.

    Returns:
      A nest of TensorSpec which describe the shape and dtype of each Tensor
      emitted by `action()`.
    """
    return self._policy_step_spec

  # TODO(kbanoop, ebrevdo): Should this be collect_data_spec to mirror agents?
  @property
  def trajectory_spec(self) -> trajectory.Trajectory:
    """Describes the Tensors written when using this policy with an environment.

    Returns:
      A `Trajectory` containing all tensor specs associated with the
      observation_spec, action_spec, policy_state_spec, and info_spec of
      this policy.
    """
    return self._trajectory_spec

  @property
  def collect_data_spec(self) -> trajectory.Trajectory:
    """Describes the Tensors written when using this policy with an environment.

    Returns:
      A nest of TensorSpec which describe the shape and dtype of each Tensor
      required to train the agent which generated this policy.
    """
    return self._trajectory_spec

  # Subclasses MAY optionally override _action.
  def _action(
      self,
      time_step: ts.TimeStep,
      policy_state: types.NestedTensor,
      seed: Optional[types.Seed] = None,
  ) -> policy_step.PolicyStep:
    """Implementation of `action`.

    Args:
      time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
      policy_state: A Tensor, or a nested dict, list or tuple of Tensors
        representing the previous policy_state.
      seed: Seed to use if action performs sampling (optional).

    Returns:
      A `PolicyStep` named tuple containing:
        `action`: An action Tensor matching the `action_spec`.
        `state`: A policy state tensor to be fed into the next call to action.
        `info`: Optional side information such as action log probabilities.
    """
    seed_stream = tfp.util.SeedStream(seed=seed, salt='tf_agents_tf_policy')
    distribution_step = self._distribution(time_step, policy_state)  # pytype: disable=wrong-arg-types
    actions = tf.nest.map_structure(
        lambda d: reparameterized_sampling.sample(d, seed=seed_stream()),
        distribution_step.action,
    )
    info = distribution_step.info
    if self.emit_log_probability:
      try:
        log_probability = tf.nest.map_structure(
            lambda a, d: d.log_prob(a), actions, distribution_step.action
        )
        info = policy_step.set_log_probability(info, log_probability)
      except:
        raise TypeError(
            '%s does not support emitting log-probabilities.'
            % type(self).__name__
        )

    return distribution_step._replace(action=actions, info=info)

  ## Subclasses MUST implement these.
  def _distribution(
      self, time_step: ts.TimeStep, policy_state: types.NestedTensorSpec
  ) -> policy_step.PolicyStep:
    """Implementation of `distribution`.

    Args:
      time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
      policy_state: A Tensor, or a nested dict, list or tuple of Tensors
        representing the previous policy_state.

    Returns:
      A `PolicyStep` named tuple containing:
        `action`: A (optionally nested) of tfp.distribution.Distribution
          capturing the distribution of next actions.
        `state`: A policy state tensor for the next call to distribution.
        `info`: Optional side information such as action log probabilities.
    """
    raise NotImplementedError()

  # Subclasses MAY optionally overwrite _get_initial_state.
  def _get_initial_state(self, batch_size: int) -> types.NestedTensor:
    """Returns the initial state of the policy network.

    Args:
      batch_size: A constant or Tensor holding the batch size. Can be None, in
        which case the state will not have a batch dimension added.

    Returns:
      A nest of zero tensors matching the spec of the policy network state.
    """
    return tensor_spec.zero_spec_nest(
        self._policy_state_spec,
        outer_dims=None if batch_size is None else [batch_size],
    )
