code: def testMasking(self):
      batch_size = 1000
      num_state_dims = 5
      num_actions = 8
      observations = tf.random.uniform([batch_size, num_state_dims])
      time_step = ts.restart(observations, batch_size=batch_size)
      input_tensor_spec = tensor_spec.TensorSpec(
          [num_state_dims], tf.float32
      )
      action_spec = tensor_spec.BoundedTensorSpec(
          [1], tf.int32, 0, num_actions - 1
      )
  
      # We create a fixed mask here for testing purposes. Normally the mask would
      # be part of the observation.
      mask = [0, 1, 0, 1, 0, 0, 1, 0]
      np_mask = np.array(mask)
      tf_mask = tf.constant([mask for _ in range(batch_size)])
      q_network = categorical_q_network.CategoricalQNetwork(
          input_tensor_spec=input_tensor_spec,
          action_spec=action_spec,
          num_atoms=3,
          fc_layer_params=[4],
      )
      policy = categorical_q_policy.CategoricalQPolicy(
          self._time_step_spec,
          action_spec,
          q_network,
          self._min_q_value,
          self._max_q_value,
          observation_and_action_constraint_splitter=(
              lambda observation: (observation, tf_mask)
          ),
      )
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
  
      # Sample from the policy 1000 times, and ensure that actions considered
      # invalid according to the mask are never chosen.
      action_step = policy.action(time_step)
      action = self.evaluate(action_step.action)
      self.assertEqual(action.shape, (batch_size,))
      self.assertAllEqual(np_mask[action], np.ones([batch_size]))
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl import flags
  import numpy
  import tensorflow
  from tf_agents.networks import categorical_q_network
  from tf_agents.networks import network
  from tf_agents.policies import categorical_q_policy
  from tf_agents.policies import policy_saver
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import test_utils
label_desc: Use image node_id 7 for calling the CategoricalQPolicyTest obj's underlying member method code with example usage: obj.testMasking() without return types