code: def testMaxExplorationProbabilityHint(self, mask):
      batch_size = 1000
      # Constructs a batch of observations such that action 1 is always the greedy
      # choice, but its gap in predicted rewards from the second-best action may
      # vary.
      features = tf.stack(
          values=[
              tf.constant(1e-2, shape=[batch_size], dtype=tf.float32),
              tf.random.uniform(
                  shape=[batch_size], maxval=1.0, dtype=tf.float32
              ),
          ],
          axis=1,
      )
      if mask is None:
          observations = features
      else:
          observations = (
              features,
              tf.constant([mask] * batch_size, dtype=tf.float32),
          )
  
      time_step = ts.restart(observations, batch_size=batch_size)
  
      # Creates a policy with a low sample count, a small
      # `exploitation_coefficient` and unset `max_exploration_probability_hint`.
      # The best allowed action is expected to be chosen the most, but not at a
      # very high probability.
      policy = self._create_policy(
          use_mask=(mask is not None),
          exploitation_coefficient=0.1,
          num_samples_list=[
              tf.compat.v2.Variable(
                  int(1), dtype=tf.int32, name="num_samples_0"
              ),
              tf.compat.v2.Variable(
                  int(1), dtype=tf.int32, name="num_samples_1"
              ),
              tf.compat.v2.Variable(
                  int(1), dtype=tf.int32, name="num_samples_2"
              ),
          ],
      )
      action_step = policy.action(time_step, seed=1)
      self.evaluate(tf.compat.v1.global_variables_initializer())
      actions = self.evaluate(action_step.action)
      if mask is None or np.sum(mask) == 3:
          # All actions are allowed, action 1 is the best.
          self.assertLess(np.sum(actions == 1), 0.5 * batch_size)
          self.assertGreater(np.sum(actions == 1), np.sum(actions == 0))
          self.assertGreater(np.sum(actions == 1), np.sum(actions == 2))
      elif np.sum(mask) == 1:
          # A single action is allowed.
          only_allowed_action = np.argwhere(mask)[0]
          self.assertEqual(
              np.sum(actions == only_allowed_action), batch_size
          )
      else:
          # Two actions are allowed.
          self.assertEqual(np.sum(mask), 2)
          action_counts = [
              np.sum(actions == action) for action in range(3)
          ]
          most_chosen_action = np.argmax(action_counts)
          # The worst action 2 is expected to never be the most chosen.
          self.assertNotEqual(most_chosen_action, 2)
          self.assertLessEqual(
              action_counts[most_chosen_action], 0.7 * batch_size
          )
  
      # Creates a policy with the same sample count and
      # `exploitation_coefficient`, but sets `max_exploration_probability_hint` to
      # 5%. The best allowed action is expected to be chosen 95% of the time.
      limited_exploration_policy = self._create_policy(
          use_mask=(mask is not None),
          exploitation_coefficient=0.1,
          max_exploration_probability_hint=0.05,
          num_samples_list=[
              tf.compat.v2.Variable(
                  int(1), dtype=tf.int32, name="num_samples_0"
              ),
              tf.compat.v2.Variable(
                  int(1), dtype=tf.int32, name="num_samples_1"
              ),
              tf.compat.v2.Variable(
                  int(1), dtype=tf.int32, name="num_samples_2"
              ),
          ],
      )
      action_step = limited_exploration_policy.action(time_step, seed=1)
      self.evaluate(tf.compat.v1.global_variables_initializer())
      actions = self.evaluate(action_step.action)
      action_counts = [np.sum(actions == action) for action in range(3)]
      most_chosen_action = np.argmax(action_counts)
      # Sets the threshold to be a bit lower than 95% for test robustness against
      # random sampling.
      self.assertGreater(
          action_counts[most_chosen_action], 0.9 * batch_size
      )
      if mask is None or np.sum(mask) == 3:
          # All actions are allowed, the greedy action (1) is expected to be chosen
          # the most.
          self.assertEqual(most_chosen_action, 1)
      elif np.sum(mask) == 1:
          # A single action is allowed, and is expected to always be chosen.
          only_allowed_action = int(np.argwhere(mask)[0])
          self.assertEqual(most_chosen_action, only_allowed_action)
          self.assertEqual(
              action_counts[only_allowed_action], batch_size
          )
      else:
          # Two actions are allowed.
          self.assertEqual(np.sum(mask), 2)
          # The greedy action is action 1 if allowed, or action 0 otherwise.
          expected_greedy_action = 1 if mask[1] > 0 else 0
          self.assertEqual(most_chosen_action, expected_greedy_action)
  
imports: from typing import List, Optional, Set
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.bandits.policies import falcon_reward_prediction_policy
  from tf_agents.networks import network
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import test_utils
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 9 for calling the FalconRewardPredictionPolicyTest obj's underlying member method code with example usage: obj.testMaxExplorationProbabilityHint(mask) without return types