class_docstr: Class to build GreedyNNPredictionPolicies.
docstr: Returns the action with largest predicted reward.
docstr_tok: ["Returns","the","action","with","largest","predicted","reward","."]
code: def _action_distribution(self, mask, predicted_rewards):
      """Returns the action with largest predicted reward."""
      # Argmax.
      batch_size = tf.shape(predicted_rewards)[0]
      if mask is not None:
          actions = policy_utilities.masked_argmax(
              predicted_rewards,
              mask,
              output_type=self.action_spec.dtype,
          )
      else:
          actions = tf.argmax(
              predicted_rewards,
              axis=-1,
              output_type=self.action_spec.dtype,
          )
  
      actions += self._action_offset
  
      bandit_policy_values = tf.fill(
          [batch_size, 1], policy_utilities.BanditPolicyType.GREEDY
      )
      return (
          tfp.distributions.Deterministic(loc=actions),
          bandit_policy_values,
      )
  
file_path: agents/tf_agents/bandits/policies/greedy_reward_prediction_policy.py
imports: import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.policies import reward_prediction_base_policy
  from tf_agents.policies import utils
label_desc: Use image node_id 1 for calling the GreedyRewardPredictionPolicy obj's underlying member method code with example usage: obj._action_distribution(mask, predicted_rewards) and returns: tfp,  bandit_policy_values