code: def testTrainWithLagrange(
      self,
      use_lagrange_cql_alpha,
      use_variable_for_cql_alpha,
      log_cql_alpha_clipping,
      expected_cql_alpha_step_one,
      expected_cql_alpha_step_two,
      expected_cql_loss_step_one,
      expected_cql_loss_step_two,
  ):
      if use_variable_for_cql_alpha:
          cql_alpha = tf.Variable(5.0)
          cql_alpha_var = cql_alpha  # Getting around type checking.
      else:
          cql_alpha = 5.0
      cql_alpha_learning_rate = 0.5
      cql_tau = 10
      num_cql_samples = 5
  
      actor_net = actor_distribution_network.ActorDistributionNetwork(
          self._obs_spec, self._action_spec, fc_layer_params=None
      )
      critic_net = critic_network.CriticNetwork(
          (self._obs_spec, self._action_spec),
          observation_fc_layer_params=(16,),
          action_fc_layer_params=(16,),
          joint_fc_layer_params=(16,),
          kernel_initializer="glorot_uniform",
          last_kernel_initializer="glorot_uniform",
      )
  
      counter = common.create_variable("test_train_counter")
      optimizer_fn = tf.compat.v1.train.AdamOptimizer
      agent = cql_sac_agent.CqlSacAgent(
          self._time_step_spec,
          self._action_spec,
          critic_network=critic_net,
          actor_network=actor_net,
          actor_optimizer=optimizer_fn(1e-3),
          critic_optimizer=optimizer_fn(1e-3),
          alpha_optimizer=optimizer_fn(1e-3),
          cql_alpha=cql_alpha,
          num_cql_samples=num_cql_samples,
          include_critic_entropy_term=False,
          use_lagrange_cql_alpha=use_lagrange_cql_alpha,
          cql_alpha_learning_rate=cql_alpha_learning_rate,
          cql_tau=cql_tau,
          random_seed=self._random_seed,
          log_cql_alpha_clipping=log_cql_alpha_clipping,
          train_step_counter=counter,
      )
  
      batch_size = 5
      observations = tf.constant(
          [[[1, 2], [3, 4]]] * batch_size, dtype=tf.float32
      )
      actions = tf.constant([[[0], [1]]] * batch_size, dtype=tf.float32)
      time_steps = ts.TimeStep(
          step_type=tf.constant([[1] * 2] * batch_size, dtype=tf.int32),
          reward=tf.constant([[1] * 2] * batch_size, dtype=tf.float32),
          discount=tf.constant(
              [[1] * 2] * batch_size, dtype=tf.float32
          ),
          observation=observations,
      )
  
      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 not tf.executing_eagerly():
          # Get experience first to make sure optimizer variables are created and
          # can be initialized.
          experience = agent.train(experience)
          with self.cached_session() as sess:
              common.initialize_uninitialized_variables(sess)
          self.assertEqual(self.evaluate(counter), 0)
          self.evaluate(experience)
          self.assertEqual(self.evaluate(counter), 1)
      else:
          # Training step one.
          self.assertEqual(self.evaluate(counter), 0)
          loss = self.evaluate(agent.train(experience))
          self.assertEqual(self.evaluate(counter), 1)
          self.assertAllClose(
              loss.extra.cql_loss, expected_cql_loss_step_one
          )
          self.assertAllClose(
              loss.extra.cql_alpha, expected_cql_alpha_step_one
          )
          if use_lagrange_cql_alpha:
              self.assertGreater(loss.extra.cql_alpha_loss, 0)
          else:
              self.assertEqual(loss.extra.cql_alpha_loss, 0)
  
          # Training step two.
          if use_variable_for_cql_alpha:
              cql_alpha_var.assign_add(1)
          loss = self.evaluate(agent.train(experience))
          self.assertEqual(self.evaluate(counter), 2)
          self.assertAllClose(
              loss.extra.cql_loss, expected_cql_loss_step_two
          )
          # GPU (V100) needs slightly increased to pass.
          if tf.test.is_gpu_available():
              self.assertAllClose(
                  loss.extra.cql_alpha,
                  expected_cql_alpha_step_two,
                  atol=4.5e-5,
                  rtol=1.5e-5,
              )
          else:
              self.assertAllClose(
                  loss.extra.cql_alpha,
                  expected_cql_alpha_step_two,
                  rtol=2e-5,
                  atol=2e-5,
              )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import tensorflow
  import tensorflow_probability
  from tf_agents.agents.cql import cql_sac_agent
  from tf_agents.agents.ddpg import critic_network
  from tf_agents.agents.ddpg import critic_rnn_network
  from tf_agents.agents.sac import tanh_normal_projection_network
  from tf_agents.networks import actor_distribution_network
  from tf_agents.networks import actor_distribution_rnn_network
  from tf_agents.networks import network
  from tf_agents.networks import utils
  from tf_agents.specs import distribution_spec
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
  from tf_agents.utils import test_utils
label_desc: Use image node_id 8 for calling the CqlSacAgentTest obj's underlying member method code with example usage: obj.testTrainWithLagrange(use_lagrange_cql_alpha, use_variable_for_cql_alpha, log_cql_alpha_clipping, expected_cql_alpha_step_one, expected_cql_alpha_step_two, expected_cql_loss_step_one, expected_cql_loss_step_two) without return types