text: # coding=utf-8
  # Copyright 2018 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
  #
  #     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.
  
  """Tests for TF Agents reinforce_agent."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from absl.testing import parameterized
  from absl.testing.absltest import mock
  
  import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import
  import tensorflow_probability as tfp
  
  from tf_agents.agents.reinforce import reinforce_agent
  from tf_agents.networks import actor_distribution_rnn_network
  from tf_agents.networks import network
  from tf_agents.networks import utils as network_utils
  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 common
  from tf_agents.utils import nest_utils
  
  from tensorflow.python.util import nest  # pylint:disable=g-direct-tensorflow-import  # TF internal
  
  
  class DummyActorNet(network.Network):
  
    def __init__(self,
                 input_tensor_spec,
                 output_tensor_spec,
                 unbounded_actions=False,
                 stateful=False):
      # When unbounded_actions=True, we skip the final tanh activation and the
      # action shift and scale. This allows us to compute the actor and critic
      # losses by hand more easily.
      # If stateful=True, the network state has the same shape as
      # `input_tensor_spec`. Otherwise it is empty.
      state_spec = (tf.TensorSpec(input_tensor_spec.shape, tf.float32)
                    if stateful else ())
      super(DummyActorNet, self).__init__(
          input_tensor_spec=input_tensor_spec,
          state_spec=state_spec,
          name='DummyActorNet')
      single_action_spec = tf.nest.flatten(output_tensor_spec)[0]
      activation_fn = None if unbounded_actions else tf.nn.tanh
      self._output_tensor_spec = output_tensor_spec
      self._layers = [
          tf.keras.layers.Dense(
              single_action_spec.shape.num_elements() * 2,
              activation=activation_fn,
              kernel_initializer=tf.compat.v1.initializers.constant(
                  [[2, 1], [1, 1]]),
              bias_initializer=tf.compat.v1.initializers.constant(5),
          ),
      ]
  
    def call(self, observations, step_type, network_state):
      del step_type
  
      states = tf.cast(tf.nest.flatten(observations)[0], tf.float32)
      for layer in self.layers:
        states = layer(states)
  
      single_action_spec = tf.nest.flatten(self._output_tensor_spec)[0]
      actions, stdevs = states[..., 0], states[..., 1]
      actions = tf.reshape(actions, [-1] + single_action_spec.shape.as_list())
      stdevs = tf.reshape(stdevs, [-1] + single_action_spec.shape.as_list())
      actions = tf.nest.pack_sequence_as(self._output_tensor_spec, [actions])
      stdevs = tf.nest.pack_sequence_as(self._output_tensor_spec, [stdevs])
  
      distribution = nest.map_structure_up_to(
          self._output_tensor_spec, tfp.distributions.Normal, actions, stdevs)
      return distribution, network_state
  
  
  class DummyValueNet(network.Network):
  
    def __init__(self, observation_spec, name=None, outer_rank=1):
      super(DummyValueNet, self).__init__(observation_spec, (), 'DummyValueNet')
      self._outer_rank = outer_rank
      self._layers.append(
          tf.keras.layers.Dense(
              1,
              kernel_initializer=tf.compat.v1.initializers.constant([2, 1]),
              bias_initializer=tf.compat.v1.initializers.constant([5])))
  
    def call(self, inputs, step_type=None, network_state=()):
      del step_type
      hidden_state = tf.cast(tf.nest.flatten(inputs), tf.float32)[0]
      batch_squash = network_utils.BatchSquash(self._outer_rank)
      hidden_state = batch_squash.flatten(hidden_state)
      for layer in self.layers:
        hidden_state = layer(hidden_state)
      value_pred = tf.squeeze(batch_squash.unflatten(hidden_state), axis=-1)
      return value_pred, network_state
  
  
  class ReinforceAgentTest(tf.test.TestCase, parameterized.TestCase):
  
    def setUp(self):
      super(ReinforceAgentTest, self).setUp()
      tf.compat.v1.enable_resource_variables()
      self._obs_spec = tensor_spec.TensorSpec([2], tf.float32)
      self._time_step_spec = ts.time_step_spec(self._obs_spec)
      self._action_spec = tensor_spec.BoundedTensorSpec([1], tf.float32, -1, 1)
  
    def testCreateAgent(self):
      reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=False),
          optimizer=None,
      )
  
    def testCreateAgentWithValueNet(self):
      reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=False),
          value_network=DummyValueNet(self._obs_spec),
          value_estimation_loss_coef=0.5,
          optimizer=None,
      )
  
    def testPolicyGradientLoss(self):
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=None,
      )
  
      observations = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
      time_steps = ts.restart(observations, batch_size=2)
      actions = tf.constant([[0], [1]], dtype=tf.float32)
      actions_distribution = agent.collect_policy.distribution(
          time_steps).action
      returns = tf.constant([1.9, 1.0], dtype=tf.float32)
  
      expected_loss = 10.983667373657227
      loss = agent.policy_gradient_loss(
          actions_distribution, actions, time_steps.is_last(), returns, 1)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_ = self.evaluate(loss)
      self.assertAllClose(loss_, expected_loss)
  
    def testPolicyGradientLossMultipleEpisodes(self):
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=None,
      )
  
      step_type = tf.constant(
          [ts.StepType.FIRST, ts.StepType.LAST, ts.StepType.FIRST,
           ts.StepType.LAST])
      reward = tf.constant([0, 0, 0, 0], dtype=tf.float32)
      discount = tf.constant([1, 1, 1, 1], dtype=tf.float32)
      observations = tf.constant(
          [[1, 2], [1, 2], [1, 2], [1, 2]], dtype=tf.float32)
      time_steps = ts.TimeStep(step_type, reward, discount, observations)
  
      actions = tf.constant([[0], [1], [2], [3]], dtype=tf.float32)
      actions_distribution = agent.collect_policy.distribution(
          time_steps).action
      returns = tf.constant([1.9, 1.9, 1.0, 1.0], dtype=tf.float32)
  
      expected_loss = 5.140229225158691
      loss = agent.policy_gradient_loss(
          actions_distribution, actions, time_steps.is_last(), returns, 2)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_ = self.evaluate(loss)
      self.assertAllClose(loss_, expected_loss)
  
    def testMaskingRewardSingleEpisodeRewardOnFirst(self):
      # Test that policy_gradient_loss reacts correctly to rewards when there are:
      #   * A single MDP episode
      #   * Returns on the tf.StepType.FIRST transitions
      #
      # F, L, M = ts.StepType.{FIRST, MID, LAST} in the chart below.
      #
      # Experience looks like this:
      # Trajectories: (F, L) -> (L, F)
      # observation : [1, 2]    [1, 2]
      # action      :   [0]       [1]
      # reward      :    3         0
      # ~is_boundary:    1         0
      # is_last     :    1         0
      # valid reward:   3*1       4*0
      #
      # The second action & reward should be masked out due to being on a
      # boundary (step_type=(L, F)) transition.
      #
      # The expected_loss is > 0.0 in this case, only LAST should be excluded.
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=None,
      )
  
      step_type = tf.constant([ts.StepType.FIRST, ts.StepType.LAST])
      reward = tf.constant([3, 4], dtype=tf.float32)
      discount = tf.constant([1, 0], dtype=tf.float32)
      observations = tf.constant([[1, 2], [1, 2]], dtype=tf.float32)
      time_steps = ts.TimeStep(step_type, reward, discount, observations)
  
      actions = tf.constant([[0], [1]], dtype=tf.float32)
      actions_distribution = agent.collect_policy.distribution(
          time_steps).action
      returns = tf.constant([3.0, 0.0], dtype=tf.float32)
  
      # Returns on the StepType.FIRST should be counted.
      expected_loss = 10.8935775757
      loss = agent.policy_gradient_loss(
          actions_distribution, actions, time_steps.is_last(), returns, 1)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_ = self.evaluate(loss)
      self.assertAllClose(loss_, expected_loss)
  
    def testMaskingReturnSingleEpisodeRewardOnLast(self):
      # Test that policy_gradient_loss reacts correctly to rewards when there are:
      #   * A single MDP episode
      #   * Returns on the tf.StepType.LAST transitions
      #
      # F, L, M = ts.StepType.{FIRST, MID, LAST} in the chart below.
      #
      # Experience looks like this:
      # Trajectories: (F, L) -> (L, F)
      # observation : [1, 2]    [1, 2]
      # action      :   [0]       [1]
      # reward      :    0         3
      # ~is_boundary:    1         0
      # is_last     :    1         0
      # valid reward:   0*1       3*0
      #
      # The second action & reward should be masked out due to being on a
      # boundary (step_type=(L, F)) transition.  The first has a 0 reward.
      #
      # The expected_loss is 0.0 in this case.
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=None,
      )
  
      step_type = tf.constant([ts.StepType.FIRST, ts.StepType.LAST])
      reward = tf.constant([0, 3], dtype=tf.float32)
      discount = tf.constant([1, 0], dtype=tf.float32)
      observations = tf.constant(
          [[1, 2], [1, 2]], dtype=tf.float32)
      time_steps = ts.TimeStep(step_type, reward, discount, observations)
  
      actions = tf.constant([[0], [1]], dtype=tf.float32)
      actions_distribution = agent.collect_policy.distribution(
          time_steps).action
      returns = tf.constant([0.0, 3.0], dtype=tf.float32)
  
      # Returns on the StepType.LAST should not be counted.
      expected_loss = 0.0
      loss = agent.policy_gradient_loss(
          actions_distribution, actions, time_steps.is_last(), returns, 1)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_ = self.evaluate(loss)
      self.assertAllClose(loss_, expected_loss)
  
    def testMaskingReturnMultipleEpisodesRewardOnFirst(self):
      # Test that policy_gradient_loss reacts correctly to rewards when there are:
      #   * Multiple MDP episodes
      #   * Returns on the tf.StepType.FIRST transitions
      #
      # F, L, M = ts.StepType.{FIRST, MID, LAST} in the chart below.
      #
      # Experience looks like this:
      # Trajectories: (F, L) -> (L, F) -> (F, L) -> (L, F)
      # observation : [1, 2]    [1, 2]    [1, 2]    [1, 2]
      # action      :   [0]       [1]       [2]       [3]
      # reward      :    3         0         4         0
      # ~is_boundary:    1         0         1         0
      # is_last     :    1         0         1         0
      # valid reward:   3*1       0*0       4*1       0*0
      #
      # The second & fourth action & reward should be masked out due to being on a
      # boundary (step_type=(L, F)) transition.
      #
      # The expected_loss is > 0.0 in this case, only LAST should be excluded.
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=None,
      )
  
      step_type = tf.constant(
          [ts.StepType.FIRST, ts.StepType.LAST, ts.StepType.FIRST,
           ts.StepType.LAST])
      reward = tf.constant([3, 0, 4, 0], dtype=tf.float32)
      discount = tf.constant([1, 0, 1, 0], dtype=tf.float32)
      observations = tf.constant(
          [[1, 2], [1, 2], [1, 2], [1, 2]], dtype=tf.float32)
      time_steps = ts.TimeStep(step_type, reward, discount, observations)
  
      actions = tf.constant([[0], [1], [2], [3]], dtype=tf.float32)
      actions_distribution = agent.collect_policy.distribution(
          time_steps).action
      returns = tf.constant([3.0, 0.0, 4.0, 0.0], dtype=tf.float32)
  
      # Returns on the StepType.FIRST should be counted.
      expected_loss = 12.2091741562
      loss = agent.policy_gradient_loss(
          actions_distribution, actions, time_steps.is_last(), returns, 2)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_ = self.evaluate(loss)
      self.assertAllClose(loss_, expected_loss)
  
    def testMaskingReturnMultipleEpisodesRewardOnLast(self):
      # Test that policy_gradient_loss reacts correctly to returns when there are:
      #   * Multiple MDP episodes
      #   * Returns on the tf.StepType.LAST transitions
      #
      # F, L, M = ts.StepType.{FIRST, MID, LAST} in the chart below.
      #
      # Experience looks like this:
      # Trajectories: (F, L) -> (L, F) -> (F, L) -> (L, F)
      # observation : [1, 2]    [1, 2]    [1, 2]    [1, 2]
      # action      :   [0]       [1]       [2]       [3]
      # reward      :    0         3         0         4
      # ~is_boundary:    1         0         1         0
      # is_last     :    1         0         1         0
      # valid reward:   0*1       3*0       0*1       4*0
      #
      # The second & fourth action & reward should be masked out due to being on a
      # boundary (step_type=(L, F)) transition.
      #
      # The expected_loss is 0.0 in this case.
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=None,
      )
  
      step_type = tf.constant(
          [ts.StepType.FIRST, ts.StepType.LAST, ts.StepType.FIRST,
           ts.StepType.LAST])
      reward = tf.constant([0, 3, 0, 4], dtype=tf.float32)
      discount = tf.constant([1, 0, 1, 0], dtype=tf.float32)
      observations = tf.constant(
          [[1, 2], [1, 2], [1, 2], [1, 2]], dtype=tf.float32)
      time_steps = ts.TimeStep(step_type, reward, discount, observations)
  
      actions = tf.constant([[0], [1], [2], [3]], dtype=tf.float32)
      actions_distribution = agent.collect_policy.distribution(
          time_steps).action
      returns = tf.constant([0.0, 3.0, 0.0, 4.0], dtype=tf.float32)
  
      # Returns on the StepType.LAST should not be counted.
      expected_loss = 0.0
      loss = agent.policy_gradient_loss(
          actions_distribution, actions, time_steps.is_last(), returns, 2)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_ = self.evaluate(loss)
      self.assertAllClose(loss_, expected_loss)
  
    @parameterized.parameters(
        ([[[0.8, 0.2]]], [1],),
        ([[[0.8, 0.2]], [[0.3, 0.7]]], [0.5, 0.5],),
    )
    def testEntropyLoss(self, probs, weights):
      probs = tf.convert_to_tensor(probs)
      distribution = tfp.distributions.Categorical(probs=probs)
      shape = probs.shape.as_list()
      action_spec = tensor_spec.TensorSpec(shape[2:-1], dtype=tf.int32)
      expected = tf.reduce_mean(
          -tf.reduce_mean(distribution.entropy()) * weights)
      actual = reinforce_agent._entropy_loss(
          distribution, action_spec, weights)
      self.assertAlmostEqual(self.evaluate(actual), self.evaluate(expected),
                             places=4)
  
    def testValueEstimationLoss(self):
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=False),
          value_network=DummyValueNet(self._obs_spec),
          value_estimation_loss_coef=0.5,
          optimizer=None,
      )
  
      observations = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
      time_steps = ts.restart(observations, batch_size=2)
      returns = tf.constant([1.9, 1.0], dtype=tf.float32)
      value_preds, _ = agent._value_network(time_steps.observation,
                                            time_steps.step_type)
  
      expected_loss = 123.20500
      loss = agent.value_estimation_loss(
          value_preds, returns, 1)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_ = self.evaluate(loss)
      self.assertAllClose(loss_, expected_loss)
  
    def testTrainMaskingRewardSingleBanditEpisode(self):
      # Test that train reacts correctly to experience when there is only a
      # single Bandit episode.  Bandit episodes are encoded differently than
      # MDP episodes.  They have only a single transition with
      # step_type=StepType.FIRST and next_step_type=StepType.LAST.
      #
      # F, L, M = ts.StepType.{FIRST, MID, LAST} in the chart below.
      #
      # Experience looks like this:
      # Trajectories: (F, L)
      # observation : [1, 2]
      # action      :   [0]
      # reward      :    3
      # ~is_boundary:    0
      # is_last     :    1
      # valid reward:   3*1
      #
      # The single bandit transition is valid and not masked.
      #
      # The expected_loss is > 0.0 in this case, matching the expected_loss of the
      # testMaskingRewardSingleEpisodeRewardOnFirst policy_gradient_loss test.
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=tf.compat.v1.train.AdamOptimizer(0.001),
          use_advantage_loss=False,
          normalize_returns=False,
      )
  
      step_type = tf.constant([ts.StepType.FIRST])
      next_step_type = tf.constant([ts.StepType.LAST])
      reward = tf.constant([3], dtype=tf.float32)
      discount = tf.constant([0], dtype=tf.float32)
      observations = tf.constant([[1, 2]], dtype=tf.float32)
      actions = tf.constant([[0]], dtype=tf.float32)
  
      experience = nest_utils.batch_nested_tensors(trajectory.Trajectory(
          step_type, observations, actions, (), next_step_type, reward, discount))
  
      # Rewards should be counted.
      expected_loss = 10.8935775757
  
      if tf.executing_eagerly():
        loss = lambda: agent.train(experience)
      else:
        loss = agent.train(experience)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_info = self.evaluate(loss)
      self.assertAllClose(loss_info.loss, expected_loss)
  
    def testTrainMaskingRewardMultipleBanditEpisodes(self):
      # Test that train reacts correctly to experience when there are multiple
      # Bandit episodes.  Bandit episodes are encoded differently than
      # MDP episodes.  They (each) have only a single transition with
      # step_type=StepType.FIRST and next_step_type=StepType.LAST.  This test
      # helps ensure that LAST->FIRST->LAST transitions are handled correctly.
      #
      # F, L, M = ts.StepType.{FIRST, MID, LAST} in the chart below.
      #
      # Experience looks like this:
      # Trajectories: (F, L) -> (F, L)
      # observation : [1, 2]    [1, 2]
      # action      :   [0]       [2]
      # reward      :    3         4
      # ~is_boundary:    0         0
      # is_last     :    1         1
      # valid reward:   3*1       4*1
      #
      # All bandit transitions are valid and none are masked.
      #
      # The expected_loss is > 0.0 in this case, matching the expected_loss of the
      # testMaskingRewardMultipleEpisodesRewardOnFirst policy_gradient_loss test.
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=tf.compat.v1.train.AdamOptimizer(0.001),
          use_advantage_loss=False,
          normalize_returns=False,
      )
  
      step_type = tf.constant([ts.StepType.FIRST, ts.StepType.FIRST])
      next_step_type = tf.constant([ts.StepType.LAST, ts.StepType.LAST])
      reward = tf.constant([3, 4], dtype=tf.float32)
      discount = tf.constant([0, 0], dtype=tf.float32)
      observations = tf.constant([[1, 2], [1, 2]], dtype=tf.float32)
      actions = tf.constant([[0], [2]], dtype=tf.float32)
  
      experience = nest_utils.batch_nested_tensors(trajectory.Trajectory(
          step_type, observations, actions, (), next_step_type, reward, discount))
  
      # Rewards on the StepType.FIRST should be counted.
      expected_loss = 12.2091741562
  
      if tf.executing_eagerly():
        loss = lambda: agent.train(experience)
      else:
        loss = agent.train(experience)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_info = self.evaluate(loss)
      self.assertAllClose(loss_info.loss, expected_loss)
  
    def testTrainMaskingRewardSingleEpisodeRewardOnFirst(self):
      # Test that train reacts correctly to experience when there are:
      #   * A single MDP episode
      #   * Rewards on the tf.StepType.FIRST transitions
      #
      # F, L, M = ts.StepType.{FIRST, MID, LAST} in the chart below.
      #
      # Experience looks like this:
      # Trajectories: (F, L) -> (L, F)
      # observation : [1, 2]    [1, 2]
      # action      :   [0]       [1]
      # reward      :    3         4
      # ~is_boundary:    1         0
      # is_last     :    1         0
      # valid reward:   3*1       4*0
      #
      # The second action & reward should be masked out due to being on a
      # boundary (step_type=(L, F)) transition.
      #
      # The expected_loss is > 0.0 in this case, matching the expected_loss of the
      # testMaskingRewardSingleEpisodeRewardOnFirst policy_gradient_loss test.
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=tf.compat.v1.train.AdamOptimizer(0.001),
          use_advantage_loss=False,
          normalize_returns=False,
      )
  
      step_type = tf.constant([ts.StepType.FIRST, ts.StepType.LAST])
      next_step_type = tf.constant([ts.StepType.LAST, ts.StepType.FIRST])
      reward = tf.constant([3, 4], dtype=tf.float32)
      discount = tf.constant([1, 0], dtype=tf.float32)
      observations = tf.constant([[1, 2], [1, 2]], dtype=tf.float32)
      actions = tf.constant([[0], [1]], dtype=tf.float32)
  
      experience = nest_utils.batch_nested_tensors(trajectory.Trajectory(
          step_type, observations, actions, (), next_step_type, reward, discount))
  
      # Rewards on the StepType.FIRST should be counted.
      expected_loss = 10.8935775757
  
      if tf.executing_eagerly():
        loss = lambda: agent.train(experience)
      else:
        loss = agent.train(experience)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_info = self.evaluate(loss)
      self.assertAllClose(loss_info.loss, expected_loss)
  
    def testTrainMaskingRewardSingleEpisodeRewardOnLast(self):
      # Test that train reacts correctly to experience when there are:
      #   * A single MDP episode
      #   * Rewards on the tf.StepType.LAST transitions
      #
      # F, L, M = ts.StepType.{FIRST, MID, LAST} in the chart below.
      #
      # Experience looks like this:
      # Trajectories: (F, L) -> (L, F)
      # observation : [1, 2]    [1, 2]
      # action      :   [0]       [1]
      # reward      :    0         3
      # ~is_boundary:    1         0
      # is_last     :    1         0
      # valid reward:   0*1       3*0
      #
      # The second action & reward should be masked out due to being on a
      # boundary (step_type=(L, F)) transition.  The first has a 0 reward.
      #
      # The expected_loss is = 0.0 in this case.
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=tf.compat.v1.train.AdamOptimizer(0.001),
          use_advantage_loss=False,
          normalize_returns=False,
      )
  
      step_type = tf.constant([ts.StepType.FIRST, ts.StepType.LAST])
      next_step_type = tf.constant([ts.StepType.LAST, ts.StepType.FIRST])
      reward = tf.constant([0, 3], dtype=tf.float32)
      discount = tf.constant([1, 0], dtype=tf.float32)
      observations = tf.constant([[1, 2], [1, 2]], dtype=tf.float32)
  
      actions = tf.constant([[0], [1]], dtype=tf.float32)
  
      experience = nest_utils.batch_nested_tensors(trajectory.Trajectory(
          step_type, observations, actions, (), next_step_type, reward, discount))
  
      # Rewards on the StepType.LAST should not be counted.
      expected_loss = 0.0
  
      if tf.executing_eagerly():
        loss = lambda: agent.train(experience)
      else:
        loss = agent.train(experience)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_info = self.evaluate(loss)
      self.assertAllClose(loss_info.loss, expected_loss)
  
    def testTrainMaskingRewardMultipleEpisodesRewardOnFirst(self):
      # Test that train reacts correctly to experience when there are:
      #   * Multiple MDP episodes
      #   * Rewards on the tf.StepType.FIRST transitions
      #
      # F, L, M = ts.StepType.{FIRST, MID, LAST} in the chart below.
      #
      # Experience looks like this:
      # Trajectories: (F, L) -> (L, F) -> (F, L) -> (L, F)
      # observation : [1, 2]    [1, 2]    [1, 2]    [1, 2]
      # action      :   [0]       [1]       [2]       [3]
      # reward      :    3         0         4         0
      # ~is_boundary:    1         0         1         0
      # is_last     :    1         0         1         0
      # valid reward:   3*1       0*0       4*1       0*0
      #
      # The second & fourth action & reward should be masked out due to being on a
      # boundary (step_type=(L, F)) transition.
      #
      # The expected_loss is > 0.0 in this case, matching the expected_loss of the
      # testMaskingRewardMultipleEpisodesRewardOnFirst policy_gradient_loss test.
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=tf.compat.v1.train.AdamOptimizer(0.001),
          use_advantage_loss=False,
          normalize_returns=False,
      )
  
      step_type = tf.constant([ts.StepType.FIRST, ts.StepType.LAST,
                               ts.StepType.FIRST, ts.StepType.LAST])
      next_step_type = tf.constant([ts.StepType.LAST, ts.StepType.FIRST,
                                    ts.StepType.LAST, ts.StepType.FIRST])
      reward = tf.constant([3, 0, 4, 0], dtype=tf.float32)
      discount = tf.constant([1, 0, 1, 0], dtype=tf.float32)
      observations = tf.constant(
          [[1, 2], [1, 2], [1, 2], [1, 2]], dtype=tf.float32)
      actions = tf.constant([[0], [1], [2], [3]], dtype=tf.float32)
  
      experience = nest_utils.batch_nested_tensors(trajectory.Trajectory(
          step_type, observations, actions, (), next_step_type, reward, discount))
  
      # Rewards on the StepType.FIRST should be counted.
      expected_loss = 12.2091741562
  
      if tf.executing_eagerly():
        loss = lambda: agent.train(experience)
      else:
        loss = agent.train(experience)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_info = self.evaluate(loss)
      self.assertAllClose(loss_info.loss, expected_loss)
  
    def testTrainMaskingPartialEpisodeMultipleEpisodesRewardOnFirst(self):
      # Test that train reacts correctly to experience when there are:
      #   * Multiple MDP episodes
      #   * Rewards on the tf.StepType.FIRST transitions
      #   * Partial episode at end of experience
      #
      # F, L, M = ts.StepType.{FIRST, MID, LAST} in the chart below.
      #
      # Experience looks like this:
      # Trajectories: (F, L) -> (L, F) -> (F, M) -> (M, M)
      # observation : [1, 2]    [1, 2]    [1, 2]    [1, 2]
      # action      :   [0]       [1]       [2]       [3]
      # reward      :    3         0         4         0
      # ~is_boundary:    1         0         1         1
      # is_last     :    1         0         0         0
      # valid reward:   3*1       0*0       4*0       0*0
      #
      # The second action & reward should be masked out due to being on a
      # boundary (step_type=(L, F)) transition.  The third & fourth transitions
      # should get masked out for everything due to it being an incomplete episode
      # (notice there is no trailing step_type=(F,L)).
      #
      # The expected_loss is > 0.0 in this case, matching the expected_loss of the
      # testMaskingRewardSingleEpisodeRewardOnFirst policy_gradient_loss test,
      # because the partial second episode should be masked out.
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=tf.compat.v1.train.AdamOptimizer(0.001),
          use_advantage_loss=False,
          normalize_returns=False,
      )
  
      step_type = tf.constant([ts.StepType.FIRST, ts.StepType.LAST,
                               ts.StepType.FIRST, ts.StepType.MID])
      next_step_type = tf.constant([ts.StepType.LAST, ts.StepType.FIRST,
                                    ts.StepType.MID, ts.StepType.MID])
      reward = tf.constant([3, 0, 4, 0], dtype=tf.float32)
      discount = tf.constant([1, 0, 1, 0], dtype=tf.float32)
      observations = tf.constant(
          [[1, 2], [1, 2], [1, 2], [1, 2]], dtype=tf.float32)
      actions = tf.constant([[0], [1], [2], [3]], dtype=tf.float32)
  
      experience = nest_utils.batch_nested_tensors(trajectory.Trajectory(
          step_type, observations, actions, (), next_step_type, reward, discount))
  
      # Rewards on the StepType.FIRST should be counted.
      expected_loss = 10.8935775757
  
      if tf.executing_eagerly():
        loss = lambda: agent.train(experience)
      else:
        loss = agent.train(experience)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_info = self.evaluate(loss)
      self.assertAllClose(loss_info.loss, expected_loss)
  
    def testTrainMaskingRewardMultipleEpisodesRewardOnLast(self):
      # Test that train reacts correctly to experience when there are:
      #   * Multiple MDP episodes
      #   * Rewards on the tf.StepType.LAST transitions
      #
      # F, L, M = ts.StepType.{FIRST, MID, LAST} in the chart below.
      #
      # Experience looks like this:
      # Trajectories: (F, L) -> (L, F) -> (F, L) -> (L, F)
      # observation : [1, 2]    [1, 2]    [1, 2]    [1, 2]
      # action      :   [0]       [1]       [2]       [3]
      # reward      :    0         3         0         4
      # ~is_boundary:    1         0         1         0
      # is_last     :    1         0         1         0
      # valid reward:   0*1       3*0       0*1       4*0
      #
      # The second & fourth action & reward should be masked out due to being on a
      # boundary (step_type=(L, F)) transition.
      #
      # The expected_loss is = 0.0 in this case.
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=True),
          optimizer=tf.compat.v1.train.AdamOptimizer(0.001),
          use_advantage_loss=False,
          normalize_returns=False,
      )
  
      step_type = tf.constant([ts.StepType.FIRST, ts.StepType.LAST,
                               ts.StepType.FIRST, ts.StepType.LAST])
      next_step_type = tf.constant([ts.StepType.LAST, ts.StepType.FIRST,
                                    ts.StepType.LAST, ts.StepType.FIRST])
      reward = tf.constant([0, 3, 0, 4], dtype=tf.float32)
      discount = tf.constant([1, 0, 1, 0], dtype=tf.float32)
      observations = tf.constant(
          [[1, 2], [1, 2], [1, 2], [1, 2]], dtype=tf.float32)
      actions = tf.constant([[0], [1], [2], [3]], dtype=tf.float32)
  
      experience = nest_utils.batch_nested_tensors(trajectory.Trajectory(
          step_type, observations, actions, (), next_step_type, reward, discount))
  
      # Rewards on the StepType.LAST should be counted.
      expected_loss = 0.0
  
      if tf.executing_eagerly():
        loss = lambda: agent.train(experience)
      else:
        loss = agent.train(experience)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      loss_info = self.evaluate(loss)
      self.assertAllClose(loss_info.loss, expected_loss)
  
    def testPolicy(self):
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=False),
          optimizer=None,
      )
      observations = tf.constant([[1, 2]], dtype=tf.float32)
      time_steps = ts.restart(observations, batch_size=2)
      actions = agent.policy.action(time_steps).action
      self.assertEqual(actions.shape.as_list(), [1, 1])
      self.evaluate(tf.compat.v1.global_variables_initializer())
      action_values = self.evaluate(actions)
      tf.nest.map_structure(
          lambda v, s: self.assertAllInRange(v, s.minimum, s.maximum),
          action_values, self._action_spec)
  
    @parameterized.parameters(
        (False,),
        (True,),
    )
    def testGetInitialPolicyState(self, stateful):
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=False,
              stateful=stateful),
          optimizer=None,
      )
      observations = tf.constant([[1, 2]], dtype=tf.float32)
      time_steps = ts.restart(observations, batch_size=3)
      initial_state = reinforce_agent._get_initial_policy_state(
          agent.collect_policy, time_steps)
      if stateful:
        self.assertAllEqual(self.evaluate(initial_state),
                            self.evaluate(tf.zeros((3, 2), dtype=tf.float32)))
      else:
        self.assertEqual(initial_state, ())
  
    def testTrainWithRnn(self):
      actor_net = actor_distribution_rnn_network.ActorDistributionRnnNetwork(
          self._obs_spec,
          self._action_spec,
          input_fc_layer_params=None,
          output_fc_layer_params=None,
          conv_layer_params=None,
          lstm_size=(40,))
  
      counter = common.create_variable('test_train_counter')
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=actor_net,
          optimizer=tf.compat.v1.train.AdamOptimizer(0.001),
          train_step_counter=counter
      )
  
      batch_size = 5
      observations = tf.constant(
          [[[1, 2], [3, 4], [5, 6]]] * batch_size, dtype=tf.float32)
      time_steps = ts.TimeStep(
          step_type=tf.constant([[1, 1, 2]] * batch_size, dtype=tf.int32),
          reward=tf.constant([[1] * 3] * batch_size, dtype=tf.float32),
          discount=tf.constant([[1] * 3] * batch_size, dtype=tf.float32),
          observation=observations)
      actions = tf.constant([[[0], [1], [1]]] * batch_size, dtype=tf.float32)
  
      experience = trajectory.Trajectory(
          time_steps.step_type, observations, actions, (),
          time_steps.step_type, time_steps.reward, time_steps.discount)
  
      # Force variable creation.
      agent.policy.variables()
  
      if tf.executing_eagerly():
        loss = lambda: agent.train(experience)
      else:
        loss = agent.train(experience)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.assertEqual(self.evaluate(counter), 0)
      self.evaluate(loss)
      self.assertEqual(self.evaluate(counter), 1)
  
    @parameterized.parameters(
        (False,), (True,)
    )
    def testWithAdvantageFn(self, with_value_network):
      advantage_fn = mock.Mock(
          side_effect=lambda returns, _: returns)
  
      value_network = (DummyValueNet(self._obs_spec) if with_value_network
                       else None)
      agent = reinforce_agent.ReinforceAgent(
          self._time_step_spec,
          self._action_spec,
          actor_network=DummyActorNet(
              self._obs_spec, self._action_spec, unbounded_actions=False),
          value_network=value_network,
          advantage_fn=advantage_fn,
          optimizer=None,
      )
  
      step_type = tf.constant([[ts.StepType.FIRST, ts.StepType.LAST,
                                ts.StepType.FIRST, ts.StepType.LAST]])
      next_step_type = tf.constant([[ts.StepType.LAST, ts.StepType.FIRST,
                                     ts.StepType.LAST, ts.StepType.FIRST]])
      reward = tf.constant([[0, 0, 0, 0]], dtype=tf.float32)
      discount = tf.constant([[1, 1, 1, 1]], dtype=tf.float32)
      observations = tf.constant(
          [[[1, 2], [1, 2], [1, 2], [1, 2]]], dtype=tf.float32)
      actions = tf.constant([[[0], [1], [2], [3]]], dtype=tf.float32)
  
      experience = trajectory.Trajectory(
          step_type, observations, actions, (), next_step_type, reward, discount)
  
      agent.total_loss(experience, reward, None)
  
      advantage_fn.assert_called_once()
  
  
  if __name__ == '__main__':
    tf.test.main()
  