code: def testLinearAgentTrainLoss(
      self, set_example_weights, distributed_train
  ):
      agent, experience, weights = _create_simple_agent_and_data(
          set_example_weights
      )
      self.evaluate(agent.initialize())
      reward = tf.squeeze(experience.reward, axis=-1)
      if distributed_train:
          train_fn = common.function_in_tf1()(
              agent._distributed_train_step
          )
          initial_loss = self.evaluate(
              train_fn(experience, weights=weights)
          )
      else:
          initial_loss = self.evaluate(
              agent.train(experience, weights=weights)
          )
      # The loss returned by the first train op is based on the initial, all-zero
      # weights, which lead to all-zero predicted rewards.
      if weights is None:
          self.assertAllClose(
              initial_loss.loss, tf.reduce_mean(tf.square(reward))
          )
      else:
          self.assertAllClose(
              initial_loss.loss,
              tf.reduce_mean(weights * tf.square(reward)),
          )
  
      # The previous train op minimizes the mse on the training data, so training
      # on the same data, we expect the loss to be minimized.
      if distributed_train:
          train_fn = common.function_in_tf1()(
              agent._distributed_train_step
          )
          minimized_loss = self.evaluate(
              train_fn(experience, weights=weights)
          )
      else:
          minimized_loss = self.evaluate(
              agent.train(experience, weights=weights)
          )
      if weights is None:
          self.assertAllClose(
              minimized_loss.loss,
              tf.reduce_mean(
                  tf.square(reward - tf.reduce_mean(reward))
              ),
          )
      else:
          self.assertAllClose(
              minimized_loss.loss,
              tf.reduce_mean(
                  weights
                  * tf.square(
                      reward
                      - tf.reduce_sum(weights * reward)
                      / tf.reduce_sum(weights)
                  )
              ),
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from typing import Optional, Tuple
  from absl import flags
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.agents import lin_ucb_agent
  from tf_agents.bandits.agents import linear_bandit_agent
  from tf_agents.bandits.agents import linear_thompson_sampling_agent
  from tf_agents.bandits.agents import utils
  from tf_agents.bandits.drivers import driver_utils
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 5 for calling the LinearBanditAgentTest obj's underlying member method code with example usage: obj.testLinearAgentTrainLoss(set_example_weights, distributed_train) without return types