code: def test_loss_and_train_output(
      test: test_utils.TestCase,
      expect_equal_loss_values: bool,
      agent: tf_agent.TFAgent,
      experience: types.NestedTensor,
      weights: Optional[types.Tensor] = None,
      **kwargs,
  ):
      """Tests that loss() and train() outputs are equivalent.
  
      Checks that the outputs have the same structures and shapes, and compares
      loss values based on `expect_equal_loss_values`.
  
      Args:
        test: An instance of `test_utils.TestCase`.
        expect_equal_loss_values: Whether to expect `LossInfo.loss` to have the same
          values for loss() and train().
        agent: An instance of `TFAgent`.
        experience: A batch of experience data in the form of a `Trajectory`.
        weights: (optional).  A `Tensor` containing weights to be used when
          calculating the total loss.
        **kwargs: Any additional data as args to `train` and `loss`.
      """
      loss_info_from_loss = agent.loss(
          experience=experience, weights=weights, **kwargs
      )
      loss_info_from_train = agent.train(
          experience=experience, weights=weights, **kwargs
      )
      if not tf.executing_eagerly():
          test.evaluate(tf.compat.v1.global_variables_initializer())
          loss_info_from_loss = test.evaluate(loss_info_from_loss)
          loss_info_from_train = test.evaluate(loss_info_from_train)
  
      test.assertIsInstance(loss_info_from_train, tf_agent.LossInfo)
      test.assertEqual(
          type(loss_info_from_train), type(loss_info_from_loss)
      )
  
      # Compare loss values.
      if expect_equal_loss_values:
          test.assertEqual(
              loss_info_from_train.loss,
              loss_info_from_loss.loss,
              msg=(
                  "Expected equal loss values, but train() has output "
                  "{loss_from_train} vs loss() output {loss_from_loss}.".format(
                      loss_from_train=loss_info_from_train.loss,
                      loss_from_loss=loss_info_from_loss.loss,
                  )
              ),
          )
      else:
          test.assertNotEqual(
              loss_info_from_train.loss,
              loss_info_from_loss.loss,
              msg=(
                  "Expected train() and loss() output to have different loss values, "
                  "but both are {loss}.".format(
                      loss=loss_info_from_train.loss
                  )
              ),
          )
  
      # Check that both `LossInfo` outputs have matching dtypes and shapes.
      nest_utils.assert_tensors_matching_dtypes_and_shapes(
          loss_info_from_train,
          loss_info_from_loss,
          test,
          "`LossInfo` from train()",
          "`LossInfo` from loss()",
      )
  
file_path: agents/tf_agents/agents/test_util.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional
  import tensorflow
  from tf_agents.agents import tf_agent
  from tf_agents.typing import types
  from tf_agents.utils import nest_utils
  from tf_agents.utils import test_utils
label_desc: Use image node_id 1 for calling a global function with example usage: test_loss_and_train_output(test, expect_equal_loss_values, agent, experience, weights) without return types