class_docstr: Class to build GreedyMultiObjectiveNeuralPolicy objects.
class_docstr_tok: ["Class","to","build","GreedyMultiObjectiveNeuralPolicy","objects","."]
code: def _distribution(
      self,
      time_step: ts.TimeStep,
      policy_state: Sequence[types.TensorSpec],
  ) -> policy_step.PolicyStep:
      observation = time_step.observation
      if self.observation_and_action_constraint_splitter is not None:
          (
              observation,
              _,
          ) = self.observation_and_action_constraint_splitter(
              observation
          )
      predicted_objective_values_tensor, policy_state = self._predict(
          observation, time_step.step_type, policy_state
      )
      scalarized_reward = scalarize_objectives(
          predicted_objective_values_tensor, self._scalarizer
      )
      # Preserve static batch size values when they are available.
      batch_size = (
          tf.compat.dimension_value(scalarized_reward.shape[0])
          or tf.shape(scalarized_reward)[0]
      )
      mask = constraints.construct_mask_from_multiple_sources(
          time_step.observation,
          self._observation_and_action_constraint_splitter,
          (),
          self._expected_num_actions,
      )
  
      # Argmax.
      if mask is not None:
          actions = policy_utilities.masked_argmax(
              scalarized_reward,
              mask,
              output_type=self.action_spec.dtype,
          )
      else:
          actions = tf.argmax(
              scalarized_reward,
              axis=-1,
              output_type=self.action_spec.dtype,
          )
  
      actions += self._action_offset
  
      bandit_policy_values = tf.fill(
          [batch_size, 1], policy_utilities.BanditPolicyType.GREEDY
      )
  
      if self._accepts_per_arm_features:
          # Saving the features for the chosen action to the policy_info.
          def gather_observation(obs):
              return tf.gather(
                  params=obs, indices=actions, batch_dims=1
              )
  
          chosen_arm_features = tf.nest.map_structure(
              gather_observation,
              observation[bandit_spec_utils.PER_ARM_FEATURE_KEY],
          )
          policy_info = policy_utilities.PerArmPolicyInfo(
              log_probability=tf.zeros([batch_size], tf.float32)
              if policy_utilities.InfoFields.LOG_PROBABILITY
              in self._emit_policy_info
              else (),
              predicted_rewards_mean=(
                  predicted_objective_values_tensor
                  if policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN
                  in self._emit_policy_info
                  else ()
              ),
              multiobjective_scalarized_predicted_rewards_mean=(
                  scalarized_reward
                  if policy_utilities.InfoFields.MULTIOBJECTIVE_SCALARIZED_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=chosen_arm_features,
          )
      else:
          policy_info = policy_utilities.PolicyInfo(
              log_probability=tf.zeros([batch_size], tf.float32)
              if policy_utilities.InfoFields.LOG_PROBABILITY
              in self._emit_policy_info
              else (),
              predicted_rewards_mean=(
                  predicted_objective_values_tensor
                  if policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN
                  in self._emit_policy_info
                  else ()
              ),
              multiobjective_scalarized_predicted_rewards_mean=(
                  scalarized_reward
                  if policy_utilities.InfoFields.MULTIOBJECTIVE_SCALARIZED_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(
          tfp.distributions.Deterministic(loc=actions),
          policy_state,
          policy_info,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  from typing import List, Optional, Sequence, Text, Tuple
  import gin
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.multi_objective import multi_objective_scalarizer
  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.networks.network import Network
  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
  from tf_agents.utils import common
label_desc: Use image node_id 5 for calling the GreedyMultiObjectiveNeuralPolicy obj's underlying member method code with example usage: obj._distribution(time_step, policy_state) and returns: policy_step, tf