class_docstr: Policy that samples actions based on the FALCON algorithm.
class_docstr_tok: ["Policy","that","samples","actions","based","on","the","FALCON","algorithm","."]
code: def _action_distribution(self, mask, predicted_rewards):
      batch_size = tf.shape(predicted_rewards)[0]
      gamma = self._compute_gamma(
          mask, predicted_rewards.dtype, batch_size
      )
      # Replace predicted rewards of masked actions with -inf.
      predictions = (
          predicted_rewards
          if mask is None
          else tf.where(
              tf.cast(mask, tf.bool),
              predicted_rewards,
              -float("Inf") * tf.ones_like(predicted_rewards),
          )
      )
  
      # Get the predicted rewards of the greedy actions.
      greedy_action_predictions = tf.reshape(
          tf.reduce_max(predictions, axis=-1), shape=[-1, 1]
      )
  
      # `other_actions_probs` is a tensor shaped as [batch_size, num_actions, d]
      # that contains valid sampling probabilities for all non-greedy actions.
      # The last dimension corresponds to different gamma parameters.
      if mask is not None:
          num_allowed_actions = tf.reshape(
              self._get_number_of_allowed_actions(mask),
              [batch_size, 1, 1],
          )
      else:
          num_allowed_actions = self._get_number_of_allowed_actions(
              mask
          )
      prediction_delta = greedy_action_predictions - predictions
      other_actions_probs = tf.math.divide_no_nan(
          1.0,
          num_allowed_actions
          + tf.matmul(
              tf.expand_dims(prediction_delta, axis=-1),
              tf.expand_dims(gamma, axis=1),
          ),
      )
      # Although `predictions` has accounted for the action mask, we still need
      # to mask the action probabilities in the case of zero gamma.
      if mask is not None:
          other_actions_probs = tf.where(
              tf.repeat(
                  input=tf.expand_dims(tf.cast(mask, tf.bool), axis=-1),
                  repeats=[tf.shape(other_actions_probs)[-1]],
                  axis=2,
              ),
              other_actions_probs,
              tf.zeros_like(other_actions_probs),
          )
  
      # Get the greedy action.
      greedy_actions = tf.reshape(
          tf.argmax(
              predictions, axis=-1, output_type=self.action_spec.dtype
          ),
          [-1, 1],
      )
  
      # Compute the probabilities of sampling the greedy actions, which is
      # 1 - (the total probability of sampling other actions),
      # shaped [batch_size, d].
      greedy_action_prob = (
          1.0
          - tf.reduce_sum(other_actions_probs, axis=1)
          + tf.squeeze(
              tf.gather(
                  other_actions_probs,
                  greedy_actions,
                  axis=1,
                  batch_dims=1,
              ),
              axis=1,
          )
      )
  
      if self._max_exploration_probability_hint is not None:
          max_exploration_prob = tf.clip_by_value(
              self._max_exploration_probability_hint,
              clip_value_min=0.0,
              clip_value_max=1.0,
          )
          (
              greedy_action_prob,
              other_actions_probs,
          ) = _find_action_probabilities(
              greedy_action_prob,
              other_actions_probs,
              max_exploration_prob,
          )
      else:
          other_actions_probs = tf.squeeze(other_actions_probs, axis=2)
  
      # Compute the sampling probabilities for all actions by combining
      # `greedy_action_prob` and `other_actions_probs`.
      greedy_action_mask = tf.equal(
          tf.tile(
              [
                  tf.range(
                      self._expected_num_actions,
                      dtype=self.action_spec.dtype,
                  )
              ],
              [batch_size, 1],
          ),
          greedy_actions,
      )
      action_probs = tf.where(
          greedy_action_mask,
          tf.tile(greedy_action_prob, [1, self._expected_num_actions]),
          other_actions_probs,
      )
  
      if self._action_offset != 0:
          distribution = shifted_categorical.ShiftedCategorical(
              probs=action_probs,
              dtype=self._action_spec.dtype,
              shift=self._action_offset,
          )
      else:
          distribution = tfp.distributions.Categorical(
              probs=action_probs, dtype=self._action_spec.dtype
          )
  
      bandit_policy_values = tf.fill(
          [batch_size, 1], policy_utilities.BanditPolicyType.FALCON
      )
      return distribution, bandit_policy_values
  
file_path: agents/tf_agents/bandits/policies/falcon_reward_prediction_policy.py
imports: from typing import Iterable, Optional, Sequence, Text, Tuple
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.policies import constraints
  from tf_agents.bandits.policies import reward_prediction_base_policy
  from tf_agents.distributions import shifted_categorical
  from tf_agents.policies import utils
  from tf_agents.typing import types
label_desc: Use image node_id 7 for calling the FalconRewardPredictionPolicy obj's underlying member method code with example usage: obj._action_distribution(mask, predicted_rewards) and returns: distribution,  bandit_policy_values