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

"""Tests for tf_agents.agents.ppo.ppo_actor_network."""

from absl import flags
from absl.testing import parameterized
import numpy as np
import tensorflow as tf
from tf_agents.agents.ppo import ppo_actor_network
from tf_agents.agents.ppo import ppo_agent
from tf_agents.agents.ppo import ppo_policy
from tf_agents.distributions import utils as distribution_utils
from tf_agents.networks import actor_distribution_network
from tf_agents.networks import value_network
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step as ts
from tf_agents.trajectories import trajectory
from tf_agents.utils import test_utils

FLAGS = flags.FLAGS


class DeterministicSeedStream(object):
  """A fake seed stream class that always generates a deterministic seed."""

  def __init__(self, seed, salt=''):
    del salt
    self._seed = seed

  def __call__(self):
    return self._seed


class PPOAgentActorDist(ppo_agent.PPOAgent):

  def __init__(self):
    observation_tensor_spec = tf.TensorSpec(shape=[1], dtype=tf.float32)
    action_tensor_spec = tensor_spec.BoundedTensorSpec([1], tf.float32, -1, 1)

    actor_net = actor_distribution_network.ActorDistributionNetwork(
        observation_tensor_spec,
        action_tensor_spec,
        fc_layer_params=(1,),
        activation_fn=tf.nn.tanh,
        kernel_initializer=tf.keras.initializers.Orthogonal(seed=1),
        seed_stream_class=DeterministicSeedStream,
        seed=1,
    )

    value_net = value_network.ValueNetwork(
        observation_tensor_spec, fc_layer_params=(1,)
    )

    super(PPOAgentActorDist, self).__init__(
        time_step_spec=ts.time_step_spec(observation_tensor_spec),
        action_spec=action_tensor_spec,
        actor_net=actor_net,
        value_net=value_net,
        # Ensures value_prediction, return and advantage are included as parts
        # of the training_data_spec.
        compute_value_and_advantage_in_train=True,
        update_normalizers_in_train=False,
        optimizer=tf.compat.v1.train.AdamOptimizer(),
    )
    # There is an artifical call on `_train` during the initialization which
    # ensures that the variables of the optimizer are initialized. This is
    # excluded from the call count.
    self.train_called_times = -1
    self.experiences = []


class PPOAgentSequential(ppo_agent.PPOAgent):

  def __init__(self):
    observation_tensor_spec = tf.TensorSpec(shape=[1], dtype=tf.float32)
    action_tensor_spec = tensor_spec.BoundedTensorSpec([1], tf.float32, -1, 1)

    actor_net_lib = ppo_actor_network.PPOActorNetwork()
    actor_net_lib.seed_stream_class = DeterministicSeedStream
    actor_net = actor_net_lib.create_sequential_actor_net(
        fc_layer_units=(1,), action_tensor_spec=action_tensor_spec, seed=1
    )
    value_net = value_network.ValueNetwork(
        observation_tensor_spec, fc_layer_params=(1,)
    )

    super(PPOAgentSequential, self).__init__(
        time_step_spec=ts.time_step_spec(observation_tensor_spec),
        action_spec=action_tensor_spec,
        actor_net=actor_net,
        value_net=value_net,
        optimizer=tf.compat.v1.train.AdamOptimizer(),
        # Ensures value_prediction, return and advantage are included as parts
        # of the training_data_spec.
        compute_value_and_advantage_in_train=True,
        update_normalizers_in_train=False,
    )
    # There is an artifical call on `_train` during the initialization which
    # ensures that the variables of the optimizer are initialized. This is
    # excluded from the call count.
    self.train_called_times = -1
    self.experiences = []


def _create_trajectories(n_time_steps, batch_size, policy_info_scale_name):
  # Observation looks like:
  # [[ 0.,  1., ... n_time_steps.],
  #  [10., 11., ... n_time_steps.],
  #  [20., 21., ... n_time_steps.],
  #  [ ...                       ],
  #  [10*batch_size., ... 10*batch_size+n_time_steps.]]
  observation_array = np.asarray(
      [np.arange(n_time_steps) + 10 * i for i in range(batch_size)]
  )
  # Adding an inner most dimension to fit the observation spec defined above.
  observation_array = np.expand_dims(observation_array, axis=2)
  observations = tf.convert_to_tensor(observation_array, dtype=tf.float32)

  default_tensor = tf.constant(
      [[1] * n_time_steps] * batch_size, dtype=tf.float32
  )
  mid_time_step_val = ts.StepType.MID.tolist()
  time_steps = ts.TimeStep(
      step_type=tf.constant(
          [[mid_time_step_val] * n_time_steps] * batch_size, dtype=tf.int32
      ),
      reward=default_tensor,
      discount=default_tensor,
      observation=observations,
  )
  actions = tf.constant([[[1]] * n_time_steps] * batch_size, dtype=tf.float32)
  policy_info = {
      'dist_params': {
          'loc': tf.constant(
              [[[1]] * n_time_steps] * batch_size, dtype=tf.float32
          ),
          policy_info_scale_name: tf.constant(
              [[[1]] * n_time_steps] * batch_size, dtype=tf.float32
          ),
      },
      'value_prediction': default_tensor,
      'return': default_tensor,
      'advantage': default_tensor,
  }
  return trajectory.Trajectory(
      time_steps.step_type,
      observations,
      actions,
      policy_info,
      time_steps.step_type,
      time_steps.reward,
      time_steps.discount,
  )


