class_docstr: Class to build Reward Prediction Policies with Boltzmann exploration.
class_docstr_tok: ["Class","to","build","Reward","Prediction","Policies","with","Boltzmann","exploration","."]
code: def _action_distribution(self, mask, predicted_rewards):
      batch_size = tf.shape(predicted_rewards)[0]
      if self._boltzmann_gumbel_exploration_constant is not None:
          logits = predicted_rewards
  
          # Apply masking if needed. Overwrite the logits for invalid actions to
          # logits.dtype.min.
          if mask is not None:
              almost_neg_inf = tf.constant(
                  logits.dtype.min, dtype=logits.dtype
              )
              logits = tf.compat.v2.where(
                  tf.cast(mask, tf.bool), logits, almost_neg_inf
              )
  
          gumbel_dist = tfp.distributions.Gumbel(loc=0.0, scale=1.0)
          gumbel_samples = gumbel_dist.sample(tf.shape(logits))
          num_samples_list_float = tf.stack(
              [
                  tf.cast(x.read_value(), tf.float32)
                  for x in self._num_samples_list
              ],
              axis=-1,
          )
          exploration_weights = tf.math.divide_no_nan(
              tf.cast(
                  self._boltzmann_gumbel_exploration_constant,
                  tf.float32,
              ),
              tf.sqrt(num_samples_list_float),
          )
          final_logits = logits + exploration_weights * gumbel_samples
          actions = tf.cast(
              tf.math.argmax(final_logits, axis=1),
              self._action_spec.dtype,
          )
          # To conform with the return type, we construct a deterministic
          # distribution here. Note that this results in the log_probability of
          # the chosen arm being 0. The true sampling probability here has no simple
          # closed-form.
          distribution = tfp.distributions.Deterministic(loc=actions)
      else:
          # Apply the temperature scaling, needed for Boltzmann exploration.
          logits = predicted_rewards / self._get_temperature_value()
  
          # Apply masking if needed. Overwrite the logits for invalid actions to
          # logits.dtype.min.
          if mask is not None:
              almost_neg_inf = tf.constant(
                  logits.dtype.min, dtype=logits.dtype
              )
              logits = tf.compat.v2.where(
                  tf.cast(mask, tf.bool), logits, almost_neg_inf
              )
  
          if self._action_offset != 0:
              distribution = shifted_categorical.ShiftedCategorical(
                  logits=logits,
                  dtype=self._action_spec.dtype,
                  shift=self._action_offset,
              )
          else:
              distribution = tfp.distributions.Categorical(
                  logits=logits, dtype=self._action_spec.dtype
              )
  
      bandit_policy_values = tf.fill(
          [batch_size, 1], policy_utilities.BanditPolicyType.BOLTZMANN
      )
      return distribution, bandit_policy_values
  
file_path: agents/tf_agents/bandits/policies/boltzmann_reward_prediction_policy.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  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 3 for calling the BoltzmannRewardPredictionPolicy obj's underlying member method code with example usage: obj._action_distribution(mask, predicted_rewards) and returns: distribution,  bandit_policy_values