code: def testCriticLossWithMaskedActions(self):
      # Observations are now a tuple of the usual observation and an action mask.
      observation_spec_with_mask = (
          self._obs_spec,
          tensor_spec.BoundedTensorSpec([2], tf.int32, 0, 1),
      )
      time_step_spec = ts.time_step_spec(observation_spec_with_mask)
      dummy_categorical_net = DummyCategoricalNet(self._obs_spec)
      agent = categorical_dqn_agent.CategoricalDqnAgent(
          time_step_spec,
          self._action_spec,
          dummy_categorical_net,
          self._optimizer,
          observation_and_action_constraint_splitter=lambda x: (
              x[0],
              x[1],
          ),
      )
  
      # For `observations`, the masks are set up so that only one action is valid
      # for each element in the batch.
      observations = (
          tf.constant([[1, 2], [3, 4]], dtype=tf.float32),
          tf.constant([[1, 0], [0, 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 the opposite actions as
      # before are valid.
      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 = test_utils.stacked_trajectory_from_transition(
          time_steps, action_steps, next_time_steps
      )
  
      # Due to the constant initialization of the DummyCategoricalNet, we can
      # expect the same loss every time. Note this is different from the loss in
      # testCriticLoss above due to previously optimal actions being masked out.
      expected_loss = 5.062895
      loss_info = agent._loss(experience)
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      evaluated_loss = self.evaluate(loss_info).loss
      self.assertAllClose(evaluated_loss, expected_loss, atol=1e-4)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow
  from tf_agents.agents.categorical_dqn import categorical_dqn_agent
  from tf_agents.networks import categorical_q_network
  from tf_agents.networks import network
  from tf_agents.networks import q_rnn_network
  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
label_desc: Use image node_id 8 for calling the CategoricalDqnAgentTest obj's underlying member method code with example usage: obj.testCriticLossWithMaskedActions() without return types