class PpoActorNetworkTest(parameterized.TestCase, test_utils.TestCase):

  def setUp(self):
    super(PpoActorNetworkTest, self).setUp()
    # Run in full eager mode in order to inspect the content of tensors.
    tf.config.experimental_run_functions_eagerly(True)

  def tearDown(self):
    tf.config.experimental_run_functions_eagerly(False)
    super(PpoActorNetworkTest, self).tearDown()

  def test_same_initialization(self):
    if not tf.executing_eagerly():
      self.skipTest('Skipping test: sequential networks not supported in TF1')
    tf.random.set_seed(111)
    agent_seq = PPOAgentSequential()
    seq_variables = agent_seq._actor_net.trainable_weights

    tf.random.set_seed(111)
    agent_act_dist = PPOAgentActorDist()
    act_dist_variables = agent_act_dist._actor_net.trainable_weights

    self.assertEqual(len(seq_variables), len(act_dist_variables))
    set_seq_weights = set()
    set_act_dist_weights = set()
    for i in range(len(seq_variables)):
      set_seq_weights.add(seq_variables[i].numpy().sum())
      set_act_dist_weights.add(act_dist_variables[i].numpy().sum())
    self.assertSetEqual(set_seq_weights, set_act_dist_weights)

  def test_no_mismatched_shape(self):
    if not tf.executing_eagerly():
      self.skipTest('Skipping test: sequential networks not supported in TF1')
    observation_tensor_spec = tf.TensorSpec(shape=[1], dtype=tf.float32)
    action_tensor_spec = tensor_spec.BoundedTensorSpec((8,), tf.float32, -1, 1)

    actor_net_lib = ppo_actor_network.PPOActorNetwork()
    actor_net_lib.seed_stream_class = DeterministicSeedStream
    actor_net = actor_net_lib.create_sequential_actor_net(
        fc_layer_units=(1,), action_tensor_spec=action_tensor_spec, seed=1
    )

    actor_output_spec = actor_net.create_variables(observation_tensor_spec)

    distribution_utils.assert_specs_are_compatible(
        actor_output_spec,
        action_tensor_spec,
        'actor_network output spec does not match action spec',
    )

  def test_same_actor_net_output(self):
    if not tf.executing_eagerly():
      self.skipTest('Skipping test: sequential networks not supported in TF1')
    observation_tensor_spec = tf.TensorSpec(shape=[1], dtype=tf.float32)
    action_tensor_spec = tensor_spec.BoundedTensorSpec((8,), tf.float32, -1, 1)

    actor_net_lib = ppo_actor_network.PPOActorNetwork()
    actor_net_lib.seed_stream_class = DeterministicSeedStream
    actor_net_sequential = actor_net_lib.create_sequential_actor_net(
        fc_layer_units=(1,), action_tensor_spec=action_tensor_spec, seed=1
    )

    actor_net_actor_dist = actor_distribution_network.ActorDistributionNetwork(
        observation_tensor_spec,
        action_tensor_spec,
        fc_layer_params=(1,),
        activation_fn=tf.nn.tanh,
        kernel_initializer=tf.keras.initializers.Orthogonal(seed=1),
        seed_stream_class=DeterministicSeedStream,
        seed=1,
    )

    sample_observation = tf.constant([[1], [2]], dtype=tf.float32)
    tf.random.set_seed(111)
    sequential_output_dist, _ = actor_net_sequential(
        sample_observation, step_type=ts.StepType.MID, network_state=()
    )
    tf.random.set_seed(111)
    actor_dist_output_dist, _ = actor_net_actor_dist(
        sample_observation, step_type=ts.StepType.MID, network_state=()
    )
    self.assertAllEqual(
        sequential_output_dist.mean(), actor_dist_output_dist.mean()
    )
    self.assertAllEqual(
        sequential_output_dist.stddev(), actor_dist_output_dist.stddev()
    )

  def test_same_policy_same_output(self):
    if not tf.executing_eagerly():
      self.skipTest('Skipping test: sequential networks not supported in TF1')
    observation_tensor_spec = tf.TensorSpec(shape=[1], dtype=tf.float32)
    action_tensor_spec = tensor_spec.BoundedTensorSpec((8,), tf.float32, -1, 1)

    value_net = value_network.ValueNetwork(
        observation_tensor_spec, fc_layer_params=(1,)
    )

    actor_net_lib = ppo_actor_network.PPOActorNetwork()
    actor_net_lib.seed_stream_class = DeterministicSeedStream
    actor_net_sequential = actor_net_lib.create_sequential_actor_net(
        fc_layer_units=(1,), action_tensor_spec=action_tensor_spec, seed=1
    )
    actor_net_actor_dist = actor_distribution_network.ActorDistributionNetwork(
        observation_tensor_spec,
        action_tensor_spec,
        fc_layer_params=(1,),
        activation_fn=tf.nn.tanh,
        kernel_initializer=tf.keras.initializers.Orthogonal(seed=1),
        seed_stream_class=DeterministicSeedStream,
        seed=1,
    )

    tf.random.set_seed(111)
    seq_policy = ppo_policy.PPOPolicy(
        ts.time_step_spec(observation_tensor_spec),
        action_tensor_spec,
        actor_net_sequential,
        value_net,
        collect=True,
    )
    tf.random.set_seed(111)
    actor_dist_policy = ppo_policy.PPOPolicy(
        ts.time_step_spec(observation_tensor_spec),
        action_tensor_spec,
        actor_net_actor_dist,
        value_net,
        collect=True,
    )

    sample_timestep = ts.TimeStep(
        step_type=tf.constant([1, 1], dtype=tf.int32),
        reward=tf.constant([1, 1], dtype=tf.float32),
        discount=tf.constant([1, 1], dtype=tf.float32),
        observation=tf.constant([[1], [2]], dtype=tf.float32),
    )
    seq_policy_step = seq_policy._distribution(sample_timestep, policy_state=())
    act_dist_policy_step = actor_dist_policy._distribution(
        sample_timestep, policy_state=()
    )

    seq_scale = seq_policy_step.info['dist_params']['scale_diag']
    act_dist_scale = act_dist_policy_step.info['dist_params']['scale']
    self.assertAllEqual(seq_scale, act_dist_scale)
    self.assertAllEqual(
        seq_policy_step.info['dist_params']['loc'],
        act_dist_policy_step.info['dist_params']['loc'],
    )

  def test_same_agent_train(self):
    if not tf.executing_eagerly():
      self.skipTest('Skipping test: sequential networks not supported in TF1')

    tf.random.set_seed(111)
    agent_seq = PPOAgentSequential()
    tf.random.set_seed(111)
    agent_act_dist = PPOAgentActorDist()
    batch_size = 2
    n_time_steps = 3

    action_distribution_parameters = {
        'loc': tf.constant(
            [[[0.0]] * n_time_steps] * batch_size, dtype=tf.float32
        ),
        'scale': tf.constant(
            [[[1.0]] * n_time_steps] * batch_size, dtype=tf.float32
        ),
    }
    observations = tf.constant(
        [
            [[1], [3], [5]],
            [[2], [4], [6]],
        ],
        dtype=tf.float32,
    )

    mid_time_step_val = ts.StepType.MID.tolist()
    time_steps = ts.TimeStep(
        step_type=tf.constant(
            [[mid_time_step_val] * n_time_steps] * batch_size, dtype=tf.int32
        ),
        reward=tf.constant([[1] * n_time_steps] * batch_size, dtype=tf.float32),
        discount=tf.constant(
            [[1] * n_time_steps] * batch_size, dtype=tf.float32
        ),
        observation=observations,
    )
    actions = tf.constant([[[0], [1], [1]], [[0], [1], [1]]], dtype=tf.float32)

    value_preds = tf.constant(
        [[9.0, 15.0, 21.0], [9.0, 15.0, 21.0]], dtype=tf.float32
    )
    policy_info = {
        'dist_params': action_distribution_parameters,
        'value_prediction': value_preds,
    }
    experience = trajectory.Trajectory(
        time_steps.step_type,
        observations,
        actions,
        policy_info,
        time_steps.step_type,
        time_steps.reward,
        time_steps.discount,
    )
    act_dist_loss = agent_act_dist.train(experience).loss

    policy_info['dist_params']['scale_diag'] = policy_info['dist_params'].pop(
        'scale'
    )
    seq_loss = agent_seq.train(experience).loss
    self.assertAllEqual(act_dist_loss, seq_loss)

    seq_variables = agent_seq._actor_net.trainable_weights
    act_dist_variables = agent_act_dist._actor_net.trainable_weights

    self.assertEqual(len(seq_variables), len(act_dist_variables))
    set_seq_weights = set()
    set_act_dist_weights = set()
    for i in range(len(seq_variables)):
      set_seq_weights.add(seq_variables[i].numpy().sum())
      set_act_dist_weights.add(act_dist_variables[i].numpy().sum())
    self.assertSetEqual(set_seq_weights, set_act_dist_weights)


if __name__ == '__main__':
  tf.test.main()
