code: def testLossWithMaskedActions(self, agent_class):
      # if 'D3qnAgent' in agent_class.__name__:
      #     self.skipTest('invalid for dueling networks')
  
      # Observations are now a tuple of the usual observation and an action mask.
      observation_spec_with_mask = (
          self._observation_spec,
          tensor_spec.BoundedTensorSpec([2], tf.int32, 0, 1),
      )
      time_step_spec = ts.time_step_spec(observation_spec_with_mask)
      q_net = DummyNet(self._observation_spec, self._action_spec)
      agent = agent_class(
          time_step_spec,
          self._action_spec,
          q_network=q_net,
          optimizer=None,
          observation_and_action_constraint_splitter=lambda x: (
              x[0],
              x[1],
          ),
      )
  
      # For `observations`, the masks are set up so that all actions are valid.
      observations = (
          tf.constant([[1, 2], [3, 4]], dtype=tf.float32),
          tf.constant([[1, 1], [1, 1]], dtype=tf.int32),
      )
      time_steps = ts.restart(observations, batch_size=2)
  
      actions = tf.constant([0, 1], dtype=tf.int32)
      action_steps = policy_step.PolicyStep(actions)
  
      rewards = tf.constant([10, 20], dtype=tf.float32)
      discounts = tf.constant([0.9, 0.9], dtype=tf.float32)
  
      # For `next_observations`, the masks are set up so that only one action is
      # valid for each element in the batch.
      next_observations = (
          tf.constant([[5, 6], [7, 8]], dtype=tf.float32),
          tf.constant([[0, 1], [1, 0]], dtype=tf.int32),
      )
      next_time_steps = ts.transition(
          next_observations, rewards, discounts
      )
  
      experience = (
          trajectories_test_utils.stacked_trajectory_from_transition(
              time_steps, action_steps, next_time_steps
          )
      )
  
      # Using the kernel initializer [[2, 1], [1, 1]] and bias initializer
      # [[1], [1]] from DummyNet above, we can calculate the following values:
      # Q-value for first observation/action pair: 2 * 1 + 1 * 2 + 1 = 5
      # Q-value for second observation/action pair: 1 * 3 + 1 * 4 + 1 = 8
      # (Here we use the second row of the kernel initializer above, since the
      # chosen action is now 1 instead of 0.)
      #
      # For target Q-values, because of the masks we only have one valid choice of
      # action for each next_observation:
      # Target Q-value for first next_observation (only action 1 is valid):
      # 1 * 5 + 1 * 6 + 1 = 12
      # Target Q-value for second next_observation (only action 0 is valid):
      # 2 * 7 + 1 * 8 + 1 = 23
      # TD targets: 10 + 0.9 * 12 = 20.8 and 20 + 0.9 * 23 = 40.7
      # TD errors: 20.8 - 5 = 20.3 and 40.7 - 8 = 32.7
      # TD loss: 19.8 and 32.2 (Huber loss subtracts 0.5)
      # Overall loss: (15.3 + 32.2) / 2 = 23.75
      expected_loss = 23.75
      if "D3qnAgent" in agent_class.__name__:
          # Using Q-values for next_observations only for D3qnAgent.
          # Q-value for first next_observation/action pair:
          # 2 * 5 + 1 * 6 + 1 = 17
          # Q-value for second next_observation/action pair:
          # 1 * 7 + 1 * 8 + 1 = 16
          # TD targets: 10 + 0.9 * 17 = 25.3 and 20 + 0.9 * 23 = 40.7
          # TD errors: 25.3 - 5 = 20.3 and 40.7 - 8 = 32.7
          # TD loss: 19.8 and 32.2 (Huber loss subtracts 0.5)
          # Overall loss: (19.8 + 32.2) / 2 = 26.0
          expected_loss = 26.0
  
      loss, _ = agent._loss(experience)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.assertAllClose(self.evaluate(loss), expected_loss)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import tensorflow
  from tf_agents.agents.dqn import dqn_agent
  from tf_agents.networks import network
  from tf_agents.networks import q_network
  from tf_agents.networks import sequential
  from tf_agents.networks import test_utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import test_utils
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import common
  from tf_agents.utils import test_utils
label_desc: Use image node_id 14 for calling the DqnAgentTest obj's underlying member method code with example usage: obj.testLossWithMaskedActions(agent_class) without return types