class_docstr: Base class to build policies based on reward predictions.
class_docstr_tok: ["Base","class","to","build","policies","based","on","reward","predictions","."]
code: def _distribution(self, time_step, policy_state):
      observation = time_step.observation
      if self.observation_and_action_constraint_splitter is not None:
          (
              observation,
              _,
          ) = self.observation_and_action_constraint_splitter(
              observation
          )
  
      predictions, policy_state = self._reward_network(
          observation, time_step.step_type, policy_state
      )
  
      if isinstance(
          self._reward_network,
          heteroscedastic_q_network.HeteroscedasticQNetwork,
      ):
          predicted_reward_values = predictions.q_value_logits
      else:
          predicted_reward_values = predictions
  
      predicted_reward_values.shape.with_rank_at_least(2)
      predicted_reward_values.shape.with_rank_at_most(3)
      if (
          predicted_reward_values.shape[-1] is not None
          and predicted_reward_values.shape[-1]
          != self._expected_num_actions
      ):
          raise ValueError(
              "The number of actions ({}) does not match the reward_network output"
              " size ({}).".format(
                  self._expected_num_actions,
                  predicted_reward_values.shape[1],
              )
          )
  
      mask = constr.construct_mask_from_multiple_sources(
          time_step.observation,
          self._observation_and_action_constraint_splitter,
          self._constraints,
          self._expected_num_actions,
      )
      (
          action_distribution,
          bandit_policy_values,
      ) = self._action_distribution(mask, predicted_reward_values)
  
      if self._accepts_per_arm_features:
          # The actual action sampling hasn't happened yet, so we leave
          # `log_probability` empty and set `chosen_arm_features` to dummy values of
          # all zeros. We need to save dummy chosen arm features to make the
          # returned policy step have the same structure as the policy state spec.
          dummy_chosen_arm_features = tf.nest.map_structure(
              lambda obs: tf.zeros_like(obs[:, 0, ...]),
              time_step.observation[
                  bandit_spec_utils.PER_ARM_FEATURE_KEY
              ],
          )
          policy_info = policy_utilities.PerArmPolicyInfo(
              log_probability=(),
              predicted_rewards_mean=(
                  predicted_reward_values
                  if policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN
                  in self._emit_policy_info
                  else ()
              ),
              bandit_policy_type=(
                  bandit_policy_values
                  if policy_utilities.InfoFields.BANDIT_POLICY_TYPE
                  in self._emit_policy_info
                  else ()
              ),
              chosen_arm_features=dummy_chosen_arm_features,
          )
      else:
          # The actual action sampling hasn't happened yet, so we leave
          # `log_probability` empty.
          policy_info = policy_utilities.PolicyInfo(
              log_probability=(),
              predicted_rewards_mean=(
                  predicted_reward_values
                  if policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN
                  in self._emit_policy_info
                  else ()
              ),
              bandit_policy_type=(
                  bandit_policy_values
                  if policy_utilities.InfoFields.BANDIT_POLICY_TYPE
                  in self._emit_policy_info
                  else ()
              ),
          )
  
      return policy_step.PolicyStep(
          action_distribution, policy_state, policy_info
      )
  
imports: import abc
  from typing import Iterable, Optional, Text, Tuple
  import gin
  import six
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.networks import heteroscedastic_q_network
  from tf_agents.bandits.policies import constraints
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import tf_policy
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
label_desc: Use image node_id 7 for calling the RewardPredictionBasePolicy obj's underlying member method code with example usage: obj._distribution(time_step, policy_state) and returns: policy_step